Fix imaginate generation

editor/src/messages/portfolio/document/node_graph/node_properties.rs

@@ -1866,8 +1866,8 @@ pub fn imaginate_properties(document_node: &DocumentNode, node_id: NodeId, conte
 
 		let mut widgets = start_widgets(document_node, node_id, resolution_index, "Resolution", FrontendGraphDataType::Number, false);
 
-		let round = |x: DVec2| {
-			let (x, y) = pick_safe_imaginate_resolution(x.into());
+		let round = |size: DVec2| {
+			let (x, y) = pick_safe_imaginate_resolution(size.into());
 			DVec2::new(x as f64, y as f64)
 		};
 

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

@@ -25,7 +25,7 @@ impl std::cmp::PartialEq for ImaginateCache {
 
 impl core::hash::Hash for ImaginateCache {
 	fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
-		self.0.lock().unwrap().hash(state);
+		self.0.try_lock().map(|g| g.hash(state));
 	}
 }
 
@@ -50,11 +50,11 @@ pub struct ImaginateController(Arc<InternalImaginateControl>);
 
 impl ImaginateController {
 	pub fn get_status(&self) -> ImaginateStatus {
-		self.0.status.lock().as_deref().cloned().unwrap_or_default()
+		self.0.status.try_lock().as_deref().cloned().unwrap_or_default()
 	}
 
 	pub fn set_status(&self, status: ImaginateStatus) {
-		if let Ok(mut lock) = self.0.status.lock() {
+		if let Ok(mut lock) = self.0.status.try_lock() {
 			*lock = status
 		}
 	}
@@ -68,13 +68,13 @@ impl ImaginateController {
 	}
 
 	pub fn request_termination(&self) {
-		if let Some(handle) = self.0.termination_sender.lock().ok().and_then(|mut lock| lock.take()) {
+		if let Some(handle) = self.0.termination_sender.try_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() {
+		if let Ok(mut lock) = self.0.termination_sender.try_lock() {
 			*lock = Some(handle)
 		}
 	}

node-graph/gstd/src/imaginate.rs

@@ -500,28 +500,31 @@ fn base64_to_image<D: AsRef<[u8]>, P: Pixel>(base64_data: D) -> Result<Image<P>,
 }
 
 pub fn pick_safe_imaginate_resolution((width, height): (f64, f64)) -> (u64, u64) {
+	const NATIVE_MODEL_RESOLUTION: f64 = 512.;
+	let size = if width * height == 0. { DVec2::splat(NATIVE_MODEL_RESOLUTION) } else { DVec2::new(width, height) };
+
 	const MAX_RESOLUTION: u64 = 1000 * 1000;
 
-	// this is the maximum width/height that can be obtained
+	// This is the maximum width/height that can be obtained
 	const MAX_DIMENSION: u64 = (MAX_RESOLUTION / 64) & !63;
 
-	// round the resolution to the nearest multiple of 64
-	let size = (DVec2::new(width, height).round().clamp(DVec2::ZERO, DVec2::splat(MAX_DIMENSION as _)).as_u64vec2() + U64Vec2::splat(32)).max(U64Vec2::splat(64)) & !U64Vec2::splat(63);
+	// Round the resolution to the nearest multiple of 64
+	let size = (size.round().clamp(DVec2::ZERO, DVec2::splat(MAX_DIMENSION as _)).as_u64vec2() + U64Vec2::splat(32)).max(U64Vec2::splat(64)) & !U64Vec2::splat(63);
 	let resolution = size.x * size.y;
 
 	if resolution > MAX_RESOLUTION {
-		// scale down the image, so it is smaller than MAX_RESOLUTION
+		// Scale down the image, so it is smaller than MAX_RESOLUTION
 		let scale = (MAX_RESOLUTION as f64 / resolution as f64).sqrt();
 		let size = size.as_dvec2() * scale;
 
 		if size.x < 64.0 {
-			// the image is extremely wide
+			// The image is extremely wide
 			(64, MAX_DIMENSION)
 		} else if size.y < 64.0 {
-			// the image is extremely high
+			// The image is extremely high
 			(MAX_DIMENSION, 64)
 		} else {
-			// round down to a multiple of 64, so that the resolution still is smaller than MAX_RESOLUTION
+			// Round down to a multiple of 64, so that the resolution still is smaller than MAX_RESOLUTION
 			(size.as_u64vec2() & !U64Vec2::splat(63)).into()
 		}
 	} else {

node-graph/gstd/src/raster.rs

@@ -480,7 +480,7 @@ macro_rules! generate_imaginate_node {
 			generation_id: G,
 			$($val: $t,)*
 			cache: std::sync::Arc<std::sync::Mutex<HashMap<u64, Image<P>>>>,
-			last_generation: u64,
+			last_generation: std::sync::atomic::AtomicU64,
 		}
 
 		impl<'e, P: Pixel, E, C, G, $($t,)*> ImaginateNode<P, E, C, G, $($t,)*>
@@ -491,7 +491,7 @@ macro_rules! generate_imaginate_node {
 		{
 			#[allow(clippy::too_many_arguments)]
 			pub fn new(editor_api: E, controller: C, $($val: $t,)*  generation_id: G ) -> Self {
-				Self { editor_api, controller, generation_id, $($val,)* cache: Default::default(), last_generation: u64::MAX }
+				Self { editor_api, controller, generation_id, $($val,)* cache: Default::default(), last_generation: std::sync::atomic::AtomicU64::new(u64::MAX) }
 			}
 		}
 
@@ -517,22 +517,39 @@ macro_rules! generate_imaginate_node {
 				Box::pin(async move {
 					let controller: ImaginateController = controller.await;
 					let generation_id = self.generation_id.eval(()).await;
-					if generation_id !=  self.last_generation {
+					if generation_id !=  self.last_generation.swap(generation_id, std::sync::atomic::Ordering::SeqCst) {
 						let image = super::imaginate::imaginate(frame.image, editor_api, controller, $($val,)*).await;
 
 						cache.lock().unwrap().insert(hash, image.clone());
 
-						return ImageFrame { image, ..frame }
+						return wrap_image_frame(image, frame.transform);
 					}
 					let image = cache.lock().unwrap().get(&hash).cloned().unwrap_or_default();
 
-					ImageFrame { image, ..frame }
+					return wrap_image_frame(image, frame.transform);
 				})
 			}
 		}
 	}
 }
 
+fn wrap_image_frame<P: Pixel>(image: Image<P>, transform: DAffine2) -> ImageFrame<P> {
+	if !transform.decompose_scale().abs_diff_eq(DVec2::ZERO, 0.00001) {
+		ImageFrame {
+			image,
+			transform,
+			alpha_blending: AlphaBlending::default(),
+		}
+	} else {
+		let resolution = DVec2::new(image.height as f64, image.width as f64);
+		ImageFrame {
+			image,
+			transform: DAffine2::from_scale_angle_translation(resolution, 0., transform.translation),
+			alpha_blending: AlphaBlending::default(),
+		}
+	}
+}
+
 #[cfg(feature = "serde")]
 generate_imaginate_node! {
 	seed: Seed: f64,