Add inpainting and outpainting to Imaginate (#864)

* Do not select layer immediatly on drag

* Add LayerReferenceInput MVP widget

* Properties Panel

* Fix dragging marker flicker

* Change mask shape to outline

* Add mask rendering

* Simplify select code

* Remove colours

* Fix inpaint/outpaint and rearrage widget UX

* Add mask blur and mask starting fill parameters

* Guard for the case when the layer is missing

* Add icon to LayerReferenceInput to finalize its UI

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

editor/src/messages/frontend/frontend_message.rs

@@ -8,7 +8,7 @@ use crate::messages::prelude::*;
 use crate::messages::tool::utility_types::HintData;
 
 use graphene::color::Color;
-use graphene::layers::imaginate_layer::{ImaginateBaseImage, ImaginateGenerationParameters};
+use graphene::layers::imaginate_layer::{ImaginateBaseImage, ImaginateGenerationParameters, ImaginateMaskFillContent, ImaginateMaskPaintMode};
 use graphene::layers::text_layer::Font;
 use graphene::LayerId;
 
@@ -60,6 +60,14 @@ pub enum FrontendMessage {
 		parameters: ImaginateGenerationParameters,
 		#[serde(rename = "baseImage")]
 		base_image: Option<ImaginateBaseImage>,
+		#[serde(rename = "maskImage")]
+		mask_image: Option<ImaginateBaseImage>,
+		#[serde(rename = "maskPaintMode")]
+		mask_paint_mode: ImaginateMaskPaintMode,
+		#[serde(rename = "maskBlurPx")]
+		mask_blur_px: u32,
+		#[serde(rename = "maskFillContent")]
+		mask_fill_content: ImaginateMaskFillContent,
 		hostname: String,
 		#[serde(rename = "refreshFrequency")]
 		refresh_frequency: f64,

editor/src/messages/layout/layout_message_handler.rs

@@ -6,8 +6,10 @@ use crate::messages::prelude::*;
 
 use graphene::color::Color;
 use graphene::layers::text_layer::Font;
+use graphene::LayerId;
 
 use serde_json::Value;
+use std::ops::Not;
 
 #[derive(Debug, Clone, Default)]
 pub struct LayoutMessageHandler {
@@ -117,6 +119,19 @@ impl<F: Fn(&MessageDiscriminant) -> Vec<KeysGroup>> MessageHandler<LayoutMessage
 						let callback_message = (invisible.on_update.callback)(&());
 						responses.push_back(callback_message);
 					}
+					Widget::LayerReferenceInput(layer_reference_input) => {
+						let update_value = value.is_null().not().then(|| {
+							value
+								.as_str()
+								.expect("LayerReferenceInput update was not of type: string")
+								.split(',')
+								.map(|id| id.parse::<LayerId>().unwrap())
+								.collect::<Vec<_>>()
+						});
+						layer_reference_input.value = update_value;
+						let callback_message = (layer_reference_input.on_update.callback)(layer_reference_input);
+						responses.push_back(callback_message);
+					}
 					Widget::NumberInput(number_input) => match value {
 						Value::Number(num) => {
 							let update_value = num.as_f64().unwrap();

editor/src/messages/layout/utility_types/layout_widget.rs

@@ -61,11 +61,12 @@ impl Layout {
 				let mut tooltip_shortcut = match &mut widget_holder.widget {
 					Widget::CheckboxInput(widget) => Some((&mut widget.tooltip, &mut widget.tooltip_shortcut)),
 					Widget::ColorInput(widget) => Some((&mut widget.tooltip, &mut widget.tooltip_shortcut)),
-					Widget::IconButton(widget) => Some((&mut widget.tooltip, &mut widget.tooltip_shortcut)),
-					Widget::OptionalInput(widget) => Some((&mut widget.tooltip, &mut widget.tooltip_shortcut)),
 					Widget::DropdownInput(widget) => Some((&mut widget.tooltip, &mut widget.tooltip_shortcut)),
 					Widget::FontInput(widget) => Some((&mut widget.tooltip, &mut widget.tooltip_shortcut)),
+					Widget::IconButton(widget) => Some((&mut widget.tooltip, &mut widget.tooltip_shortcut)),
+					Widget::LayerReferenceInput(widget) => Some((&mut widget.tooltip, &mut widget.tooltip_shortcut)),
 					Widget::NumberInput(widget) => Some((&mut widget.tooltip, &mut widget.tooltip_shortcut)),
+					Widget::OptionalInput(widget) => Some((&mut widget.tooltip, &mut widget.tooltip_shortcut)),
 					Widget::PopoverButton(widget) => Some((&mut widget.tooltip, &mut widget.tooltip_shortcut)),
 					Widget::TextButton(widget) => Some((&mut widget.tooltip, &mut widget.tooltip_shortcut)),
 					Widget::IconLabel(_)
@@ -290,6 +291,7 @@ pub enum Widget {
 	IconButton(IconButton),
 	IconLabel(IconLabel),
 	InvisibleStandinInput(InvisibleStandinInput),
+	LayerReferenceInput(LayerReferenceInput),
 	NumberInput(NumberInput),
 	OptionalInput(OptionalInput),
 	PivotAssist(PivotAssist),

editor/src/messages/layout/utility_types/widgets/input_widgets.rs

@@ -1,7 +1,7 @@
 use crate::messages::input_mapper::utility_types::misc::ActionKeys;
 use crate::messages::layout::utility_types::layout_widget::WidgetCallback;
 
-use graphene::color::Color;
+use graphene::{color::Color, layers::layer_info::LayerDataTypeDiscriminant, LayerId};
 
 use derivative::*;
 use serde::{Deserialize, Serialize};
@@ -150,6 +150,34 @@ pub struct InvisibleStandinInput {
 	pub on_update: WidgetCallback<()>,
 }
 
+#[derive(Clone, Serialize, Deserialize, Derivative)]
+#[derivative(Debug, PartialEq, Default)]
+pub struct LayerReferenceInput {
+	pub value: Option<Vec<LayerId>>,
+
+	#[serde(rename = "layerName")]
+	pub layer_name: Option<String>,
+
+	#[serde(rename = "layerType")]
+	pub layer_type: Option<LayerDataTypeDiscriminant>,
+
+	pub disabled: bool,
+
+	pub tooltip: String,
+
+	#[serde(skip)]
+	pub tooltip_shortcut: Option<ActionKeys>,
+
+	// Styling
+	#[serde(rename = "minWidth")]
+	pub min_width: u32,
+
+	// Callbacks
+	#[serde(skip)]
+	#[derivative(Debug = "ignore", PartialEq = "ignore")]
+	pub on_update: WidgetCallback<LayerReferenceInput>,
+}
+
 #[derive(Clone, Serialize, Deserialize, Derivative)]
 #[derivative(Debug, PartialEq, Default)]
 pub struct NumberInput {

editor/src/messages/portfolio/document/document_message_handler.rs

@@ -963,23 +963,55 @@ impl DocumentMessageHandler {
 			restore_faces: imaginate_layer.restore_faces,
 			tiling: imaginate_layer.tiling,
 		};
-		let base_image = if imaginate_layer.use_img2img {
+		let mask_paint_mode = imaginate_layer.mask_paint_mode;
+		let mask_blur_px = imaginate_layer.mask_blur_px;
+		let mask_fill_content = imaginate_layer.mask_fill_content;
+		let (base_image, mask_image) = if imaginate_layer.use_img2img {
+			let mask = imaginate_layer.mask_layer_ref.clone();
+
 			// Calculate the size of the region to be exported
 			let size = DVec2::new(transform.transform_vector2(DVec2::new(1., 0.)).length(), transform.transform_vector2(DVec2::new(0., 1.)).length());
 
 			let old_transforms = self.remove_document_transform();
 			let svg = self.render_document(size, transform.inverse(), persistent_data, DocumentRenderMode::OnlyBelowLayerInFolder(&layer_path));
-			self.restore_document_transform(old_transforms);
 
-			Some(ImaginateBaseImage { svg, size })
+			let mask_image = mask.and_then(|mask_layer_path| match self.graphene_document.layer(&mask_layer_path) {
+				Ok(_) => {
+					let svg = self.render_document(size, transform.inverse(), persistent_data, DocumentRenderMode::LayerCutout(&mask_layer_path, Color::WHITE));
+
+					Some(ImaginateBaseImage { svg, size })
+				}
+				Err(_) => None,
+			});
+
+			if mask_image.is_none() {
+				return Some(
+					DialogMessage::DisplayDialogError {
+						title: "Masking layer is missing".into(),
+						description: "
+							It may have been deleted or moved. Please drag a new layer reference\n\
+							into the 'Masking Layer' parameter input, then generate again."
+							.trim()
+							.into(),
+					}
+					.into(),
+				);
+			}
+
+			self.restore_document_transform(old_transforms);
+			(Some(ImaginateBaseImage { svg, size }), mask_image)
 		} else {
-			None
+			(None, None)
 		};
 
 		Some(
 			FrontendMessage::TriggerImaginateGenerate {
 				parameters,
 				base_image,
+				mask_image,
+				mask_paint_mode,
+				mask_blur_px,
+				mask_fill_content,
 				hostname: preferences.imaginate_server_hostname.clone(),
 				refresh_frequency: preferences.imaginate_refresh_frequency,
 				document_id,
@@ -1042,13 +1074,17 @@ impl DocumentMessageHandler {
 
 		let render_data = RenderData::new(ViewMode::Normal, &persistent_data.font_cache, None);
 
-		let artwork = match render_mode {
-			DocumentRenderMode::Root => self.graphene_document.render_root(render_data),
-			DocumentRenderMode::OnlyBelowLayerInFolder(below_layer_path) => self.graphene_document.render_layers_below(below_layer_path, render_data).unwrap(),
+		let (artwork, outside) = match render_mode {
+			DocumentRenderMode::Root => (self.graphene_document.render_root(render_data), None),
+			DocumentRenderMode::OnlyBelowLayerInFolder(below_layer_path) => (self.graphene_document.render_layers_below(below_layer_path, render_data).unwrap(), None),
+			DocumentRenderMode::LayerCutout(layer_path, background) => (self.graphene_document.render_layer(layer_path, render_data).unwrap(), Some(background)),
 		};
 		let artboards = self.artboard_message_handler.artboards_graphene_document.render_root(render_data);
-		let outside_artboards_color = if self.artboard_message_handler.artboard_ids.is_empty() { "#ffffff" } else { "#222222" };
-		let outside_artboards = format!(r#"<rect x="0" y="0" width="100%" height="100%" fill="{}" />"#, outside_artboards_color);
+		let outside_artboards_color = outside.map_or_else(
+			|| if self.artboard_message_handler.artboard_ids.is_empty() { "ffffff" } else { "222222" }.to_string(),
+			|col| col.rgba_hex(),
+		);
+		let outside_artboards = format!(r##"<rect x="0" y="0" width="100%" height="100%" fill="#{}" />"##, outside_artboards_color);
 		let matrix = transform
 			.to_cols_array()
 			.iter()

editor/src/messages/portfolio/document/properties_panel/properties_panel_message.rs

@@ -3,7 +3,7 @@ use crate::messages::layout::utility_types::widgets::assist_widgets::PivotPositi
 use crate::messages::portfolio::document::utility_types::misc::TargetDocument;
 use crate::messages::prelude::*;
 
-use graphene::layers::imaginate_layer::ImaginateSamplingMethod;
+use graphene::layers::imaginate_layer::{ImaginateMaskFillContent, ImaginateMaskPaintMode, ImaginateSamplingMethod};
 use graphene::layers::style::{Fill, Stroke};
 use graphene::LayerId;
 
@@ -29,6 +29,10 @@ pub enum PropertiesPanelMessage {
 	SetActiveLayers { paths: Vec<Vec<LayerId>>, document: TargetDocument },
 	SetImaginateCfgScale { cfg_scale: f64 },
 	SetImaginateDenoisingStrength { denoising_strength: f64 },
+	SetImaginateLayerPath { layer_path: Option<Vec<LayerId>> },
+	SetImaginateMaskBlurPx { mask_blur_px: u32 },
+	SetImaginateMaskFillContent { mode: ImaginateMaskFillContent },
+	SetImaginateMaskPaintMode { paint: ImaginateMaskPaintMode },
 	SetImaginateNegativePrompt { negative_prompt: String },
 	SetImaginatePrompt { prompt: String },
 	SetImaginateRestoreFaces { restore_faces: bool },

editor/src/messages/portfolio/document/properties_panel/properties_panel_message_handler.rs

@@ -140,10 +140,11 @@ impl<'a> MessageHandler<PropertiesPanelMessage, (&PersistentData, PropertiesPane
 			}
 			ResendActiveProperties => {
 				if let Some((path, target_document)) = self.active_selection.clone() {
-					let layer = get_document(target_document).layer(&path).unwrap();
+					let document = get_document(target_document);
+					let layer = document.layer(&path).unwrap();
 					match target_document {
 						TargetDocument::Artboard => register_artboard_layer_properties(layer, responses, persistent_data),
-						TargetDocument::Artwork => register_artwork_layer_properties(path, layer, responses, persistent_data, node_graph_message_handler),
+						TargetDocument::Artwork => register_artwork_layer_properties(document, path, layer, responses, persistent_data, node_graph_message_handler),
 					}
 				}
 			}
@@ -166,6 +167,22 @@ impl<'a> MessageHandler<PropertiesPanelMessage, (&PersistentData, PropertiesPane
 				let (path, _) = self.active_selection.clone().expect("Received update for properties panel with no active layer");
 				responses.push_back(Operation::ImaginateSetDenoisingStrength { path, denoising_strength }.into());
 			}
+			SetImaginateLayerPath { layer_path } => {
+				let (path, _) = self.active_selection.clone().expect("Received update for properties panel with no active layer");
+				responses.push_back(Operation::ImaginateSetLayerPath { path, layer_path }.into());
+			}
+			SetImaginateMaskBlurPx { mask_blur_px } => {
+				let (path, _) = self.active_selection.clone().expect("Received update for properties panel with no active layer");
+				responses.push_back(Operation::ImaginateSetMaskBlurPx { path, mask_blur_px }.into());
+			}
+			SetImaginateMaskFillContent { mode } => {
+				let (path, _) = self.active_selection.clone().expect("Received update for properties panel with no active layer");
+				responses.push_back(Operation::ImaginateSetMaskFillContent { path, mode }.into());
+			}
+			SetImaginateMaskPaintMode { paint } => {
+				let (path, _) = self.active_selection.clone().expect("Received update for properties panel with no active layer");
+				responses.push_back(Operation::ImaginateSetMaskPaintMode { path, paint }.into());
+			}
 			SetImaginateSamples { samples } => {
 				let (path, _) = self.active_selection.clone().expect("Received update for properties panel with no active layer");
 				responses.push_back(Operation::ImaginateSetSamples { path, samples }.into());

editor/src/messages/portfolio/document/utility_types/layer_panel.rs

@@ -61,12 +61,12 @@ impl LayerPanelEntry {
 	pub fn new(layer_metadata: &LayerMetadata, transform: DAffine2, layer: &Layer, path: Vec<LayerId>, font_cache: &FontCache) -> Self {
 		let name = layer.name.clone().unwrap_or_else(|| String::from(""));
 
-		let tooltip = if cfg!(debug_assertions) {
-			let joined = &path.iter().map(|id| id.to_string()).collect::<Vec<_>>().join(" / ");
-			name.clone() + "\nLayer Path: " + joined.as_str()
-		} else {
-			name.clone()
-		};
+		let mut tooltip = name.clone();
+		if cfg!(debug_assertions) {
+			tooltip += "\nLayer Path: ";
+			tooltip += &path.iter().map(|id| id.to_string()).collect::<Vec<_>>().join(" / ");
+			tooltip = tooltip.trim().to_string();
+		}
 
 		let arr = layer.data.bounding_box(transform, font_cache).unwrap_or([DVec2::ZERO, DVec2::ZERO]);
 		let arr = arr.iter().map(|x| (*x).into()).collect::<Vec<(f64, f64)>>();

editor/src/messages/portfolio/document/utility_types/misc.rs

@@ -1,5 +1,6 @@
 pub use super::layer_panel::{LayerMetadata, LayerPanelEntry};
 
+use graphene::color::Color;
 use graphene::document::Document as GrapheneDocument;
 use graphene::LayerId;
 
@@ -65,4 +66,5 @@ impl DocumentMode {
 pub enum DocumentRenderMode<'a> {
 	Root,
 	OnlyBelowLayerInFolder(&'a [LayerId]),
+	LayerCutout(&'a [LayerId], Color),
 }