Fix Imaginate by porting its JS roundtrip code to graph-based async execution in Rust (#1250)

* Create asynchronous rust imaginate node

* Make a first imaginate request via rust

* Implement parsing of imaginate API result image

* Stop refresh timer from affecting imaginate progress requests

* Add cargo-about clarification for rustls-webpki

* Delete imaginate.ts and all uses of its functions

* Add imaginate img2img feature

* Fix imaginate random seed button

* Fix imaginate ui inferring non-custom resolutions

* Fix the imaginate progress indicator

* Remove ImaginatePreferences from being compiled into node graph

* Regenerate imaginate only when hitting button

* Add ability to terminate imaginate requests

* Add imaginate server check feature

* Do not compile wasm_bindgen bindings in graphite_editor for tests

* Address some review suggestions

- move wasm futures dependency in editor to the future-executor crate
- guard wasm-bindgen in editor behind a `wasm` feature flag
- dont make seed number input a slider
- remove poll_server_check from process_message function beginning
- guard wasm related code behind `cfg(target_arch = "wasm32")` instead
  of `cfg(test)`
- Call the imaginate idle states "Ready" and "Done" instead of "Nothing
  to do"
- Call the imaginate uploading state "Uploading Image" instead of
  "Uploading Input Image"
- Remove the EvalSyncNode

* Fix imaginate host name being restored between graphite instances

also change the progress status texts a bit.

---------

Co-authored-by: Keavon Chambers <keavon@keavon.com>

node-graph/graph-craft/src/document.rs

@@ -441,7 +441,7 @@ impl NodeNetwork {
 	}
 
 	/// Check if the specified node id is connected to the output
-	pub fn connected_to_output(&self, target_node_id: NodeId, ignore_imaginate: bool) -> bool {
+	pub fn connected_to_output(&self, target_node_id: NodeId) -> bool {
 		// If the node is the output then return true
 		if self.outputs.iter().any(|&NodeOutput { node_id, .. }| node_id == target_node_id) {
 			return true;
@@ -454,11 +454,6 @@ impl NodeNetwork {
 		already_visited.extend(self.outputs.iter().map(|output| output.node_id));
 
 		while let Some(node) = stack.pop() {
-			// Skip the imaginate node inputs
-			if ignore_imaginate && node.name == "Imaginate" {
-				continue;
-			}
-
 			for input in &node.inputs {
 				if let &NodeInput::Node { node_id: ref_id, .. } = input {
 					// Skip if already viewed
@@ -680,7 +675,7 @@ impl NodeNetwork {
 
 			let mut dummy_input = NodeInput::ShortCircut(concrete!(()));
 			std::mem::swap(&mut dummy_input, input);
-			if let NodeInput::Value { tagged_value, exposed } = dummy_input {
+			if let NodeInput::Value { mut tagged_value, exposed } = dummy_input {
 				let value_node_id = gen_id();
 				let merged_node_id = map_ids(id, value_node_id);
 				let path = if let Some(mut new_path) = node.path.clone() {

node-graph/graph-craft/src/document/value.rs

@@ -1,6 +1,6 @@
 use super::DocumentNode;
 use crate::graphene_compiler::Any;
-pub use crate::imaginate_input::{ImaginateMaskStartingFill, ImaginateSamplingMethod, ImaginateStatus};
+pub use crate::imaginate_input::{ImaginateCache, ImaginateController, ImaginateMaskStartingFill, ImaginateSamplingMethod};
 use crate::proto::{Any as DAny, FutureAny};
 
 use graphene_core::raster::brush_cache::BrushCache;
@@ -27,7 +27,7 @@ pub enum TaggedValue {
 	OptionalDVec2(Option<DVec2>),
 	DAffine2(DAffine2),
 	Image(graphene_core::raster::Image<Color>),
-	RcImage(Option<Arc<graphene_core::raster::Image<Color>>>),
+	ImaginateCache(ImaginateCache),
 	ImageFrame(graphene_core::raster::ImageFrame<Color>),
 	Color(graphene_core::raster::color::Color),
 	Subpaths(Vec<bezier_rs::Subpath<graphene_core::uuid::ManipulatorGroupId>>),
@@ -36,7 +36,7 @@ pub enum TaggedValue {
 	LuminanceCalculation(LuminanceCalculation),
 	ImaginateSamplingMethod(ImaginateSamplingMethod),
 	ImaginateMaskStartingFill(ImaginateMaskStartingFill),
-	ImaginateStatus(ImaginateStatus),
+	ImaginateController(ImaginateController),
 	LayerPath(Option<Vec<u64>>),
 	VectorData(graphene_core::vector::VectorData),
 	Fill(graphene_core::vector::style::Fill),
@@ -83,7 +83,7 @@ impl Hash for TaggedValue {
 			}
 			Self::DAffine2(m) => m.to_cols_array().iter().for_each(|x| x.to_bits().hash(state)),
 			Self::Image(i) => i.hash(state),
-			Self::RcImage(i) => i.hash(state),
+			Self::ImaginateCache(i) => i.hash(state),
 			Self::Color(c) => c.hash(state),
 			Self::Subpaths(s) => s.iter().for_each(|subpath| subpath.hash(state)),
 			Self::RcSubpath(s) => s.hash(state),
@@ -91,7 +91,7 @@ impl Hash for TaggedValue {
 			Self::LuminanceCalculation(l) => l.hash(state),
 			Self::ImaginateSamplingMethod(m) => m.hash(state),
 			Self::ImaginateMaskStartingFill(f) => f.hash(state),
-			Self::ImaginateStatus(s) => s.hash(state),
+			Self::ImaginateController(s) => s.hash(state),
 			Self::LayerPath(p) => p.hash(state),
 			Self::ImageFrame(i) => i.hash(state),
 			Self::VectorData(vector_data) => vector_data.hash(state),
@@ -146,7 +146,7 @@ impl<'a> TaggedValue {
 			TaggedValue::OptionalDVec2(x) => Box::new(x),
 			TaggedValue::DAffine2(x) => Box::new(x),
 			TaggedValue::Image(x) => Box::new(x),
-			TaggedValue::RcImage(x) => Box::new(x),
+			TaggedValue::ImaginateCache(x) => Box::new(x),
 			TaggedValue::ImageFrame(x) => Box::new(x),
 			TaggedValue::Color(x) => Box::new(x),
 			TaggedValue::Subpaths(x) => Box::new(x),
@@ -155,7 +155,7 @@ impl<'a> TaggedValue {
 			TaggedValue::LuminanceCalculation(x) => Box::new(x),
 			TaggedValue::ImaginateSamplingMethod(x) => Box::new(x),
 			TaggedValue::ImaginateMaskStartingFill(x) => Box::new(x),
-			TaggedValue::ImaginateStatus(x) => Box::new(x),
+			TaggedValue::ImaginateController(x) => Box::new(x),
 			TaggedValue::LayerPath(x) => Box::new(x),
 			TaggedValue::VectorData(x) => Box::new(x),
 			TaggedValue::Fill(x) => Box::new(x),
@@ -210,7 +210,7 @@ impl<'a> TaggedValue {
 			TaggedValue::DVec2(_) => concrete!(DVec2),
 			TaggedValue::OptionalDVec2(_) => concrete!(Option<DVec2>),
 			TaggedValue::Image(_) => concrete!(graphene_core::raster::Image<Color>),
-			TaggedValue::RcImage(_) => concrete!(Option<Arc<graphene_core::raster::Image<Color>>>),
+			TaggedValue::ImaginateCache(_) => concrete!(ImaginateCache),
 			TaggedValue::ImageFrame(_) => concrete!(graphene_core::raster::ImageFrame<Color>),
 			TaggedValue::Color(_) => concrete!(graphene_core::raster::Color),
 			TaggedValue::Subpaths(_) => concrete!(Vec<bezier_rs::Subpath<graphene_core::uuid::ManipulatorGroupId>>),
@@ -218,7 +218,7 @@ impl<'a> TaggedValue {
 			TaggedValue::BlendMode(_) => concrete!(BlendMode),
 			TaggedValue::ImaginateSamplingMethod(_) => concrete!(ImaginateSamplingMethod),
 			TaggedValue::ImaginateMaskStartingFill(_) => concrete!(ImaginateMaskStartingFill),
-			TaggedValue::ImaginateStatus(_) => concrete!(ImaginateStatus),
+			TaggedValue::ImaginateController(_) => concrete!(ImaginateController),
 			TaggedValue::LayerPath(_) => concrete!(Option<Vec<u64>>),
 			TaggedValue::DAffine2(_) => concrete!(DAffine2),
 			TaggedValue::LuminanceCalculation(_) => concrete!(LuminanceCalculation),
@@ -263,7 +263,7 @@ impl<'a> TaggedValue {
 			x if x == TypeId::of::<DVec2>() => Ok(TaggedValue::DVec2(*downcast(input).unwrap())),
 			x if x == TypeId::of::<Option<DVec2>>() => Ok(TaggedValue::OptionalDVec2(*downcast(input).unwrap())),
 			x if x == TypeId::of::<graphene_core::raster::Image<Color>>() => Ok(TaggedValue::Image(*downcast(input).unwrap())),
-			x if x == TypeId::of::<Option<Arc<graphene_core::raster::Image<Color>>>>() => Ok(TaggedValue::RcImage(*downcast(input).unwrap())),
+			x if x == TypeId::of::<ImaginateCache>() => Ok(TaggedValue::ImaginateCache(*downcast(input).unwrap())),
 			x if x == TypeId::of::<graphene_core::raster::ImageFrame<Color>>() => Ok(TaggedValue::ImageFrame(*downcast(input).unwrap())),
 			x if x == TypeId::of::<graphene_core::raster::Color>() => Ok(TaggedValue::Color(*downcast(input).unwrap())),
 			x if x == TypeId::of::<Vec<bezier_rs::Subpath<graphene_core::uuid::ManipulatorGroupId>>>() => Ok(TaggedValue::Subpaths(*downcast(input).unwrap())),
@@ -271,7 +271,7 @@ impl<'a> TaggedValue {
 			x if x == TypeId::of::<BlendMode>() => Ok(TaggedValue::BlendMode(*downcast(input).unwrap())),
 			x if x == TypeId::of::<ImaginateSamplingMethod>() => Ok(TaggedValue::ImaginateSamplingMethod(*downcast(input).unwrap())),
 			x if x == TypeId::of::<ImaginateMaskStartingFill>() => Ok(TaggedValue::ImaginateMaskStartingFill(*downcast(input).unwrap())),
-			x if x == TypeId::of::<ImaginateStatus>() => Ok(TaggedValue::ImaginateStatus(*downcast(input).unwrap())),
+			x if x == TypeId::of::<ImaginateController>() => Ok(TaggedValue::ImaginateController(*downcast(input).unwrap())),
 			x if x == TypeId::of::<Option<Vec<u64>>>() => Ok(TaggedValue::LayerPath(*downcast(input).unwrap())),
 			x if x == TypeId::of::<DAffine2>() => Ok(TaggedValue::DAffine2(*downcast(input).unwrap())),
 			x if x == TypeId::of::<LuminanceCalculation>() => Ok(TaggedValue::LuminanceCalculation(*downcast(input).unwrap())),

node-graph/graph-craft/src/imaginate_input.rs

@@ -1,50 +1,155 @@
 use dyn_any::{DynAny, StaticType};
-use glam::DVec2;
+use graphene_core::Color;
+use std::borrow::Cow;
 use std::fmt::Debug;
+use std::sync::{
+	atomic::{AtomicBool, Ordering},
+	Arc, Mutex,
+};
 
-#[derive(Default, Debug, Clone, Copy, PartialEq, DynAny, specta::Type)]
+#[derive(Default, Debug, Clone, DynAny, specta::Type)]
+#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
+pub struct ImaginateCache(Arc<Mutex<graphene_core::raster::Image<Color>>>);
+
+impl ImaginateCache {
+	pub fn into_inner(self) -> Arc<Mutex<graphene_core::raster::Image<Color>>> {
+		self.0
+	}
+}
+
+impl std::cmp::PartialEq for ImaginateCache {
+	fn eq(&self, other: &Self) -> bool {
+		Arc::ptr_eq(&self.0, &other.0)
+	}
+}
+
+impl core::hash::Hash for ImaginateCache {
+	fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
+		self.0.lock().unwrap().hash(state);
+	}
+}
+
+pub trait ImaginateTerminationHandle: Debug + Send + Sync + 'static {
+	fn terminate(&self);
+}
+
+#[derive(Default, Debug, specta::Type)]
+#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
+struct InternalImaginateControl {
+	status: Mutex<ImaginateStatus>,
+	trigger_regenerate: AtomicBool,
+	#[serde(skip)]
+	termination_sender: Mutex<Option<Box<dyn ImaginateTerminationHandle>>>,
+}
+
+#[derive(Debug, Default, Clone, DynAny, specta::Type)]
+#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
+pub struct ImaginateController(Arc<InternalImaginateControl>);
+
+impl ImaginateController {
+	pub fn get_status(&self) -> ImaginateStatus {
+		self.0.status.lock().as_deref().cloned().unwrap_or_default()
+	}
+
+	pub fn set_status(&self, status: ImaginateStatus) {
+		if let Ok(mut lock) = self.0.status.lock() {
+			*lock = status
+		}
+	}
+
+	pub fn take_regenerate_trigger(&self) -> bool {
+		self.0.trigger_regenerate.swap(false, Ordering::SeqCst)
+	}
+
+	pub fn trigger_regenerate(&self) {
+		self.0.trigger_regenerate.store(true, Ordering::SeqCst)
+	}
+
+	pub fn request_termination(&self) {
+		if let Some(handle) = self.0.termination_sender.lock().ok().and_then(|mut lock| lock.take()) {
+			handle.terminate()
+		}
+	}
+
+	pub fn set_termination_handle<H: ImaginateTerminationHandle>(&self, handle: Box<H>) {
+		if let Ok(mut lock) = self.0.termination_sender.lock() {
+			*lock = Some(handle)
+		}
+	}
+}
+
+impl std::cmp::PartialEq for ImaginateController {
+	fn eq(&self, other: &Self) -> bool {
+		Arc::ptr_eq(&self.0, &other.0)
+	}
+}
+
+impl core::hash::Hash for ImaginateController {
+	fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
+		core::ptr::hash(Arc::as_ptr(&self.0), state)
+	}
+}
+
+#[derive(Default, Debug, Clone, PartialEq, DynAny, specta::Type)]
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
 pub enum ImaginateStatus {
 	#[default]
-	Idle,
+	Ready,
+	ReadyDone,
 	Beginning,
-	Uploading(f64),
-	Generating,
+	Uploading,
+	Generating(f64),
 	Terminating,
 	Terminated,
+	Failed(String),
+}
+
+impl ImaginateStatus {
+	pub fn to_text(&self) -> Cow<'static, str> {
+		match self {
+			Self::Ready => Cow::Borrowed("Ready"),
+			Self::ReadyDone => Cow::Borrowed("Done"),
+			Self::Beginning => Cow::Borrowed("Beginning…"),
+			Self::Uploading => Cow::Borrowed("Downloading Image…"),
+			Self::Generating(percent) => Cow::Owned(format!("Generating {percent:.0}%")),
+			Self::Terminating => Cow::Owned(format!("Terminating…")),
+			Self::Terminated => Cow::Owned(format!("Terminated")),
+			Self::Failed(err) => Cow::Owned(format!("Failed: {err}")),
+		}
+	}
 }
 
 #[allow(clippy::derived_hash_with_manual_eq)]
 impl core::hash::Hash for ImaginateStatus {
 	fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
+		core::mem::discriminant(self).hash(state);
 		match self {
-			Self::Idle => 0.hash(state),
-			Self::Beginning => 1.hash(state),
-			Self::Uploading(f) => {
-				2.hash(state);
-				f.to_bits().hash(state);
-			}
-			Self::Generating => 3.hash(state),
-			Self::Terminating => 4.hash(state),
-			Self::Terminated => 5.hash(state),
+			Self::Ready | Self::ReadyDone | Self::Beginning | Self::Uploading | Self::Terminating | Self::Terminated => (),
+			Self::Generating(f) => f.to_bits().hash(state),
+			Self::Failed(err) => err.hash(state),
 		}
 	}
 }
 
-#[derive(Debug, Clone, PartialEq, specta::Type)]
-#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
-pub struct ImaginateBaseImage {
-	pub mime: String,
-	#[cfg_attr(feature = "serde", serde(rename = "imageData"))]
-	pub image_data: Vec<u8>,
-	pub size: DVec2,
+#[derive(PartialEq, Eq, Clone, Default, Debug)]
+pub enum ImaginateServerStatus {
+	#[default]
+	Unknown,
+	Checking,
+	Connected,
+	Failed(String),
+	Unavailable,
 }
 
-#[derive(Debug, Clone, PartialEq, specta::Type)]
-#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
-pub struct ImaginateMaskImage {
-	pub svg: String,
-	pub size: DVec2,
+impl ImaginateServerStatus {
+	pub fn to_text(&self) -> Cow<'static, str> {
+		match self {
+			Self::Unknown | Self::Checking => Cow::Borrowed("Checking..."),
+			Self::Connected => Cow::Borrowed("Connected"),
+			Self::Failed(err) => Cow::Owned(err.clone()),
+			Self::Unavailable => Cow::Borrowed("Unavailable"),
+		}
+	}
 }
 
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
@@ -180,24 +285,26 @@ impl std::fmt::Display for ImaginateSamplingMethod {
 	}
 }
 
-#[derive(Debug, Clone, PartialEq, specta::Type)]
+#[derive(Clone, Debug, PartialEq, Hash, specta::Type)]
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
-pub struct ImaginateGenerationParameters {
-	pub seed: u64,
-	pub samples: u32,
-	/// Use `ImaginateSamplingMethod::api_value()` to generate this string
-	#[cfg_attr(feature = "serde", serde(rename = "samplingMethod"))]
-	pub sampling_method: String,
-	#[cfg_attr(feature = "serde", serde(rename = "denoisingStrength"))]
-	pub image_creativity: Option<f64>,
-	#[cfg_attr(feature = "serde", serde(rename = "cfgScale"))]
-	pub text_guidance: f64,
-	#[cfg_attr(feature = "serde", serde(rename = "prompt"))]
-	pub text_prompt: String,
-	#[cfg_attr(feature = "serde", serde(rename = "negativePrompt"))]
-	pub negative_prompt: String,
-	pub resolution: (u32, u32),
-	#[cfg_attr(feature = "serde", serde(rename = "restoreFaces"))]
-	pub restore_faces: bool,
-	pub tiling: bool,
+pub struct ImaginatePreferences {
+	pub host_name: String,
+}
+
+impl graphene_core::application_io::GetImaginatePreferences for ImaginatePreferences {
+	fn get_host_name(&self) -> &str {
+		&self.host_name
+	}
+}
+
+impl Default for ImaginatePreferences {
+	fn default() -> Self {
+		Self {
+			host_name: "http://localhost:7860/".into(),
+		}
+	}
+}
+
+unsafe impl dyn_any::StaticType for ImaginatePreferences {
+	type Static = ImaginatePreferences;
 }