diff --git a/editor/src/messages/portfolio/document/node_graph/node_graph_message_handler/document_node_types.rs b/editor/src/messages/portfolio/document/node_graph/node_graph_message_handler/document_node_types.rs
index 3317b694b..573c7f042 100644
--- a/editor/src/messages/portfolio/document/node_graph/node_graph_message_handler/document_node_types.rs
+++ b/editor/src/messages/portfolio/document/node_graph/node_graph_message_handler/document_node_types.rs
@@ -1761,7 +1761,9 @@ impl DocumentNodeType {
 
 		let inner_network = match &self.identifier {
 			NodeImplementation::DocumentNode(network) => network.clone(),
-			NodeImplementation::ProtoNode(ident) => {
+
+			NodeImplementation::ProtoNode(ident) => return DocumentNodeImplementation::Unresolved(ident.clone()),
+			/*
 				NodeNetwork {
 					inputs: (0..num_inputs).map(|_| 0).collect(),
 					outputs: vec![NodeOutput::new(0, 0)],
@@ -1779,23 +1781,10 @@ impl DocumentNodeType {
 					.collect(),
 					..Default::default()
 				}
+
 			}
-			NodeImplementation::Extract => NodeNetwork {
-				inputs: (0..num_inputs).map(|_| 0).collect(),
-				outputs: vec![NodeOutput::new(0, 0)],
-				nodes: [(
-					0,
-					DocumentNode {
-						name: "ExtractNode".to_string(),
-						implementation: DocumentNodeImplementation::Extract,
-						inputs: self.inputs.iter().map(|i| NodeInput::Network(i.default.ty())).collect(),
-						..Default::default()
-					},
-				)]
-				.into_iter()
-				.collect(),
-				..Default::default()
-			},
+			*/
+			NodeImplementation::Extract => return DocumentNodeImplementation::Extract,
 		};
 
 		DocumentNodeImplementation::Network(inner_network)
@@ -1824,9 +1813,11 @@ impl DocumentNodeType {
 }
 
 pub fn wrap_network_in_scope(mut network: NodeNetwork) -> NodeNetwork {
-	let node_ids = network.nodes.keys().copied().collect::<Vec<_>>();
 
 	network.generate_node_paths(&[]);
+	network.resolve_extract_nodes();
+
+	let node_ids = network.nodes.keys().copied().collect::<Vec<_>>();
 	for id in node_ids {
 		network.flatten(id);
 	}
diff --git a/node-graph/gcore/src/raster/adjustments.rs b/node-graph/gcore/src/raster/adjustments.rs
index a91da0cac..1e26ffc47 100644
--- a/node-graph/gcore/src/raster/adjustments.rs
+++ b/node-graph/gcore/src/raster/adjustments.rs
@@ -321,13 +321,6 @@ fn grayscale_color_node(color: Color, tint: Color, reds: f64, yellows: f64, gree
 	color.to_linear_srgb()
 }
 
-#[cfg(not(target_arch = "spirv"))]
-pub use hue_shift::HueSaturationNode;
-
-// TODO: Make this work on GPU so it can be removed from the wrapper module that excludes GPU (it doesn't work because of the modulo)
-#[cfg(not(target_arch = "spirv"))]
-mod hue_shift {
-	use super::*;
 
 	#[derive(Debug)]
 	pub struct HueSaturationNode<Hue, Saturation, Lightness> {
@@ -353,7 +346,6 @@ mod hue_shift {
 
 		color.to_linear_srgb()
 	}
-}
 
 #[derive(Debug, Clone, Copy)]
 pub struct InvertRGBNode;
diff --git a/node-graph/graph-craft/src/document.rs b/node-graph/graph-craft/src/document.rs
index bb0d32982..51f2a279a 100644
--- a/node-graph/graph-craft/src/document.rs
+++ b/node-graph/graph-craft/src/document.rs
@@ -99,7 +99,7 @@ impl DocumentNode {
 				document_node_path: self.path.unwrap_or(Vec::new()),
 			}
 		} else {
-			unreachable!("tried to resolve not flattened node on resolved node");
+			unreachable!("tried to resolve not flattened node on resolved node {:?}", self);
 		}
 	}
 
@@ -665,6 +665,7 @@ impl NodeNetwork {
 			self.nodes.insert(id, node);
 			return;
 		}
+		log::debug!("Flattening node {:?}", &node);
 
 		// replace value inputs with value nodes
 		for input in &mut node.inputs {
@@ -706,6 +707,8 @@ impl NodeNetwork {
 		}
 
 		if let DocumentNodeImplementation::Network(mut inner_network) = node.implementation {
+			// Resolve all extract nodes in the inner network
+			inner_network.resolve_extract_nodes();
 			// Connect all network inputs to either the parent network nodes, or newly created value nodes.
 			inner_network.map_ids(|inner_id| map_ids(id, inner_id));
 			let new_nodes = inner_network.nodes.keys().cloned().collect::<Vec<_>>();
@@ -717,8 +720,10 @@ impl NodeNetwork {
 			assert_eq!(
 				node.inputs.len(),
 				inner_network.inputs.len(),
-				"The number of inputs to the node and the inner network must be the same {}",
-				node.name
+				"The number of inputs to the node and the inner network must be the same for {}. The node has {:?} inputs, the network has {:?} inputs.",
+				node.name,
+				node.inputs,
+				inner_network.inputs
 			);
 			// Match the document node input and the inputs of the inner network
 			for (document_input, network_input) in node.inputs.into_iter().zip(inner_network.inputs.iter()) {
@@ -829,21 +834,27 @@ impl NodeNetwork {
 		self.nodes.retain(|_, node| !matches!(node.implementation, DocumentNodeImplementation::Extract));
 
 		for (_, node) in &mut extraction_nodes {
+			log::info!("extraction network: {:#?}", &self);
 			if let DocumentNodeImplementation::Extract = node.implementation {
 				assert_eq!(node.inputs.len(), 1);
+				log::debug!("Resolving extract node {:?}", node);
 				let NodeInput::Node { node_id, output_index, .. } = node.inputs.pop().unwrap() else {
-					panic!("Extract node has no input");
+					panic!("Extract node has no input, inputs: {:?}", node.inputs);
 				};
 				assert_eq!(output_index, 0);
 				// TODO: check if we can readd lambda checking
 				let mut input_node = self.nodes.remove(&node_id).unwrap();
 				node.implementation = DocumentNodeImplementation::Unresolved("graphene_core::value::ValueNode".into());
+				if let Some(input) = input_node.inputs.get_mut(0) {
+					*input = NodeInput::Network(input.ty());
+				}
+
 				for input in input_node.inputs.iter_mut() {
-					match input {
-						NodeInput::Node { .. } | NodeInput::Value { .. } => *input = NodeInput::Network(generic!(T)),
-						_ => (),
+					if let NodeInput::Node { .. } = input {
+						*input = NodeInput::Network(generic!(T))
 					}
 				}
+				log::debug!("Extract node {:?} resolved to {:?}", node, input_node);
 				node.inputs = vec![NodeInput::value(TaggedValue::DocumentNode(input_node), false)];
 			}
 		}
diff --git a/node-graph/graph-craft/src/graphene_compiler.rs b/node-graph/graph-craft/src/graphene_compiler.rs
index fbdeb373e..d7417098a 100644
--- a/node-graph/graph-craft/src/graphene_compiler.rs
+++ b/node-graph/graph-craft/src/graphene_compiler.rs
@@ -11,11 +11,11 @@ impl Compiler {
 	pub fn compile(&self, mut network: NodeNetwork, resolve_inputs: bool) -> impl Iterator<Item = ProtoNetwork> {
 		let node_ids = network.nodes.keys().copied().collect::<Vec<_>>();
 		println!("flattening");
+		network.resolve_extract_nodes();
 		for id in node_ids {
 			network.flatten(id);
 		}
 		network.remove_redundant_id_nodes();
-		network.resolve_extract_nodes();
 		network.remove_dead_nodes();
 		let proto_networks = network.into_proto_networks();
 		proto_networks.map(move |mut proto_network| {
diff --git a/node-graph/wgpu-executor/src/lib.rs b/node-graph/wgpu-executor/src/lib.rs
index e61b77abf..65256f24c 100644
--- a/node-graph/wgpu-executor/src/lib.rs
+++ b/node-graph/wgpu-executor/src/lib.rs
@@ -287,7 +287,7 @@ impl gpu_executor::GpuExecutor for WgpuExecutor {
 			log::warn!("No surface formats available");
 			//return Ok(());
 		}
-		let config = self.surface_config.take().unwrap();
+		let Some(config) = self.surface_config.take() else {return Ok(())};
 		let new_config = config.clone();
 		self.surface_config.replace(Some(config));
 		let output = match result {