From 953de99fb0e7856b28e3930af5f868bb2e5618ba Mon Sep 17 00:00:00 2001
From: Kasisnu <5044120+kasisnu@users.noreply.github.com>
Date: Tue, 27 Feb 2024 14:08:21 +0530
Subject: [PATCH] some updates

---
 image-generator-demo/src/app/app.element.ts | 44 +++++++++++++--------
 studio/src/wasm/screenshot.rs               | 26 +++++++-----
 2 files changed, 44 insertions(+), 26 deletions(-)

diff --git a/image-generator-demo/src/app/app.element.ts b/image-generator-demo/src/app/app.element.ts
index 96c0ae7..4cf3465 100644
--- a/image-generator-demo/src/app/app.element.ts
+++ b/image-generator-demo/src/app/app.element.ts
@@ -75,6 +75,8 @@ export class AppElement extends LitElement {
 	@state() realtimeCapture = true;
 	@state() lcmLiveStatus = LCMLiveStatus.DISCONNECTED;
 	@state() websocket: WebSocket | null = null;
+	@state() websocketUUID: string | null = null;
+	@state() latestBlob: Blob | null = null;
 
 	params = new StudioURLParams();
 	viewer = createRef<ViewerElement>();
@@ -108,23 +110,27 @@ export class AppElement extends LitElement {
 	@on("screenshot-ready")
 	onScreenshotReady({ detail: screenshot }: ScreenshotReadyEvent) {
 		const prevUrl = this.imageCaptureUrl;
-		const blob = new Blob([screenshot.data.buffer], { type: "image/png" });
+		// const blob = new Blob([screenshot.data.buffer], { type: "image/png" });
+		const blob = new Blob([screenshot.data.buffer], {type : 'image/jpeg'});
+		this.latestBlob = blob;
+		// const blob = new Blob (["iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNk+A8AAQUBAScY42YAAAAASUVORK5CYII="]);
 
-		const innerSend = (blob: Blob | { [key: string]: any }) : void => {
-			if (this.websocket && this.websocket.readyState === WebSocket.OPEN) {
-				if (blob instanceof Blob) {
-					this.websocket.send(blob);
-				} else {
-					this.websocket.send(JSON.stringify(blob));
-				}
-			} else {
-				console.log("WebSocket not connected durign innerSend with state" + this.lcmLiveStatus + " and readyState " + this.websocket?.readyState + " and blob " + blob + " and websocket " + this.websocket + " and websocket.readyState " + this.websocket?.readyState);
-			}
-		};
+		// const innerSend = (blob: Blob | { [key: string]: any }) : void => {
+		// 	if (this.websocket && this.websocket.readyState === WebSocket.OPEN) {
+		// 		if (blob instanceof Blob) {
+		// 			this.websocket.send(blob);
+		// 		} else {
+		// 			this.websocket.send(JSON.stringify(blob));
+		// 		}
+		// 	} else {
+		// 		console.log("WebSocket not connected durign innerSend with state" + this.lcmLiveStatus + " and readyState " + this.websocket?.readyState + " and blob " + blob + " and websocket " + this.websocket + " and websocket.readyState " + this.websocket?.readyState);
+		// 	}
+		// };
 
 		// if websocket is disconnected, connect
 		if (this.lcmLiveStatus === LCMLiveStatus.DISCONNECTED) {
 			const userId = crypto.randomUUID();
+			this.websocketUUID = userId;
 			this.websocket = new WebSocket("wss://b3000.pancake-chimaera.ts.net/api/ws/" + userId);
 			this.websocket.onopen = () => {
 				console.log("Marking connection to websocket as open for user " + userId);
@@ -150,12 +156,18 @@ export class AppElement extends LitElement {
 							seed: 12345,
 							guidance_scale: 1.2,
 						});
-						const streamData = blob;
+						const streamData = [this.latestBlob];
 						console.log("Sending next frame with streamData" + streamData + " and websocket " + this.websocket + " and websocket.readyState " + this.websocket?.readyState)
 						this.websocket?.send(JSON.stringify({ status: "next_frame" }));
-						// for (const d of streamData) {
-							innerSend(streamData);
-						// }
+						this.websocket?.send(JSON.stringify({ prompt: default_prompt }));
+
+						for (const d of streamData) {
+							if (d instanceof Blob) {
+								this.websocket?.send(d);
+							} else {
+								this.websocket?.send(JSON.stringify(d));
+							}
+						}
 						break;
 					case "wait":
 						this.lcmLiveStatus = LCMLiveStatus.WAIT;
diff --git a/studio/src/wasm/screenshot.rs b/studio/src/wasm/screenshot.rs
index b4141c4..b3979de 100644
--- a/studio/src/wasm/screenshot.rs
+++ b/studio/src/wasm/screenshot.rs
@@ -10,7 +10,7 @@
 //!    to the frontend by converting the `Vec<u8>` to a [`js_sys::Uint8Array`]
 
 use bevy::{prelude::*, render::view::screenshot::ScreenshotManager, window::PrimaryWindow};
-use image::{codecs::png::PngEncoder, ColorType, ImageEncoder};
+use image::{codecs::{jpeg::JpegEncoder, png::PngEncoder}, ColorType, ImageEncoder};
 use wasm_bindgen::prelude::{wasm_bindgen, JsValue};
 
 use super::{AsJsValue, StaticBuffer, WasmMessageBuffer};
@@ -111,20 +111,26 @@ fn take_screenshots(
                 }
             };
             let src_buffer = &img.as_raw()[..buffer_len];
+            let mut jpeg_bytes = vec![];
 
-            let mut png_bytes = vec![];
-            if let Err(err) = PngEncoder::new(&mut png_bytes).write_image(
-                src_buffer,
-                img.width(),
-                img.height(),
-                ColorType::Rgb8,
-            ) {
+            if let Err(err) = JpegEncoder::new(&mut jpeg_bytes)
+                .encode(src_buffer, img.width(), img.height(), ColorType::Rgb8)
+                .map_err(|err| format!("{err}")) {
                 error!("{err}");
                 return;
-            }
+                };
+            // if let Err(err) = PngEncoder::new(&mut png_bytes).write_image(
+            //     src_buffer,
+            //     img.width(),
+            //     img.height(),
+            //     ColorType::Rgb8,
+            // ) {
+            //     error!("{err}");
+            //     return;
+            // }
 
             SCREENSHOT.write_message(ScreenshotNative {
-                data: png_bytes,
+                data: jpeg_bytes,
                 width: img.width() as usize,
                 height: img.height() as usize,
             });