From dab40c993160914bfa9144fe72e9776c46186863 Mon Sep 17 00:00:00 2001 From: Rich Date: Wed, 29 Nov 2023 16:32:54 -0800 Subject: [PATCH] wip: PoC of inline bot. --- Cargo.lock | 756 ++++++++++++++++++ crates/stable-diffusion-bot/Cargo.toml | 2 + .../src/bot/handlers/image.rs | 108 ++- crates/stable-diffusion-bot/src/bot/mod.rs | 62 +- crates/stable-diffusion-bot/src/main.rs | 13 + 5 files changed, 926 insertions(+), 15 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 555b8ab..9a03120 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -183,6 +183,384 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa" +[[package]] +name = "aws-config" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "80c950a809d39bc9480207cb1cfc879ace88ea7e3a4392a8e9999e45d6e5692e" +dependencies = [ + "aws-credential-types", + "aws-http", + "aws-runtime", + "aws-sdk-sso", + "aws-sdk-ssooidc", + "aws-sdk-sts", + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-json", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-types", + "aws-types", + "bytes", + "fastrand", + "hex", + "http", + "hyper", + "ring", + "time", + "tokio", + "tracing", + "zeroize", +] + +[[package]] +name = "aws-credential-types" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8c1317e1a3514b103cf7d5828bbab3b4d30f56bd22d684f8568bc51b6cfbbb1c" +dependencies = [ + "aws-smithy-async", + "aws-smithy-runtime-api", + "aws-smithy-types", + "zeroize", +] + +[[package]] +name = "aws-http" +version = "0.60.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "361c4310fdce94328cc2d1ca0c8a48c13f43009c61d3367585685a50ca8c66b6" +dependencies = [ + "aws-smithy-runtime-api", + "aws-smithy-types", + "aws-types", + "bytes", + "http", + "http-body", + "pin-project-lite", + "tracing", +] + +[[package]] +name = "aws-runtime" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ed7ef604a15fd0d4d9e43701295161ea6b504b63c44990ead352afea2bc15e9" +dependencies = [ + "aws-credential-types", + "aws-http", + "aws-sigv4", + "aws-smithy-async", + "aws-smithy-eventstream", + "aws-smithy-http", + "aws-smithy-runtime-api", + "aws-smithy-types", + "aws-types", + "fastrand", + "http", + "percent-encoding", + "tracing", + "uuid", +] + +[[package]] +name = "aws-sdk-s3" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dcafc2fe52cc30b2d56685e2fa6a879ba50d79704594852112337a472ddbd24" +dependencies = [ + "aws-credential-types", + "aws-http", + "aws-runtime", + "aws-sigv4", + "aws-smithy-async", + "aws-smithy-checksums", + "aws-smithy-eventstream", + "aws-smithy-http", + "aws-smithy-json", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-types", + "aws-smithy-xml", + "aws-types", + "bytes", + "http", + "http-body", + "once_cell", + "percent-encoding", + "regex", + "tracing", + "url", +] + +[[package]] +name = "aws-sdk-sso" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0619ab97a5ca8982e7de073cdc66f93e5f6a1b05afc09e696bec1cb3607cd4df" +dependencies = [ + "aws-credential-types", + "aws-http", + "aws-runtime", + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-json", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-types", + "aws-types", + "bytes", + "http", + "regex", + "tracing", +] + +[[package]] +name = "aws-sdk-ssooidc" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f04b9f5474cc0f35d829510b2ec8c21e352309b46bf9633c5a81fb9321e9b1c7" +dependencies = [ + "aws-credential-types", + "aws-http", + "aws-runtime", + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-json", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-types", + "aws-types", + "bytes", + "http", + "regex", + "tracing", +] + +[[package]] +name = "aws-sdk-sts" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5700da387716ccfc30b27f44b008f457e1baca5b0f05b6b95455778005e3432a" +dependencies = [ + "aws-credential-types", + "aws-http", + "aws-runtime", + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-json", + "aws-smithy-query", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-types", + "aws-smithy-xml", + "aws-types", + "http", + "regex", + "tracing", +] + +[[package]] +name = "aws-sigv4" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "380adcc8134ad8bbdfeb2ace7626a869914ee266322965276cbc54066186d236" +dependencies = [ + "aws-credential-types", + "aws-smithy-eventstream", + "aws-smithy-http", + "aws-smithy-runtime-api", + "aws-smithy-types", + "bytes", + "crypto-bigint 0.5.5", + "form_urlencoded", + "hex", + "hmac", + "http", + "once_cell", + "p256", + "percent-encoding", + "regex", + "ring", + "sha2", + "subtle", + "time", + "tracing", + "zeroize", +] + +[[package]] +name = "aws-smithy-async" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3e37ca17d25fe1e210b6d4bdf59b81caebfe99f986201a1228cb5061233b4b13" +dependencies = [ + "futures-util", + "pin-project-lite", + "tokio", +] + +[[package]] +name = "aws-smithy-checksums" +version = "0.60.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c5a373ec01aede3dd066ec018c1bc4e8f5dd11b2c11c59c8eef1a5c68101f397" +dependencies = [ + "aws-smithy-http", + "aws-smithy-types", + "bytes", + "crc32c", + "crc32fast", + "hex", + "http", + "http-body", + "md-5", + "pin-project-lite", + "sha1", + "sha2", + "tracing", +] + +[[package]] +name = "aws-smithy-eventstream" +version = "0.60.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c669e1e5fc0d79561bf7a122b118bd50c898758354fe2c53eb8f2d31507cbc3" +dependencies = [ + "aws-smithy-types", + "bytes", + "crc32fast", +] + +[[package]] +name = "aws-smithy-http" +version = "0.60.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b1de8aee22f67de467b2e3d0dd0fb30859dc53f579a63bd5381766b987db644" +dependencies = [ + "aws-smithy-eventstream", + "aws-smithy-runtime-api", + "aws-smithy-types", + "bytes", + "bytes-utils", + "futures-core", + "http", + "http-body", + "once_cell", + "percent-encoding", + "pin-project-lite", + "pin-utils", + "tracing", +] + +[[package]] +name = "aws-smithy-json" +version = "0.60.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a46dd338dc9576d6a6a5b5a19bd678dcad018ececee11cf28ecd7588bd1a55c" +dependencies = [ + "aws-smithy-types", +] + +[[package]] +name = "aws-smithy-query" +version = "0.60.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "feb5b8c7a86d4b6399169670723b7e6f21a39fc833a30f5c5a2f997608178129" +dependencies = [ + "aws-smithy-types", + "urlencoding", +] + +[[package]] +name = "aws-smithy-runtime" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "273479291efc55e7b0bce985b139d86b6031adb8e50f65c1f712f20ba38f6388" +dependencies = [ + "aws-smithy-async", + "aws-smithy-http", + "aws-smithy-runtime-api", + "aws-smithy-types", + "bytes", + "fastrand", + "h2", + "http", + "http-body", + "hyper", + "hyper-rustls", + "once_cell", + "pin-project-lite", + "pin-utils", + "rustls", + "tokio", + "tracing", +] + +[[package]] +name = "aws-smithy-runtime-api" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6cebff0d977b6b6feed2fd07db52aac58ba3ccaf26cdd49f1af4add5061bef9" +dependencies = [ + "aws-smithy-async", + "aws-smithy-types", + "bytes", + "http", + "pin-project-lite", + "tokio", + "tracing", + "zeroize", +] + +[[package]] +name = "aws-smithy-types" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7f48b3f27ddb40ab19892a5abda331f403e3cb877965e4e51171447807104af" +dependencies = [ + "base64-simd", + "bytes", + "bytes-utils", + "futures-core", + "http", + "http-body", + "itoa", + "num-integer", + "pin-project-lite", + "pin-utils", + "ryu", + "serde", + "time", + "tokio", + "tokio-util", +] + +[[package]] +name = "aws-smithy-xml" +version = "0.60.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ec40d74a67fd395bc3f6b4ccbdf1543672622d905ef3f979689aea5b730cb95" +dependencies = [ + "xmlparser", +] + +[[package]] +name = "aws-types" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8403fc56b1f3761e8efe45771ddc1165e47ec3417c68e68a4519b5cb030159ca" +dependencies = [ + "aws-credential-types", + "aws-smithy-async", + "aws-smithy-runtime-api", + "aws-smithy-types", + "http", + "rustc_version", + "tracing", +] + [[package]] name = "backtrace" version = "0.3.69" @@ -198,6 +576,12 @@ dependencies = [ "rustc-demangle", ] +[[package]] +name = "base16ct" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "349a06037c7bf932dd7e7d1f653678b2038b9ad46a74102f1fc7bd7872678cce" + [[package]] name = "base64" version = "0.13.1" @@ -210,6 +594,22 @@ version = "0.21.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "35636a1494ede3b646cc98f74f8e62c773a38a659ebc777a2cf26b9b74171df9" +[[package]] +name = "base64-simd" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "339abbe78e73178762e23bea9dfd08e697eb3f3301cd4be981c0f78ba5859195" +dependencies = [ + "outref", + "vsimd", +] + +[[package]] +name = "base64ct" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8c3c1a368f70d6cf7302d78f8f7093da241fb8e8807c05cc9e51a125895a6d5b" + [[package]] name = "bitflags" version = "1.3.2" @@ -255,6 +655,16 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a2bd12c1caf447e69cd4528f47f94d203fd2582878ecb9e9465484c4148a8223" +[[package]] +name = "bytes-utils" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7dafe3a8757b027e2be6e4e5601ed563c55989fcf1546e933c66c8eb3a058d35" +dependencies = [ + "bytes", + "either", +] + [[package]] name = "cc" version = "1.0.83" @@ -329,6 +739,12 @@ version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "acbf1af155f9b9ef647e42cdc158db4b64a1b61f743629225fde6f3e0be2a7c7" +[[package]] +name = "const-oid" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "28c122c3980598d243d63d9a704629a2d748d101f278052ff068be5a4423ab6f" + [[package]] name = "convert_case" version = "0.4.0" @@ -360,6 +776,24 @@ dependencies = [ "libc", ] +[[package]] +name = "crc32c" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8f48d60e5b4d2c53d5c2b1d8a58c849a70ae5e5509b08a48d047e3b65714a74" +dependencies = [ + "rustc_version", +] + +[[package]] +name = "crc32fast" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b540bd8bc810d3885c6ea91e2018302f68baba2129ab3e88f32389ee9370880d" +dependencies = [ + "cfg-if", +] + [[package]] name = "crossbeam-queue" version = "0.3.8" @@ -379,6 +813,28 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "crypto-bigint" +version = "0.4.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ef2b4b23cddf68b89b8f8069890e8c270d54e2d5fe1b143820234805e4cb17ef" +dependencies = [ + "generic-array", + "rand_core", + "subtle", + "zeroize", +] + +[[package]] +name = "crypto-bigint" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0dc92fb57ca44df6db8059111ab3af99a63d5d0f8375d9972e319a379c6bab76" +dependencies = [ + "rand_core", + "subtle", +] + [[package]] name = "crypto-common" version = "0.1.6" @@ -459,6 +915,16 @@ dependencies = [ "syn 2.0.38", ] +[[package]] +name = "der" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1a467a65c5e759bce6e65eaf91cc29f466cdc57cb65777bd646872a8a1fd4de" +dependencies = [ + "const-oid", + "zeroize", +] + [[package]] name = "deranged" version = "0.3.9" @@ -508,12 +974,44 @@ dependencies = [ "futures", ] +[[package]] +name = "ecdsa" +version = "0.14.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413301934810f597c1d19ca71c8710e99a3f1ba28a0d2ebc01551a2daeea3c5c" +dependencies = [ + "der", + "elliptic-curve", + "rfc6979", + "signature", +] + [[package]] name = "either" version = "1.9.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a26ae43d7bcc3b814de94796a5e736d4029efb0ee900c12e2d54c993ad1a1e07" +[[package]] +name = "elliptic-curve" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7bb888ab5300a19b8e5bceef25ac745ad065f3c9f7efc6de1b91958110891d3" +dependencies = [ + "base16ct", + "crypto-bigint 0.4.9", + "der", + "digest", + "ff", + "generic-array", + "group", + "pkcs8", + "rand_core", + "sec1", + "subtle", + "zeroize", +] + [[package]] name = "encoding_rs" version = "0.8.33" @@ -561,6 +1059,16 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "25cbce373ec4653f1a01a31e8a5e5ec0c622dc27ff9c4e6606eefef5cbbed4a5" +[[package]] +name = "ff" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d013fc25338cc558c5c2cfbad646908fb23591e2404481826742b651c9af7160" +dependencies = [ + "rand_core", + "subtle", +] + [[package]] name = "figment" version = "0.10.11" @@ -750,6 +1258,17 @@ version = "0.28.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6fb8d784f27acf97159b40fc4db5ecd8aa23b9ad5ef69cdd136d3bc80665f0c0" +[[package]] +name = "group" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5dfbfb3a6cfbd390d5c9564ab283a0349b9b9fcd46a706c1eb10e0db70bfbac7" +dependencies = [ + "ff", + "rand_core", + "subtle", +] + [[package]] name = "h2" version = "0.3.21" @@ -882,6 +1401,22 @@ dependencies = [ "want", ] +[[package]] +name = "hyper-rustls" +version = "0.24.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec3efd23720e2049821a693cbc7e65ea87c72f1c58ff2f9522ff332b1491e590" +dependencies = [ + "futures-util", + "http", + "hyper", + "log", + "rustls", + "rustls-native-certs", + "tokio", + "tokio-rustls", +] + [[package]] name = "hyper-tls" version = "0.5.0" @@ -1081,6 +1616,16 @@ dependencies = [ "regex-automata 0.1.10", ] +[[package]] +name = "md-5" +version = "0.10.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d89e7ee0cfbedfc4da3340218492196241d89eefb6dab27de5df917a6d2e78cf" +dependencies = [ + "cfg-if", + "digest", +] + [[package]] name = "memchr" version = "2.6.4" @@ -1194,6 +1739,16 @@ dependencies = [ "winapi", ] +[[package]] +name = "num-integer" +version = "0.1.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "225d3389fb3509a24c93f5c29eb6bde2586b98d9f016636dff58d7c6f7569cd9" +dependencies = [ + "autocfg", + "num-traits", +] + [[package]] name = "num-traits" version = "0.2.17" @@ -1272,12 +1827,29 @@ dependencies = [ "vcpkg", ] +[[package]] +name = "outref" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4030760ffd992bef45b0ae3f10ce1aba99e33464c90d14dd7c039884963ddc7a" + [[package]] name = "overload" version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39" +[[package]] +name = "p256" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51f44edd08f51e2ade572f141051021c5af22677e42b7dd28a88155151c33594" +dependencies = [ + "ecdsa", + "elliptic-curve", + "sha2", +] + [[package]] name = "parking_lot" version = "0.11.2" @@ -1370,6 +1942,16 @@ version = "0.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" +[[package]] +name = "pkcs8" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9eca2c590a5f85da82668fa685c09ce2888b9430e83299debf1f34b65fd4a4ba" +dependencies = [ + "der", + "spki", +] + [[package]] name = "pkg-config" version = "0.3.27" @@ -1437,6 +2019,15 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom", +] + [[package]] name = "rc-box" version = "1.2.0" @@ -1549,6 +2140,31 @@ dependencies = [ "winreg", ] +[[package]] +name = "rfc6979" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7743f17af12fa0b03b803ba12cd6a8d9483a587e89c69445e3909655c0b9fabb" +dependencies = [ + "crypto-bigint 0.4.9", + "hmac", + "zeroize", +] + +[[package]] +name = "ring" +version = "0.17.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "684d5e6e18f669ccebf64a92236bb7db9a34f07be010e3627368182027180866" +dependencies = [ + "cc", + "getrandom", + "libc", + "spin", + "untrusted", + "windows-sys", +] + [[package]] name = "rustc-demangle" version = "0.1.23" @@ -1577,6 +2193,49 @@ dependencies = [ "windows-sys", ] +[[package]] +name = "rustls" +version = "0.21.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "629648aced5775d558af50b2b4c7b02983a04b312126d45eeead26e7caa498b9" +dependencies = [ + "log", + "ring", + "rustls-webpki", + "sct", +] + +[[package]] +name = "rustls-native-certs" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a9aace74cb666635c918e9c12bc0d348266037aa8eb599b5cba565709a8dff00" +dependencies = [ + "openssl-probe", + "rustls-pemfile", + "schannel", + "security-framework", +] + +[[package]] +name = "rustls-pemfile" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c74cae0a4cf6ccbbf5f359f08efdf8ee7e1dc532573bf0db71968cb56b1448c" +dependencies = [ + "base64 0.21.5", +] + +[[package]] +name = "rustls-webpki" +version = "0.101.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b6275d1ee7a1cd780b64aca7726599a1dbc893b1e64144529e55c3c2f745765" +dependencies = [ + "ring", + "untrusted", +] + [[package]] name = "ryu" version = "1.0.15" @@ -1598,6 +2257,30 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +[[package]] +name = "sct" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da046153aa2352493d6cb7da4b6e5c0c057d8a1d0a9aa8560baffdd945acd414" +dependencies = [ + "ring", + "untrusted", +] + +[[package]] +name = "sec1" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3be24c1842290c45df0a7bf069e0c268a747ad05a192f2fd7dcfdbc1cba40928" +dependencies = [ + "base16ct", + "der", + "generic-array", + "pkcs8", + "subtle", + "zeroize", +] + [[package]] name = "security-framework" version = "2.9.2" @@ -1719,6 +2402,17 @@ dependencies = [ "syn 2.0.38", ] +[[package]] +name = "sha1" +version = "0.10.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + [[package]] name = "sha2" version = "0.10.8" @@ -1748,6 +2442,16 @@ dependencies = [ "libc", ] +[[package]] +name = "signature" +version = "1.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74233d3b3b2f6d4b006dc19dee745e73e2a6bfb6f93607cd3b02bd5b00797d7c" +dependencies = [ + "digest", + "rand_core", +] + [[package]] name = "slab" version = "0.4.9" @@ -1792,6 +2496,16 @@ dependencies = [ "lock_api", ] +[[package]] +name = "spki" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67cf02bbac7a337dc36e4f5a693db6c21e7863f45070f7064577eb4367a3212b" +dependencies = [ + "base64ct", + "der", +] + [[package]] name = "sqlformat" version = "0.2.2" @@ -1901,6 +2615,8 @@ name = "stable-diffusion-bot" version = "0.1.0" dependencies = [ "anyhow", + "aws-config", + "aws-sdk-s3", "base64 0.21.5", "bytes", "clap", @@ -2195,6 +2911,16 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-rustls" +version = "0.24.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c28327cf380ac148141087fbfb9de9d7bd4e84ab5d2c28fbc911d753de8a7081" +dependencies = [ + "rustls", + "tokio", +] + [[package]] name = "tokio-stream" version = "0.1.14" @@ -2408,6 +3134,12 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" +[[package]] +name = "untrusted" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" + [[package]] name = "url" version = "2.4.1" @@ -2420,6 +3152,12 @@ dependencies = [ "serde", ] +[[package]] +name = "urlencoding" +version = "2.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "daf8dba3b7eb870caf1ddeed7bc9d2a049f3cfdfae7cb521b087cc33ae4c49da" + [[package]] name = "utf8parse" version = "0.2.1" @@ -2454,6 +3192,12 @@ version = "0.9.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "49874b5167b65d7193b8aba1567f5c7d93d001cafc34600cee003eda787e483f" +[[package]] +name = "vsimd" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c3082ca00d5a5ef149bb8b555a72ae84c9c59f7250f013ac822ac2e49b19c64" + [[package]] name = "want" version = "0.3.1" @@ -2674,6 +3418,12 @@ dependencies = [ "windows-sys", ] +[[package]] +name = "xmlparser" +version = "0.13.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "66fee0b777b0f5ac1c69bb06d361268faafa61cd4682ae064a171c16c433e9e4" + [[package]] name = "yansi" version = "1.0.0-rc.1" @@ -2699,3 +3449,9 @@ dependencies = [ "quote", "syn 2.0.38", ] + +[[package]] +name = "zeroize" +version = "1.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "525b4ec142c6b68a2d10f01f7bbf6755599ca3f81ea53b8431b7dd348f5fdb2d" diff --git a/crates/stable-diffusion-bot/Cargo.toml b/crates/stable-diffusion-bot/Cargo.toml index 2a64dff..4c5db08 100644 --- a/crates/stable-diffusion-bot/Cargo.toml +++ b/crates/stable-diffusion-bot/Cargo.toml @@ -22,6 +22,8 @@ tracing-subscriber = { version = "0.3.16", features = ["env-filter"] } stable-diffusion-api = { path = "../stable-diffusion-api" } clap = { version = "4.4.7", features = ["derive"] } tracing-journald = "0.3.0" +aws-config = { version = "1.0.1", features = ["behavior-version-latest"] } +aws-sdk-s3 = "1.4.0" [target.'cfg(target_os = "linux")'.dependencies] libsystemd = "0.7.0" diff --git a/crates/stable-diffusion-bot/src/bot/handlers/image.rs b/crates/stable-diffusion-bot/src/bot/handlers/image.rs index 0e79db6..d737930 100644 --- a/crates/stable-diffusion-bot/src/bot/handlers/image.rs +++ b/crates/stable-diffusion-bot/src/bot/handlers/image.rs @@ -1,4 +1,7 @@ +use std::time::Duration; + use anyhow::{anyhow, Context}; +use futures::{stream::FuturesUnordered, StreamExt}; use stable_diffusion_api::{Img2ImgRequest, ImgInfo, ImgResponse, Txt2ImgRequest}; use teloxide::{ dispatching::UpdateHandler, @@ -7,11 +10,12 @@ use teloxide::{ payloads::setters::*, prelude::*, types::{ - ChatAction, InlineKeyboardButton, InlineKeyboardMarkup, InputFile, InputMedia, - InputMediaPhoto, MessageId, PhotoSize, + ChatAction, InlineKeyboardButton, InlineKeyboardMarkup, InlineQueryResult, + InlineQueryResultPhoto, InputFile, InputMedia, InputMediaPhoto, MessageId, PhotoSize, }, }; -use tracing::info; +use tokio::join; +use tracing::{info, warn}; use crate::{ bot::{helpers, State}, @@ -324,6 +328,99 @@ async fn handle_prompt( Ok(()) } +pub(crate) async fn handle_inline_query( + bot: Bot, + cfg: ConfigParameters, + InlineQuery { id, query, .. }: InlineQuery, + client: Option, +) -> anyhow::Result<()> { + let (client, bucket_id) = match (client, &cfg.aws_bucket_id) { + (Some(client), Some(bucket_id)) => (client, bucket_id), + _ => { + bot.answer_inline_query(id, vec![]).await?; + return Ok(()); + } + }; + + let mut txt2img = cfg.txt2img_defaults.clone(); + txt2img.with_batch_size(2).with_steps(1); + + let resp = do_txt2img(query.as_str(), &cfg, &mut txt2img).await?; + let info = resp.info().unwrap(); + let images = Photo::album(resp.images)?; + + let presigning_config = aws_sdk_s3::presigning::PresigningConfig::builder() + .expires_in(Duration::from_secs(300)) + .build()?; + + let upload = match images { + Photo::Album(images) => images.into_iter().enumerate().map(|(i, image)| { + client + .put_object() + .key(info.all_seeds.as_ref().unwrap()[i].to_string() + ".png") + .bucket(bucket_id.clone()) + .body(image.into()) + .send() + }), + _ => return Ok(()), + } + .collect::>() + .collect::>(); + + let urls = info + .all_seeds + .as_ref() + .unwrap() + .iter() + .map(|s| { + client + .get_object() + .key(s.to_string() + ".png") + .bucket(bucket_id.clone()) + .presigned(presigning_config.clone()) + }) + .collect::>() + .collect::>(); + + let (upload_result, url_result) = join!(upload, urls); + + let output = upload_result + .into_iter() + .zip(url_result.into_iter()) + .filter_map(|result| match result { + (Ok(_), Ok(url)) => Some(url), + (Ok(_), Err(e)) => { + warn!("Error getting presigned url: {}", e); + None + } + (Err(e), Ok(_)) => { + warn!("Error during upload: {}", e); + None + } + (Err(e1), Err(e2)) => { + warn!("Error during upload: {}", e1); + warn!("Error getting presigned url: {}", e2); + None + } + }); + + let photos = output.into_iter().enumerate().filter_map(|(i, resp)| { + let url = match reqwest::Url::parse(resp.uri()).context("Failed to parse URL") { + Ok(url) => url, + Err(_) => return None, + }; + Some(InlineQueryResult::Photo(InlineQueryResultPhoto::new( + info.all_seeds.as_ref().unwrap()[i].to_string(), + url.clone(), + url, + ))) + }); + + bot.answer_inline_query(id, photos).await?; + + Ok(()) +} + fn keyboard(seed: i64) -> InlineKeyboardMarkup { let seed_button = if seed == -1 { InlineKeyboardButton::callback("🎲 Seed", "reuse/-1") @@ -531,8 +628,13 @@ pub(crate) fn image_schema() -> UpdateHandler { .endpoint(handle_rerun), ); + let inline_query_handler = Update::filter_inline_query() + .chain(filter_map_settings()) + .endpoint(handle_inline_query); + dptree::entry() .branch(gen_command_handler) .branch(message_handler) .branch(callback_handler) + .branch(inline_query_handler) } diff --git a/crates/stable-diffusion-bot/src/bot/mod.rs b/crates/stable-diffusion-bot/src/bot/mod.rs index 20ab7c4..422d7fc 100644 --- a/crates/stable-diffusion-bot/src/bot/mod.rs +++ b/crates/stable-diffusion-bot/src/bot/mod.rs @@ -103,6 +103,7 @@ pub struct StableDiffusionBot { bot: Bot, storage: DialogueStorage, config: ConfigParameters, + aws_client: Option, } impl StableDiffusionBot { @@ -159,6 +160,7 @@ impl StableDiffusionBot { bot, storage, config, + aws_client, } = self; let mut commands = UnauthenticatedCommands::bot_commands(); @@ -167,19 +169,25 @@ impl StableDiffusionBot { bot.set_my_commands(commands) .scope(teloxide::types::BotCommandScope::Default) .await?; + let inline_query_handler = Update::filter_inline_query().endpoint(handle_inline_query); - Dispatcher::builder(bot, Self::schema()) - .dependencies(dptree::deps![config, storage]) - .default_handler(|upd| async move { - warn!("Unhandled update: {:?}", upd); - }) - .error_handler(LoggingErrorHandler::with_custom_text( - "An error has occurred in the dispatcher", - )) - .enable_ctrlc_handler() - .build() - .dispatch() - .await; + Dispatcher::builder( + bot, + dptree::entry() + .branch(Self::schema()) + .branch(inline_query_handler), + ) + .dependencies(dptree::deps![config, storage, aws_client]) + .default_handler(|upd| async move { + warn!("Unhandled update: {:?}", upd); + }) + .error_handler(LoggingErrorHandler::with_custom_text( + "An error has occurred in the dispatcher", + )) + .enable_ctrlc_handler() + .build() + .dispatch() + .await; Ok(()) } @@ -192,6 +200,7 @@ pub(crate) struct ConfigParameters { txt2img_defaults: Txt2ImgRequest, img2img_defaults: Img2ImgRequest, allow_all_users: bool, + aws_bucket_id: Option, } impl ConfigParameters { @@ -210,6 +219,8 @@ pub struct StableDiffusionBotBuilder { txt2img_defaults: Option, img2img_defaults: Option, allow_all_users: bool, + aws_client: Option, + aws_bucket_id: Option, } impl StableDiffusionBotBuilder { @@ -228,6 +239,8 @@ impl StableDiffusionBotBuilder { txt2img_defaults: None, img2img_defaults: None, allow_all_users, + aws_client: None, + aws_bucket_id: None, } } @@ -320,6 +333,16 @@ impl StableDiffusionBotBuilder { self } + pub fn aws_client(mut self, client: aws_sdk_s3::Client) -> Self { + self.aws_client = Some(client); + self + } + + pub fn aws_bucket_id(mut self, bucket_id: String) -> Self { + self.aws_bucket_id = Some(bucket_id); + self + } + /// Consumes the builder and builds a `StableDiffusionBot` instance. /// /// # Examples @@ -355,18 +378,33 @@ impl StableDiffusionBotBuilder { let api = Api::new_with_client_and_url(client, self.sd_api_url.clone()) .context("Failed to initialize sd api")?; + let (aws_client, aws_bucket_id) = match (self.aws_client, self.aws_bucket_id) { + (Some(client), Some(bucket_id)) => (Some(client), Some(bucket_id)), + (Some(_), None) => { + warn!("AWS client provided but no bucket ID"); + (None, None) + } + (None, Some(_)) => { + warn!("AWS bucket ID provided but no client"); + (None, None) + } + (None, None) => (None, None), + }; + let parameters = ConfigParameters { allowed_users, api, txt2img_defaults: default_txt2img(self.txt2img_defaults.unwrap_or_default()), img2img_defaults: default_img2img(self.img2img_defaults.unwrap_or_default()), allow_all_users: self.allow_all_users, + aws_bucket_id, }; Ok(StableDiffusionBot { bot, storage, config: parameters, + aws_client, }) } } diff --git a/crates/stable-diffusion-bot/src/main.rs b/crates/stable-diffusion-bot/src/main.rs index 1e93fed..107cf8e 100644 --- a/crates/stable-diffusion-bot/src/main.rs +++ b/crates/stable-diffusion-bot/src/main.rs @@ -36,6 +36,8 @@ struct Config { txt2img: Option, img2img: Option, allow_all_users: Option, + aws_endpoint_url: Option, + aws_bucket_id: Option, } #[tokio::main] @@ -74,6 +76,15 @@ async fn main() -> anyhow::Result<()> { .extract() .context("Invalid configuration")?; + let sdk_config = aws_config::from_env() + .endpoint_url(config.aws_endpoint_url.unwrap()) + .load() + .await; + let mut s3_config_builder: aws_sdk_s3::config::Builder = (&sdk_config).into(); + s3_config_builder.set_force_path_style(Some(true)); + let s3_config = s3_config_builder.build(); + let client = aws_sdk_s3::Client::from_conf(s3_config); + StableDiffusionBotBuilder::new( config.api_key, config.allowed_users, @@ -83,6 +94,8 @@ async fn main() -> anyhow::Result<()> { .db_path(config.db_path) .txt2img_defaults(config.txt2img.unwrap_or_default()) .img2img_defaults(config.img2img.unwrap_or_default()) + .aws_client(client) + .aws_bucket_id(config.aws_bucket_id.unwrap()) .build() .await? .run()