Thumbnails for the layer node (#1210)

* Thumbnails for the layer node * Raster node graph frames * Downscale to a random resolution * Cleanup and bug fixes * Generate paths before duplicating outputs * Fix stable id test * Code review changes * Code review pass with minor changes --------- Co-authored-by: Keavon Chambers <keavon@keavon.com>
2026-09-17 15:28:04 +08:00 · 2023-05-18 05:12:50 +01:00
parent d76bd36447
commit 88c4713e75
24 changed files with 609 additions and 131 deletions
--- a/editor/src/node_graph_executor.rs
+++ b/editor/src/node_graph_executor.rs
@@ -7,10 +7,13 @@ use crate::messages::prelude::*;
 use document_legacy::{document::pick_safe_imaginate_resolution, layers::layer_info::LayerDataType};
 use document_legacy::{LayerId, Operation};
 use dyn_any::DynAny;
-use graph_craft::document::{generate_uuid, NodeId, NodeInput, NodeNetwork, NodeOutput};
+use graph_craft::document::{generate_uuid, DocumentNodeImplementation, NodeId, NodeInput, NodeNetwork, NodeOutput};
 use graph_craft::executor::Compiler;
+use graph_craft::imaginate_input::*;
 use graph_craft::{concrete, Type, TypeDescriptor};
 use graphene_core::raster::{Image, ImageFrame};
+use graphene_core::renderer::{SvgSegment, SvgSegmentList};
+use graphene_core::vector::style::ViewMode;
 use graphene_core::vector::VectorData;
 use graphene_core::{Color, EditorApi};
 use interpreted_executor::executor::DynamicExecutor;
@@ -24,16 +27,33 @@ pub struct NodeGraphExecutor {
 	pub(crate) executor: DynamicExecutor,
 	// TODO: This is a memory leak since layers are never removed
 	pub(crate) last_output_type: HashMap<Vec<LayerId>, Option<Type>>,
+	pub(crate) thumbnails: HashMap<LayerId, HashMap<NodeId, SvgSegmentList>>,
+}
+
+fn get_imaginate_index(name: &str) -> usize {
+	use crate::messages::portfolio::document::node_graph::IMAGINATE_NODE;
+	IMAGINATE_NODE.inputs.iter().position(|input| input.name == name).unwrap_or_else(|| panic!("Input {name} not found"))
 }

 impl NodeGraphExecutor {
-	/// Execute the network by flattening it and creating a borrow stack.
-	fn execute_network<'a>(&'a mut self, network: NodeNetwork, editor_api: EditorApi<'a>) -> Result<Box<dyn dyn_any::DynAny + 'a>, String> {
+	/// Wraps a network in a scope and returns the new network and the paths to the monitor nodes.
+	fn wrap_network(network: NodeNetwork) -> (NodeNetwork, Vec<Vec<NodeId>>) {
 		let mut scoped_network = wrap_network_in_scope(network);

+		scoped_network.generate_node_paths(&[]);
+		let monitor_nodes = scoped_network
+			.recursive_nodes()
+			.filter(|(node, _, _)| node.implementation == DocumentNodeImplementation::proto("graphene_std::memo::MonitorNode<_>"))
+			.map(|(_, _, path)| path)
+			.collect();
 		scoped_network.duplicate_outputs(&mut generate_uuid);
 		scoped_network.remove_dead_nodes();

+		(scoped_network, monitor_nodes)
+	}
+
+	/// Executes the network by flattening it and creating a borrow stack.
+	fn execute_network<'a>(&'a mut self, scoped_network: NodeNetwork, editor_api: EditorApi<'a>) -> Result<Box<dyn dyn_any::DynAny + 'a>, String> {
 		// We assume only one output
 		assert_eq!(scoped_network.outputs.len(), 1, "Graph with multiple outputs not yet handled");
 		let c = Compiler {};
@@ -100,6 +120,7 @@ impl NodeGraphExecutor {
 			}
 		}

+		let (network, _) = Self::wrap_network(network);
 		let boxed = self.execute_network(network, editor_api.into_owned())?;

 		dyn_any::downcast::<T>(boxed).map(|v| *v)
@@ -125,73 +146,61 @@ impl NodeGraphExecutor {
 		Ok::<_, String>((image_data, size))
 	}

-	fn generate_imaginate(
-		&mut self,
-		network: NodeNetwork,
-		imaginate_node_path: Vec<NodeId>,
-		(document, document_id): (&mut DocumentMessageHandler, u64),
-		layer_path: Vec<LayerId>,
-		mut editor_api: EditorApi<'_>,
-		(preferences, persistent_data): (&PreferencesMessageHandler, &PersistentData),
-	) -> Result<Message, String> {
-		use crate::messages::portfolio::document::node_graph::IMAGINATE_NODE;
-		use graph_craft::imaginate_input::*;
-
-		let image = editor_api.image_frame.take();
-
-		let get = |name: &str| IMAGINATE_NODE.inputs.iter().position(|input| input.name == name).unwrap_or_else(|| panic!("Input {name} not found"));
-
-		// Get the node graph layer
-		let layer = document.document_legacy.layer(&layer_path).map_err(|e| format!("No layer: {e:?}"))?;
-		let transform = layer.transform;
-
-		let resolution: Option<glam::DVec2> = self.compute_input(&network, &imaginate_node_path, get("Resolution"), Cow::Borrowed(&editor_api))?;
-		let resolution = resolution.unwrap_or_else(|| {
-			let (x, y) = pick_safe_imaginate_resolution((transform.transform_vector2(DVec2::new(1., 0.)).length(), transform.transform_vector2(DVec2::new(0., 1.)).length()));
-			DVec2::new(x as f64, y as f64)
-		});
-
-		let parameters = ImaginateGenerationParameters {
-			seed: self.compute_input::<f64>(&network, &imaginate_node_path, get("Seed"), Cow::Borrowed(&editor_api))? as u64,
+	fn imaginate_parameters(&mut self, network: &NodeNetwork, node_path: &[LayerId], resolution: DVec2, editor_api: &EditorApi) -> Result<ImaginateGenerationParameters, String> {
+		let get = get_imaginate_index;
+		Ok(ImaginateGenerationParameters {
+			seed: self.compute_input::<f64>(network, node_path, get("Seed"), Cow::Borrowed(editor_api))? as u64,
 			resolution: resolution.as_uvec2().into(),
-			samples: self.compute_input::<f64>(&network, &imaginate_node_path, get("Samples"), Cow::Borrowed(&editor_api))? as u32,
+			samples: self.compute_input::<f64>(network, node_path, get("Samples"), Cow::Borrowed(editor_api))? as u32,
 			sampling_method: self
-				.compute_input::<ImaginateSamplingMethod>(&network, &imaginate_node_path, get("Sampling Method"), Cow::Borrowed(&editor_api))?
+				.compute_input::<ImaginateSamplingMethod>(network, node_path, get("Sampling Method"), Cow::Borrowed(editor_api))?
 				.api_value()
 				.to_string(),
-			text_guidance: self.compute_input(&network, &imaginate_node_path, get("Prompt Guidance"), Cow::Borrowed(&editor_api))?,
-			text_prompt: self.compute_input(&network, &imaginate_node_path, get("Prompt"), Cow::Borrowed(&editor_api))?,
-			negative_prompt: self.compute_input(&network, &imaginate_node_path, get("Negative Prompt"), Cow::Borrowed(&editor_api))?,
-			image_creativity: Some(self.compute_input::<f64>(&network, &imaginate_node_path, get("Image Creativity"), Cow::Borrowed(&editor_api))? / 100.),
-			restore_faces: self.compute_input(&network, &imaginate_node_path, get("Improve Faces"), Cow::Borrowed(&editor_api))?,
-			tiling: self.compute_input(&network, &imaginate_node_path, get("Tiling"), Cow::Borrowed(&editor_api))?,
-		};
-		let use_base_image = self.compute_input::<bool>(&network, &imaginate_node_path, get("Adapt Input Image"), Cow::Borrowed(&editor_api))?;
+			text_guidance: self.compute_input(network, node_path, get("Prompt Guidance"), Cow::Borrowed(editor_api))?,
+			text_prompt: self.compute_input(network, node_path, get("Prompt"), Cow::Borrowed(editor_api))?,
+			negative_prompt: self.compute_input(network, node_path, get("Negative Prompt"), Cow::Borrowed(editor_api))?,
+			image_creativity: Some(self.compute_input::<f64>(network, node_path, get("Image Creativity"), Cow::Borrowed(editor_api))? / 100.),
+			restore_faces: self.compute_input(network, node_path, get("Improve Faces"), Cow::Borrowed(editor_api))?,
+			tiling: self.compute_input(network, node_path, get("Tiling"), Cow::Borrowed(editor_api))?,
+		})
+	}

-		editor_api.image_frame = image;
+	fn imaginate_base_image(&mut self, network: &NodeNetwork, imaginate_node_path: &[LayerId], resolution: DVec2, editor_api: &EditorApi) -> Result<Option<(ImaginateBaseImage, DAffine2)>, String> {
+		let use_base_image = self.compute_input::<bool>(&network, &imaginate_node_path, get_imaginate_index("Adapt Input Image"), Cow::Borrowed(editor_api))?;
 		let input_image_frame: Option<ImageFrame<Color>> = if use_base_image {
-			Some(self.compute_input::<ImageFrame<Color>>(&network, &imaginate_node_path, get("Input Image"), Cow::Borrowed(&editor_api))?)
+			Some(self.compute_input::<ImageFrame<Color>>(&network, &imaginate_node_path, get_imaginate_index("Input Image"), Cow::Borrowed(editor_api))?)
 		} else {
 			None
 		};

-		let image_transform = input_image_frame.as_ref().map(|image_frame| image_frame.transform);
-		let base_image = if let Some(ImageFrame { image, .. }) = input_image_frame {
+		let base_image = if let Some(ImageFrame { image, transform }) = input_image_frame {
 			// Only use if has size
 			if image.width > 0 && image.height > 0 {
 				let (image_data, size) = Self::encode_img(image, Some(resolution), image::ImageOutputFormat::Png)?;
 				let size = DVec2::new(size.0 as f64, size.1 as f64);
 				let mime = "image/png".to_string();
-				Some(ImaginateBaseImage { image_data, size, mime })
+				Some((ImaginateBaseImage { image_data, size, mime }, transform))
 			} else {
+				info!("Base image is input but has no size.");
 				None
 			}
 		} else {
 			None
 		};
+		Ok(base_image)
+	}

-		let mask_image = if let Some(transform) = image_transform {
-			let mask_path: Option<Vec<LayerId>> = self.compute_input(&network, &imaginate_node_path, get("Masking Layer"), Cow::Borrowed(&editor_api))?;
+	fn imaginate_mask_image(
+		&mut self,
+		network: &NodeNetwork,
+		node_path: &[LayerId],
+		editor_api: &EditorApi<'_>,
+		image_transform: Option<DAffine2>,
+		document: &mut DocumentMessageHandler,
+		persistent_data: &PersistentData,
+	) -> Result<Option<ImaginateMaskImage>, String> {
+		if let Some(transform) = image_transform {
+			let mask_path: Option<Vec<LayerId>> = self.compute_input(&network, &node_path, get_imaginate_index("Masking Layer"), Cow::Borrowed(&editor_api))?;

 			// Calculate the size of the frame
 			let size = DVec2::new(transform.transform_vector2(DVec2::new(1., 0.)).length(), transform.transform_vector2(DVec2::new(0., 1.)).length());
@@ -214,22 +223,53 @@ impl NodeGraphExecutor {
 			}

 			document.restore_document_transform(old_transforms);
-			mask_image
+			Ok(mask_image)
 		} else {
-			None
-		};
+			Ok(None)
+		}
+	}
+
+	fn generate_imaginate(
+		&mut self,
+		network: NodeNetwork,
+		imaginate_node_path: Vec<NodeId>,
+		(document, document_id): (&mut DocumentMessageHandler, u64),
+		layer_path: Vec<LayerId>,
+		mut editor_api: EditorApi<'_>,
+		(preferences, persistent_data): (&PreferencesMessageHandler, &PersistentData),
+	) -> Result<Message, String> {
+		let image = editor_api.image_frame.take();
+
+		// Get the node graph layer
+		let layer = document.document_legacy.layer(&layer_path).map_err(|e| format!("No layer: {e:?}"))?;
+		let transform = layer.transform;
+
+		let resolution: Option<glam::DVec2> = self.compute_input(&network, &imaginate_node_path, get_imaginate_index("Resolution"), Cow::Borrowed(&editor_api))?;
+		let resolution = resolution.unwrap_or_else(|| {
+			let (x, y) = pick_safe_imaginate_resolution((transform.transform_vector2(DVec2::new(1., 0.)).length(), transform.transform_vector2(DVec2::new(0., 1.)).length()));
+			DVec2::new(x as f64, y as f64)
+		});
+
+		let parameters = self.imaginate_parameters(&network, &imaginate_node_path, resolution, &editor_api)?;
+
+		editor_api.image_frame = image;
+		let base = self.imaginate_base_image(&network, &imaginate_node_path, resolution, &editor_api)?;
+		let image_transform = base.as_ref().map(|base| base.1);
+		let base_image = base.map(|base| base.0);
+
+		let mask_image = self.imaginate_mask_image(&network, &imaginate_node_path, &editor_api, image_transform, document, persistent_data)?;

 		Ok(FrontendMessage::TriggerImaginateGenerate {
 			parameters: Box::new(parameters),
 			base_image: base_image.map(Box::new),
 			mask_image: mask_image.map(Box::new),
-			mask_paint_mode: if self.compute_input::<bool>(&network, &imaginate_node_path, get("Inpaint"), Cow::Borrowed(&editor_api))? {
+			mask_paint_mode: if self.compute_input::<bool>(&network, &imaginate_node_path, get_imaginate_index("Inpaint"), Cow::Borrowed(&editor_api))? {
 				ImaginateMaskPaintMode::Inpaint
 			} else {
 				ImaginateMaskPaintMode::Outpaint
 			},
-			mask_blur_px: self.compute_input::<f64>(&network, &imaginate_node_path, get("Mask Blur"), Cow::Borrowed(&editor_api))? as u32,
-			imaginate_mask_starting_fill: self.compute_input(&network, &imaginate_node_path, get("Mask Starting Fill"), Cow::Borrowed(&editor_api))?,
+			mask_blur_px: self.compute_input::<f64>(&network, &imaginate_node_path, get_imaginate_index("Mask Blur"), Cow::Borrowed(&editor_api))? as u32,
+			imaginate_mask_starting_fill: self.compute_input(&network, &imaginate_node_path, get_imaginate_index("Mask Starting Fill"), Cow::Borrowed(&editor_api))?,
 			hostname: preferences.imaginate_server_hostname.clone(),
 			refresh_frequency: preferences.imaginate_refresh_frequency,
 			document_id,
@@ -239,6 +279,22 @@ impl NodeGraphExecutor {
 		.into())
 	}

+	/// Generate a new [`FrontendImageData`] from the [`Image`].
+	fn to_frontend_image_data(image: Image<Color>, transform: Option<[f64; 6]>, layer_path: &[LayerId], node_id: Option<u64>, resize: Option<DVec2>) -> Result<FrontendImageData, String> {
+		let (image_data, _size) = Self::encode_img(image, resize, image::ImageOutputFormat::Bmp)?;
+
+		let mime = "image/bmp".to_string();
+		let image_data = std::sync::Arc::new(image_data);
+
+		Ok(FrontendImageData {
+			path: layer_path.to_vec(),
+			node_id,
+			image_data,
+			mime,
+			transform,
+		})
+	}
+
 	/// Evaluates a node graph, computing either the Imaginate node or the entire graph
 	pub fn evaluate_node_graph(
 		&mut self,
@@ -276,6 +332,7 @@ impl NodeGraphExecutor {
 			return Ok(());
 		}
 		// Execute the node graph
+		let (network, monitor_nodes) = Self::wrap_network(network);
 		let boxed_node_graph_output = self.execute_network(network, editor_api)?;

 		// Check if the output is vector data
@@ -303,29 +360,74 @@ impl NodeGraphExecutor {
 					responses.add(Operation::SetLayerTransform { path: layer_path.clone(), transform });
 				}
 			} else {
-				// Update the image data
-				let (image_data, _size) = Self::encode_img(image, None, image::ImageOutputFormat::Bmp)?;
-
-				let mime = "image/bmp".to_string();
-				let image_data = std::sync::Arc::new(image_data);
-				let image_data = vec![FrontendImageData {
-					path: layer_path.clone(),
-					image_data,
-					mime,
-					transform,
-				}];
+				let image_data = vec![Self::to_frontend_image_data(image, transform, &layer_path, None, None)?];
 				responses.add(FrontendMessage::UpdateImageData { document_id, image_data });
 			}
 		} else if core::any::TypeId::of::<graphene_core::Artboard>() == DynAny::type_id(boxed_node_graph_output.as_ref()) {
 			let artboard: graphene_core::Artboard = dyn_any::downcast(boxed_node_graph_output).map(|artboard| *artboard)?;
 			info!("{artboard:#?}");
+			self.update_thumbnails(&layer_path, monitor_nodes, responses);
+
 			return Err(format!("Artboard (see console)"));
 		} else if core::any::TypeId::of::<graphene_core::GraphicGroup>() == DynAny::type_id(boxed_node_graph_output.as_ref()) {
 			let graphic_group: graphene_core::GraphicGroup = dyn_any::downcast(boxed_node_graph_output).map(|graphic| *graphic)?;
 			info!("{graphic_group:#?}");
+			self.update_thumbnails(&layer_path, monitor_nodes, responses);
+
 			return Err(format!("Graphic group (see console)"));
 		}

 		Ok(())
 	}
+
+	/// Recomputes the thumbnails for the layers in the graph, modifying the state and updating the UI.
+	pub fn update_thumbnails(&mut self, layer_path: &[LayerId], monitor_nodes: Vec<Vec<u64>>, responses: &mut VecDeque<Message>) {
+		let mut thumbnails_changed: bool = false;
+		let mut image_data: Vec<_> = Vec::new();
+		for node_path in monitor_nodes {
+			let Some(value) = self.executor.introspect(&node_path).flatten() else {
+				warn!("No introspect");
+				continue;
+			};
+			let Some(graphic_group) = value.downcast_ref::<graphene_core::GraphicGroup>() else {
+				warn!("Not graphic");
+				continue;
+			};
+			use graphene_core::renderer::*;
+			let bounds = graphic_group.bounding_box(DAffine2::IDENTITY);
+			let render_params = RenderParams::new(ViewMode::Normal, bounds, true);
+			let mut render = SvgRender::new();
+			graphic_group.render_svg(&mut render, &render_params);
+			let [min, max] = bounds.unwrap_or_default();
+			render.format_svg(min, max);
+			info!("SVG {}", render.svg);
+
+			if let (Some(layer_id), Some(node_id)) = (layer_path.last().copied(), node_path.get(node_path.len() - 2).copied()) {
+				let old_thumbnail = self.thumbnails.entry(layer_id).or_default().entry(node_id).or_default();
+				if *old_thumbnail != render.svg {
+					*old_thumbnail = render.svg;
+					thumbnails_changed = true;
+				}
+			}
+			let resize = Some(DVec2::splat(100.));
+			let create_image_data = |(node_id, image)| Self::to_frontend_image_data(image, None, layer_path, Some(node_id), resize).ok();
+			image_data.extend(render.image_data.into_iter().filter_map(create_image_data))
+		}
+		if !image_data.is_empty() {
+			responses.add(FrontendMessage::UpdateImageData { document_id: 0, image_data });
+		} else if thumbnails_changed {
+			responses.add(NodeGraphMessage::SendGraph { should_rerender: false });
+		}
+	}
+
+	/// When a blob url for a thumbnail is loaded, update the state and the UI.
+	pub fn insert_thumbnail_blob_url(&mut self, blob_url: String, layer_id: LayerId, node_id: NodeId, responses: &mut VecDeque<Message>) {
+		if let Some(layer) = self.thumbnails.get_mut(&layer_id) {
+			if let Some(segment) = layer.values_mut().flat_map(|segments| segments.iter_mut()).find(|segment| **segment == SvgSegment::BlobUrl(node_id)) {
+				*segment = SvgSegment::String(blob_url);
+				responses.add(NodeGraphMessage::SendGraph { should_rerender: false });
+				return;
+			}
+		}
+	}
 }