From bc1d1a0574de1d86ec2332f1467383c53b53d3d5 Mon Sep 17 00:00:00 2001 From: Kevin Flansburg Date: Thu, 29 Jan 2026 14:09:46 -0800 Subject: [PATCH] Implement `run_bytes` method on AI binding for models that return binary data. The method returns a `ByteStream` which implements `Stream`. This can be passed directly to the response or collected to `Vec`, as shown in the docstring. Also added an example worker to demonstrate / test this behavior. Closes #783 --- Cargo.lock | 9 + examples/ai-image/.gitignore | 2 + examples/ai-image/Cargo.lock | 759 ++++++++++++++++++++++++++++++++ examples/ai-image/Cargo.toml | 15 + examples/ai-image/src/lib.rs | 84 ++++ examples/ai-image/wrangler.toml | 9 + worker/src/ai.rs | 106 +++++ worker/src/streams.rs | 8 + 8 files changed, 992 insertions(+) create mode 100644 examples/ai-image/.gitignore create mode 100644 examples/ai-image/Cargo.lock create mode 100644 examples/ai-image/Cargo.toml create mode 100644 examples/ai-image/src/lib.rs create mode 100644 examples/ai-image/wrangler.toml diff --git a/Cargo.lock b/Cargo.lock index e4372f752..dc1c1e236 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -28,6 +28,15 @@ dependencies = [ "memchr", ] +[[package]] +name = "ai-image-on-workers" +version = "0.1.0" +dependencies = [ + "futures-util", + "serde", + "worker 0.7.4", +] + [[package]] name = "android_system_properties" version = "0.1.5" diff --git a/examples/ai-image/.gitignore b/examples/ai-image/.gitignore new file mode 100644 index 000000000..22a146e87 --- /dev/null +++ b/examples/ai-image/.gitignore @@ -0,0 +1,2 @@ +target +dist \ No newline at end of file diff --git a/examples/ai-image/Cargo.lock b/examples/ai-image/Cargo.lock new file mode 100644 index 000000000..1b8444a01 --- /dev/null +++ b/examples/ai-image/Cargo.lock @@ -0,0 +1,759 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "ai-image-on-workers" +version = "0.1.0" +dependencies = [ + "futures-util", + "serde", + "worker", +] + +[[package]] +name = "async-trait" +version = "0.1.89" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "autocfg" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" + +[[package]] +name = "bumpalo" +version = "3.19.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5dd9dc738b7a8311c7ade152424974d8115f2cdad61e8dab8dac9f2362298510" + +[[package]] +name = "bytes" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b35204fbdc0b3f4446b89fc1ac2cf84a8a68971995d0bf2e925ec7cd960f9cb3" + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "chrono" +version = "0.4.43" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fac4744fb15ae8337dc853fee7fb3f4e48c0fbaa23d0afe49c447b4fab126118" +dependencies = [ + "js-sys", + "num-traits", + "wasm-bindgen", +] + +[[package]] +name = "displaydoc" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "form_urlencoded" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb4cb245038516f5f85277875cdaa4f7d2c9a0fa0468de06ed190163b1581fcf" +dependencies = [ + "percent-encoding", +] + +[[package]] +name = "futures-channel" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2dff15bf788c671c1934e366d07e30c1814a8ef514e1af724a602e8a2fbe1b10" +dependencies = [ + "futures-core", +] + +[[package]] +name = "futures-core" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05f29059c0c2090612e8d742178b0580d2dc940c837851ad723096f87af6663e" + +[[package]] +name = "futures-io" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e5c1b78ca4aae1ac06c48a526a655760685149f0d465d21f37abfe57ce075c6" + +[[package]] +name = "futures-macro" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "162ee34ebcb7c64a8abebc059ce0fee27c2262618d7b60ed8faf72fef13c3650" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "futures-sink" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e575fab7d1e0dcb8d0c7bcf9a63ee213816ab51902e6d244a95819acacf1d4f7" + +[[package]] +name = "futures-task" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f90f7dce0722e95104fcb095585910c0977252f286e354b5e3bd38902cd99988" + +[[package]] +name = "futures-util" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fa08315bb612088cc391249efdc3bc77536f16c91f6cf495e6fbe85b20a4a81" +dependencies = [ + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "pin-utils", + "slab", +] + +[[package]] +name = "http" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3ba2a386d7f85a81f119ad7498ebe444d2e22c2af0b86b069416ace48b3311a" +dependencies = [ + "bytes", + "itoa", +] + +[[package]] +name = "http-body" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1efedce1fb8e6913f23e0c92de8e62cd5b772a67e7b3946df930a62566c93184" +dependencies = [ + "bytes", + "http", +] + +[[package]] +name = "icu_collections" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c6b649701667bbe825c3b7e6388cb521c23d88644678e83c0c4d0a621a34b43" +dependencies = [ + "displaydoc", + "potential_utf", + "yoke", + "zerofrom", + "zerovec", +] + +[[package]] +name = "icu_locale_core" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "edba7861004dd3714265b4db54a3c390e880ab658fec5f7db895fae2046b5bb6" +dependencies = [ + "displaydoc", + "litemap", + "tinystr", + "writeable", + "zerovec", +] + +[[package]] +name = "icu_normalizer" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f6c8828b67bf8908d82127b2054ea1b4427ff0230ee9141c54251934ab1b599" +dependencies = [ + "icu_collections", + "icu_normalizer_data", + "icu_properties", + "icu_provider", + "smallvec", + "zerovec", +] + +[[package]] +name = "icu_normalizer_data" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7aedcccd01fc5fe81e6b489c15b247b8b0690feb23304303a9e560f37efc560a" + +[[package]] +name = "icu_properties" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "020bfc02fe870ec3a66d93e677ccca0562506e5872c650f893269e08615d74ec" +dependencies = [ + "icu_collections", + "icu_locale_core", + "icu_properties_data", + "icu_provider", + "zerotrie", + "zerovec", +] + +[[package]] +name = "icu_properties_data" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "616c294cf8d725c6afcd8f55abc17c56464ef6211f9ed59cccffe534129c77af" + +[[package]] +name = "icu_provider" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85962cf0ce02e1e0a629cc34e7ca3e373ce20dda4c4d7294bbd0bf1fdb59e614" +dependencies = [ + "displaydoc", + "icu_locale_core", + "writeable", + "yoke", + "zerofrom", + "zerotrie", + "zerovec", +] + +[[package]] +name = "idna" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b0875f23caa03898994f6ddc501886a45c7d3d62d04d2d90788d47be1b1e4de" +dependencies = [ + "idna_adapter", + "smallvec", + "utf8_iter", +] + +[[package]] +name = "idna_adapter" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3acae9609540aa318d1bc588455225fb2085b9ed0c4f6bd0d9d5bcd86f1a0344" +dependencies = [ + "icu_normalizer", + "icu_properties", +] + +[[package]] +name = "itoa" +version = "1.0.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92ecc6618181def0457392ccd0ee51198e065e016d1d527a7ac1b6dc7c1f09d2" + +[[package]] +name = "js-sys" +version = "0.3.85" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8c942ebf8e95485ca0d52d97da7c5a2c387d0e7f0ba4c35e93bfcaee045955b3" +dependencies = [ + "once_cell", + "wasm-bindgen", +] + +[[package]] +name = "litemap" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6373607a59f0be73a39b6fe456b8192fcc3585f602af20751600e974dd455e77" + +[[package]] +name = "matchit" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e7465ac9959cc2b1404e8e2367b43684a6d13790fe23056cc8c6c5a6b7bcb94" + +[[package]] +name = "memchr" +version = "2.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f52b00d39961fc5b2736ea853c9cc86238e165017a493d1d5c8eac6bdc4cc273" + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + +[[package]] +name = "once_cell" +version = "1.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" + +[[package]] +name = "percent-encoding" +version = "2.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" + +[[package]] +name = "pin-project" +version = "1.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "677f1add503faace112b9f1373e43e9e054bfdd22ff1a63c1bc485eaec6a6a8a" +dependencies = [ + "pin-project-internal", +] + +[[package]] +name = "pin-project-internal" +version = "1.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e918e4ff8c4549eb882f14b3a4bc8c8bc93de829416eacf579f1207a8fbf861" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b3cff922bd51709b605d9ead9aa71031d81447142d828eb4a6eba76fe619f9b" + +[[package]] +name = "pin-utils" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" + +[[package]] +name = "potential_utf" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b73949432f5e2a09657003c25bca5e19a0e9c84f8058ca374f49e0ebe605af77" +dependencies = [ + "zerovec", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "21b2ebcf727b7760c461f091f9f0f539b77b8e87f2fd88131e7f1b433b3cece4" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "rustversion" +version = "1.0.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" + +[[package]] +name = "ryu" +version = "1.0.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a50f4cf475b65d88e057964e0e9bb1f0aa9bbb2036dc65c64596b42932536984" + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde-wasm-bindgen" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8302e169f0eddcc139c70f139d19d6467353af16f9fce27e8c30158036a1e16b" +dependencies = [ + "js-sys", + "serde", + "wasm-bindgen", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.149" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "serde_urlencoded" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3491c14715ca2294c4d6a88f15e84739788c1d030eed8c110436aafdaa2f3fd" +dependencies = [ + "form_urlencoded", + "itoa", + "ryu", + "serde", +] + +[[package]] +name = "slab" +version = "0.4.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a2ae44ef20feb57a68b23d846850f861394c2e02dc425a50098ae8c90267589" + +[[package]] +name = "smallvec" +version = "1.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" + +[[package]] +name = "stable_deref_trait" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" + +[[package]] +name = "syn" +version = "2.0.114" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d4d107df263a3013ef9b1879b0df87d706ff80f65a86ea879bd9c31f9b307c2a" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "synstructure" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tinystr" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42d3e9c45c09de15d06dd8acf5f4e0e399e85927b7f00711024eb7ae10fa4869" +dependencies = [ + "displaydoc", + "zerovec", +] + +[[package]] +name = "tokio" +version = "1.49.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72a2903cd7736441aac9df9d7688bd0ce48edccaadf181c3b90be801e81d3d86" +dependencies = [ + "pin-project-lite", +] + +[[package]] +name = "unicode-ident" +version = "1.0.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9312f7c4f6ff9069b165498234ce8be658059c6728633667c526e27dc2cf1df5" + +[[package]] +name = "url" +version = "2.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff67a8a4397373c3ef660812acab3268222035010ab8680ec4215f38ba3d0eed" +dependencies = [ + "form_urlencoded", + "idna", + "percent-encoding", + "serde", +] + +[[package]] +name = "utf8_iter" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" + +[[package]] +name = "wasm-bindgen" +version = "0.2.108" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "64024a30ec1e37399cf85a7ffefebdb72205ca1c972291c51512360d90bd8566" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-futures" +version = "0.4.58" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "70a6e77fd0ae8029c9ea0063f87c46fde723e7d887703d74ad2616d792e51e6f" +dependencies = [ + "cfg-if", + "futures-util", + "js-sys", + "once_cell", + "wasm-bindgen", + "web-sys", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.108" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "008b239d9c740232e71bd39e8ef6429d27097518b6b30bdf9086833bd5b6d608" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.108" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5256bae2d58f54820e6490f9839c49780dff84c65aeab9e772f15d5f0e913a55" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.108" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f01b580c9ac74c8d8f0c0e4afb04eeef2acf145458e52c03845ee9cd23e3d12" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "wasm-streams" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d1ec4f6517c9e11ae630e200b2b65d193279042e28edd4a2cda233e46670bbb" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + +[[package]] +name = "web-sys" +version = "0.3.85" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "312e32e551d92129218ea9a2452120f4aabc03529ef03e4d0d82fb2780608598" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "worker" +version = "0.7.4" +dependencies = [ + "async-trait", + "bytes", + "chrono", + "futures-channel", + "futures-util", + "http", + "http-body", + "js-sys", + "matchit", + "pin-project", + "serde", + "serde-wasm-bindgen", + "serde_json", + "serde_urlencoded", + "tokio", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "wasm-streams", + "web-sys", + "worker-macros", + "worker-sys", +] + +[[package]] +name = "worker-macros" +version = "0.7.4" +dependencies = [ + "async-trait", + "proc-macro2", + "quote", + "syn", + "wasm-bindgen", + "wasm-bindgen-futures", + "wasm-bindgen-macro-support", + "worker-sys", +] + +[[package]] +name = "worker-sys" +version = "0.7.4" +dependencies = [ + "cfg-if", + "js-sys", + "wasm-bindgen", + "web-sys", +] + +[[package]] +name = "writeable" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9edde0db4769d2dc68579893f2306b26c6ecfbe0ef499b013d731b7b9247e0b9" + +[[package]] +name = "yoke" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72d6e5c6afb84d73944e5cedb052c4680d5657337201555f9f2a16b7406d4954" +dependencies = [ + "stable_deref_trait", + "yoke-derive", + "zerofrom", +] + +[[package]] +name = "yoke-derive" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b659052874eb698efe5b9e8cf382204678a0086ebf46982b79d6ca3182927e5d" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zerofrom" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50cc42e0333e05660c3587f3bf9d0478688e15d870fab3346451ce7f8c9fbea5" +dependencies = [ + "zerofrom-derive", +] + +[[package]] +name = "zerofrom-derive" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d71e5d6e06ab090c67b5e44993ec16b72dcbaabc526db883a360057678b48502" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zerotrie" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a59c17a5562d507e4b54960e8569ebee33bee890c70aa3fe7b97e85a9fd7851" +dependencies = [ + "displaydoc", + "yoke", + "zerofrom", +] + +[[package]] +name = "zerovec" +version = "0.11.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c28719294829477f525be0186d13efa9a3c602f7ec202ca9e353d310fb9a002" +dependencies = [ + "yoke", + "zerofrom", + "zerovec-derive", +] + +[[package]] +name = "zerovec-derive" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eadce39539ca5cb3985590102671f2567e659fca9666581ad3411d59207951f3" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zmij" +version = "1.0.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "02aae0f83f69aafc94776e879363e9771d7ecbffe2c7fbb6c14c5e00dfe88439" diff --git a/examples/ai-image/Cargo.toml b/examples/ai-image/Cargo.toml new file mode 100644 index 000000000..c494c407d --- /dev/null +++ b/examples/ai-image/Cargo.toml @@ -0,0 +1,15 @@ +[package] +name = "ai-image-on-workers" +version = "0.1.0" +edition = "2021" + +[package.metadata.release] +release = false + +[lib] +crate-type = ["cdylib", "rlib"] + +[dependencies] +futures-util = "0.3" +serde = { version = "1.0", features = ["derive"] } +worker = { path = "../../worker" } diff --git a/examples/ai-image/src/lib.rs b/examples/ai-image/src/lib.rs new file mode 100644 index 000000000..c2ce41865 --- /dev/null +++ b/examples/ai-image/src/lib.rs @@ -0,0 +1,84 @@ +//! Example demonstrating AI image generation using Workers AI. +//! +//! This example shows how to use the `run_bytes` method to generate images +//! using Stable Diffusion and return them as a PNG response. + +use futures_util::StreamExt; +use serde::Serialize; +use worker::{event, Env, Request, Response, Result, Router, Url}; + +#[derive(Serialize)] +struct ImageGenRequest { + prompt: String, +} + +#[event(fetch)] +async fn main(req: Request, env: Env, _ctx: worker::Context) -> Result { + let router = Router::new(); + + router + .get_async("/stream", |req, ctx| async move { + // This endpoint streams the image directly without buffering + generate_image_stream(req, ctx.env).await + }) + .get_async("/buffered", |req, ctx| async move { + // This endpoint buffers the entire image before responding + generate_image_buffered(req, ctx.env).await + }) + .run(req, env) + .await +} + +/// Stream the image directly to the response without buffering. +/// This is more memory-efficient for large responses. +async fn generate_image_stream(req: Request, env: Env) -> Result { + let url = req.url()?; + let prompt = get_prompt_from_url(&url).unwrap_or_else(|| "a beautiful sunset".to_string()); + + let ai = env.ai("AI")?; + let request = ImageGenRequest { prompt }; + + // Get the ByteStream from the AI model + let stream = ai + .run_bytes("@cf/stabilityai/stable-diffusion-xl-base-1.0", &request) + .await?; + + // Stream directly to the response without buffering in memory + let mut response = Response::from_stream(stream)?; + response.headers_mut().set("Content-Type", "image/png")?; + + Ok(response) +} + +/// Buffer the entire image before responding. +/// This allows you to inspect or modify the bytes before sending. +async fn generate_image_buffered(req: Request, env: Env) -> Result { + let url = req.url()?; + let prompt = get_prompt_from_url(&url).unwrap_or_else(|| "a beautiful sunset".to_string()); + + let ai = env.ai("AI")?; + let request = ImageGenRequest { prompt }; + + // Get the ByteStream from the AI model + let mut stream = ai + .run_bytes("@cf/stabilityai/stable-diffusion-xl-base-1.0", &request) + .await?; + + // Collect all chunks into a Vec + let mut image_bytes = Vec::new(); + while let Some(chunk) = stream.next().await { + image_bytes.extend_from_slice(&chunk?); + } + + // Return the image with appropriate content type + let mut response = Response::from_bytes(image_bytes)?; + response.headers_mut().set("Content-Type", "image/png")?; + + Ok(response) +} + +fn get_prompt_from_url(url: &Url) -> Option { + url.query_pairs() + .find(|(key, _)| key == "prompt") + .map(|(_, value)| value.into_owned()) +} diff --git a/examples/ai-image/wrangler.toml b/examples/ai-image/wrangler.toml new file mode 100644 index 000000000..591bfe6f7 --- /dev/null +++ b/examples/ai-image/wrangler.toml @@ -0,0 +1,9 @@ +name = "ai-image-worker" +main = "build/worker/shim.mjs" +compatibility_date = "2026-01-01" + +[build] +command = "cargo install worker-build@^0.7 && worker-build --release" + +[ai] +binding = "AI" \ No newline at end of file diff --git a/worker/src/ai.rs b/worker/src/ai.rs index da2bdbdff..1679272f4 100644 --- a/worker/src/ai.rs +++ b/worker/src/ai.rs @@ -1,9 +1,11 @@ +use crate::streams::ByteStream; use crate::{env::EnvBinding, send::SendFuture}; use crate::{Error, Result}; use serde::de::DeserializeOwned; use serde::Serialize; use wasm_bindgen::{JsCast, JsValue}; use wasm_bindgen_futures::JsFuture; +use web_sys::ReadableStream; use worker_sys::Ai as AiSys; /// Enables access to Workers AI functionality. @@ -28,6 +30,110 @@ impl Ai { Err(err) => Err(Error::from(err)), } } + + /// Execute a Workers AI operation that returns binary data as a [`ByteStream`]. + /// + /// This method is designed for AI models that return raw bytes, such as: + /// - Image generation models (e.g., Stable Diffusion) + /// - Text-to-speech models + /// - Any other model that returns binary output + /// + /// The returned [`ByteStream`] implements [`Stream`](futures_util::Stream) and can be: + /// - Streamed directly to a [`Response`] using [`Response::from_stream`] + /// - Collected into a `Vec` by iterating over the chunks + /// + /// # Examples + /// + /// ## Streaming directly to a response (recommended) + /// + /// This approach is more memory-efficient as it doesn't buffer the entire + /// response in memory: + /// + /// ```ignore + /// use worker::*; + /// use serde::Serialize; + /// + /// #[derive(Serialize)] + /// struct ImageGenRequest { + /// prompt: String, + /// } + /// + /// async fn generate_image(env: &Env) -> Result { + /// let ai = env.ai("AI")?; + /// let request = ImageGenRequest { + /// prompt: "a beautiful sunset".to_string(), + /// }; + /// let stream = ai.run_bytes( + /// "@cf/stabilityai/stable-diffusion-xl-base-1.0", + /// &request + /// ).await?; + /// + /// // Stream directly to the response + /// let mut response = Response::from_stream(stream)?; + /// response.headers_mut().set("Content-Type", "image/png")?; + /// Ok(response) + /// } + /// ``` + /// + /// ## Collecting into bytes + /// + /// Use this approach if you need to inspect or modify the bytes before sending: + /// + /// ```ignore + /// use worker::*; + /// use serde::Serialize; + /// use futures_util::StreamExt; + /// + /// #[derive(Serialize)] + /// struct ImageGenRequest { + /// prompt: String, + /// } + /// + /// async fn generate_image(env: &Env) -> Result { + /// let ai = env.ai("AI")?; + /// let request = ImageGenRequest { + /// prompt: "a beautiful sunset".to_string(), + /// }; + /// let mut stream = ai.run_bytes( + /// "@cf/stabilityai/stable-diffusion-xl-base-1.0", + /// &request + /// ).await?; + /// + /// // Collect all chunks into a Vec + /// let mut bytes = Vec::new(); + /// while let Some(chunk) = stream.next().await { + /// bytes.extend_from_slice(&chunk?); + /// } + /// + /// let mut response = Response::from_bytes(bytes)?; + /// response.headers_mut().set("Content-Type", "image/png")?; + /// Ok(response) + /// } + /// ``` + pub async fn run_bytes( + &self, + model: impl AsRef, + input: T, + ) -> Result { + let fut = SendFuture::new(JsFuture::from( + self.0 + .run(model.as_ref(), serde_wasm_bindgen::to_value(&input)?), + )); + match fut.await { + Ok(output) => { + if output.is_instance_of::() { + let stream = ReadableStream::unchecked_from_js(output); + Ok(ByteStream::from(stream)) + } else { + Err(Error::RustError( + "AI model did not return binary data. Use run() for non-binary responses." + .into(), + )) + } + } + Err(err) => Err(Error::from(err)), + } + } } unsafe impl Sync for Ai {} diff --git a/worker/src/streams.rs b/worker/src/streams.rs index 90f982fdd..388adba46 100644 --- a/worker/src/streams.rs +++ b/worker/src/streams.rs @@ -20,6 +20,14 @@ pub struct ByteStream { pub(crate) inner: IntoStream<'static>, } +impl From for ByteStream { + fn from(stream: ReadableStream) -> Self { + Self { + inner: wasm_streams::ReadableStream::from_raw(stream.unchecked_into()).into_stream(), + } + } +} + impl Stream for ByteStream { type Item = Result>;