Graphene CLI + quantization research (#1320)

* Implement skeleton for graphene-cli

* Configure gpu surface on non wasm32 targets

* Create window with full hd size

* Create window using the graphen-cli

* Use window size for surface creation

* Reuse surface configuration

* Reduce window size for native applications to 800x600

* Add compute pipeline test

* Poll wgpu execution externally

* Remove cache node after texture upload

* Add profiling instructions

* Add more debug markers

* Evaluate extract node before flattening the network

* Reenable hue saturation node for compilation

* Make hue saturation node work on the gpu + make f32 default for user inputs

* Add version of test files without caching

* Only dispatch each workgroup not pixel

* ICE

* Add quantization to gpu code

* Fix quantization

* Load images at graph runtime

* Fix quantization calculation

* Feature gate quantization

* Use git version of autoquant

* Add license to `graphene-cli`

* Fix graphene-cli test case

* Ignore tests on non unix platforms

* Fix flattening test

node-graph/gstd/src/gpu_nodes.rs

@@ -1,13 +1,17 @@
+use dyn_any::{StaticType, StaticTypeSized};
 use glam::{DAffine2, DVec2, Mat2, Vec2};
 use gpu_executor::{Bindgroup, ComputePassDimensions, PipelineLayout, StorageBufferOptions};
 use gpu_executor::{GpuExecutor, ShaderIO, ShaderInput};
 use graph_craft::document::value::TaggedValue;
 use graph_craft::document::*;
 use graph_craft::proto::*;
+use graphene_core::quantization::{PackedPixel, QuantizationChannels};
 use graphene_core::raster::*;
 use graphene_core::*;
 use wgpu_executor::WgpuExecutor;
 
+use std::cell::RefCell;
+use std::collections::HashMap;
 use std::sync::Arc;
 
 use crate::wasm_application_io::WasmApplicationIo;
@@ -42,18 +46,110 @@ async fn compile_gpu(node: &'input DocumentNode, mut typing_context: TypingConte
 pub struct MapGpuNode<Node, EditorApi> {
 	node: Node,
 	editor_api: EditorApi,
+	cache: RefCell<HashMap<String, ComputePass<WgpuExecutor>>>,
 }
 
-#[node_macro::node_fn(MapGpuNode)]
+struct ComputePass<T: GpuExecutor> {
+	pipeline_layout: PipelineLayout<T>,
+	readback_buffer: Option<Arc<ShaderInput<T>>>,
+}
+
+impl<T: GpuExecutor> Clone for ComputePass<T> {
+	fn clone(&self) -> Self {
+		Self {
+			pipeline_layout: self.pipeline_layout.clone(),
+			readback_buffer: self.readback_buffer.clone(),
+		}
+	}
+}
+
+#[node_macro::node_impl(MapGpuNode)]
 async fn map_gpu<'a: 'input>(image: ImageFrame<Color>, node: DocumentNode, editor_api: graphene_core::application_io::EditorApi<'a, WasmApplicationIo>) -> ImageFrame<Color> {
 	log::debug!("Executing gpu node");
+	let executor = &editor_api.application_io.gpu_executor.as_ref().unwrap();
+
+	#[cfg(feature = "quantization")]
+	let quantization = crate::quantization::generate_quantization_from_image_frame(&image);
+	#[cfg(not(feature = "quantization"))]
+	let quantization = QuantizationChannels::default();
+	log::debug!("quantization: {:?}", quantization);
+
+	#[cfg(feature = "quantization")]
+	let image = ImageFrame {
+		image: Image {
+			data: image.image.data.iter().map(|c| quantization::quantize_color(*c, quantization)).collect(),
+			width: image.image.width,
+			height: image.image.height,
+		},
+		transform: image.transform,
+	};
+	// TODO: The cache should be based on the network topology not the node name
+	let compute_pass_descriptor = if self.cache.borrow().contains_key(&node.name) {
+		self.cache.borrow().get(&node.name).unwrap().clone()
+	} else {
+		let name = node.name.clone();
+		let compute_pass_descriptor = create_compute_pass_descriptor(node, &image, executor, quantization).await;
+		self.cache.borrow_mut().insert(name, compute_pass_descriptor.clone());
+		log::error!("created compute pass");
+		compute_pass_descriptor
+	};
+
+	let compute_pass = executor
+		.create_compute_pass(
+			&compute_pass_descriptor.pipeline_layout,
+			compute_pass_descriptor.readback_buffer.clone(),
+			ComputePassDimensions::XY(image.image.width / 12 + 1, image.image.height / 8 + 1),
+		)
+		.unwrap();
+	executor.execute_compute_pipeline(compute_pass).unwrap();
+	log::debug!("executed pipeline");
+	log::debug!("reading buffer");
+	let result = executor.read_output_buffer(compute_pass_descriptor.readback_buffer.clone().unwrap()).await.unwrap();
+	#[cfg(feature = "quantization")]
+	let colors = bytemuck::pod_collect_to_vec::<u8, PackedPixel>(result.as_slice());
+	#[cfg(feature = "quantization")]
+	log::debug!("first color: {:b}", colors[0].0);
+	#[cfg(feature = "quantization")]
+	let colors: Vec<_> = colors.iter().map(|c| quantization::dequantize_color(*c, quantization)).collect();
+	#[cfg(not(feature = "quantization"))]
+	let colors = bytemuck::pod_collect_to_vec::<u8, Color>(result.as_slice());
+	log::debug!("first color: {:?}", colors[0]);
+	ImageFrame {
+		image: Image {
+			data: colors,
+			width: image.image.width,
+			height: image.image.height,
+		},
+		transform: image.transform,
+	}
+}
+
+impl<Node, EditorApi> MapGpuNode<Node, EditorApi> {
+	pub fn new(node: Node, editor_api: EditorApi) -> Self {
+		Self {
+			node,
+			editor_api,
+			cache: RefCell::new(HashMap::new()),
+		}
+	}
+}
+
+async fn create_compute_pass_descriptor<T: Clone + Pixel + StaticTypeSized>(
+	node: DocumentNode,
+	image: &ImageFrame<T>,
+	executor: &&WgpuExecutor,
+	quantization: QuantizationChannels,
+) -> ComputePass<WgpuExecutor> {
 	let compiler = graph_craft::graphene_compiler::Compiler {};
 	let inner_network = NodeNetwork::value_network(node);
 
 	log::debug!("inner_network: {:?}", inner_network);
 	let network = NodeNetwork {
-		inputs: vec![], //vec![0, 1],
-		outputs: vec![NodeOutput::new(1, 0)],
+		inputs: vec![2, 1], //vec![0, 1],
+		#[cfg(feature = "quantization")]
+		outputs: vec![NodeOutput::new(5, 0)],
+		#[cfg(not(feature = "quantization"))]
+		outputs: vec![NodeOutput::new(3, 0)],
 		nodes: [
 			DocumentNode {
 				name: "Slice".into(),
@@ -61,6 +157,18 @@ async fn map_gpu<'a: 'input>(image: ImageFrame<Color>, node: DocumentNode, edito
 				implementation: DocumentNodeImplementation::Unresolved("graphene_core::value::CopiedNode".into()),
 				..Default::default()
 			},
+			DocumentNode {
+				name: "Quantization".into(),
+				inputs: vec![NodeInput::Network(concrete!(quantization::Quantization))],
+				implementation: DocumentNodeImplementation::Unresolved("graphene_core::ops::IdNode".into()),
+				..Default::default()
+			},
+			DocumentNode {
+				name: "Width".into(),
+				inputs: vec![NodeInput::Network(concrete!(u32))],
+				implementation: DocumentNodeImplementation::Unresolved("graphene_core::ops::IdNode".into()),
+				..Default::default()
+			},
 			/*DocumentNode {
 				name: "Index".into(),
 				//inputs: vec![NodeInput::Network(concrete!(UVec3))],
@@ -68,30 +176,48 @@ async fn map_gpu<'a: 'input>(image: ImageFrame<Color>, node: DocumentNode, edito
 				implementation: DocumentNodeImplementation::Unresolved("graphene_core::value::CopiedNode".into()),
 				..Default::default()
 			},*/
-			/*
+				/*
 			DocumentNode {
 				name: "GetNode".into(),
 				inputs: vec![NodeInput::node(1, 0), NodeInput::node(0, 0)],
 				implementation: DocumentNodeImplementation::Unresolved("graphene_core::storage::GetNode".into()),
 				..Default::default()
 			},*/
+			#[cfg(feature = "quantization")]
+			DocumentNode {
+				name: "Dequantize".into(),
+				inputs: vec![NodeInput::node(0, 0), NodeInput::node(1, 0)],
+				implementation: DocumentNodeImplementation::proto("graphene_core::quantization::DeQuantizeNode"),
+				..Default::default()
+			},
 			DocumentNode {
 				name: "MapNode".into(),
+				#[cfg(feature = "quantization")]
+				inputs: vec![NodeInput::node(3, 0)],
+				#[cfg(not(feature = "quantization"))]
 				inputs: vec![NodeInput::node(0, 0)],
 				implementation: DocumentNodeImplementation::Network(inner_network),
 				..Default::default()
 			},
+			#[cfg(feature = "quantization")]
+			DocumentNode {
+				name: "Quantize".into(),
+				inputs: vec![NodeInput::node(4, 0), NodeInput::node(1, 0)],
+				implementation: DocumentNodeImplementation::proto("graphene_core::quantization::QuantizeNode"),
+				..Default::default()
+			},
 			/*
 			DocumentNode {
 				name: "SaveNode".into(),
 				inputs: vec![
-					//NodeInput::node(0, 0),
+					NodeInput::node(5, 0),
 					NodeInput::Inline(InlineRust::new(
-						"o0[_global_index.x as usize] = i0[_global_index.x as usize]".into(),
-						Type::Fn(Box::new(concrete!(Color)), Box::new(concrete!(()))),
+						"|x| o0[(_global_index.y * i1 + _global_index.x) as usize] = x".into(),
+						//"|x|()".into(),
+						Type::Fn(Box::new(concrete!(PackedPixel)), Box::new(concrete!(()))),
 					)),
 				],
-				implementation: DocumentNodeImplementation::Unresolved("graphene_core::value::ValueNode".into()),
+				implementation: DocumentNodeImplementation::Unresolved("graphene_core::generic::FnMutNode".into()),
 				..Default::default()
 			},
 			*/
@@ -110,12 +236,23 @@ async fn map_gpu<'a: 'input>(image: ImageFrame<Color>, node: DocumentNode, edito
 		vec![concrete!(u32), concrete!(Color)], //, concrete!(u32)],
 		vec![concrete!(Color)],
 		ShaderIO {
+			#[cfg(feature = "quantization")]
+			inputs: vec![
+				ShaderInput::UniformBuffer((), concrete!(u32)),
+				ShaderInput::StorageBuffer((), concrete!(PackedPixel)),
+				ShaderInput::UniformBuffer((), concrete!(quantization::QuantizationChannels)),
+				//ShaderInput::Constant(gpu_executor::GPUConstant::GlobalInvocationId),
+				ShaderInput::OutputBuffer((), concrete!(PackedPixel)),
+			],
+			#[cfg(not(feature = "quantization"))]
 			inputs: vec![
 				ShaderInput::UniformBuffer((), concrete!(u32)),
 				ShaderInput::StorageBuffer((), concrete!(Color)),
-				//ShaderInput::Constant(gpu_executor::GPUConstant::GlobalInvocationId),
 				ShaderInput::OutputBuffer((), concrete!(Color)),
 			],
+			#[cfg(feature = "quantization")]
+			output: ShaderInput::OutputBuffer((), concrete!(PackedPixel)),
+			#[cfg(not(feature = "quantization"))]
 			output: ShaderInput::OutputBuffer((), concrete!(Color)),
 		},
 	)
@@ -124,8 +261,6 @@ async fn map_gpu<'a: 'input>(image: ImageFrame<Color>, node: DocumentNode, edito
 	//return ImageFrame::empty();
 	let len: usize = image.image.data.len();
 
-	let executor = &editor_api.application_io.gpu_executor.as_ref().unwrap();
-
 	/*
 	let canvas = editor_api.application_io.create_surface();
 
@@ -144,6 +279,7 @@ async fn map_gpu<'a: 'input>(image: ImageFrame<Color>, node: DocumentNode, edito
 	return frame;*/
 	log::debug!("creating buffer");
 	let width_uniform = executor.create_uniform_buffer(image.image.width).unwrap();
+	let quantization_uniform = executor.create_uniform_buffer(quantization).unwrap();
 	let storage_buffer = executor
 		.create_storage_buffer(
 			image.image.data.clone(),
@@ -156,6 +292,7 @@ async fn map_gpu<'a: 'input>(image: ImageFrame<Color>, node: DocumentNode, edito
 		)
 		.unwrap();
 	let width_uniform = Arc::new(width_uniform);
+	let quantization_uniform = Arc::new(quantization_uniform);
 	let storage_buffer = Arc::new(storage_buffer);
 	let output_buffer = executor.create_output_buffer(len, concrete!(Color), false).unwrap();
 	let output_buffer = Arc::new(output_buffer);
@@ -163,6 +300,9 @@ async fn map_gpu<'a: 'input>(image: ImageFrame<Color>, node: DocumentNode, edito
 	let readback_buffer = Arc::new(readback_buffer);
 	log::debug!("created buffer");
 	let bind_group = Bindgroup {
+		#[cfg(feature = "quantization")]
+		buffers: vec![width_uniform.clone(), storage_buffer.clone(), quantization_uniform.clone()],
+		#[cfg(not(feature = "quantization"))]
 		buffers: vec![width_uniform.clone(), storage_buffer.clone()],
 	};
 
@@ -175,36 +315,18 @@ async fn map_gpu<'a: 'input>(image: ImageFrame<Color>, node: DocumentNode, edito
 	let shader = executor.load_shader(shader).unwrap();
 	log::debug!("loaded shader");
 	let pipeline = PipelineLayout {
-		shader,
+		shader: shader.into(),
 		entry_point: "eval".to_string(),
-		bind_group,
+		bind_group: bind_group.into(),
 		output_buffer: output_buffer.clone(),
 	};
 	log::debug!("created pipeline");
-	let compute_pass = executor
-		.create_compute_pass(&pipeline, Some(readback_buffer.clone()), ComputePassDimensions::XY(image.image.width, image.image.height))
-		.unwrap();
-	executor.execute_compute_pipeline(compute_pass).unwrap();
-	log::debug!("executed pipeline");
-	log::debug!("reading buffer");
-	let result = executor.read_output_buffer(readback_buffer).await.unwrap();
-	let colors = bytemuck::pod_collect_to_vec::<u8, Color>(result.as_slice());
-	ImageFrame {
-		image: Image {
-			data: colors,
-			width: image.image.width,
-			height: image.image.height,
-		},
-		transform: image.transform,
-	}
 
-	/*
-	let executor: GpuExecutor = GpuExecutor::new(Context::new().await.unwrap(), shader.into(), "gpu::eval".into()).unwrap();
-	let data: Vec<_> = input.into_iter().collect();
-	let result = executor.execute(Box::new(data)).unwrap();
-	let result = dyn_any::downcast::<Vec<_O>>(result).unwrap();
-	*result
-	*/
+	let compute_pass_descriptor = ComputePass {
+		pipeline_layout: pipeline,
+		readback_buffer: Some(readback_buffer.clone()),
+	};
+	compute_pass_descriptor
 }
 /*
 #[node_macro::node_fn(MapGpuNode)]
@@ -414,9 +536,9 @@ async fn blend_gpu_image(foreground: ImageFrame<Color>, background: ImageFrame<C
 	let shader = executor.load_shader(shader).unwrap();
 	log::debug!("loaded shader");
 	let pipeline = PipelineLayout {
-		shader,
+		shader: shader.into(),
 		entry_point: "eval".to_string(),
-		bind_group,
+		bind_group: bind_group.into(),
 		output_buffer: output_buffer.clone(),
 	};
 	log::debug!("created pipeline");

node-graph/gstd/src/quantization.rs

@@ -1,3 +1,4 @@
+use autoquant::packing::ErrorFunction;
 use dyn_any::{DynAny, StaticType};
 use graphene_core::quantization::*;
 use graphene_core::raster::{Color, ImageFrame};
@@ -13,47 +14,93 @@ pub struct GenerateQuantizationNode<N, M> {
 
 #[node_macro::node_fn(GenerateQuantizationNode)]
 fn generate_quantization_fn(image_frame: ImageFrame<Color>, samples: u32, function: u32) -> [Quantization; 4] {
-	let image = image_frame.image;
+	generate_quantization_from_image_frame(&image_frame)
+}
+
+pub fn generate_quantization_from_image_frame(image_frame: &ImageFrame<Color>) -> [Quantization; 4] {
+	let image = &image_frame.image;
 
 	let len = image.data.len().min(10000);
-	let mut channels: Vec<_> = (0..4).map(|_| Vec::with_capacity(image.data.len())).collect();
-	image
+	let data = image
 		.data
 		.iter()
 		.enumerate()
 		.filter(|(i, _)| i % (image.data.len() / len) == 0)
-		.map(|(_, x)| vec![x.r() as f64, x.g() as f64, x.b() as f64, x.a() as f64])
-		.for_each(|x| x.into_iter().enumerate().for_each(|(i, value)| channels[i].push(value)));
-	let quantization: Vec<Quantization> = channels.into_iter().map(|x| generate_quantization_per_channel(x, samples)).collect();
-	core::array::from_fn(|i| quantization[i].clone())
+		.flat_map(|(_, x)| vec![x.r() as f64, x.g() as f64, x.b() as f64, x.a() as f64])
+		.collect::<Vec<_>>();
+	generate_quantization(data, len)
+}
+fn generate_quantization(data: Vec<f64>, samples: usize) -> [Quantization; 4] {
+	let red = create_distribution(data.clone(), samples, 0);
+	let green = create_distribution(data.clone(), samples, 1);
+	let blue = create_distribution(data.clone(), samples, 2);
+	let alpha = create_distribution(data, samples, 3);
+
+	let fit_red = autoquant::calculate_error_function(&red, 1, &red);
+	let fit_green = autoquant::calculate_error_function(&green, 1, &green);
+	let fit_blue = autoquant::calculate_error_function(&blue, 1, &blue);
+	let fit_alpha = autoquant::calculate_error_function(&alpha, 1, &alpha);
+	let red_error: ErrorFunction<10> = autoquant::packing::ErrorFunction::new(fit_red.as_slice());
+	let green_error: ErrorFunction<10> = autoquant::packing::ErrorFunction::new(fit_green.as_slice());
+	let blue_error: ErrorFunction<10> = autoquant::packing::ErrorFunction::new(fit_blue.as_slice());
+	let alpha_error: ErrorFunction<10> = autoquant::packing::ErrorFunction::new(fit_alpha.as_slice());
+	let merged: ErrorFunction<20> = autoquant::packing::merge_error_functions(&red_error, &green_error);
+	let merged: ErrorFunction<30> = autoquant::packing::merge_error_functions(&merged, &blue_error);
+	let merged: ErrorFunction<40> = autoquant::packing::merge_error_functions(&merged, &alpha_error);
+
+	let bin_size = 32;
+	let mut distributions = [red, green, blue, alpha].into_iter();
+
+	let bits = &merged.bits[bin_size];
+
+	core::array::from_fn(|i| {
+		let fit = autoquant::models::OptimizedLin::new(distributions.next().unwrap(), (1 << bits[i]) - 1);
+		let parameters = fit.parameters();
+		Quantization::new(parameters[0] as f32, parameters[1] as f32, bits[i] as u32)
+	})
 }
 
-fn generate_quantization_per_channel(data: Vec<f64>, samples: u32) -> Quantization {
+/*
+// TODO: make this work with generic size parameters
+fn generate_quantization<const N: usize>(data: Vec<f64>, samples: usize, channels: usize) -> [Quantization; N] {
+	let mut quantizations = Vec::new();
+	let mut merged_error: Option<ErrorFunction<10>> = None;
+	let bin_size = 32;
+
+	for i in 0..channels {
+		let channel_data = create_distribution(data.clone(), samples, i);
+
+		let fit = autoquant::calculate_error_function(&channel_data, 0, &channel_data);
+		let error: ErrorFunction<10> = autoquant::packing::ErrorFunction::new(fit.as_slice());
+
+		// Merge current error function with previous ones
+		merged_error = match merged_error {
+			Some(prev_error) => Some(autoquant::packing::merge_error_functions(&prev_error, &error)),
+			None => Some(error.clone()),
+		};
+
+		println!("Merged: {:?}", merged_error);
+
+		let bits = merged_error.as_ref().unwrap().bits.iter().map(|x| x[i]).collect::<Vec<_>>();
+		let model_fit = autoquant::models::OptimizedLin::new(channel_data, 1 << bits[bin_size]);
+		let parameters = model_fit.parameters();
+		let quantization = Quantization::new(parameters[0] as f32, parameters[1] as u32, bits[bin_size] as u32);
+
+		quantizations.push(quantization);
+	}
+
+	core::array::from_fn(|x| quantizations[x])
+}*/
+
+fn create_distribution(data: Vec<f64>, samples: usize, channel: usize) -> Vec<(f64, f64)> {
+	let data: Vec<f64> = data.chunks(4 * (data.len() / (4 * samples.min(data.len() / 4)))).map(|x| x[channel] as f64).collect();
+	let max = *data.iter().max_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)).unwrap();
+	let data: Vec<f64> = data.iter().map(|x| x / max).collect();
+	dbg!(max);
+	//let data = autoquant::generate_normal_distribution(3.0, 1.1, 1000);
+	//data.iter_mut().for_each(|x| *x = x.abs());
 	let mut dist = autoquant::integrate_distribution(data);
 	autoquant::drop_duplicates(&mut dist);
 	let dist = autoquant::normalize_distribution(dist.as_slice());
-	let max = dist.last().unwrap().0;
-	/*let linear = Box::new(autoquant::SimpleFitFn {
-		function: move |x| x / max,
-		inverse: move |x| x * max,
-		name: "identity",
-	});*/
-
-	let linear = Quantization {
-		fn_index: 0,
-		a: max as f32,
-		b: 0.,
-		c: 0.,
-		d: 0.,
-	};
-	let log_fit = autoquant::models::OptimizedLog::new(dist, samples as u64);
-	let parameters = log_fit.parameters();
-	let log_fit = Quantization {
-		fn_index: 1,
-		a: parameters[0] as f32,
-		b: parameters[1] as f32,
-		c: parameters[2] as f32,
-		d: parameters[3] as f32,
-	};
-	log_fit
+	dist
 }

node-graph/gstd/src/value.rs

@@ -7,8 +7,7 @@ use dyn_any::DynAny;
 pub struct AnyRefNode<'n, N: Node<'n>>(N, PhantomData<&'n ()>);
 
 impl<'n, N: Node<'n, Output = &'n O>, O: DynAny<'n> + 'n> Node<'n> for AnyRefNode<'n, N> {
-	type Output = &'n (dyn DynAny<'n>);
-	fn eval(&'n self) -> Self::Output {
+	fn eval(&'n self) -> &'n (dyn DynAny<'n>) {
 		let value: &O = self.0.eval();
 		value
 	}
@@ -22,8 +21,7 @@ impl<'n, N: Node<'n, Output = &'n O>, O: 'n + ?Sized> AnyRefNode<'n, N> {
 pub struct StorageNode<'n>(&'n dyn Node<'n, Output = &'n dyn DynAny<'n>>);
 
 impl<'n> Node<'n> for StorageNode<'n> {
-	type Output = &'n (dyn DynAny<'n>);
-	fn eval(&'n self) -> Self::Output {
+	fn eval(&'n self) -> &'n (dyn DynAny<'n>) {
 		self.0.eval()
 	}
 }
@@ -36,7 +34,6 @@ impl<'n> StorageNode<'n> {
 #[derive(Default)]
 pub struct AnyValueNode<'n, T>(T, PhantomData<&'n ()>);
 impl<'n, T: 'n + DynAny<'n>> Node<'n> for AnyValueNode<'n, T> {
-	type Output = &'n dyn DynAny<'n>;
 	fn eval(&'n self) -> &'n dyn DynAny<'n> {
 		&self.0
 	}

node-graph/gstd/src/wasm_application_io.rs

@@ -1,13 +1,21 @@
+use std::any::Any;
 use std::cell::RefCell;
 
+use core::future::Future;
 use dyn_any::StaticType;
-use graphene_core::application_io::{ApplicationIo, SurfaceHandle, SurfaceHandleFrame, SurfaceId};
+use graphene_core::application_io::{ApplicationError, ApplicationIo, SurfaceHandle, SurfaceHandleFrame, SurfaceId};
+use graphene_core::raster::Image;
+use graphene_core::Color;
 use graphene_core::{
 	raster::{color::SRGBA8, ImageFrame},
 	Node,
 };
 use js_sys::{Object, Reflect};
+use std::collections::HashMap;
+use std::pin::Pin;
 use std::sync::Arc;
+#[cfg(feature = "tokio")]
+use tokio::io::AsyncReadExt;
 use wasm_bindgen::{Clamped, JsCast, JsValue};
 use web_sys::{window, CanvasRenderingContext2d, HtmlCanvasElement};
 #[cfg(feature = "wgpu")]
@@ -20,15 +28,23 @@ pub struct WasmApplicationIo {
 	ids: RefCell<u64>,
 	#[cfg(feature = "wgpu")]
 	pub(crate) gpu_executor: Option<WgpuExecutor>,
+	#[cfg(not(target_arch = "wasm32"))]
+	windows: RefCell<Vec<Arc<winit::window::Window>>>,
+	pub resources: HashMap<String, Arc<[u8]>>,
 }
 
 impl WasmApplicationIo {
 	pub async fn new() -> Self {
-		Self {
+		let mut io = Self {
 			ids: RefCell::new(0),
 			#[cfg(feature = "wgpu")]
 			gpu_executor: WgpuExecutor::new().await,
-		}
+			#[cfg(not(target_arch = "wasm32"))]
+			windows: RefCell::new(Vec::new()),
+			resources: HashMap::new(),
+		};
+		io.resources.insert("null".to_string(), Arc::from(include_bytes!("null.png").to_vec()));
+		io
 	}
 }
 
@@ -51,12 +67,16 @@ impl<'a> From<&'a WasmApplicationIo> for &'a WgpuExecutor {
 pub type WasmEditorApi<'a> = graphene_core::application_io::EditorApi<'a, WasmApplicationIo>;
 
 impl ApplicationIo for WasmApplicationIo {
+	#[cfg(target_arch = "wasm32")]
 	type Surface = HtmlCanvasElement;
+	#[cfg(not(target_arch = "wasm32"))]
+	type Surface = Arc<winit::window::Window>;
 	#[cfg(feature = "wgpu")]
 	type Executor = WgpuExecutor;
 	#[cfg(not(feature = "wgpu"))]
 	type Executor = ();
 
+	#[cfg(target_arch = "wasm32")]
 	fn create_surface(&self) -> SurfaceHandle<Self::Surface> {
 		let wrapper = || {
 			let document = window().expect("should have a window in this context").document().expect("window should have a document");
@@ -90,7 +110,29 @@ impl ApplicationIo for WasmApplicationIo {
 
 		wrapper().expect("should be able to set canvas in global scope")
 	}
+	#[cfg(not(target_arch = "wasm32"))]
+	fn create_surface(&self) -> SurfaceHandle<Self::Surface> {
+		#[cfg(feature = "wayland")]
+		use winit::platform::wayland::EventLoopBuilderExtWayland;
 
+		#[cfg(feature = "wayland")]
+		let event_loop = winit::event_loop::EventLoopBuilder::new().with_any_thread(true).build();
+		#[cfg(not(feature = "wayland"))]
+		let event_loop = winit::event_loop::EventLoop::new();
+		let window = winit::window::WindowBuilder::new()
+			.with_title("Graphite")
+			.with_inner_size(winit::dpi::PhysicalSize::new(800, 600))
+			.build(&event_loop)
+			.unwrap();
+		let window = Arc::new(window);
+		self.windows.borrow_mut().push(window.clone());
+		SurfaceHandle {
+			surface_id: SurfaceId(window.id().into()),
+			surface: window,
+		}
+	}
+
+	#[cfg(target_arch = "wasm32")]
 	fn destroy_surface(&self, surface_id: SurfaceId) {
 		let window = window().expect("should have a window in this context");
 		let window = Object::from(window);
@@ -111,10 +153,50 @@ impl ApplicationIo for WasmApplicationIo {
 		wrapper().expect("should be able to set canvas in global scope")
 	}
 
+	#[cfg(not(target_arch = "wasm32"))]
+	fn destroy_surface(&self, _surface_id: SurfaceId) {}
+
 	#[cfg(feature = "wgpu")]
 	fn gpu_executor(&self) -> Option<&Self::Executor> {
 		self.gpu_executor.as_ref()
 	}
+
+	fn load_resource(&self, url: impl AsRef<str>) -> Result<Pin<Box<dyn Future<Output = Result<Arc<[u8]>, ApplicationError>>>>, ApplicationError> {
+		let url = url::Url::parse(url.as_ref()).map_err(|_| ApplicationError::InvalidUrl)?;
+		log::info!("Loading resource: {:?}", url);
+		match url.scheme() {
+			#[cfg(feature = "tokio")]
+			"file" => {
+				let path = url.to_file_path().map_err(|_| ApplicationError::NotFound)?;
+				let path = path.to_str().ok_or(ApplicationError::NotFound)?;
+				let path = path.to_owned();
+				Ok(Box::pin(async move {
+					let file = tokio::fs::File::open(path).await.map_err(|_| ApplicationError::NotFound)?;
+					let mut reader = tokio::io::BufReader::new(file);
+					let mut data = Vec::new();
+					reader.read_to_end(&mut data).await.map_err(|_| ApplicationError::NotFound)?;
+					Ok(Arc::from(data))
+				}) as Pin<Box<dyn Future<Output = Result<Arc<[u8]>, _>>>>)
+			}
+			"http" | "https" => {
+				let url = url.to_string();
+				Ok(Box::pin(async move {
+					let client = reqwest::Client::new();
+					let response = client.get(url).send().await.map_err(|_| ApplicationError::NotFound)?;
+					let data = response.bytes().await.map_err(|_| ApplicationError::NotFound)?;
+					Ok(Arc::from(data.to_vec()))
+				}) as Pin<Box<dyn Future<Output = Result<Arc<[u8]>, _>>>>)
+			}
+			"graphite" => {
+				let path = url.path();
+				let path = path.to_owned();
+				log::info!("Loading resource: {}", path);
+				let data = self.resources.get(&path).ok_or(ApplicationError::NotFound)?.clone();
+				Ok(Box::pin(async move { Ok(data.clone()) }) as Pin<Box<dyn Future<Output = Result<Arc<[u8]>, _>>>>)
+			}
+			_ => Err(ApplicationError::NotFound),
+		}
+	}
 }
 
 pub type WasmSurfaceHandle = SurfaceHandle<HtmlCanvasElement>;
@@ -123,7 +205,7 @@ pub type WasmSurfaceHandleFrame = SurfaceHandleFrame<HtmlCanvasElement>;
 pub struct CreateSurfaceNode {}
 
 #[node_macro::node_fn(CreateSurfaceNode)]
-async fn create_surface_node<'a: 'input>(editor: WasmEditorApi<'a>) -> Arc<SurfaceHandle<HtmlCanvasElement>> {
+async fn create_surface_node<'a: 'input>(editor: WasmEditorApi<'a>) -> Arc<SurfaceHandle<<WasmApplicationIo as ApplicationIo>::Surface>> {
 	editor.application_io.create_surface().into()
 }
 
@@ -149,3 +231,29 @@ async fn draw_image_frame_node<'a: 'input>(image: ImageFrame<SRGBA8>, surface_ha
 		transform: image.transform,
 	}
 }
+
+pub struct LoadResourceNode<Url> {
+	url: Url,
+}
+
+#[node_macro::node_fn(LoadResourceNode)]
+async fn load_resource_node<'a: 'input>(editor: WasmEditorApi<'a>, url: String) -> Arc<[u8]> {
+	editor.application_io.load_resource(url).unwrap().await.unwrap()
+}
+
+pub struct DecodeImageNode;
+
+#[node_macro::node_fn(DecodeImageNode)]
+fn decode_image_node<'a: 'input>(data: Arc<[u8]>) -> ImageFrame<Color> {
+	let image = image::load_from_memory(data.as_ref()).expect("Failed to decode image");
+	let image = image.to_rgba32f();
+	let image = ImageFrame {
+		image: Image {
+			data: image.chunks(4).map(|pixel| Color::from_unassociated_alpha(pixel[0], pixel[1], pixel[2], pixel[3])).collect(),
+			width: image.width(),
+			height: image.height(),
+		},
+		transform: glam::DAffine2::IDENTITY,
+	};
+	image
+}