WIP websocket integration

image-generator-demo/src/app/app.element.ts

@@ -12,49 +12,74 @@ import { customElement, state } from "lit/decorators.js";
 import { styleMap } from "lit/directives/style-map.js";
 import { classMap } from "lit/directives/class-map.js";
 
+import { ImageInfo, LcmClient } from "./lcm-client";
+
 import "@storyteller/studio-web/canvas-provider";
 import "@storyteller/studio-web/scene-hierarchy";
 import "@storyteller/studio-web/studio";
 import "@storyteller/studio-web/transform-toolbar";
 
 import "./app.element.scss"
+import { createRef, ref } from "lit/directives/ref.js";
 
 @customElement("app-root")
 export class AppElement extends LitElement {
-	@state() imageCaptureUrl?: string;
-
 	@state() hierarchyWidth = 256;
 	@state() resizing = false;
 
 	@state() loading = true;
 
+	@state() blob?: Blob;
 	@state() prompt = "";
 	@state() negPrompt = "";
 	@state() seed = randomSeed();
 
+	@state() inputUrl?: string;
+	@state() outputUrl?: string;
+
+	#lcmClient = new LcmClient(crypto.randomUUID());
+	#outputImageRef = createRef<HTMLImageElement>();
+
+	constructor () {
+		super();
+
+		this.#lcmClient.addEventListener("lcm-image-ready", event => {
+			this.outputUrl = event.detail;
+			this.requestUpdate();
+		});
+	}
+
 	protected override createRenderRoot = () => this;
 
 	override disconnectedCallback(): void {
-		if (this.imageCaptureUrl != null)
-			URL.revokeObjectURL(this.imageCaptureUrl);
+		if (this.inputUrl)
+			URL.revokeObjectURL(this.inputUrl);
 
 		super.disconnectedCallback();
 	}
 
 	protected override willUpdate(changes: PropertyValues<this>): void {
+		let updates: Partial<ImageInfo> = {};
+
+		if (changes.has("blob"))
+			updates.blob = this.blob;
+
+		if (changes.has("prompt"))
+			updates.prompt = this.prompt;
+
+		if (changes.has("negPrompt"))
+			updates.negativePrompt = this.negPrompt;
+
+		if (changes.has("seed"))
+			updates.seed = this.seed;
+
 		if (
-			changes.has("prompt")
+			changes.has("blob")
+			|| changes.has("prompt")
 			|| changes.has("negPrompt")
 			|| changes.has("seed")
-			|| changes.has("imageCaptureUrl")
 		) {
-			// TODO: Send parameters to model
-			console.log("changes:", {
-				prompt: this.prompt,
-				negPrompt: this.negPrompt,
-				seed: this.seed,
-				imageCaptureUrl: this.imageCaptureUrl,
-			});
+			this.#lcmClient.update(updates);
 		}
 
 		super.willUpdate(changes);
@@ -85,10 +110,10 @@ export class AppElement extends LitElement {
 
 	@on("screenshot-ready")
 	onScreenshotReady({ detail: screenshot }: ScreenshotReadyEvent) {
-		const prevUrl = this.imageCaptureUrl;
+		this.blob = new Blob([screenshot.data.buffer], { type: "image/jpg" });
 
-		const blob = new Blob([screenshot.data.buffer], { type: "image/jpg" });
-		this.imageCaptureUrl = URL.createObjectURL(blob);
+		const prevUrl = this.inputUrl;
+		this.inputUrl = URL.createObjectURL(this.blob);
 
 		if (prevUrl != null)
 			URL.revokeObjectURL(prevUrl);
@@ -165,10 +190,11 @@ export class AppElement extends LitElement {
 					></sts-transform-toolbar>
 				</sts-studio>
 
-				${this.imageCaptureUrl != null ? html`
+				${this.inputUrl || this.outputUrl ? html`
 					<img
+						${ref(this.#outputImageRef)}
 						class="image-capture"
-						src=${this.imageCaptureUrl}
+						src=${(this.outputUrl ?? this.inputUrl)!}
 						alt="Image capture"
 					/>
 				` : html`

image-generator-demo/src/app/lcm-client.ts

@@ -0,0 +1,134 @@
+export enum LcmLiveStatus {
+	None = "NONE",
+	Connected = "connected",
+	Disconnected = "disconnected",
+	Wait = "wait",
+	SendFrame = "send_frame",
+	Timeout = "timeout",
+}
+
+export interface ImageInfo {
+	blob: Blob;
+	prompt: string;
+	negativePrompt: string;
+	seed: number;
+	guidanceScale: number;
+}
+
+interface LcmMessageData {
+	status: LcmLiveStatus | "error";
+	message?: string;
+}
+
+export class LcmClient extends WebSocket {
+	#userId: string;
+
+	get status() { return this.#status; }
+	#status = LcmLiveStatus.None;
+
+	#waiting = false;
+
+	#blob?: Blob;
+	#prompt?: string;
+	#negativePrompt?: string;
+	#seed?: number;
+	#guidanceScale = 0.8;
+
+	constructor (userId: string) {
+		super(`wss://b3000.pancake-chimaera.ts.net/api/ws/${userId}`);
+		this.#userId = userId;
+
+		this.addEventListener("open", () => this.#onOpen());
+		this.addEventListener("close", () => this.#onClose());
+		this.addEventListener("message", event => this.#onMessage(event))
+	}
+
+	update(data: Partial<ImageInfo>): void {
+		if ("blob" in data)
+			this.#blob = data.blob;
+
+		if ("prompt" in data)
+			this.#prompt = data.prompt;
+
+		if ("negativePrompt" in data)
+			this.#negativePrompt = data.negativePrompt;
+
+		if ("seed" in data)
+			this.#seed = data.seed;
+
+		this.#update();
+	}
+
+	#onOpen(): void {
+		console.log(`Marking connection to websocket as open for user ${this.#userId}`);
+		this.#status = LcmLiveStatus.Connected;
+	}
+
+	#onClose(): void {
+		console.log(`Marking connection to websocket as closed for user ${this.#userId}`);
+		this.#status = LcmLiveStatus.Disconnected;
+	}
+
+	#onMessage(event: MessageEvent<string>): void {
+		console.log("onMessage:", event);
+		const data: LcmMessageData = JSON.parse(event.data);
+
+		if (data.status === "error") {
+			this.#status = LcmLiveStatus.Disconnected;
+			console.error(data.message);
+
+			return;
+		}
+
+		this.#status = data.status;
+
+		if (this.#status === LcmLiveStatus.SendFrame && this.#waiting)
+			this.#sendFrame();
+	}
+
+	#update(): void {
+		if (this.#status === LcmLiveStatus.SendFrame) {
+			this.#sendFrame();
+		} else {
+			this.#waiting = true;
+		}
+	}
+
+	#sendFrame(): void {
+		if (!this.#blob || !this.#prompt || !this.#negativePrompt || this.#seed == null) {
+			this.#waiting = true;
+			return;
+		}
+
+		console.log("Sending next frame...");
+
+		this.send(JSON.stringify({ status: "next_frame" }));
+		this.send(JSON.stringify({
+			prompt: this.#prompt,
+			negativePrompt: this.#negativePrompt,
+			seed: this.#seed,
+			guidanceScale: this.#guidanceScale,
+		}));
+		this.send(this.#blob);
+
+		this.#waiting = false;
+
+		setTimeout(() => {
+			this.dispatchEvent(new CustomEvent("lcm-image-ready", {
+				detail: `https://b3000.pancake-chimaera.ts.net/api/stream/${this.#userId}`
+			}));
+		});
+	}
+}
+
+export interface LcmClientEventMap extends WebSocketEventMap {
+	"lcm-image-ready": CustomEvent<string>;
+}
+
+export interface LcmClient {
+	addEventListener<K extends keyof LcmClientEventMap>(
+		type: K,
+		listener: (this: LcmClient, ev: LcmClientEventMap[K]) => any,
+		options?: boolean | AddEventListenerOptions
+	): void;
+}