migrate from wsi to flux

.vscode/settings.json

@@ -5,5 +5,9 @@
     "[toml]": {
         "editor.defaultFormatter": "tamasfe.even-better-toml"
     },
-    "rust-analyzer.check.command": "clippy"
+    "rust-analyzer.check.command": "clippy",
+    "terminal.integrated.env.linux": {
+        "GTK_PATH": null,
+        "GIO_MODULE_DIR": null,
+    },
 }

Cargo.toml

@@ -12,4 +12,5 @@ members = [
     "assyst-database",
     "assyst-proc-macro",
 ]
+exclude = ["flux"]
 resolver = "2"

assyst-common/Cargo.toml

@@ -23,3 +23,4 @@ tracing-subscriber = { version = "0.3.16", features = ["time", "env-filter"] }
 time = { version = "0.3.31", features = ["macros"] }
 num_cpus = "1.16.0"
 rayon = "1.8.1"
+rand = "0.8.5"

assyst-common/src/config/config.rs

@@ -17,7 +17,6 @@ pub struct Urls {
     pub proxy: Vec<String>,
     pub filer: String,
     pub eval: String,
-    pub wsi: String,
     pub bad_translation: String,
 }
 
@@ -100,5 +99,5 @@ pub struct DevAttributes {
     pub dev_guild: u64,
     pub dev_channel: u64,
     pub dev_message: bool,
-    pub wsi_retry_limit: u64,
+    pub flux_workspace_root_path_override: String,
 }

assyst-common/src/util/mod.rs

@@ -1,6 +1,8 @@
 use std::borrow::Cow;
+use std::hash::{DefaultHasher, Hash, Hasher};
 use std::time::Duration;
 
+use rand::Rng;
 use time::format_description;
 use tracing::info;
 use tracing_subscriber::fmt::time::UtcTime;
@@ -53,9 +55,7 @@ pub fn parse_to_millis(input: &str) -> Result<u64, ParseToMillisError> {
             .parse::<u64>()
             .map_err(|_| ParseToMillisError::ParseIntError)?;
 
-        let unit: u64 = unit_to_ms(&current[2])
-            .try_into()
-            .map_err(|_| ParseToMillisError::Overflow)?;
+        let unit: u64 = unit_to_ms(&current[2]);
 
         let ms = amount.checked_mul(unit).ok_or(ParseToMillisError::Overflow)?;
 
@@ -139,3 +139,11 @@ pub fn string_from_likely_utf8(bytes: Vec<u8>) -> String {
         String::from_utf8_lossy(err.as_bytes()).into_owned()
     })
 }
+
+/// Hashes a buffer. Appends a random string.
+pub fn hash_buffer(buf: &[u8]) -> String {
+    let mut body_hasher = DefaultHasher::new();
+    buf.hash(&mut body_hasher);
+    let rand = rand::thread_rng().gen::<usize>();
+    format!("{:x}{:x}", body_hasher.finish(), rand)
+}

assyst-core/Cargo.toml

@@ -31,7 +31,6 @@ futures-util = "0.3.30"
 serde_json = "1.0.113"
 moka = { version = "0.12.3", features = ["sync"] }
 human_bytes = { version = "0.4", default-features = false }
-shared = { git = "https://github.com/jacherr/wsi-shared" }
 bincode = "1.3.3"
 paste = "1.0.14"
 rand = "0.8.5"

assyst-core/src/assyst.rs

@@ -1,10 +1,10 @@
 use crate::command_ratelimits::CommandRatelimits;
+use crate::flux_handler::FluxHandler;
 use crate::persistent_cache_handler::PersistentCacheHandler;
 use crate::replies::Replies;
 use crate::rest::patreon::Patron;
 use crate::rest::rest_cache_handler::RestCacheHandler;
 use crate::task::Task;
-use crate::wsi_handler::WsiHandler;
 use assyst_common::config::CONFIG;
 use assyst_common::metrics_handler::MetricsHandler;
 use assyst_common::pipe::CACHE_PIPE_PATH;
@@ -26,7 +26,7 @@ pub struct Assyst {
     /// Handler for the Assyst database. RwLocked to allow concurrent reads.
     pub database_handler: Arc<DatabaseHandler>,
     /// Handler for WSI.
-    pub wsi_handler: WsiHandler,
+    pub wsi_handler: FluxHandler,
     /// Handler for the REST cache.
     pub rest_cache_handler: RestCacheHandler,
     /// HTTP client for Discord. Handles all HTTP requests to Discord, storing stateful information
@@ -70,7 +70,7 @@ impl Assyst {
             tasks: Mutex::new(vec![]),
             shard_count,
             replies: Replies::new(),
-            wsi_handler: WsiHandler::new(database_handler.clone(), premium_users.clone()),
+            wsi_handler: FluxHandler::new(database_handler.clone(), premium_users.clone()),
             rest_cache_handler: RestCacheHandler::new(http_client.clone()),
             command_ratelimits: CommandRatelimits::new(),
         })

assyst-core/src/command/flags.rs

@@ -208,6 +208,46 @@ impl FlagDecode for DownloadFlags {
 }
 flag_parse_argument! { DownloadFlags }
 
+#[derive(Default)]
+pub struct BloomFlags {
+    pub radius: Option<u64>,
+    pub brightness: Option<u64>,
+    pub sharpness: Option<u64>,
+}
+impl FlagDecode for BloomFlags {
+    fn from_str(input: &str) -> anyhow::Result<Self> {
+        let mut valid_flags = HashMap::new();
+        valid_flags.insert("radius", FlagType::WithValue);
+        valid_flags.insert("sharpness", FlagType::WithValue);
+        valid_flags.insert("brightness", FlagType::WithValue);
+
+        let raw_decode = flags_from_str(input, valid_flags)?;
+        let result = Self {
+            radius: raw_decode
+                .get("radius")
+                .unwrap_or(&None)
+                .clone()
+                .map(|x| x.parse().context("Provided radius is invalid"))
+                .transpose()?,
+            sharpness: raw_decode
+                .get("sharpness")
+                .unwrap_or(&None)
+                .clone()
+                .map(|x| x.parse().context("Provided sharpness is invalid"))
+                .transpose()?,
+            brightness: raw_decode
+                .get("brightness")
+                .unwrap_or(&None)
+                .clone()
+                .map(|x| x.parse().context("Provided brightness is invalid"))
+                .transpose()?,
+        };
+
+        Ok(result)
+    }
+}
+flag_parse_argument! { BloomFlags }
+
 pub enum FlagType {
     WithValue,
     NoValue,

assyst-core/src/command/wsi/mod.rs → assyst-core/src/command/image/mod.rs

@@ -4,13 +4,14 @@ use anyhow::Context;
 use assyst_proc_macro::command;
 
 use super::arguments::{Image, Rest, Word};
+use super::flags::BloomFlags;
 use crate::command::{Availability, Category, CommandCtxt};
 
 #[command(
     description = "ah shit here we go again",
     cooldown = Duration::from_secs(2),
     access = Availability::Public,
-    category = Category::Wsi,
+    category = Category::Image,
     usage = "[image]",
     examples = ["https://link.to.my/image.png"],
     send_processing = true
@@ -27,16 +28,13 @@ pub async fn ahshit(ctxt: CommandCtxt<'_>, source: Image) -> anyhow::Result<()>
     description = "april fools!!!!",
     cooldown = Duration::from_secs(2),
     access = Availability::Public,
-    category = Category::Wsi,
+    category = Category::Image,
     usage = "[image]",
     examples = ["https://link.to.my/image.png"],
     send_processing = true
 )]
 pub async fn aprilfools(ctxt: CommandCtxt<'_>, source: Image) -> anyhow::Result<()> {
-    let result = ctxt
-        .wsi_handler()
-        .aprilfools(source.0, ctxt.data.author.id.get())
-        .await?;
+    let result = ctxt.wsi_handler().aprilfools(source.0).await?;
 
     ctxt.reply(result).await?;
 
@@ -47,15 +45,26 @@ pub async fn aprilfools(ctxt: CommandCtxt<'_>, source: Image) -> anyhow::Result<
     description = "bloom an image",
     cooldown = Duration::from_secs(2),
     access = Availability::Public,
-    category = Category::Wsi,
-    usage = "[image] <radius>",
-    examples = ["https://link.to.my/image.png"],
-    send_processing = true
+    category = Category::Image,
+    usage = "[image] <flags>",
+    examples = ["https://link.to.my/image.png", "https://link.to.my/image.png --brightness 100 --sharpness 25 --radius 10"],
+    send_processing = true,
+    flag_descriptions = [
+        ("radius", "Bloom radius as a number"),
+        ("brightness", "Bloom brightness as a number"),
+        ("sharpness", "Bloom sharpness as a number"),
+    ]
 )]
-pub async fn bloom(ctxt: CommandCtxt<'_>, source: Image, radius: Option<u64>) -> anyhow::Result<()> {
+pub async fn bloom(ctxt: CommandCtxt<'_>, source: Image, flags: BloomFlags) -> anyhow::Result<()> {
     let result = ctxt
         .wsi_handler()
-        .bloom(source.0, radius.unwrap_or(5) as usize, ctxt.data.author.id.get())
+        .bloom(
+            source.0,
+            flags.radius,
+            flags.sharpness,
+            flags.brightness,
+            ctxt.data.author.id.get(),
+        )
         .await?;
 
     ctxt.reply(result).await?;
@@ -67,15 +76,15 @@ pub async fn bloom(ctxt: CommandCtxt<'_>, source: Image, radius: Option<u64>) ->
     description = "blur an image",
     cooldown = Duration::from_secs(2),
     access = Availability::Public,
-    category = Category::Wsi,
+    category = Category::Image,
     usage = "[image] <radius>",
     examples = ["https://link.to.my/image.png"],
     send_processing = true
 )]
 pub async fn blur(ctxt: CommandCtxt<'_>, source: Image, strength: Option<f32>) -> anyhow::Result<()> {
     let result = ctxt
         .wsi_handler()
-        .blur(source.0, strength.unwrap_or(1.0), ctxt.data.author.id.get())
+        .blur(source.0, strength, ctxt.data.author.id.get())
         .await?;
 
     ctxt.reply(result).await?;
@@ -87,7 +96,7 @@ pub async fn blur(ctxt: CommandCtxt<'_>, source: Image, strength: Option<f32>) -
     description = "add a caption to an image",
     cooldown = Duration::from_secs(2),
     access = Availability::Public,
-    category = Category::Wsi,
+    category = Category::Image,
     usage = "[image] [caption]",
     examples = ["https://link.to.my/image.png hello there"],
     send_processing = true
@@ -107,7 +116,7 @@ pub async fn caption(ctxt: CommandCtxt<'_>, source: Image, text: Rest) -> anyhow
     description = "resize an image based on scale or WxH (default 2x)",
     cooldown = Duration::from_secs(2),
     access = Availability::Public,
-    category = Category::Wsi,
+    category = Category::Image,
     usage = "[image] <scale>",
     examples = ["https://link.to.my/image.png", "https://link.to.my/image.png 128x128", "https://link.to.my/image.png 2"],
     send_processing = true

assyst-core/src/command/mod.rs

@@ -34,7 +34,7 @@ use std::time::{Duration, Instant};
 
 use super::gateway_handler::reply as gateway_reply;
 use crate::assyst::ThreadSafeAssyst;
-use crate::wsi_handler::WsiHandler;
+use crate::flux_handler::FluxHandler;
 use assyst_common::config::CONFIG;
 use async_trait::async_trait;
 use errors::TagParseError;
@@ -56,12 +56,12 @@ pub mod errors;
 pub mod flags;
 pub mod fun;
 pub mod group;
+pub mod image;
 pub mod messagebuilder;
 pub mod misc;
 pub mod registry;
 pub mod services;
 pub mod source;
-pub mod wsi;
 
 /// Defines who can use a command in a server.
 #[derive(Clone, Copy, Debug, PartialEq)]
@@ -142,7 +142,7 @@ pub struct CommandInteractionInfo {
 pub enum Category {
     Fun,
     Makesweet,
-    Wsi,
+    Image,
     Misc,
     Services,
     None(String),
@@ -156,7 +156,7 @@ impl Display for Category {
             match self {
                 Self::Fun => "fun",
                 Self::Makesweet => "makesweet",
-                Self::Wsi => "wsi",
+                Self::Image => "image",
                 Self::Misc => "misc",
                 Self::Services => "services",
                 Self::None(t) => &**t,
@@ -170,7 +170,7 @@ impl From<String> for Category {
         match &*v {
             "fun" => Category::Fun,
             "misc" => Category::Misc,
-            "wsi" => Category::Wsi,
+            "image" => Category::Image,
             "makesweet" => Category::Makesweet,
             t => Category::None(t.to_string()),
         }
@@ -382,7 +382,7 @@ impl<'a> CommandCtxt<'a> {
         self.data.assyst
     }
 
-    pub fn wsi_handler(&self) -> &'a WsiHandler {
+    pub fn wsi_handler(&self) -> &'a FluxHandler {
         &self.data.assyst.wsi_handler
     }
 }

assyst-core/src/command/registry.rs

@@ -7,7 +7,7 @@ use twilight_model::application::command::Command as InteractionCommand;
 use crate::assyst::ThreadSafeAssyst;
 use crate::command::CommandMetadata;
 
-use super::{fun, misc, services, wsi, TCommand};
+use super::{fun, image, misc, services, TCommand};
 
 macro_rules! declare_commands {
     ($($name:path),*) => {
@@ -36,12 +36,12 @@ declare_commands!(
     services::cooltext_command,
     services::download_command,
     services::r34_command,
-    wsi::ahshit_command,
-    wsi::aprilfools_command,
-    wsi::bloom_command,
-    wsi::blur_command,
-    wsi::caption_command,
-    wsi::resize_command
+    image::ahshit_command,
+    image::aprilfools_command,
+    image::bloom_command,
+    image::blur_command,
+    image::caption_command,
+    image::resize_command
 );
 
 static COMMANDS: OnceLock<HashMap<&'static str, TCommand>> = OnceLock::new();

assyst-core/src/flux_handler/flux_request.rs

@@ -0,0 +1,46 @@
+use std::collections::HashMap;
+
+use super::limits::LimitData;
+
+/// A step in a Flux execution.
+pub enum FluxStep {
+    /// Input file. Saves the file and passes to Flux as `-i path`. Input must be the first step.
+    Input(Vec<u8>),
+    /// Operation. Passes to Flux as `-o operation[k=v]`
+    Operation((String, HashMap<String, String>)),
+    /// Output. Passes to Flux as `path` at the end. Output must be the last step.
+    Output,
+    ImagePageLimit(u64),
+    ResolutionLimit((u64, u64)),
+}
+
+pub struct FluxRequest(pub Vec<FluxStep>);
+impl FluxRequest {
+    pub fn new() -> Self {
+        Self(vec![])
+    }
+
+    pub fn new_with_input_and_limits(input: Vec<u8>, limits: &LimitData) -> Self {
+        let mut new = Self(vec![]);
+        new.input(input);
+        new.limits(limits);
+        new
+    }
+
+    pub fn input(&mut self, input: Vec<u8>) {
+        self.0.push(FluxStep::Input(input));
+    }
+
+    pub fn operation(&mut self, name: String, options: HashMap<String, String>) {
+        self.0.push(FluxStep::Operation((name, options)));
+    }
+
+    pub fn output(&mut self) {
+        self.0.push(FluxStep::Output);
+    }
+
+    pub fn limits(&mut self, limits: &LimitData) {
+        self.0.push(FluxStep::ImagePageLimit(limits.frames));
+        self.0.push(FluxStep::ResolutionLimit((limits.size, limits.size)));
+    }
+}