diff --git a/image-generator-demo/src/app/app.element.ts b/image-generator-demo/src/app/app.element.ts index 0b997a3..15f89ff 100644 --- a/image-generator-demo/src/app/app.element.ts +++ b/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(); + + 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): void { + let updates: Partial = {}; + + 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 { > - ${this.imageCaptureUrl != null ? html` + ${this.inputUrl || this.outputUrl ? html` Image capture ` : html` diff --git a/image-generator-demo/src/app/lcm-client.ts b/image-generator-demo/src/app/lcm-client.ts new file mode 100644 index 0000000..96c419e --- /dev/null +++ b/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): 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): 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; +} + +export interface LcmClient { + addEventListener( + type: K, + listener: (this: LcmClient, ev: LcmClientEventMap[K]) => any, + options?: boolean | AddEventListenerOptions + ): void; +}