controlnet

app.py

@@ -3,7 +3,7 @@ from fastapi import FastAPI
 from config import args
 from device import device, torch_dtype
 from app_init import init_app
-from user_queue import user_queue_map
+from user_queue import user_data_events
 from util import get_pipeline_class
 
 
@@ -11,4 +11,4 @@ app = FastAPI()
 
 pipeline_class = get_pipeline_class(args.pipeline)
 pipeline = pipeline_class(args, device, torch_dtype)
-init_app(app, user_queue_map, args, pipeline)
+init_app(app, user_data_events, args, pipeline)

app_init.py

@@ -6,15 +6,16 @@ from fastapi.staticfiles import StaticFiles
 import logging
 import traceback
 from config import Args
-from user_queue import UserQueueDict
+from user_queue import UserDataEventMap, UserDataEvent
 import uuid
-import asyncio
+from asyncio import Event, sleep
 import time
 from PIL import Image
 import io
+from types import SimpleNamespace
 
 
-def init_app(app: FastAPI, user_queue_map: UserQueueDict, args: Args, pipeline):
+def init_app(app: FastAPI, user_data_events: UserDataEventMap, args: Args, pipeline):
     app.add_middleware(
         CORSMiddleware,
         allow_origins=["*"],
@@ -27,19 +28,20 @@ def init_app(app: FastAPI, user_queue_map: UserQueueDict, args: Args, pipeline):
     @app.websocket("/ws")
     async def websocket_endpoint(websocket: WebSocket):
         await websocket.accept()
-        if args.max_queue_size > 0 and len(user_queue_map) >= args.max_queue_size:
+        if args.max_queue_size > 0 and len(user_data_events) >= args.max_queue_size:
             print("Server is full")
             await websocket.send_json({"status": "error", "message": "Server is full"})
             await websocket.close()
             return
 
         try:
-            uid = uuid.uuid4()
+            uid = str(uuid.uuid4())
             print(f"New user connected: {uid}")
             await websocket.send_json(
                 {"status": "success", "message": "Connected", "userId": uid}
             )
-            user_queue_map[uid] = {"queue": asyncio.Queue()}
+            user_data_events[uid] = UserDataEvent()
+            print(f"User data events: {user_data_events}")
             await websocket.send_json(
                 {"status": "start", "message": "Start Streaming", "userId": uid}
             )
@@ -49,40 +51,27 @@ def init_app(app: FastAPI, user_queue_map: UserQueueDict, args: Args, pipeline):
             traceback.print_exc()
         finally:
             print(f"User disconnected: {uid}")
-            queue_value = user_queue_map.pop(uid, None)
-            queue = queue_value.get("queue", None)
-            if queue:
-                while not queue.empty():
-                    try:
-                        queue.get_nowait()
-                    except asyncio.QueueEmpty:
-                        continue
+            del user_data_events[uid]
 
     @app.get("/queue_size")
     async def get_queue_size():
-        queue_size = len(user_queue_map)
+        queue_size = len(user_data_events)
         return JSONResponse({"queue_size": queue_size})
 
     @app.get("/stream/{user_id}")
     async def stream(user_id: uuid.UUID):
-        uid = user_id
+        uid = str(user_id)
         try:
-            user_queue = user_queue_map[uid]
-            queue = user_queue["queue"]
 
             async def generate():
                 last_prompt: str = None
                 while True:
-                    data = await queue.get()
-                    input_image = data["image"]
+                    data = await user_data_events[uid].wait_for_data()
                     params = data["params"]
-                    if input_image is None:
-                        continue
-
-                    image = pipeline.predict(
-                        input_image,
-                        params,
-                    )
+                    # input_image = data["image"]
+                    # if input_image is None:
+                    # continue
+                    image = pipeline.predict(params)
                     if image is None:
                         continue
                     frame_data = io.BytesIO()
@@ -91,36 +80,31 @@ def init_app(app: FastAPI, user_queue_map: UserQueueDict, args: Args, pipeline):
                     if frame_data is not None and len(frame_data) > 0:
                         yield b"--frame\r\nContent-Type: image/jpeg\r\n\r\n" + frame_data + b"\r\n"
 
-                    await asyncio.sleep(1.0 / 120.0)
+                    await sleep(1.0 / 120.0)
 
             return StreamingResponse(
                 generate(), media_type="multipart/x-mixed-replace;boundary=frame"
             )
         except Exception as e:
-            logging.error(f"Streaming Error: {e}, {user_queue_map}")
+            logging.error(f"Streaming Error: {e}, {user_data_events}")
             traceback.print_exc()
             return HTTPException(status_code=404, detail="User not found")
 
     async def handle_websocket_data(websocket: WebSocket, user_id: uuid.UUID):
-        uid = user_id
-        user_queue = user_queue_map[uid]
-        queue = user_queue["queue"]
-        if not queue:
+        uid = str(user_id)
+        if uid not in user_data_events:
             return HTTPException(status_code=404, detail="User not found")
         last_time = time.time()
         try:
             while True:
-                data = await websocket.receive_bytes()
                 params = await websocket.receive_json()
                 params = pipeline.InputParams(**params)
-                pil_image = Image.open(io.BytesIO(data))
-
-                while not queue.empty():
-                    try:
-                        queue.get_nowait()
-                    except asyncio.QueueEmpty:
-                        continue
-                await queue.put({"image": pil_image, "params": params})
+                params = SimpleNamespace(**params.dict())
+                if hasattr(params, "image"):
+                    image_data = await websocket.receive_bytes()
+                    pil_image = Image.open(io.BytesIO(image_data))
+                    params.image = pil_image
+                user_data_events[uid].update_data({"params": params})
                 if args.timeout > 0 and time.time() - last_time > args.timeout:
                     await websocket.send_json(
                         {

frontend/package-lock.json

@@ -7,6 +7,9 @@
     "": {
       "name": "frontend",
       "version": "0.0.1",
+      "dependencies": {
+        "rvfc-polyfill": "^1.0.7"
+      },
       "devDependencies": {
         "@sveltejs/adapter-auto": "^2.0.0",
         "@sveltejs/adapter-static": "^2.0.3",
@@ -3035,6 +3038,11 @@
         "queue-microtask": "^1.2.2"
       }
     },
+    "node_modules/rvfc-polyfill": {
+      "version": "1.0.7",
+      "resolved": "https://registry.npmjs.org/rvfc-polyfill/-/rvfc-polyfill-1.0.7.tgz",
+      "integrity": "sha512-seBl7J1J3/k0LuzW2T9fG6JIOpni5AbU+/87LA+zTYKgTVhsfShmS8K/yOo1eeEjGJHnAdkVAUUM+PEjN9Mpkw=="
+    },
     "node_modules/sade": {
       "version": "1.8.1",
       "resolved": "https://registry.npmjs.org/sade/-/sade-1.8.1.tgz",

frontend/package.json

@@ -33,5 +33,8 @@
     "typescript": "^5.0.0",
     "vite": "^4.4.2"
   },
-  "type": "module"
+  "type": "module",
+  "dependencies": {
+    "rvfc-polyfill": "^1.0.7"
+  }
 }

frontend/src/lib/components/Button.svelte

@@ -1,8 +1,9 @@
 <script lang="ts">
   export let classList: string = '';
+  export let disabled: boolean = false;
 </script>
 
-<button class="button {classList}" on:click>
+<button class="button {classList}" on:click {disabled}>
   <slot />
 </button>
 

frontend/src/lib/components/Checkbox.svelte

@@ -0,0 +1,10 @@
+<script lang="ts">
+  import type { FieldProps } from '$lib/types';
+  export let value = false;
+  export let params: FieldProps;
+</script>
+
+<div class="grid max-w-md grid-cols-4 items-center justify-items-start gap-3">
+  <label class="text-sm font-medium" for={params.id}>{params?.title}</label>
+  <input bind:checked={value} type="checkbox" id={params.id} class="cursor-pointer" />
+</div>

frontend/src/lib/components/ImagePlayer.svelte

@@ -1,12 +1,18 @@
 <script lang="ts">
+  import { isLCMRunning, lcmLiveState, lcmLiveActions } from '$lib/lcmLive';
+  import { onFrameChangeStore } from '$lib/mediaStream';
+  import { PUBLIC_BASE_URL } from '$env/static/public';
+
+  $: streamId = $lcmLiveState.streamId;
 </script>
 
 <div class="relative overflow-hidden rounded-lg border border-slate-300">
   <!-- svelte-ignore a11y-missing-attribute -->
-  <img
-    class="aspect-square w-full rounded-lg"
-    src="data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII="
-  />
+  {#if $isLCMRunning}
+    <img class="aspect-square w-full rounded-lg" src={PUBLIC_BASE_URL + '/stream/' + streamId} />
+  {:else}
+    <div class="aspect-square w-full rounded-lg" />
+  {/if}
   <div class="absolute left-0 top-0 aspect-square w-1/4">
     <div class="relative z-10 aspect-square w-full object-cover">
       <slot />

frontend/src/lib/components/InputRange.svelte

@@ -8,14 +8,14 @@
   });
 </script>
 
-<div class="grid max-w-md grid-cols-4 items-center gap-3">
-  <label class="text-sm font-medium" for="guidance-scale">{params?.title}</label>
+<div class="grid grid-cols-4 items-center gap-3">
+  <label class="text-sm font-medium" for={params.id}>{params?.title}</label>
   <input
-    class="col-span-2"
+    class="col-span-2 h-2 w-full cursor-pointer appearance-none rounded-lg bg-gray-300 dark:bg-gray-500"
     bind:value
     type="range"
-    id="guidance-scale"
-    name="guidance-scale"
+    id={params.id}
+    name={params.id}
     min={params?.min}
     max={params?.max}
     step={params?.step ?? 1}
@@ -24,6 +24,27 @@
     type="number"
     step={params?.step ?? 1}
     bind:value
-    class="rounded-md border border-gray-700 px-1 py-1 text-center text-xs font-bold dark:text-black"
+    class="rounded-md border px-1 py-1 text-center text-xs font-bold dark:text-black"
   />
 </div>
+<!-- 
+<style lang="postcss" scoped>
+  input[type='range']::-webkit-slider-runnable-track {
+    @apply h-2 cursor-pointer rounded-lg dark:bg-gray-50;
+  }
+  input[type='range']::-webkit-slider-thumb {
+    @apply cursor-pointer rounded-lg dark:bg-gray-50;
+  }
+  input[type='range']::-moz-range-track {
+    @apply cursor-pointer rounded-lg dark:bg-gray-50;
+  }
+  input[type='range']::-moz-range-thumb {
+    @apply cursor-pointer rounded-lg dark:bg-gray-50;
+  }
+  input[type='range']::-ms-track {
+    @apply cursor-pointer rounded-lg dark:bg-gray-50;
+  }
+  input[type='range']::-ms-thumb {
+    @apply cursor-pointer rounded-lg dark:bg-gray-50;
+  }
+</style> -->

frontend/src/lib/components/PipelineOptions.svelte

@@ -5,6 +5,7 @@
   import InputRange from './InputRange.svelte';
   import SeedInput from './SeedInput.svelte';
   import TextArea from './TextArea.svelte';
+  import Checkbox from './Checkbox.svelte';
 
   export let pipelineParams: FieldProps[];
   export let pipelineValues = {} as any;
@@ -17,11 +18,13 @@
   {#if featuredOptions}
     {#each featuredOptions as params}
       {#if params.field === FieldType.range}
-        <InputRange {params} bind:value={pipelineValues[params.title]}></InputRange>
+        <InputRange {params} bind:value={pipelineValues[params.id]}></InputRange>
       {:else if params.field === FieldType.seed}
-        <SeedInput bind:value={pipelineValues[params.title]}></SeedInput>
+        <SeedInput bind:value={pipelineValues[params.id]}></SeedInput>
       {:else if params.field === FieldType.textarea}
-        <TextArea {params} bind:value={pipelineValues[params.title]}></TextArea>
+        <TextArea {params} bind:value={pipelineValues[params.id]}></TextArea>
+      {:else if params.field === FieldType.checkbox}
+        <Checkbox {params} bind:value={pipelineValues[params.id]}></Checkbox>
       {/if}
     {/each}
   {/if}
@@ -29,15 +32,17 @@
 
 <details open>
   <summary class="cursor-pointer font-medium">Advanced Options</summary>
-  <div class="flex flex-col gap-3 py-3">
+  <div class="grid grid-cols-1 items-center gap-3 sm:grid-cols-2">
     {#if advanceOptions}
       {#each advanceOptions as params}
         {#if params.field === FieldType.range}
-          <InputRange {params} bind:value={pipelineValues[params.title]}></InputRange>
+          <InputRange {params} bind:value={pipelineValues[params.id]}></InputRange>
         {:else if params.field === FieldType.seed}
-          <SeedInput bind:value={pipelineValues[params.title]}></SeedInput>
+          <SeedInput bind:value={pipelineValues[params.id]}></SeedInput>
         {:else if params.field === FieldType.textarea}
-          <TextArea {params} bind:value={pipelineValues[params.title]}></TextArea>
+          <TextArea {params} bind:value={pipelineValues[params.id]}></TextArea>
+        {:else if params.field === FieldType.checkbox}
+          <Checkbox {params} bind:value={pipelineValues[params.id]}></Checkbox>
         {/if}
       {/each}
     {/if}

frontend/src/lib/components/SeedInput.svelte

@@ -16,5 +16,5 @@
     name="seed"
     class="col-span-2 rounded-md border border-gray-700 p-2 text-right font-light dark:text-black"
   />
-  <Button on:click={randomize}>Randomize</Button>
+  <Button on:click={randomize}>Rand</Button>
 </div>

frontend/src/lib/components/VideoInput.svelte

@@ -1,4 +1,73 @@
 <script lang="ts">
+  import 'rvfc-polyfill';
+  import { onMount, onDestroy } from 'svelte';
+  import {
+    mediaStreamState,
+    mediaStreamActions,
+    isMediaStreaming,
+    MediaStreamStatus,
+    onFrameChangeStore
+  } from '$lib/mediaStream';
+
+  $: mediaStream = $mediaStreamState.mediaStream;
+
+  let videoEl: HTMLVideoElement;
+  let videoFrameCallbackId: number;
+  const WIDTH = 512;
+  const HEIGHT = 512;
+
+  onDestroy(() => {
+    if (videoFrameCallbackId) videoEl.cancelVideoFrameCallback(videoFrameCallbackId);
+  });
+
+  function srcObject(node: HTMLVideoElement, stream: MediaStream) {
+    node.srcObject = stream;
+    return {
+      update(newStream: MediaStream) {
+        if (node.srcObject != newStream) {
+          node.srcObject = newStream;
+        }
+      }
+    };
+  }
+  async function onFrameChange(now: DOMHighResTimeStamp, metadata: VideoFrameCallbackMetadata) {
+    const blob = await grapBlobImg();
+    onFrameChangeStore.set({ now, metadata, blob });
+    videoFrameCallbackId = videoEl.requestVideoFrameCallback(onFrameChange);
+  }
+
+  $: if ($isMediaStreaming == MediaStreamStatus.CONNECTED) {
+    videoFrameCallbackId = videoEl.requestVideoFrameCallback(onFrameChange);
+  }
+  async function grapBlobImg() {
+    const canvas = new OffscreenCanvas(WIDTH, HEIGHT);
+    const videoW = videoEl.videoWidth;
+    const videoH = videoEl.videoHeight;
+    const aspectRatio = WIDTH / HEIGHT;
+
+    const ctx = canvas.getContext('2d') as OffscreenCanvasRenderingContext2D;
+    ctx.drawImage(
+      videoEl,
+      videoW / 2 - (videoH * aspectRatio) / 2,
+      0,
+      videoH * aspectRatio,
+      videoH,
+      0,
+      0,
+      WIDTH,
+      HEIGHT
+    );
+    const blob = await canvas.convertToBlob({ type: 'image/jpeg', quality: 1 });
+    return blob;
+  }
 </script>
 
-<video playsinline autoplay muted loop />
+<video
+  class="aspect-square w-full object-cover"
+  bind:this={videoEl}
+  playsinline
+  autoplay
+  muted
+  loop
+  use:srcObject={mediaStream}
+></video>

frontend/src/lib/lcmLive.ts

@@ -0,0 +1,99 @@
+import { writable } from 'svelte/store';
+import { PUBLIC_BASE_URL, PUBLIC_WSS_URL } from '$env/static/public';
+
+export const isStreaming = writable(false);
+export const isLCMRunning = writable(false);
+
+
+export enum LCMLiveStatus {
+    INIT = "init",
+    CONNECTED = "connected",
+    DISCONNECTED = "disconnected",
+}
+
+interface lcmLive {
+    streamId: string | null;
+    status: LCMLiveStatus
+}
+
+const initialState: lcmLive = {
+    streamId: null,
+    status: LCMLiveStatus.INIT
+};
+
+export const lcmLiveState = writable(initialState);
+
+let websocket: WebSocket | null = null;
+export const lcmLiveActions = {
+    async start() {
+
+        isLCMRunning.set(true);
+        try {
+            const websocketURL = PUBLIC_WSS_URL ? PUBLIC_WSS_URL : `${window.location.protocol === "https:" ? "wss" : "ws"
+                }:${window.location.host}/ws`;
+
+            websocket = new WebSocket(websocketURL);
+            websocket.onopen = () => {
+                console.log("Connected to websocket");
+            };
+            websocket.onclose = () => {
+                lcmLiveState.update((state) => ({
+                    ...state,
+                    status: LCMLiveStatus.DISCONNECTED
+                }));
+                console.log("Disconnected from websocket");
+                isLCMRunning.set(false);
+            };
+            websocket.onerror = (err) => {
+                console.error(err);
+            };
+            websocket.onmessage = (event) => {
+                const data = JSON.parse(event.data);
+                console.log("WS: ", data);
+                switch (data.status) {
+                    case "success":
+                        break;
+                    case "start":
+                        const streamId = data.userId;
+                        lcmLiveState.update((state) => ({
+                            ...state,
+                            status: LCMLiveStatus.CONNECTED,
+                            streamId: streamId,
+                        }));
+                        break;
+                    case "timeout":
+                        console.log("timeout");
+                    case "error":
+                        console.log(data.message);
+                        isLCMRunning.set(false);
+                }
+            };
+            lcmLiveState.update((state) => ({
+                ...state,
+            }));
+        } catch (err) {
+            console.error(err);
+            isLCMRunning.set(false);
+        }
+    },
+    send(data: Blob | { [key: string]: any }) {
+        if (websocket && websocket.readyState === WebSocket.OPEN) {
+            if (data instanceof Blob) {
+                websocket.send(data);
+            } else {
+                websocket.send(JSON.stringify(data));
+            }
+        } else {
+            console.log("WebSocket not connected");
+        }
+    },
+    async stop() {
+
+        if (websocket) {
+            websocket.close();
+        }
+        websocket = null;
+        lcmLiveState.set({ status: LCMLiveStatus.DISCONNECTED, streamId: null });
+        isLCMRunning.set(false)
+    },
+};

frontend/src/lib/mediaStream.ts

@@ -0,0 +1,93 @@
+import { writable, type Writable } from 'svelte/store';
+
+export enum MediaStreamStatus {
+    INIT = "init",
+    CONNECTED = "connected",
+    DISCONNECTED = "disconnected",
+}
+export const onFrameChangeStore: Writable<{ now: Number, metadata: VideoFrameCallbackMetadata, blob: Blob }> = writable();
+export const isMediaStreaming = writable(MediaStreamStatus.INIT);
+
+interface mediaStream {
+    mediaStream: MediaStream | null;
+    status: MediaStreamStatus
+    devices: MediaDeviceInfo[];
+}
+
+const initialState: mediaStream = {
+    mediaStream: null,
+    status: MediaStreamStatus.INIT,
+    devices: [],
+};
+
+export const mediaStreamState = writable(initialState);
+
+export const mediaStreamActions = {
+    async enumerateDevices() {
+        console.log("Enumerating devices");
+        await navigator.mediaDevices.enumerateDevices()
+            .then(devices => {
+                const cameras = devices.filter(device => device.kind === 'videoinput');
+                console.log("Cameras: ", cameras);
+                mediaStreamState.update((state) => ({
+                    ...state,
+                    devices: cameras,
+                }));
+            })
+            .catch(err => {
+                console.error(err);
+            });
+    },
+    async start(mediaDevicedID?: string) {
+        const constraints = {
+            audio: false,
+            video: {
+                width: 1024, height: 1024, deviceId: mediaDevicedID
+            }
+        };
+
+        await navigator.mediaDevices
+            .getUserMedia(constraints)
+            .then((mediaStream) => {
+                mediaStreamState.update((state) => ({
+                    ...state,
+                    mediaStream: mediaStream,
+                    status: MediaStreamStatus.CONNECTED,
+                }));
+                isMediaStreaming.set(MediaStreamStatus.CONNECTED);
+            })
+            .catch((err) => {
+                console.error(`${err.name}: ${err.message}`);
+                isMediaStreaming.set(MediaStreamStatus.DISCONNECTED);
+            });
+    },
+    async switchCamera(mediaDevicedID: string) {
+        const constraints = {
+            audio: false,
+            video: { width: 1024, height: 1024, deviceId: mediaDevicedID }
+        };
+        await navigator.mediaDevices
+            .getUserMedia(constraints)
+            .then((mediaStream) => {
+                mediaStreamState.update((state) => ({
+                    ...state,
+                    mediaStream: mediaStream,
+                    status: MediaStreamStatus.CONNECTED,
+                }));
+            })
+            .catch((err) => {
+                console.error(`${err.name}: ${err.message}`);
+            });
+    },
+    async stop() {
+        navigator.mediaDevices.getUserMedia({ video: true }).then((mediaStream) => {
+            mediaStream.getTracks().forEach((track) => track.stop());
+        });
+        mediaStreamState.update((state) => ({
+            ...state,
+            mediaStream: null,
+            status: MediaStreamStatus.DISCONNECTED,
+        }));
+        isMediaStreaming.set(MediaStreamStatus.DISCONNECTED);
+    },
+};

frontend/src/lib/types.ts

@@ -2,6 +2,7 @@ export const enum FieldType {
     range = "range",
     seed = "seed",
     textarea = "textarea",
+    checkbox = "checkbox",
 }
 
 export interface FieldProps {
@@ -13,6 +14,7 @@ export interface FieldProps {
     step?: number;
     disabled?: boolean;
     hide?: boolean;
+    id: string;
 }
 export interface PipelineInfo {
     name: string;

frontend/src/lib/utils.ts

@@ -0,0 +1,145 @@
+export function LCMLive(webcamVideo, liveImage) {
+    let websocket: WebSocket;
+
+    async function start() {
+        return new Promise((resolve, reject) => {
+            const websocketURL = `${window.location.protocol === "https:" ? "wss" : "ws"
+                }:${window.location.host}/ws`;
+
+            const socket = new WebSocket(websocketURL);
+            socket.onopen = () => {
+                console.log("Connected to websocket");
+            };
+            socket.onclose = () => {
+                console.log("Disconnected from websocket");
+                stop();
+                resolve({ "status": "disconnected" });
+            };
+            socket.onerror = (err) => {
+                console.error(err);
+                reject(err);
+            };
+            socket.onmessage = (event) => {
+                const data = JSON.parse(event.data);
+                switch (data.status) {
+                    case "success":
+                        break;
+                    case "start":
+                        const userId = data.userId;
+                        initVideoStream(userId);
+                        break;
+                    case "timeout":
+                        stop();
+                        resolve({ "status": "timeout" });
+                    case "error":
+                        stop();
+                        reject(data.message);
+
+                }
+            };
+            websocket = socket;
+        })
+    }
+    function switchCamera() {
+        const constraints = {
+            audio: false,
+            video: { width: 1024, height: 1024, deviceId: mediaDevices[webcamsEl.value].deviceId }
+        };
+        navigator.mediaDevices
+            .getUserMedia(constraints)
+            .then((mediaStream) => {
+                webcamVideo.removeEventListener("timeupdate", videoTimeUpdateHandler);
+                webcamVideo.srcObject = mediaStream;
+                webcamVideo.onloadedmetadata = () => {
+                    webcamVideo.play();
+                    webcamVideo.addEventListener("timeupdate", videoTimeUpdateHandler);
+                };
+            })
+            .catch((err) => {
+                console.error(`${err.name}: ${err.message}`);
+            });
+    }
+
+    async function videoTimeUpdateHandler() {
+        const dimension = getValue("input[name=dimension]:checked");
+        const [WIDTH, HEIGHT] = JSON.parse(dimension);
+
+        const canvas = new OffscreenCanvas(WIDTH, HEIGHT);
+        const videoW = webcamVideo.videoWidth;
+        const videoH = webcamVideo.videoHeight;
+        const aspectRatio = WIDTH / HEIGHT;
+
+        const ctx = canvas.getContext("2d");
+        ctx.drawImage(webcamVideo, videoW / 2 - videoH * aspectRatio / 2, 0, videoH * aspectRatio, videoH, 0, 0, WIDTH, HEIGHT)
+        const blob = await canvas.convertToBlob({ type: "image/jpeg", quality: 1 });
+        websocket.send(blob);
+        websocket.send(JSON.stringify({
+            "seed": getValue("#seed"),
+            "prompt": getValue("#prompt"),
+            "guidance_scale": getValue("#guidance-scale"),
+            "strength": getValue("#strength"),
+            "steps": getValue("#steps"),
+            "lcm_steps": getValue("#lcm_steps"),
+            "width": WIDTH,
+            "height": HEIGHT,
+            "controlnet_scale": getValue("#controlnet_scale"),
+            "controlnet_start": getValue("#controlnet_start"),
+            "controlnet_end": getValue("#controlnet_end"),
+            "canny_low_threshold": getValue("#canny_low_threshold"),
+            "canny_high_threshold": getValue("#canny_high_threshold"),
+            "debug_canny": getValue("#debug_canny")
+        }));
+    }
+    let mediaDevices = [];
+    async function initVideoStream(userId) {
+        liveImage.src = `/stream/${userId}`;
+        await navigator.mediaDevices.enumerateDevices()
+            .then(devices => {
+                const cameras = devices.filter(device => device.kind === 'videoinput');
+                mediaDevices = cameras;
+                webcamsEl.innerHTML = "";
+                cameras.forEach((camera, index) => {
+                    const option = document.createElement("option");
+                    option.value = index;
+                    option.innerText = camera.label;
+                    webcamsEl.appendChild(option);
+                    option.selected = index === 0;
+                });
+                webcamsEl.addEventListener("change", switchCamera);
+            })
+            .catch(err => {
+                console.error(err);
+            });
+        const constraints = {
+            audio: false,
+            video: { width: 1024, height: 1024, deviceId: mediaDevices[0].deviceId }
+        };
+        navigator.mediaDevices
+            .getUserMedia(constraints)
+            .then((mediaStream) => {
+                webcamVideo.srcObject = mediaStream;
+                webcamVideo.onloadedmetadata = () => {
+                    webcamVideo.play();
+                    webcamVideo.addEventListener("timeupdate", videoTimeUpdateHandler);
+                };
+            })
+            .catch((err) => {
+                console.error(`${err.name}: ${err.message}`);
+            });
+    }
+
+
+    async function stop() {
+        websocket.close();
+        navigator.mediaDevices.getUserMedia({ video: true }).then((mediaStream) => {
+            mediaStream.getTracks().forEach((track) => track.stop());
+        });
+        webcamVideo.removeEventListener("timeupdate", videoTimeUpdateHandler);
+        webcamsEl.removeEventListener("change", switchCamera);
+        webcamVideo.srcObject = null;
+    }
+    return {
+        start,
+        stop
+    }
+}

frontend/src/routes/+page.svelte

@@ -7,6 +7,13 @@
   import Button from '$lib/components/Button.svelte';
   import PipelineOptions from '$lib/components/PipelineOptions.svelte';
   import Spinner from '$lib/icons/spinner.svelte';
+  import { isLCMRunning, lcmLiveState, lcmLiveActions, LCMLiveStatus } from '$lib/lcmLive';
+  import {
+    mediaStreamState,
+    mediaStreamActions,
+    isMediaStreaming,
+    onFrameChangeStore
+  } from '$lib/mediaStream';
 
   let pipelineParams: FieldProps[];
   let pipelineInfo: PipelineInfo;
@@ -21,11 +28,58 @@
     pipelineParams = Object.values(settings.input_params.properties);
     pipelineInfo = settings.info.properties;
     pipelineParams = pipelineParams.filter((e) => e?.disabled !== true);
+    console.log('PARAMS', pipelineParams);
     console.log('SETTINGS', pipelineInfo);
   }
 
-  $: {
-    console.log('PARENT', pipelineValues);
+  // $: {
+  //   console.log('isLCMRunning', $isLCMRunning);
+  // }
+  // $: {
+  //   console.log('lcmLiveState', $lcmLiveState);
+  // }
+  // $: {
+  //   console.log('mediaStreamState', $mediaStreamState);
+  // }
+  // $: if ($lcmLiveState.status === LCMLiveStatus.CONNECTED) {
+  //   lcmLiveActions.send(pipelineValues);
+  // }
+  onFrameChangeStore.subscribe(async (frame) => {
+    if ($lcmLiveState.status === LCMLiveStatus.CONNECTED) {
+      lcmLiveActions.send(pipelineValues);
+      lcmLiveActions.send(frame.blob);
+    }
+  });
+  let startBt: Button;
+  let stopBt: Button;
+  let snapShotBt: Button;
+
+  async function toggleLcmLive() {
+    if (!$isLCMRunning) {
+      await mediaStreamActions.enumerateDevices();
+      await mediaStreamActions.start();
+      lcmLiveActions.start();
+    } else {
+      mediaStreamActions.stop();
+      lcmLiveActions.stop();
+    }
+  }
+  async function startLcmLive() {
+    try {
+      $isLCMRunning = true;
+      // const res = await lcmLive.start();
+      $isLCMRunning = false;
+      // if (res.status === "timeout")
+      // toggleMessage("success")
+    } catch (err) {
+      console.log(err);
+      // toggleMessage("error")
+      $isLCMRunning = false;
+    }
+  }
+  async function stopLcmLive() {
+    // await lcmLive.stop();
+    $isLCMRunning = false;
   }
 </script>
 
@@ -58,19 +112,26 @@
     </p>
   </article>
   {#if pipelineParams}
-    <h2 class="font-medium">Prompt</h2>
-    <p class="text-sm text-gray-500">
-      Change the prompt to generate different images, accepts <a
-        href="https://github.com/damian0815/compel/blob/main/doc/syntax.md"
-        target="_blank"
-        class="text-blue-500 underline hover:no-underline">Compel</a
-      > syntax.
-    </p>
+    <header>
+      <h2 class="font-medium">Prompt</h2>
+      <p class="text-sm text-gray-500">
+        Change the prompt to generate different images, accepts <a
+          href="https://github.com/damian0815/compel/blob/main/doc/syntax.md"
+          target="_blank"
+          class="text-blue-500 underline hover:no-underline">Compel</a
+        > syntax.
+      </p>
+    </header>
     <PipelineOptions {pipelineParams} bind:pipelineValues></PipelineOptions>
     <div class="flex gap-3">
-      <Button>Start</Button>
-      <Button>Stop</Button>
-      <Button classList={'ml-auto'}>Snapshot</Button>
+      <Button on:click={toggleLcmLive}>
+        {#if $isLCMRunning}
+          Stop
+        {:else}
+          Start
+        {/if}
+      </Button>
+      <Button disabled={$isLCMRunning} classList={'ml-auto'}>Snapshot</Button>
     </div>
 
     <ImagePlayer>

latent_consistency_controlnet.py

@@ -1,1100 +0,0 @@
-# from https://github.com/taabata/LCM_Inpaint_Outpaint_Comfy/blob/main/LCM/pipeline_cn.py
-# Copyright 2023 Stanford University Team and The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion
-# and https://github.com/hojonathanho/diffusion
-
-import math
-from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple, Union
-
-import numpy as np
-import torch
-from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
-
-from diffusers import (
-    AutoencoderKL,
-    ConfigMixin,
-    DiffusionPipeline,
-    SchedulerMixin,
-    UNet2DConditionModel,
-    ControlNetModel,
-    logging,
-)
-from diffusers.configuration_utils import register_to_config
-from diffusers.image_processor import VaeImageProcessor, PipelineImageInput
-from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
-from diffusers.pipelines.stable_diffusion.safety_checker import (
-    StableDiffusionSafetyChecker,
-)
-from diffusers.utils import BaseOutput
-
-from diffusers.utils.torch_utils import randn_tensor, is_compiled_module
-from diffusers.pipelines.controlnet.multicontrolnet import MultiControlNetModel
-
-
-import PIL.Image
-
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-
-# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
-def retrieve_latents(encoder_output, generator):
-    if hasattr(encoder_output, "latent_dist"):
-        return encoder_output.latent_dist.sample(generator)
-    elif hasattr(encoder_output, "latents"):
-        return encoder_output.latents
-    else:
-        raise AttributeError("Could not access latents of provided encoder_output")
-
-
-class LatentConsistencyModelPipeline_controlnet(DiffusionPipeline):
-    _optional_components = ["scheduler"]
-
-    def __init__(
-        self,
-        vae: AutoencoderKL,
-        text_encoder: CLIPTextModel,
-        tokenizer: CLIPTokenizer,
-        controlnet: Union[
-            ControlNetModel,
-            List[ControlNetModel],
-            Tuple[ControlNetModel],
-            MultiControlNetModel,
-        ],
-        unet: UNet2DConditionModel,
-        scheduler: "LCMScheduler",
-        safety_checker: StableDiffusionSafetyChecker,
-        feature_extractor: CLIPImageProcessor,
-        requires_safety_checker: bool = True,
-    ):
-        super().__init__()
-
-        scheduler = (
-            scheduler
-            if scheduler is not None
-            else LCMScheduler_X(
-                beta_start=0.00085,
-                beta_end=0.0120,
-                beta_schedule="scaled_linear",
-                prediction_type="epsilon",
-            )
-        )
-
-        self.register_modules(
-            vae=vae,
-            text_encoder=text_encoder,
-            tokenizer=tokenizer,
-            unet=unet,
-            controlnet=controlnet,
-            scheduler=scheduler,
-            safety_checker=safety_checker,
-            feature_extractor=feature_extractor,
-        )
-        self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
-        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
-        self.control_image_processor = VaeImageProcessor(
-            vae_scale_factor=self.vae_scale_factor,
-            do_convert_rgb=True,
-            do_normalize=False,
-        )
-
-    def _encode_prompt(
-        self,
-        prompt,
-        device,
-        num_images_per_prompt,
-        prompt_embeds: None,
-    ):
-        r"""
-        Encodes the prompt into text encoder hidden states.
-        Args:
-            prompt (`str` or `List[str]`, *optional*):
-                prompt to be encoded
-            device: (`torch.device`):
-                torch device
-            num_images_per_prompt (`int`):
-                number of images that should be generated per prompt
-            prompt_embeds (`torch.FloatTensor`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
-                provided, text embeddings will be generated from `prompt` input argument.
-        """
-
-        if prompt is not None and isinstance(prompt, str):
-            pass
-        elif prompt is not None and isinstance(prompt, list):
-            len(prompt)
-        else:
-            prompt_embeds.shape[0]
-
-        if prompt_embeds is None:
-            text_inputs = self.tokenizer(
-                prompt,
-                padding="max_length",
-                max_length=self.tokenizer.model_max_length,
-                truncation=True,
-                return_tensors="pt",
-            )
-            text_input_ids = text_inputs.input_ids
-            untruncated_ids = self.tokenizer(
-                prompt, padding="longest", return_tensors="pt"
-            ).input_ids
-
-            if untruncated_ids.shape[-1] >= text_input_ids.shape[
-                -1
-            ] and not torch.equal(text_input_ids, untruncated_ids):
-                removed_text = self.tokenizer.batch_decode(
-                    untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]
-                )
-                logger.warning(
-                    "The following part of your input was truncated because CLIP can only handle sequences up to"
-                    f" {self.tokenizer.model_max_length} tokens: {removed_text}"
-                )
-
-            if (
-                hasattr(self.text_encoder.config, "use_attention_mask")
-                and self.text_encoder.config.use_attention_mask
-            ):
-                attention_mask = text_inputs.attention_mask.to(device)
-            else:
-                attention_mask = None
-
-            prompt_embeds = self.text_encoder(
-                text_input_ids.to(device),
-                attention_mask=attention_mask,
-            )
-            prompt_embeds = prompt_embeds[0]
-
-        if self.text_encoder is not None:
-            prompt_embeds_dtype = self.text_encoder.dtype
-        elif self.unet is not None:
-            prompt_embeds_dtype = self.unet.dtype
-        else:
-            prompt_embeds_dtype = prompt_embeds.dtype
-
-        prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
-
-        bs_embed, seq_len, _ = prompt_embeds.shape
-        # duplicate text embeddings for each generation per prompt, using mps friendly method
-        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
-        prompt_embeds = prompt_embeds.view(
-            bs_embed * num_images_per_prompt, seq_len, -1
-        )
-
-        # Don't need to get uncond prompt embedding because of LCM Guided Distillation
-        return prompt_embeds
-
-    def run_safety_checker(self, image, device, dtype):
-        if self.safety_checker is None:
-            has_nsfw_concept = None
-        else:
-            if torch.is_tensor(image):
-                feature_extractor_input = self.image_processor.postprocess(
-                    image, output_type="pil"
-                )
-            else:
-                feature_extractor_input = self.image_processor.numpy_to_pil(image)
-            safety_checker_input = self.feature_extractor(
-                feature_extractor_input, return_tensors="pt"
-            ).to(device)
-            image, has_nsfw_concept = self.safety_checker(
-                images=image, clip_input=safety_checker_input.pixel_values.to(dtype)
-            )
-        return image, has_nsfw_concept
-
-    def prepare_control_image(
-        self,
-        image,
-        width,
-        height,
-        batch_size,
-        num_images_per_prompt,
-        device,
-        dtype,
-        do_classifier_free_guidance=False,
-        guess_mode=False,
-    ):
-        image = self.control_image_processor.preprocess(
-            image, height=height, width=width
-        ).to(dtype=dtype)
-        image_batch_size = image.shape[0]
-
-        if image_batch_size == 1:
-            repeat_by = batch_size
-        else:
-            # image batch size is the same as prompt batch size
-            repeat_by = num_images_per_prompt
-
-        image = image.repeat_interleave(repeat_by, dim=0)
-
-        image = image.to(device=device, dtype=dtype)
-
-        if do_classifier_free_guidance and not guess_mode:
-            image = torch.cat([image] * 2)
-
-        return image
-
-    def prepare_latents(
-        self,
-        image,
-        timestep,
-        batch_size,
-        num_channels_latents,
-        height,
-        width,
-        dtype,
-        device,
-        latents=None,
-        generator=None,
-    ):
-        shape = (
-            batch_size,
-            num_channels_latents,
-            height // self.vae_scale_factor,
-            width // self.vae_scale_factor,
-        )
-
-        if not isinstance(image, (torch.Tensor, PIL.Image.Image, list)):
-            raise ValueError(
-                f"`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(image)}"
-            )
-
-        image = image.to(device=device, dtype=dtype)
-
-        # batch_size = batch_size * num_images_per_prompt
-
-        if image.shape[1] == 4:
-            init_latents = image
-
-        else:
-            if isinstance(generator, list) and len(generator) != batch_size:
-                raise ValueError(
-                    f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
-                    f" size of {batch_size}. Make sure the batch size matches the length of the generators."
-                )
-
-            elif isinstance(generator, list):
-                init_latents = [
-                    retrieve_latents(
-                        self.vae.encode(image[i : i + 1]), generator=generator[i]
-                    )
-                    for i in range(batch_size)
-                ]
-                init_latents = torch.cat(init_latents, dim=0)
-            else:
-                init_latents = retrieve_latents(
-                    self.vae.encode(image), generator=generator
-                )
-
-            init_latents = self.vae.config.scaling_factor * init_latents
-
-        if (
-            batch_size > init_latents.shape[0]
-            and batch_size % init_latents.shape[0] == 0
-        ):
-            # expand init_latents for batch_size
-            deprecation_message = (
-                f"You have passed {batch_size} text prompts (`prompt`), but only {init_latents.shape[0]} initial"
-                " images (`image`). Initial images are now duplicating to match the number of text prompts. Note"
-                " that this behavior is deprecated and will be removed in a version 1.0.0. Please make sure to update"
-                " your script to pass as many initial images as text prompts to suppress this warning."
-            )
-            # deprecate("len(prompt) != len(image)", "1.0.0", deprecation_message, standard_warn=False)
-            additional_image_per_prompt = batch_size // init_latents.shape[0]
-            init_latents = torch.cat(
-                [init_latents] * additional_image_per_prompt, dim=0
-            )
-        elif (
-            batch_size > init_latents.shape[0]
-            and batch_size % init_latents.shape[0] != 0
-        ):
-            raise ValueError(
-                f"Cannot duplicate `image` of batch size {init_latents.shape[0]} to {batch_size} text prompts."
-            )
-        else:
-            init_latents = torch.cat([init_latents], dim=0)
-
-        shape = init_latents.shape
-        noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
-
-        # get latents
-        init_latents = self.scheduler.add_noise(init_latents, noise, timestep)
-        latents = init_latents
-
-        return latents
-
-        if latents is None:
-            latents = torch.randn(shape, dtype=dtype).to(device)
-        else:
-            latents = latents.to(device)
-        # scale the initial noise by the standard deviation required by the scheduler
-        latents = latents * self.scheduler.init_noise_sigma
-        return latents
-
-    def get_w_embedding(self, w, embedding_dim=512, dtype=torch.float32):
-        """
-        see https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298
-        Args:
-        timesteps: torch.Tensor: generate embedding vectors at these timesteps
-        embedding_dim: int: dimension of the embeddings to generate
-        dtype: data type of the generated embeddings
-        Returns:
-        embedding vectors with shape `(len(timesteps), embedding_dim)`
-        """
-        assert len(w.shape) == 1
-        w = w * 1000.0
-
-        half_dim = embedding_dim // 2
-        emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
-        emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb)
-        emb = w.to(dtype)[:, None] * emb[None, :]
-        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
-        if embedding_dim % 2 == 1:  # zero pad
-            emb = torch.nn.functional.pad(emb, (0, 1))
-        assert emb.shape == (w.shape[0], embedding_dim)
-        return emb
-
-    def get_timesteps(self, num_inference_steps, strength, device):
-        # get the original timestep using init_timestep
-        init_timestep = min(int(num_inference_steps * strength), num_inference_steps)
-
-        t_start = max(num_inference_steps - init_timestep, 0)
-        timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :]
-
-        return timesteps, num_inference_steps - t_start
-
-    @torch.no_grad()
-    def __call__(
-        self,
-        prompt: Union[str, List[str]] = None,
-        image: PipelineImageInput = None,
-        control_image: PipelineImageInput = None,
-        strength: float = 0.8,
-        height: Optional[int] = 768,
-        width: Optional[int] = 768,
-        guidance_scale: float = 7.5,
-        num_images_per_prompt: Optional[int] = 1,
-        latents: Optional[torch.FloatTensor] = None,
-        generator: Optional[torch.Generator] = None,
-        num_inference_steps: int = 4,
-        lcm_origin_steps: int = 50,
-        prompt_embeds: Optional[torch.FloatTensor] = None,
-        output_type: Optional[str] = "pil",
-        return_dict: bool = True,
-        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
-        controlnet_conditioning_scale: Union[float, List[float]] = 0.8,
-        guess_mode: bool = True,
-        control_guidance_start: Union[float, List[float]] = 0.0,
-        control_guidance_end: Union[float, List[float]] = 1.0,
-    ):
-        controlnet = (
-            self.controlnet._orig_mod
-            if is_compiled_module(self.controlnet)
-            else self.controlnet
-        )
-        # 0. Default height and width to unet
-        height = height or self.unet.config.sample_size * self.vae_scale_factor
-        width = width or self.unet.config.sample_size * self.vae_scale_factor
-        if not isinstance(control_guidance_start, list) and isinstance(
-            control_guidance_end, list
-        ):
-            control_guidance_start = len(control_guidance_end) * [
-                control_guidance_start
-            ]
-        elif not isinstance(control_guidance_end, list) and isinstance(
-            control_guidance_start, list
-        ):
-            control_guidance_end = len(control_guidance_start) * [control_guidance_end]
-        elif not isinstance(control_guidance_start, list) and not isinstance(
-            control_guidance_end, list
-        ):
-            mult = (
-                len(controlnet.nets)
-                if isinstance(controlnet, MultiControlNetModel)
-                else 1
-            )
-            control_guidance_start, control_guidance_end = mult * [
-                control_guidance_start
-            ], mult * [control_guidance_end]
-        # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        device = self._execution_device
-        # do_classifier_free_guidance = guidance_scale > 0.0  # In LCM Implementation:  cfg_noise = noise_cond + cfg_scale * (noise_cond - noise_uncond) , (cfg_scale > 0.0 using CFG)
-        global_pool_conditions = (
-            controlnet.config.global_pool_conditions
-            if isinstance(controlnet, ControlNetModel)
-            else controlnet.nets[0].config.global_pool_conditions
-        )
-        guess_mode = guess_mode or global_pool_conditions
-        # 3. Encode input prompt
-        prompt_embeds = self._encode_prompt(
-            prompt,
-            device,
-            num_images_per_prompt,
-            prompt_embeds=prompt_embeds,
-        )
-
-        # 3.5 encode image
-        image = self.image_processor.preprocess(image)
-
-        if isinstance(controlnet, ControlNetModel):
-            control_image = self.prepare_control_image(
-                image=control_image,
-                width=width,
-                height=height,
-                batch_size=batch_size * num_images_per_prompt,
-                num_images_per_prompt=num_images_per_prompt,
-                device=device,
-                dtype=controlnet.dtype,
-                guess_mode=guess_mode,
-            )
-        elif isinstance(controlnet, MultiControlNetModel):
-            control_images = []
-
-            for control_image_ in control_image:
-                control_image_ = self.prepare_control_image(
-                    image=control_image_,
-                    width=width,
-                    height=height,
-                    batch_size=batch_size * num_images_per_prompt,
-                    num_images_per_prompt=num_images_per_prompt,
-                    device=device,
-                    dtype=controlnet.dtype,
-                    do_classifier_free_guidance=do_classifier_free_guidance,
-                    guess_mode=guess_mode,
-                )
-
-                control_images.append(control_image_)
-
-            control_image = control_images
-        else:
-            assert False
-
-        # 4. Prepare timesteps
-        self.scheduler.set_timesteps(strength, num_inference_steps, lcm_origin_steps)
-        # timesteps = self.scheduler.timesteps
-        # timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, 1.0, device)
-        timesteps = self.scheduler.timesteps
-        latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)
-
-        print("timesteps: ", timesteps)
-
-        # 5. Prepare latent variable
-        num_channels_latents = self.unet.config.in_channels
-        latents = self.prepare_latents(
-            image,
-            latent_timestep,
-            batch_size * num_images_per_prompt,
-            num_channels_latents,
-            height,
-            width,
-            prompt_embeds.dtype,
-            device,
-            latents,
-        )
-        bs = batch_size * num_images_per_prompt
-
-        # 6. Get Guidance Scale Embedding
-        w = torch.tensor(guidance_scale).repeat(bs)
-        w_embedding = self.get_w_embedding(w, embedding_dim=256).to(
-            device=device, dtype=latents.dtype
-        )
-        controlnet_keep = []
-        for i in range(len(timesteps)):
-            keeps = [
-                1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e)
-                for s, e in zip(control_guidance_start, control_guidance_end)
-            ]
-            controlnet_keep.append(
-                keeps[0] if isinstance(controlnet, ControlNetModel) else keeps
-            )
-        # 7. LCM MultiStep Sampling Loop:
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
-                ts = torch.full((bs,), t, device=device, dtype=torch.long)
-                latents = latents.to(prompt_embeds.dtype)
-                if guess_mode:
-                    # Infer ControlNet only for the conditional batch.
-                    control_model_input = latents
-                    control_model_input = self.scheduler.scale_model_input(
-                        control_model_input, ts
-                    )
-                    controlnet_prompt_embeds = prompt_embeds
-                else:
-                    control_model_input = latents
-                    controlnet_prompt_embeds = prompt_embeds
-                if isinstance(controlnet_keep[i], list):
-                    cond_scale = [
-                        c * s
-                        for c, s in zip(
-                            controlnet_conditioning_scale, controlnet_keep[i]
-                        )
-                    ]
-                else:
-                    controlnet_cond_scale = controlnet_conditioning_scale
-                    if isinstance(controlnet_cond_scale, list):
-                        controlnet_cond_scale = controlnet_cond_scale[0]
-                    cond_scale = controlnet_cond_scale * controlnet_keep[i]
-
-                down_block_res_samples, mid_block_res_sample = self.controlnet(
-                    control_model_input,
-                    ts,
-                    encoder_hidden_states=controlnet_prompt_embeds,
-                    controlnet_cond=control_image,
-                    conditioning_scale=cond_scale,
-                    guess_mode=guess_mode,
-                    return_dict=False,
-                )
-                # model prediction (v-prediction, eps, x)
-                model_pred = self.unet(
-                    latents,
-                    ts,
-                    timestep_cond=w_embedding,
-                    encoder_hidden_states=prompt_embeds,
-                    cross_attention_kwargs=cross_attention_kwargs,
-                    down_block_additional_residuals=down_block_res_samples,
-                    mid_block_additional_residual=mid_block_res_sample,
-                    return_dict=False,
-                )[0]
-
-                # compute the previous noisy sample x_t -> x_t-1
-                latents, denoised = self.scheduler.step(
-                    model_pred, i, t, latents, return_dict=False
-                )
-
-                # # call the callback, if provided
-                # if i == len(timesteps) - 1:
-                progress_bar.update()
-
-        denoised = denoised.to(prompt_embeds.dtype)
-        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
-            self.unet.to("cpu")
-            self.controlnet.to("cpu")
-            torch.cuda.empty_cache()
-        if not output_type == "latent":
-            image = self.vae.decode(
-                denoised / self.vae.config.scaling_factor, return_dict=False
-            )[0]
-            image, has_nsfw_concept = self.run_safety_checker(
-                image, device, prompt_embeds.dtype
-            )
-        else:
-            image = denoised
-            has_nsfw_concept = None
-
-        if has_nsfw_concept is None:
-            do_denormalize = [True] * image.shape[0]
-        else:
-            do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]
-
-        image = self.image_processor.postprocess(
-            image, output_type=output_type, do_denormalize=do_denormalize
-        )
-
-        if not return_dict:
-            return (image, has_nsfw_concept)
-
-        return StableDiffusionPipelineOutput(
-            images=image, nsfw_content_detected=has_nsfw_concept
-        )
-
-
-@dataclass
-# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->DDIM
-class LCMSchedulerOutput(BaseOutput):
-    """
-    Output class for the scheduler's `step` function output.
-    Args:
-        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
-            Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
-            denoising loop.
-        pred_original_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
-            The predicted denoised sample `(x_{0})` based on the model output from the current timestep.
-            `pred_original_sample` can be used to preview progress or for guidance.
-    """
-
-    prev_sample: torch.FloatTensor
-    denoised: Optional[torch.FloatTensor] = None
-
-
-# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar
-def betas_for_alpha_bar(
-    num_diffusion_timesteps,
-    max_beta=0.999,
-    alpha_transform_type="cosine",
-):
-    """
-    Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of
-    (1-beta) over time from t = [0,1].
-    Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up
-    to that part of the diffusion process.
-    Args:
-        num_diffusion_timesteps (`int`): the number of betas to produce.
-        max_beta (`float`): the maximum beta to use; use values lower than 1 to
-                     prevent singularities.
-        alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar.
-                     Choose from `cosine` or `exp`
-    Returns:
-        betas (`np.ndarray`): the betas used by the scheduler to step the model outputs
-    """
-    if alpha_transform_type == "cosine":
-
-        def alpha_bar_fn(t):
-            return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2
-
-    elif alpha_transform_type == "exp":
-
-        def alpha_bar_fn(t):
-            return math.exp(t * -12.0)
-
-    else:
-        raise ValueError(f"Unsupported alpha_tranform_type: {alpha_transform_type}")
-
-    betas = []
-    for i in range(num_diffusion_timesteps):
-        t1 = i / num_diffusion_timesteps
-        t2 = (i + 1) / num_diffusion_timesteps
-        betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta))
-    return torch.tensor(betas, dtype=torch.float32)
-
-
-def rescale_zero_terminal_snr(betas):
-    """
-    Rescales betas to have zero terminal SNR Based on https://arxiv.org/pdf/2305.08891.pdf (Algorithm 1)
-    Args:
-        betas (`torch.FloatTensor`):
-            the betas that the scheduler is being initialized with.
-    Returns:
-        `torch.FloatTensor`: rescaled betas with zero terminal SNR
-    """
-    # Convert betas to alphas_bar_sqrt
-    alphas = 1.0 - betas
-    alphas_cumprod = torch.cumprod(alphas, dim=0)
-    alphas_bar_sqrt = alphas_cumprod.sqrt()
-
-    # Store old values.
-    alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone()
-    alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone()
-
-    # Shift so the last timestep is zero.
-    alphas_bar_sqrt -= alphas_bar_sqrt_T
-
-    # Scale so the first timestep is back to the old value.
-    alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T)
-
-    # Convert alphas_bar_sqrt to betas
-    alphas_bar = alphas_bar_sqrt**2  # Revert sqrt
-    alphas = alphas_bar[1:] / alphas_bar[:-1]  # Revert cumprod
-    alphas = torch.cat([alphas_bar[0:1], alphas])
-    betas = 1 - alphas
-
-    return betas
-
-
-class LCMScheduler_X(SchedulerMixin, ConfigMixin):
-    """
-    `LCMScheduler` extends the denoising procedure introduced in denoising diffusion probabilistic models (DDPMs) with
-    non-Markovian guidance.
-    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
-    methods the library implements for all schedulers such as loading and saving.
-    Args:
-        num_train_timesteps (`int`, defaults to 1000):
-            The number of diffusion steps to train the model.
-        beta_start (`float`, defaults to 0.0001):
-            The starting `beta` value of inference.
-        beta_end (`float`, defaults to 0.02):
-            The final `beta` value.
-        beta_schedule (`str`, defaults to `"linear"`):
-            The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from
-            `linear`, `scaled_linear`, or `squaredcos_cap_v2`.
-        trained_betas (`np.ndarray`, *optional*):
-            Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`.
-        clip_sample (`bool`, defaults to `True`):
-            Clip the predicted sample for numerical stability.
-        clip_sample_range (`float`, defaults to 1.0):
-            The maximum magnitude for sample clipping. Valid only when `clip_sample=True`.
-        set_alpha_to_one (`bool`, defaults to `True`):
-            Each diffusion step uses the alphas product value at that step and at the previous one. For the final step
-            there is no previous alpha. When this option is `True` the previous alpha product is fixed to `1`,
-            otherwise it uses the alpha value at step 0.
-        steps_offset (`int`, defaults to 0):
-            An offset added to the inference steps. You can use a combination of `offset=1` and
-            `set_alpha_to_one=False` to make the last step use step 0 for the previous alpha product like in Stable
-            Diffusion.
-        prediction_type (`str`, defaults to `epsilon`, *optional*):
-            Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process),
-            `sample` (directly predicts the noisy sample`) or `v_prediction` (see section 2.4 of [Imagen
-            Video](https://imagen.research.google/video/paper.pdf) paper).
-        thresholding (`bool`, defaults to `False`):
-            Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such
-            as Stable Diffusion.
-        dynamic_thresholding_ratio (`float`, defaults to 0.995):
-            The ratio for the dynamic thresholding method. Valid only when `thresholding=True`.
-        sample_max_value (`float`, defaults to 1.0):
-            The threshold value for dynamic thresholding. Valid only when `thresholding=True`.
-        timestep_spacing (`str`, defaults to `"leading"`):
-            The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
-            Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
-        rescale_betas_zero_snr (`bool`, defaults to `False`):
-            Whether to rescale the betas to have zero terminal SNR. This enables the model to generate very bright and
-            dark samples instead of limiting it to samples with medium brightness. Loosely related to
-            [`--offset_noise`](https://github.com/huggingface/diffusers/blob/74fd735eb073eb1d774b1ab4154a0876eb82f055/examples/dreambooth/train_dreambooth.py#L506).
-    """
-
-    # _compatibles = [e.name for e in KarrasDiffusionSchedulers]
-    order = 1
-
-    @register_to_config
-    def __init__(
-        self,
-        num_train_timesteps: int = 1000,
-        beta_start: float = 0.0001,
-        beta_end: float = 0.02,
-        beta_schedule: str = "linear",
-        trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
-        clip_sample: bool = True,
-        set_alpha_to_one: bool = True,
-        steps_offset: int = 0,
-        prediction_type: str = "epsilon",
-        thresholding: bool = False,
-        dynamic_thresholding_ratio: float = 0.995,
-        clip_sample_range: float = 1.0,
-        sample_max_value: float = 1.0,
-        timestep_spacing: str = "leading",
-        rescale_betas_zero_snr: bool = False,
-    ):
-        if trained_betas is not None:
-            self.betas = torch.tensor(trained_betas, dtype=torch.float32)
-        elif beta_schedule == "linear":
-            self.betas = torch.linspace(
-                beta_start, beta_end, num_train_timesteps, dtype=torch.float32
-            )
-        elif beta_schedule == "scaled_linear":
-            # this schedule is very specific to the latent diffusion model.
-            self.betas = (
-                torch.linspace(
-                    beta_start**0.5,
-                    beta_end**0.5,
-                    num_train_timesteps,
-                    dtype=torch.float32,
-                )
-                ** 2
-            )
-        elif beta_schedule == "squaredcos_cap_v2":
-            # Glide cosine schedule
-            self.betas = betas_for_alpha_bar(num_train_timesteps)
-        else:
-            raise NotImplementedError(
-                f"{beta_schedule} does is not implemented for {self.__class__}"
-            )
-
-        # Rescale for zero SNR
-        if rescale_betas_zero_snr:
-            self.betas = rescale_zero_terminal_snr(self.betas)
-
-        self.alphas = 1.0 - self.betas
-        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
-
-        # At every step in ddim, we are looking into the previous alphas_cumprod
-        # For the final step, there is no previous alphas_cumprod because we are already at 0
-        # `set_alpha_to_one` decides whether we set this parameter simply to one or
-        # whether we use the final alpha of the "non-previous" one.
-        self.final_alpha_cumprod = (
-            torch.tensor(1.0) if set_alpha_to_one else self.alphas_cumprod[0]
-        )
-
-        # standard deviation of the initial noise distribution
-        self.init_noise_sigma = 1.0
-
-        # setable values
-        self.num_inference_steps = None
-        self.timesteps = torch.from_numpy(
-            np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64)
-        )
-
-    def scale_model_input(
-        self, sample: torch.FloatTensor, timestep: Optional[int] = None
-    ) -> torch.FloatTensor:
-        """
-        Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
-        current timestep.
-        Args:
-            sample (`torch.FloatTensor`):
-                The input sample.
-            timestep (`int`, *optional*):
-                The current timestep in the diffusion chain.
-        Returns:
-            `torch.FloatTensor`:
-                A scaled input sample.
-        """
-        return sample
-
-    def _get_variance(self, timestep, prev_timestep):
-        alpha_prod_t = self.alphas_cumprod[timestep]
-        alpha_prod_t_prev = (
-            self.alphas_cumprod[prev_timestep]
-            if prev_timestep >= 0
-            else self.final_alpha_cumprod
-        )
-        beta_prod_t = 1 - alpha_prod_t
-        beta_prod_t_prev = 1 - alpha_prod_t_prev
-
-        variance = (beta_prod_t_prev / beta_prod_t) * (
-            1 - alpha_prod_t / alpha_prod_t_prev
-        )
-
-        return variance
-
-    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample
-    def _threshold_sample(self, sample: torch.FloatTensor) -> torch.FloatTensor:
-        """
-        "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the
-        prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by
-        s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing
-        pixels from saturation at each step. We find that dynamic thresholding results in significantly better
-        photorealism as well as better image-text alignment, especially when using very large guidance weights."
-        https://arxiv.org/abs/2205.11487
-        """
-        dtype = sample.dtype
-        batch_size, channels, height, width = sample.shape
-
-        if dtype not in (torch.float32, torch.float64):
-            sample = (
-                sample.float()
-            )  # upcast for quantile calculation, and clamp not implemented for cpu half
-
-        # Flatten sample for doing quantile calculation along each image
-        sample = sample.reshape(batch_size, channels * height * width)
-
-        abs_sample = sample.abs()  # "a certain percentile absolute pixel value"
-
-        s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1)
-        s = torch.clamp(
-            s, min=1, max=self.config.sample_max_value
-        )  # When clamped to min=1, equivalent to standard clipping to [-1, 1]
-
-        s = s.unsqueeze(1)  # (batch_size, 1) because clamp will broadcast along dim=0
-        sample = (
-            torch.clamp(sample, -s, s) / s
-        )  # "we threshold xt0 to the range [-s, s] and then divide by s"
-
-        sample = sample.reshape(batch_size, channels, height, width)
-        sample = sample.to(dtype)
-
-        return sample
-
-    def set_timesteps(
-        self,
-        stength,
-        num_inference_steps: int,
-        lcm_origin_steps: int,
-        device: Union[str, torch.device] = None,
-    ):
-        """
-        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
-        Args:
-            num_inference_steps (`int`):
-                The number of diffusion steps used when generating samples with a pre-trained model.
-        """
-
-        if num_inference_steps > self.config.num_train_timesteps:
-            raise ValueError(
-                f"`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:"
-                f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle"
-                f" maximal {self.config.num_train_timesteps} timesteps."
-            )
-
-        self.num_inference_steps = num_inference_steps
-
-        # LCM Timesteps Setting:  # Linear Spacing
-        c = self.config.num_train_timesteps // lcm_origin_steps
-        lcm_origin_timesteps = (
-            np.asarray(list(range(1, int(lcm_origin_steps * stength) + 1))) * c - 1
-        )  # LCM Training  Steps Schedule
-        skipping_step = max(len(lcm_origin_timesteps) // num_inference_steps, 1)
-        timesteps = lcm_origin_timesteps[::-skipping_step][
-            :num_inference_steps
-        ]  # LCM Inference Steps Schedule
-
-        self.timesteps = torch.from_numpy(timesteps.copy()).to(device)
-
-    def get_scalings_for_boundary_condition_discrete(self, t):
-        self.sigma_data = 0.5  # Default: 0.5
-
-        # By dividing 0.1: This is almost a delta function at t=0.
-        c_skip = self.sigma_data**2 / ((t / 0.1) ** 2 + self.sigma_data**2)
-        c_out = (t / 0.1) / ((t / 0.1) ** 2 + self.sigma_data**2) ** 0.5
-        return c_skip, c_out
-
-    def step(
-        self,
-        model_output: torch.FloatTensor,
-        timeindex: int,
-        timestep: int,
-        sample: torch.FloatTensor,
-        eta: float = 0.0,
-        use_clipped_model_output: bool = False,
-        generator=None,
-        variance_noise: Optional[torch.FloatTensor] = None,
-        return_dict: bool = True,
-    ) -> Union[LCMSchedulerOutput, Tuple]:
-        """
-        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
-        process from the learned model outputs (most often the predicted noise).
-        Args:
-            model_output (`torch.FloatTensor`):
-                The direct output from learned diffusion model.
-            timestep (`float`):
-                The current discrete timestep in the diffusion chain.
-            sample (`torch.FloatTensor`):
-                A current instance of a sample created by the diffusion process.
-            eta (`float`):
-                The weight of noise for added noise in diffusion step.
-            use_clipped_model_output (`bool`, defaults to `False`):
-                If `True`, computes "corrected" `model_output` from the clipped predicted original sample. Necessary
-                because predicted original sample is clipped to [-1, 1] when `self.config.clip_sample` is `True`. If no
-                clipping has happened, "corrected" `model_output` would coincide with the one provided as input and
-                `use_clipped_model_output` has no effect.
-            generator (`torch.Generator`, *optional*):
-                A random number generator.
-            variance_noise (`torch.FloatTensor`):
-                Alternative to generating noise with `generator` by directly providing the noise for the variance
-                itself. Useful for methods such as [`CycleDiffusion`].
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] or `tuple`.
-        Returns:
-            [`~schedulers.scheduling_utils.LCMSchedulerOutput`] or `tuple`:
-                If return_dict is `True`, [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] is returned, otherwise a
-                tuple is returned where the first element is the sample tensor.
-        """
-        if self.num_inference_steps is None:
-            raise ValueError(
-                "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
-            )
-
-        # 1. get previous step value
-        prev_timeindex = timeindex + 1
-        if prev_timeindex < len(self.timesteps):
-            prev_timestep = self.timesteps[prev_timeindex]
-        else:
-            prev_timestep = timestep
-
-        # 2. compute alphas, betas
-        alpha_prod_t = self.alphas_cumprod[timestep]
-        alpha_prod_t_prev = (
-            self.alphas_cumprod[prev_timestep]
-            if prev_timestep >= 0
-            else self.final_alpha_cumprod
-        )
-
-        beta_prod_t = 1 - alpha_prod_t
-        beta_prod_t_prev = 1 - alpha_prod_t_prev
-
-        # 3. Get scalings for boundary conditions
-        c_skip, c_out = self.get_scalings_for_boundary_condition_discrete(timestep)
-
-        # 4. Different Parameterization:
-        parameterization = self.config.prediction_type
-
-        if parameterization == "epsilon":  # noise-prediction
-            pred_x0 = (sample - beta_prod_t.sqrt() * model_output) / alpha_prod_t.sqrt()
-
-        elif parameterization == "sample":  # x-prediction
-            pred_x0 = model_output
-
-        elif parameterization == "v_prediction":  # v-prediction
-            pred_x0 = alpha_prod_t.sqrt() * sample - beta_prod_t.sqrt() * model_output
-
-        # 4. Denoise model output using boundary conditions
-        denoised = c_out * pred_x0 + c_skip * sample
-
-        # 5. Sample z ~ N(0, I), For MultiStep Inference
-        # Noise is not used for one-step sampling.
-        if len(self.timesteps) > 1:
-            noise = torch.randn(model_output.shape).to(model_output.device)
-            prev_sample = (
-                alpha_prod_t_prev.sqrt() * denoised + beta_prod_t_prev.sqrt() * noise
-            )
-        else:
-            prev_sample = denoised
-
-        if not return_dict:
-            return (prev_sample, denoised)
-
-        return LCMSchedulerOutput(prev_sample=prev_sample, denoised=denoised)
-
-    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noise
-    def add_noise(
-        self,
-        original_samples: torch.FloatTensor,
-        noise: torch.FloatTensor,
-        timesteps: torch.IntTensor,
-    ) -> torch.FloatTensor:
-        # Make sure alphas_cumprod and timestep have same device and dtype as original_samples
-        alphas_cumprod = self.alphas_cumprod.to(
-            device=original_samples.device, dtype=original_samples.dtype
-        )
-        timesteps = timesteps.to(original_samples.device)
-
-        sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
-        sqrt_alpha_prod = sqrt_alpha_prod.flatten()
-        while len(sqrt_alpha_prod.shape) < len(original_samples.shape):
-            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
-
-        sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
-        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
-        while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape):
-            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
-
-        noisy_samples = (
-            sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
-        )
-        return noisy_samples
-
-    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.get_velocity
-    def get_velocity(
-        self,
-        sample: torch.FloatTensor,
-        noise: torch.FloatTensor,
-        timesteps: torch.IntTensor,
-    ) -> torch.FloatTensor:
-        # Make sure alphas_cumprod and timestep have same device and dtype as sample
-        alphas_cumprod = self.alphas_cumprod.to(
-            device=sample.device, dtype=sample.dtype
-        )
-        timesteps = timesteps.to(sample.device)
-
-        sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
-        sqrt_alpha_prod = sqrt_alpha_prod.flatten()
-        while len(sqrt_alpha_prod.shape) < len(sample.shape):
-            sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
-
-        sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
-        sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
-        while len(sqrt_one_minus_alpha_prod.shape) < len(sample.shape):
-            sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
-
-        velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample
-        return velocity
-
-    def __len__(self):
-        return self.config.num_train_timesteps

pipelines/controlnet.py

@@ -1,8 +1,11 @@
-from diffusers import DiffusionPipeline, AutoencoderTiny
-from latent_consistency_controlnet import LatentConsistencyModelPipeline_controlnet
-
+from diffusers import (
+    StableDiffusionControlNetImg2ImgPipeline,
+    AutoencoderTiny,
+    ControlNetModel,
+)
 from compel import Compel
 import torch
+from pipelines.utils.canny_gpu import SobelOperator
 
 try:
     import intel_extension_for_pytorch as ipex  # type: ignore
@@ -11,80 +14,202 @@ except:
 
 import psutil
 from config import Args
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 from PIL import Image
-from typing import Callable
 
 base_model = "SimianLuo/LCM_Dreamshaper_v7"
-WIDTH = 512
-HEIGHT = 512
+taesd_model = "madebyollin/taesd"
+controlnet_model = "lllyasviel/control_v11p_sd15_canny"
+
+default_prompt = "Portrait of The Terminator with , glare pose, detailed, intricate, full of colour, cinematic lighting, trending on artstation, 8k, hyperrealistic, focused, extreme details, unreal engine 5 cinematic, masterpiece"
 
 
 class Pipeline:
+    class Info(BaseModel):
+        name: str = "txt2img"
+        description: str = "Generates an image from a text prompt"
+
     class InputParams(BaseModel):
-        seed: int = 2159232
-        prompt: str
-        guidance_scale: float = 8.0
-        strength: float = 0.5
-        steps: int = 4
-        lcm_steps: int = 50
-        width: int = WIDTH
-        height: int = HEIGHT
-
-    @staticmethod
-    def create_pipeline(
-        args: Args, device: torch.device, torch_dtype: torch.dtype
-    ) -> Callable[["Pipeline.InputParams"], Image.Image]:
+        prompt: str = Field(
+            default_prompt,
+            title="Prompt",
+            field="textarea",
+            id="prompt",
+        )
+        seed: int = Field(
+            2159232, min=0, title="Seed", field="seed", hide=True, id="seed"
+        )
+        steps: int = Field(
+            4, min=2, max=15, title="Steps", field="range", hide=True, id="steps"
+        )
+        width: int = Field(
+            512, min=2, max=15, title="Width", disabled=True, hide=True, id="width"
+        )
+        height: int = Field(
+            512, min=2, max=15, title="Height", disabled=True, hide=True, id="height"
+        )
+        guidance_scale: float = Field(
+            0.2,
+            min=0,
+            max=2,
+            step=0.001,
+            title="Guidance Scale",
+            field="range",
+            hide=True,
+            id="guidance_scale",
+        )
+        strength: float = Field(
+            0.5,
+            min=0.25,
+            max=1.0,
+            step=0.001,
+            title="Strength",
+            field="range",
+            hide=True,
+            id="strength",
+        )
+        controlnet_scale: float = Field(
+            0.8,
+            min=0,
+            max=1.0,
+            step=0.001,
+            title="Controlnet Scale",
+            field="range",
+            hide=True,
+            id="controlnet_scale",
+        )
+        controlnet_start: float = Field(
+            0.0,
+            min=0,
+            max=1.0,
+            step=0.001,
+            title="Controlnet Start",
+            field="range",
+            hide=True,
+            id="controlnet_start",
+        )
+        controlnet_end: float = Field(
+            1.0,
+            min=0,
+            max=1.0,
+            step=0.001,
+            title="Controlnet End",
+            field="range",
+            hide=True,
+            id="controlnet_end",
+        )
+        canny_low_threshold: float = Field(
+            0.31,
+            min=0,
+            max=1.0,
+            step=0.001,
+            title="Canny Low Threshold",
+            field="range",
+            hide=True,
+            id="canny_low_threshold",
+        )
+        canny_high_threshold: float = Field(
+            0.125,
+            min=0,
+            max=1.0,
+            step=0.001,
+            title="Canny High Threshold",
+            field="range",
+            hide=True,
+            id="canny_high_threshold",
+        )
+        debug_canny: bool = Field(
+            False,
+            title="Debug Canny",
+            field="checkbox",
+            hide=True,
+            id="debug_canny",
+        )
+        image: bool = True
+
+    def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype):
+        controlnet_canny = ControlNetModel.from_pretrained(
+            controlnet_model, torch_dtype=torch_dtype
+        ).to(device)
         if args.safety_checker:
-            pipe = DiffusionPipeline.from_pretrained(base_model)
+            self.pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
+                base_model, controlnet=controlnet_canny
+            )
         else:
-            pipe = DiffusionPipeline.from_pretrained(base_model, safety_checker=None)
+            self.pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
+                base_model,
+                safety_checker=None,
+                controlnet=controlnet_canny,
+            )
         if args.use_taesd:
-            pipe.vae = AutoencoderTiny.from_pretrained(
-                "madebyollin/taesd", torch_dtype=torch_dtype, use_safetensors=True
+            self.pipe.vae = AutoencoderTiny.from_pretrained(
+                taesd_model, torch_dtype=torch_dtype, use_safetensors=True
             )
-
-        pipe.set_progress_bar_config(disable=True)
-        pipe.to(device=device, dtype=torch_dtype)
-        pipe.unet.to(memory_format=torch.channels_last)
+        self.canny_torch = SobelOperator(device=device)
+        self.pipe.set_progress_bar_config(disable=True)
+        self.pipe.to(device=device, dtype=torch_dtype)
+        self.pipe.unet.to(memory_format=torch.channels_last)
 
         # check if computer has less than 64GB of RAM using sys or os
         if psutil.virtual_memory().total < 64 * 1024**3:
-            pipe.enable_attention_slicing()
+            self.pipe.enable_attention_slicing()
 
         if args.torch_compile:
-            pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
-            pipe.vae = torch.compile(pipe.vae, mode="reduce-overhead", fullgraph=True)
+            self.pipe.unet = torch.compile(
+                self.pipe.unet, mode="reduce-overhead", fullgraph=True
+            )
+            self.pipe.vae = torch.compile(
+                self.pipe.vae, mode="reduce-overhead", fullgraph=True
+            )
 
-            pipe(prompt="warmup", num_inference_steps=1, guidance_scale=8.0)
+            self.pipe(
+                prompt="warmup",
+                image=[Image.new("RGB", (768, 768))],
+                control_image=[Image.new("RGB", (768, 768))],
+            )
 
-        compel_proc = Compel(
-            tokenizer=pipe.tokenizer,
-            text_encoder=pipe.text_encoder,
+        self.compel_proc = Compel(
+            tokenizer=self.pipe.tokenizer,
+            text_encoder=self.pipe.text_encoder,
             truncate_long_prompts=False,
         )
 
-        def predict(params: "Pipeline.InputParams") -> Image.Image:
-            generator = torch.manual_seed(params.seed)
-            prompt_embeds = compel_proc(params.prompt)
-            # Can be set to 1~50 steps. LCM support fast inference even <= 4 steps. Recommend: 1~8 steps.
-            results = pipe(
-                prompt_embeds=prompt_embeds,
-                generator=generator,
-                num_inference_steps=params.steps,
-                guidance_scale=params.guidance_scale,
-                width=params.width,
-                height=params.height,
-                original_inference_steps=params.lcm_steps,
-                output_type="pil",
-            )
-            nsfw_content_detected = (
-                results.nsfw_content_detected[0]
-                if "nsfw_content_detected" in results
-                else False
-            )
-            if nsfw_content_detected:
-                return None
-            return results.images[0]
+    def predict(self, params: "Pipeline.InputParams") -> Image.Image:
+        generator = torch.manual_seed(params.seed)
+        prompt_embeds = self.compel_proc(params.prompt)
+        control_image = self.canny_torch(
+            params.image, params.canny_low_threshold, params.canny_high_threshold
+        )
+
+        results = self.pipe(
+            image=params.image,
+            control_image=control_image,
+            prompt_embeds=prompt_embeds,
+            generator=generator,
+            strength=params.strength,
+            num_inference_steps=params.steps,
+            guidance_scale=params.guidance_scale,
+            width=params.width,
+            height=params.height,
+            output_type="pil",
+            controlnet_conditioning_scale=params.controlnet_scale,
+            control_guidance_start=params.controlnet_start,
+            control_guidance_end=params.controlnet_end,
+        )
+
+        nsfw_content_detected = (
+            results.nsfw_content_detected[0]
+            if "nsfw_content_detected" in results
+            else False
+        )
+        if nsfw_content_detected:
+            return None
+        result_image = results.images[0]
+        if params.debug_canny:
+            # paste control_image on top of result_image
+            w0, h0 = (200, 200)
+            control_image = control_image.resize((w0, h0))
+            w1, h1 = result_image.size
+            result_image.paste(control_image, (w1 - w0, h1 - h0))
 
-        return predict
+        return result_image

pipelines/txt2img.py

@@ -11,7 +11,6 @@ import psutil
 from config import Args
 from pydantic import BaseModel, Field
 from PIL import Image
-from typing import Callable
 
 base_model = "SimianLuo/LCM_Dreamshaper_v7"
 taesd_model = "madebyollin/taesd"
@@ -29,22 +28,19 @@ class Pipeline:
             default_prompt,
             title="Prompt",
             field="textarea",
+            id="prompt",
         )
-        seed: int = Field(2159232, min=0, title="Seed", field="seed", hide=True)
-        strength: float = Field(
-            0.5,
-            min=0,
-            max=1,
-            step=0.001,
-            title="Strength",
-            field="range",
-            hide=True,
+        seed: int = Field(
+            2159232, min=0, title="Seed", field="seed", hide=True, id="seed"
+        )
+        steps: int = Field(
+            4, min=2, max=15, title="Steps", field="range", hide=True, id="steps"
+        )
+        width: int = Field(
+            512, min=2, max=15, title="Width", disabled=True, hide=True, id="width"
         )
-
-        steps: int = Field(4, min=2, max=15, title="Steps", field="range", hide=True)
-        width: int = Field(512, min=2, max=15, title="Width", disabled=True, hide=True)
         height: int = Field(
-            512, min=2, max=15, title="Height", disabled=True, hide=True
+            512, min=2, max=15, title="Height", disabled=True, hide=True, id="height"
         )
         guidance_scale: float = Field(
             8.0,
@@ -54,6 +50,10 @@ class Pipeline:
             title="Guidance Scale",
             field="range",
             hide=True,
+            id="guidance_scale",
+        )
+        image: bool = Field(
+            True, title="Image", field="checkbox", hide=True, id="image"
         )
 
     def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype):

requirements.txt

@@ -1,4 +1,4 @@
-diffusers==0.23.0
+git+https://github.com/huggingface/diffusers@c697f524761abd2314c030221a3ad2f7791eab4e
 transformers==4.34.1
 gradio==3.50.2
 --extra-index-url https://download.pytorch.org/whl/cu121;

user_queue.py

@@ -1,18 +1,29 @@
 from typing import Dict, Union
 from uuid import UUID
-from asyncio import Queue
+import asyncio
 from PIL import Image
-from typing import Tuple, Union
-from uuid import UUID
-from asyncio import Queue
+from typing import Dict, Union
 from PIL import Image
 
+InputParams = dict
 UserId = UUID
+EventDataContent = Dict[str, InputParams]
 
-InputParams = dict
 
-QueueContent = Dict[str, Union[Image.Image, InputParams]]
+class UserDataEvent:
+    def __init__(self):
+        self.data_event = asyncio.Event()
+        self.data_content: EventDataContent = {}
+
+    def update_data(self, new_data: EventDataContent):
+        self.data_content = new_data
+        self.data_event.set()
+
+    async def wait_for_data(self) -> EventDataContent:
+        await self.data_event.wait()
+        self.data_event.clear()
+        return self.data_content
 
-UserQueueDict = Dict[UserId, Queue[QueueContent]]
 
-user_queue_map: UserQueueDict = {}
+UserDataEventMap = Dict[UserId, UserDataEvent]
+user_data_events: UserDataEventMap = {}