Integrate Stable Diffusion with the Imaginate layer (#784)

* Add AI Artist layer * WIP add a button to download the rendered folder under an AI Artist layer * Successfully download the correct image * Break out image downloading JS into helper function * Change file download from using data URLs to blob URLs * WIP rasterize to blob * Remove dimensions from AI Artist layer * Successfully draw rasterized image on layer after calculation * Working txt2img generation based on user prompt * Add img2img and the main parameters * Fix ability to rasterize multi-depth documents with blob URL images by switching them to base64 * Fix test * Rasterize with artboard background color * Allow aspect ratio stretch of AI Artist images * Add automatic resolution choosing * Add a terminate button, and make the lifecycle more robust * Add negative prompt * Add range bounds for parameter inputs * Add seed * Add tiling and restore faces * Add server status check, server hostname customization, and resizing layer to fit AI Artist resolution * Fix background color of infinite canvas rasterization * Escape prompt text sent in the JSON * Revoke blob URLs when cleared/replaced to reduce memory leak * Fix welcome screen logo color * Add PreferencesMessageHandler * Add persistent storage of preferences * Fix crash introduced in previous commit when moving mouse on page load * Add tooltips to the AI Artist layer properties * Integrate AI Artist tool into the raster section of the tool shelf * Add a refresh button to the connection status * Fix crash when generating and switching to a different document tab * Add persistent image storage to AI Artist layers and fix duplication bugs * Add a generate with random seed button * Simplify and standardize message names * Majorly improve robustness of networking code * Fix race condition causing default server hostname to show disconnected when app loads with AI Artist layer selected (probably, not confirmed fixed) * Clean up messages and function calls by changing arguments into structs * Update API to more recent server commit * Add support for picking the sampling method * Add machinery for filtering selected layers with type * Replace placeholder button icons * Improve the random icon by tilting the dice * Use selected_layers() instead of repeating that code * Fix borrow error * Change message flow in progress towards fixing #797 * Allow loading image on non-active document (fixes #797) * Reduce code duplication with rasterization * Add AI Artist tool and layer icons, and remove ugly node layer icon style * Rename "AI Artist" codename to "Imaginate" feature name Co-authored-by: otdavies <oliver@psyfer.io> Co-authored-by: 0hypercube <0hypercube@gmail.com>
2026-10-01 07:28:12 +08:00 · 2022-10-18 22:33:27 -07:00
parent 562217015d
commit fe1a03fac7
118 changed files with 3767 additions and 678 deletions
@@ -0,0 +1,286 @@
+use super::base64_serde;
+use super::layer_info::LayerData;
+use super::style::{RenderData, ViewMode};
+use crate::intersection::{intersect_quad_bez_path, Quad};
+use crate::layers::text_layer::FontCache;
+use crate::LayerId;
+
+use glam::{DAffine2, DMat2, DVec2};
+use kurbo::{Affine, BezPath, Shape as KurboShape};
+use serde::{Deserialize, Serialize};
+use std::fmt::Write;
+
+#[derive(Clone, PartialEq, Deserialize, Serialize)]
+pub struct ImaginateLayer {
+	// User-configurable layer parameters
+	pub seed: u64,
+	pub samples: u32,
+	pub sampling_method: ImaginateSamplingMethod,
+	pub use_img2img: bool,
+	pub denoising_strength: f64,
+	pub cfg_scale: f64,
+	pub prompt: String,
+	pub negative_prompt: String,
+	pub restore_faces: bool,
+	pub tiling: bool,
+
+	// Image stored in layer after generation completes
+	pub image_data: Option<ImaginateImageData>,
+	pub mime: String,
+	/// 0 is not started, 100 is complete.
+	pub percent_complete: f64,
+
+	// TODO: Have the browser dispose of this blob URL when this is dropped (like when the layer is deleted)
+	#[serde(skip)]
+	pub blob_url: Option<String>,
+	#[serde(skip)]
+	pub status: ImaginateStatus,
+	#[serde(skip)]
+	pub dimensions: DVec2,
+}
+
+#[derive(Default, Debug, Clone, PartialEq, Deserialize, Serialize)]
+pub enum ImaginateStatus {
+	#[default]
+	Idle,
+	Beginning,
+	Uploading(f64),
+	Generating,
+	Terminating,
+	Terminated,
+}
+
+#[derive(Debug, Clone, Eq, PartialEq, Deserialize, Serialize)]
+pub struct ImaginateImageData {
+	#[serde(serialize_with = "base64_serde::as_base64", deserialize_with = "base64_serde::from_base64")]
+	pub image_data: Vec<u8>,
+}
+
+#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
+pub struct ImaginateBaseImage {
+	pub svg: String,
+	pub size: DVec2,
+}
+
+#[derive(Debug, Default, Copy, Clone, Eq, PartialEq, Deserialize, Serialize)]
+pub enum ImaginateSamplingMethod {
+	#[default]
+	EulerA,
+	Euler,
+	LMS,
+	Heun,
+	DPM2,
+	DPM2A,
+	DPMFast,
+	DPMAdaptive,
+	LMSKarras,
+	DPM2Karras,
+	DPM2AKarras,
+	DDIM,
+	PLMS,
+}
+
+impl ImaginateSamplingMethod {
+	pub fn api_value(&self) -> &str {
+		match self {
+			ImaginateSamplingMethod::EulerA => "Euler a",
+			ImaginateSamplingMethod::Euler => "Euler",
+			ImaginateSamplingMethod::LMS => "LMS",
+			ImaginateSamplingMethod::Heun => "Heun",
+			ImaginateSamplingMethod::DPM2 => "DPM2",
+			ImaginateSamplingMethod::DPM2A => "DPM2 a",
+			ImaginateSamplingMethod::DPMFast => "DPM fast",
+			ImaginateSamplingMethod::DPMAdaptive => "DPM adaptive",
+			ImaginateSamplingMethod::LMSKarras => "LMS Karras",
+			ImaginateSamplingMethod::DPM2Karras => "DPM2 Karras",
+			ImaginateSamplingMethod::DPM2AKarras => "DPM2 a Karras",
+			ImaginateSamplingMethod::DDIM => "DDIM",
+			ImaginateSamplingMethod::PLMS => "PLMS",
+		}
+	}
+
+	pub fn list() -> [ImaginateSamplingMethod; 13] {
+		[
+			ImaginateSamplingMethod::EulerA,
+			ImaginateSamplingMethod::Euler,
+			ImaginateSamplingMethod::LMS,
+			ImaginateSamplingMethod::Heun,
+			ImaginateSamplingMethod::DPM2,
+			ImaginateSamplingMethod::DPM2A,
+			ImaginateSamplingMethod::DPMFast,
+			ImaginateSamplingMethod::DPMAdaptive,
+			ImaginateSamplingMethod::LMSKarras,
+			ImaginateSamplingMethod::DPM2Karras,
+			ImaginateSamplingMethod::DPM2AKarras,
+			ImaginateSamplingMethod::DDIM,
+			ImaginateSamplingMethod::PLMS,
+		]
+	}
+}
+
+impl std::fmt::Display for ImaginateSamplingMethod {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		match self {
+			ImaginateSamplingMethod::EulerA => write!(f, "Euler A (Recommended)"),
+			ImaginateSamplingMethod::Euler => write!(f, "Euler"),
+			ImaginateSamplingMethod::LMS => write!(f, "LMS"),
+			ImaginateSamplingMethod::Heun => write!(f, "Heun"),
+			ImaginateSamplingMethod::DPM2 => write!(f, "DPM2"),
+			ImaginateSamplingMethod::DPM2A => write!(f, "DPM2 A"),
+			ImaginateSamplingMethod::DPMFast => write!(f, "DPM Fast"),
+			ImaginateSamplingMethod::DPMAdaptive => write!(f, "DPM Adaptive"),
+			ImaginateSamplingMethod::LMSKarras => write!(f, "LMS Karras"),
+			ImaginateSamplingMethod::DPM2Karras => write!(f, "DPM2 Karras"),
+			ImaginateSamplingMethod::DPM2AKarras => write!(f, "DPM2 A Karras"),
+			ImaginateSamplingMethod::DDIM => write!(f, "DDIM"),
+			ImaginateSamplingMethod::PLMS => write!(f, "PLMS"),
+		}
+	}
+}
+
+#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
+pub struct ImaginateGenerationParameters {
+	pub seed: u64,
+	pub samples: u32,
+	/// Use `ImaginateSamplingMethod::api_value()` to generate this string
+	#[serde(rename = "samplingMethod")]
+	pub sampling_method: String,
+	#[serde(rename = "denoisingStrength")]
+	pub denoising_strength: Option<f64>,
+	#[serde(rename = "cfgScale")]
+	pub cfg_scale: f64,
+	pub prompt: String,
+	#[serde(rename = "negativePrompt")]
+	pub negative_prompt: String,
+	pub resolution: (u64, u64),
+	#[serde(rename = "restoreFaces")]
+	pub restore_faces: bool,
+	pub tiling: bool,
+}
+
+impl Default for ImaginateLayer {
+	fn default() -> Self {
+		Self {
+			seed: 0,
+			samples: 30,
+			sampling_method: Default::default(),
+			use_img2img: false,
+			denoising_strength: 0.66,
+			cfg_scale: 10.,
+			prompt: "".into(),
+			negative_prompt: "".into(),
+			restore_faces: false,
+			tiling: false,
+
+			image_data: None,
+			mime: "image/png".into(),
+
+			blob_url: None,
+			percent_complete: 0.,
+			status: Default::default(),
+			dimensions: Default::default(),
+		}
+	}
+}
+
+impl LayerData for ImaginateLayer {
+	fn render(&mut self, svg: &mut String, _svg_defs: &mut String, transforms: &mut Vec<DAffine2>, render_data: RenderData) {
+		let transform = self.transform(transforms, render_data.view_mode);
+		let inverse = transform.inverse();
+
+		let matrix_values = transform.matrix2.to_cols_array();
+		let (width, height) = (matrix_values[0], matrix_values[3]);
+
+		if !inverse.is_finite() {
+			let _ = write!(svg, "<!-- SVG shape has an invalid transform -->");
+			return;
+		}
+
+		let _ = writeln!(svg, r#"<g transform="matrix("#);
+		inverse.to_cols_array().iter().enumerate().for_each(|(i, entry)| {
+			let _ = svg.write_str(&(entry.to_string() + if i == 5 { "" } else { "," }));
+		});
+		let _ = svg.write_str(r#")">"#);
+
+		if let Some(blob_url) = &self.blob_url {
+			let _ = write!(
+				svg,
+				r#"<image width="{}" height="{}" transform="matrix(1,0,0,1,{},{})" preserveAspectRatio="none" href="{}"/>"#,
+				width.abs(),
+				height.abs(),
+				if width >= 0. { transform.translation.x } else { transform.translation.x + width },
+				if height >= 0. { transform.translation.y } else { transform.translation.y + height },
+				blob_url
+			);
+		} else {
+			let _ = write!(
+				svg,
+				r#"<rect width="{}" height="{}" transform="matrix(1,0,0,1,{},{})" fill="none" stroke="var(--color-data-raster)" stroke-width="3" stroke-dasharray="8"/>"#,
+				width.abs(),
+				height.abs(),
+				if width >= 0. { transform.translation.x } else { transform.translation.x + width },
+				if height >= 0. { transform.translation.y } else { transform.translation.y + height },
+			);
+		}
+
+		let _ = svg.write_str("</g>");
+	}
+
+	fn bounding_box(&self, transform: glam::DAffine2, _font_cache: &FontCache) -> Option<[DVec2; 2]> {
+		let mut path = self.bounds();
+
+		if transform.matrix2 == DMat2::ZERO {
+			return None;
+		}
+		path.apply_affine(glam_to_kurbo(transform));
+
+		let kurbo::Rect { x0, y0, x1, y1 } = path.bounding_box();
+		Some([(x0, y0).into(), (x1, y1).into()])
+	}
+
+	fn intersects_quad(&self, quad: Quad, path: &mut Vec<LayerId>, intersections: &mut Vec<Vec<LayerId>>, _font_cache: &FontCache) {
+		if intersect_quad_bez_path(quad, &self.bounds(), true) {
+			intersections.push(path.clone());
+		}
+	}
+}
+
+impl ImaginateLayer {
+	pub fn transform(&self, transforms: &[DAffine2], mode: ViewMode) -> DAffine2 {
+		let start = match mode {
+			ViewMode::Outline => 0,
+			_ => (transforms.len() as i32 - 1).max(0) as usize,
+		};
+		transforms.iter().skip(start).cloned().reduce(|a, b| a * b).unwrap_or(DAffine2::IDENTITY)
+	}
+
+	fn bounds(&self) -> BezPath {
+		kurbo::Rect::from_origin_size(kurbo::Point::ZERO, kurbo::Size::new(1., 1.)).to_path(0.)
+	}
+}
+
+fn glam_to_kurbo(transform: DAffine2) -> Affine {
+	Affine::new(transform.to_cols_array())
+}
+
+impl std::fmt::Debug for ImaginateLayer {
+	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+		f.debug_struct("ImaginateLayer")
+			.field("seed", &self.seed)
+			.field("samples", &self.samples)
+			.field("use_img2img", &self.use_img2img)
+			.field("denoising_strength", &self.denoising_strength)
+			.field("cfg_scale", &self.cfg_scale)
+			.field("prompt", &self.prompt)
+			.field("negative_prompt", &self.negative_prompt)
+			.field("restore_faces", &self.restore_faces)
+			.field("tiling", &self.tiling)
+			.field("image_data", &self.image_data.as_ref().map(|_| "..."))
+			.field("mime", &self.mime)
+			.field("percent_complete", &self.percent_complete)
+			.field("blob_url", &self.blob_url)
+			.field("status", &self.status)
+			.field("dimensions", &self.dimensions)
+			.finish()
+	}
+}