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>

editor/src/messages/tool/tool_messages/artboard_tool.rs

@@ -132,7 +132,7 @@ impl Fsm for ArtboardToolFsmState {
 		self,
 		event: ToolMessage,
 		tool_data: &mut Self::ToolData,
-		(document, _global_tool_data, input, font_cache): ToolActionHandlerData,
+		(document, _document_id, _global_tool_data, input, font_cache): ToolActionHandlerData,
 		_tool_options: &Self::ToolOptions,
 		responses: &mut VecDeque<Message>,
 	) -> Self {

editor/src/messages/tool/tool_messages/ellipse_tool.rs

@@ -122,7 +122,7 @@ impl Fsm for EllipseToolFsmState {
 		self,
 		event: ToolMessage,
 		tool_data: &mut Self::ToolData,
-		(document, global_tool_data, input, font_cache): ToolActionHandlerData,
+		(document, _document_id, global_tool_data, input, font_cache): ToolActionHandlerData,
 		_tool_options: &Self::ToolOptions,
 		responses: &mut VecDeque<Message>,
 	) -> Self {

editor/src/messages/tool/tool_messages/eyedropper_tool.rs

@@ -104,7 +104,7 @@ impl Fsm for EyedropperToolFsmState {
 		self,
 		event: ToolMessage,
 		_tool_data: &mut Self::ToolData,
-		(document, _global_tool_data, input, font_cache): ToolActionHandlerData,
+		(document, _document_id, _global_tool_data, input, font_cache): ToolActionHandlerData,
 		_tool_options: &Self::ToolOptions,
 		responses: &mut VecDeque<Message>,
 	) -> Self {

editor/src/messages/tool/tool_messages/fill_tool.rs

@@ -105,7 +105,7 @@ impl Fsm for FillToolFsmState {
 		self,
 		event: ToolMessage,
 		_tool_data: &mut Self::ToolData,
-		(document, global_tool_data, input, font_cache): ToolActionHandlerData,
+		(document, _document_id, global_tool_data, input, font_cache): ToolActionHandlerData,
 		_tool_options: &Self::ToolOptions,
 		responses: &mut VecDeque<Message>,
 	) -> Self {

editor/src/messages/tool/tool_messages/freehand_tool.rs

@@ -161,7 +161,7 @@ impl Fsm for FreehandToolFsmState {
 		self,
 		event: ToolMessage,
 		tool_data: &mut Self::ToolData,
-		(document, global_tool_data, input, _font_cache): ToolActionHandlerData,
+		(document, _document_id, global_tool_data, input, _font_cache): ToolActionHandlerData,
 		tool_options: &Self::ToolOptions,
 		responses: &mut VecDeque<Message>,
 	) -> Self {

editor/src/messages/tool/tool_messages/gradient_tool.rs

@@ -335,7 +335,7 @@ impl Fsm for GradientToolFsmState {
 		self,
 		event: ToolMessage,
 		tool_data: &mut Self::ToolData,
-		(document, global_tool_data, input, font_cache): ToolActionHandlerData,
+		(document, _document_id, global_tool_data, input, font_cache): ToolActionHandlerData,
 		tool_options: &Self::ToolOptions,
 		responses: &mut VecDeque<Message>,
 	) -> Self {

editor/src/messages/tool/tool_messages/imaginate_tool.rs

@@ -0,0 +1,231 @@
+use crate::consts::DRAG_THRESHOLD;
+use crate::messages::frontend::utility_types::MouseCursorIcon;
+use crate::messages::input_mapper::utility_types::input_keyboard::{Key, KeysGroup, MouseMotion};
+use crate::messages::layout::utility_types::layout_widget::PropertyHolder;
+use crate::messages::prelude::*;
+use crate::messages::tool::common_functionality::resize::Resize;
+use crate::messages::tool::utility_types::{EventToMessageMap, Fsm, ToolActionHandlerData, ToolMetadata, ToolTransition, ToolType};
+use crate::messages::tool::utility_types::{HintData, HintGroup, HintInfo};
+
+use graphene::Operation;
+
+use glam::DAffine2;
+use serde::{Deserialize, Serialize};
+
+#[derive(Default)]
+pub struct ImaginateTool {
+	fsm_state: ImaginateToolFsmState,
+	tool_data: ImaginateToolData,
+}
+
+#[remain::sorted]
+#[impl_message(Message, ToolMessage, Imaginate)]
+#[derive(PartialEq, Eq, Clone, Debug, Hash, Serialize, Deserialize)]
+pub enum ImaginateToolMessage {
+	// Standard messages
+	#[remain::unsorted]
+	Abort,
+
+	// Tool-specific messages
+	DragStart,
+	DragStop,
+	Resize {
+		center: Key,
+		lock_ratio: Key,
+	},
+}
+
+impl PropertyHolder for ImaginateTool {}
+
+impl<'a> MessageHandler<ToolMessage, ToolActionHandlerData<'a>> for ImaginateTool {
+	fn process_message(&mut self, message: ToolMessage, tool_data: ToolActionHandlerData<'a>, responses: &mut VecDeque<Message>) {
+		if message == ToolMessage::UpdateHints {
+			self.fsm_state.update_hints(responses);
+			return;
+		}
+
+		if message == ToolMessage::UpdateCursor {
+			self.fsm_state.update_cursor(responses);
+			return;
+		}
+
+		let new_state = self.fsm_state.transition(message, &mut self.tool_data, tool_data, &(), responses);
+
+		if self.fsm_state != new_state {
+			self.fsm_state = new_state;
+			self.fsm_state.update_hints(responses);
+			self.fsm_state.update_cursor(responses);
+		}
+	}
+
+	fn actions(&self) -> ActionList {
+		use ImaginateToolFsmState::*;
+
+		match self.fsm_state {
+			Ready => actions!(ImaginateToolMessageDiscriminant;
+				DragStart,
+			),
+			Drawing => actions!(ImaginateToolMessageDiscriminant;
+				DragStop,
+				Abort,
+				Resize,
+			),
+		}
+	}
+}
+
+impl ToolMetadata for ImaginateTool {
+	fn icon_name(&self) -> String {
+		"RasterImaginateTool".into()
+	}
+	fn tooltip(&self) -> String {
+		"Imaginate Tool".into()
+	}
+	fn tool_type(&self) -> crate::messages::tool::utility_types::ToolType {
+		ToolType::Imaginate
+	}
+}
+
+impl ToolTransition for ImaginateTool {
+	fn event_to_message_map(&self) -> EventToMessageMap {
+		EventToMessageMap {
+			document_dirty: None,
+			tool_abort: Some(ImaginateToolMessage::Abort.into()),
+			selection_changed: None,
+		}
+	}
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+enum ImaginateToolFsmState {
+	Ready,
+	Drawing,
+}
+
+impl Default for ImaginateToolFsmState {
+	fn default() -> Self {
+		ImaginateToolFsmState::Ready
+	}
+}
+#[derive(Clone, Debug, Default)]
+struct ImaginateToolData {
+	data: Resize,
+}
+
+impl Fsm for ImaginateToolFsmState {
+	type ToolData = ImaginateToolData;
+	type ToolOptions = ();
+
+	fn transition(
+		self,
+		event: ToolMessage,
+		tool_data: &mut Self::ToolData,
+		(document, _document_id, _global_tool_data, input, font_cache): ToolActionHandlerData,
+		_tool_options: &Self::ToolOptions,
+		responses: &mut VecDeque<Message>,
+	) -> Self {
+		use ImaginateToolFsmState::*;
+		use ImaginateToolMessage::*;
+
+		let mut shape_data = &mut tool_data.data;
+
+		if let ToolMessage::Imaginate(event) = event {
+			match (self, event) {
+				(Ready, DragStart) => {
+					shape_data.start(responses, document, input.mouse.position, font_cache);
+					responses.push_back(DocumentMessage::StartTransaction.into());
+					shape_data.path = Some(document.get_path_for_new_layer());
+					responses.push_back(DocumentMessage::DeselectAllLayers.into());
+
+					responses.push_back(
+						Operation::AddImaginateFrame {
+							path: shape_data.path.clone().unwrap(),
+							insert_index: -1,
+							transform: DAffine2::ZERO.to_cols_array(),
+						}
+						.into(),
+					);
+
+					Drawing
+				}
+				(state, Resize { center, lock_ratio }) => {
+					if let Some(message) = shape_data.calculate_transform(responses, document, center, lock_ratio, input) {
+						responses.push_back(message);
+					}
+
+					state
+				}
+				(Drawing, DragStop) => {
+					match shape_data.drag_start.distance(input.mouse.position) <= DRAG_THRESHOLD {
+						true => responses.push_back(DocumentMessage::AbortTransaction.into()),
+						false => responses.push_back(DocumentMessage::CommitTransaction.into()),
+					}
+
+					shape_data.cleanup(responses);
+
+					Ready
+				}
+				(Drawing, Abort) => {
+					responses.push_back(DocumentMessage::AbortTransaction.into());
+
+					shape_data.cleanup(responses);
+
+					Ready
+				}
+				_ => self,
+			}
+		} else {
+			self
+		}
+	}
+
+	fn update_hints(&self, responses: &mut VecDeque<Message>) {
+		let hint_data = match self {
+			ImaginateToolFsmState::Ready => HintData(vec![HintGroup(vec![
+				HintInfo {
+					key_groups: vec![],
+					key_groups_mac: None,
+					mouse: Some(MouseMotion::LmbDrag),
+					label: String::from("Draw Repaint Frame"),
+					plus: false,
+				},
+				HintInfo {
+					key_groups: vec![KeysGroup(vec![Key::Shift])],
+					key_groups_mac: None,
+					mouse: None,
+					label: String::from("Constrain Square"),
+					plus: true,
+				},
+				HintInfo {
+					key_groups: vec![KeysGroup(vec![Key::Alt])],
+					key_groups_mac: None,
+					mouse: None,
+					label: String::from("From Center"),
+					plus: true,
+				},
+			])]),
+			ImaginateToolFsmState::Drawing => HintData(vec![HintGroup(vec![
+				HintInfo {
+					key_groups: vec![KeysGroup(vec![Key::Shift])],
+					key_groups_mac: None,
+					mouse: None,
+					label: String::from("Constrain Square"),
+					plus: false,
+				},
+				HintInfo {
+					key_groups: vec![KeysGroup(vec![Key::Alt])],
+					key_groups_mac: None,
+					mouse: None,
+					label: String::from("From Center"),
+					plus: false,
+				},
+			])]),
+		};
+
+		responses.push_back(FrontendMessage::UpdateInputHints { hint_data }.into());
+	}
+
+	fn update_cursor(&self, responses: &mut VecDeque<Message>) {
+		responses.push_back(FrontendMessage::UpdateMouseCursor { cursor: MouseCursorIcon::Crosshair }.into());
+	}
+}

editor/src/messages/tool/tool_messages/line_tool.rs

@@ -170,7 +170,7 @@ impl Fsm for LineToolFsmState {
 		self,
 		event: ToolMessage,
 		tool_data: &mut Self::ToolData,
-		(document, global_tool_data, input, font_cache): ToolActionHandlerData,
+		(document, _document_id, global_tool_data, input, font_cache): ToolActionHandlerData,
 		tool_options: &Self::ToolOptions,
 		responses: &mut VecDeque<Message>,
 	) -> Self {

editor/src/messages/tool/tool_messages/mod.rs

@@ -4,6 +4,7 @@ pub mod eyedropper_tool;
 pub mod fill_tool;
 pub mod freehand_tool;
 pub mod gradient_tool;
+pub mod imaginate_tool;
 pub mod line_tool;
 pub mod navigate_tool;
 pub mod path_tool;

editor/src/messages/tool/tool_messages/navigate_tool.rs

@@ -126,7 +126,7 @@ impl Fsm for NavigateToolFsmState {
 		self,
 		message: ToolMessage,
 		tool_data: &mut Self::ToolData,
-		(_document, _global_tool_data, input, _font_cache): ToolActionHandlerData,
+		(_document, _document_id, _global_tool_data, input, _font_cache): ToolActionHandlerData,
 		_tool_options: &Self::ToolOptions,
 		messages: &mut VecDeque<Message>,
 	) -> Self {

editor/src/messages/tool/tool_messages/path_tool.rs

@@ -142,7 +142,7 @@ impl Fsm for PathToolFsmState {
 		self,
 		event: ToolMessage,
 		tool_data: &mut Self::ToolData,
-		(document, _global_tool_data, input, font_cache): ToolActionHandlerData,
+		(document, _document_id, _global_tool_data, input, font_cache): ToolActionHandlerData,
 		_tool_options: &Self::ToolOptions,
 		responses: &mut VecDeque<Message>,
 	) -> Self {

editor/src/messages/tool/tool_messages/pen_tool.rs

@@ -182,7 +182,7 @@ impl Fsm for PenToolFsmState {
 		self,
 		event: ToolMessage,
 		tool_data: &mut Self::ToolData,
-		(document, global_tool_data, input, font_cache): ToolActionHandlerData,
+		(document, _document_id, global_tool_data, input, font_cache): ToolActionHandlerData,
 		tool_options: &Self::ToolOptions,
 		responses: &mut VecDeque<Message>,
 	) -> Self {

editor/src/messages/tool/tool_messages/rectangle_tool.rs

@@ -121,7 +121,7 @@ impl Fsm for RectangleToolFsmState {
 		self,
 		event: ToolMessage,
 		tool_data: &mut Self::ToolData,
-		(document, global_tool_data, input, font_cache): ToolActionHandlerData,
+		(document, _document_id, global_tool_data, input, font_cache): ToolActionHandlerData,
 		_tool_options: &Self::ToolOptions,
 		responses: &mut VecDeque<Message>,
 	) -> Self {

editor/src/messages/tool/tool_messages/select_tool.rs

@@ -1,6 +1,6 @@
 use crate::application::generate_uuid;
 use crate::consts::{ROTATE_SNAP_ANGLE, SELECTION_TOLERANCE};
-use crate::messages::frontend::utility_types::MouseCursorIcon;
+use crate::messages::frontend::utility_types::{FrontendImageData, MouseCursorIcon};
 use crate::messages::input_mapper::utility_types::input_keyboard::{Key, KeysGroup, MouseMotion};
 use crate::messages::input_mapper::utility_types::input_mouse::ViewportPosition;
 use crate::messages::layout::utility_types::layout_widget::{Layout, LayoutGroup, PropertyHolder, Widget, WidgetCallback, WidgetHolder, WidgetLayout};
@@ -374,7 +374,7 @@ impl Fsm for SelectToolFsmState {
 		self,
 		event: ToolMessage,
 		tool_data: &mut Self::ToolData,
-		(document, _global_tool_data, input, font_cache): ToolActionHandlerData,
+		(document, document_id, _global_tool_data, input, font_cache): ToolActionHandlerData,
 		_tool_options: &Self::ToolOptions,
 		responses: &mut VecDeque<Message>,
 	) -> Self {
@@ -550,12 +550,6 @@ impl Fsm for SelectToolFsmState {
 						.flat_map(snapping::expand_bounds)
 						.collect();
 
-					if input.keyboard.get(duplicate as usize) && tool_data.not_duplicated_layers.is_none() {
-						tool_data.start_duplicates(document, responses);
-					} else if !input.keyboard.get(duplicate as usize) && tool_data.not_duplicated_layers.is_some() {
-						tool_data.stop_duplicates(responses);
-					}
-
 					let closest_move = tool_data.snap_manager.snap_layers(responses, document, snap, mouse_delta);
 					// TODO: Cache the result of `shallowest_unique_layers` to avoid this heavy computation every frame of movement, see https://github.com/GraphiteEditor/Graphite/pull/481
 					for path in Document::shallowest_unique_layers(tool_data.layers_dragging.iter()) {
@@ -568,6 +562,13 @@ impl Fsm for SelectToolFsmState {
 						);
 					}
 					tool_data.drag_current = mouse_position + closest_move;
+
+					if input.keyboard.get(duplicate as usize) && tool_data.not_duplicated_layers.is_none() {
+						tool_data.start_duplicates(document, document_id, responses);
+					} else if !input.keyboard.get(duplicate as usize) && tool_data.not_duplicated_layers.is_some() {
+						tool_data.stop_duplicates(responses);
+					}
+
 					Dragging
 				}
 				(ResizingBounds, PointerMove { axis_align, center, .. }) => {
@@ -919,8 +920,8 @@ impl Fsm for SelectToolFsmState {
 }
 
 impl SelectToolData {
-	/// Duplicates the currently dragging layers. Called when alt is pressed and the layers have not yet been duplicated.
-	fn start_duplicates(&mut self, document: &DocumentMessageHandler, responses: &mut VecDeque<Message>) {
+	/// Duplicates the currently dragging layers. Called when Alt is pressed and the layers have not yet been duplicated.
+	fn start_duplicates(&mut self, document: &DocumentMessageHandler, document_id: u64, responses: &mut VecDeque<Message>) {
 		responses.push_back(DocumentMessage::DeselectAllLayers.into());
 
 		self.not_duplicated_layers = Some(self.layers_dragging.clone());
@@ -938,19 +939,34 @@ impl SelectToolData {
 
 			// Copy the layers.
 			// Not using the Copy message allows us to retrieve the ids of the new layers to initialize the drag.
-			let layer = match document.graphene_document.layer(layer_path) {
+			let mut layer = match document.graphene_document.layer(layer_path) {
 				Ok(layer) => layer.clone(),
 				Err(e) => {
 					warn!("Could not access selected layer {:?}: {:?}", layer_path, e);
 					continue;
 				}
 			};
+
 			let layer_metadata = *document.layer_metadata(layer_path);
 			*layer_path.last_mut().unwrap() = generate_uuid();
 
+			let image_data = if let LayerDataType::Imaginate(imaginate) = &mut layer.data {
+				imaginate.blob_url = None;
+
+				imaginate.image_data.as_ref().map(|data| {
+					vec![FrontendImageData {
+						path: layer_path.clone(),
+						image_data: data.image_data.clone(),
+						mime: imaginate.mime.clone(),
+					}]
+				})
+			} else {
+				None
+			};
+
 			responses.push_back(
 				Operation::InsertLayer {
-					layer,
+					layer: Box::new(layer),
 					destination_path: layer_path.clone(),
 					insert_index: -1,
 				}
@@ -964,10 +980,14 @@ impl SelectToolData {
 				}
 				.into(),
 			);
+
+			if let Some(image_data) = image_data {
+				responses.push_back(FrontendMessage::UpdateImageData { image_data, document_id }.into());
+			}
 		}
 	}
 
-	/// Removes the duplicated layers. Called when alt is released and the layers have been duplicated.
+	/// Removes the duplicated layers. Called when Alt is released and the layers have been duplicated.
 	fn stop_duplicates(&mut self, responses: &mut VecDeque<Message>) {
 		let originals = match self.not_duplicated_layers.take() {
 			Some(x) => x,

editor/src/messages/tool/tool_messages/shape_tool.rs

@@ -162,7 +162,7 @@ impl Fsm for ShapeToolFsmState {
 		self,
 		event: ToolMessage,
 		tool_data: &mut Self::ToolData,
-		(document, global_tool_data, input, font_cache): ToolActionHandlerData,
+		(document, _document_id, global_tool_data, input, font_cache): ToolActionHandlerData,
 		tool_options: &Self::ToolOptions,
 		responses: &mut VecDeque<Message>,
 	) -> Self {

editor/src/messages/tool/tool_messages/spline_tool.rs

@@ -170,7 +170,7 @@ impl Fsm for SplineToolFsmState {
 		self,
 		event: ToolMessage,
 		tool_data: &mut Self::ToolData,
-		(document, global_tool_data, input, font_cache): ToolActionHandlerData,
+		(document, _document_id, global_tool_data, input, font_cache): ToolActionHandlerData,
 		tool_options: &Self::ToolOptions,
 		responses: &mut VecDeque<Message>,
 	) -> Self {

editor/src/messages/tool/tool_messages/text_tool.rs

@@ -274,7 +274,7 @@ impl Fsm for TextToolFsmState {
 		self,
 		event: ToolMessage,
 		tool_data: &mut Self::ToolData,
-		(document, global_tool_data, input, font_cache): ToolActionHandlerData,
+		(document, _document_id, global_tool_data, input, font_cache): ToolActionHandlerData,
 		tool_options: &Self::ToolOptions,
 		responses: &mut VecDeque<Message>,
 	) -> Self {
@@ -302,7 +302,7 @@ impl Fsm for TextToolFsmState {
 					{
 						if state == TextToolFsmState::Editing {
 							responses.push_back(
-								DocumentMessage::SetTexboxEditability {
+								DocumentMessage::SetTextboxEditability {
 									path: tool_data.path.clone(),
 									editable: false,
 								}
@@ -313,7 +313,7 @@ impl Fsm for TextToolFsmState {
 						tool_data.path = l.clone();
 
 						responses.push_back(
-							DocumentMessage::SetTexboxEditability {
+							DocumentMessage::SetTextboxEditability {
 								path: tool_data.path.clone(),
 								editable: true,
 							}
@@ -358,7 +358,7 @@ impl Fsm for TextToolFsmState {
 						);
 
 						responses.push_back(
-							DocumentMessage::SetTexboxEditability {
+							DocumentMessage::SetTextboxEditability {
 								path: tool_data.path.clone(),
 								editable: true,
 							}
@@ -376,7 +376,7 @@ impl Fsm for TextToolFsmState {
 					} else {
 						// Removing old text as editable
 						responses.push_back(
-							DocumentMessage::SetTexboxEditability {
+							DocumentMessage::SetTextboxEditability {
 								path: tool_data.path.clone(),
 								editable: false,
 							}
@@ -393,7 +393,7 @@ impl Fsm for TextToolFsmState {
 				(state, Abort) => {
 					if state == TextToolFsmState::Editing {
 						responses.push_back(
-							DocumentMessage::SetTexboxEditability {
+							DocumentMessage::SetTextboxEditability {
 								path: tool_data.path.clone(),
 								editable: false,
 							}
@@ -420,7 +420,7 @@ impl Fsm for TextToolFsmState {
 					);
 
 					responses.push_back(
-						DocumentMessage::SetTexboxEditability {
+						DocumentMessage::SetTextboxEditability {
 							path: tool_data.path.clone(),
 							editable: false,
 						}