img2img

app_init.py

@@ -36,10 +36,16 @@ def init_app(app: FastAPI, user_data: UserData, args: Args, pipeline):
         try:
             user_id = uuid.uuid4()
             print(f"New user connected: {user_id}")
+
             await user_data.create_user(user_id, websocket)
             await websocket.send_json(
                 {"status": "connected", "message": "Connected", "userId": str(user_id)}
             )
+            await websocket.send_json(
+                {
+                    "status": "send_frame",
+                }
+            )
             await handle_websocket_data(user_id, websocket)
         except WebSocketDisconnect as e:
             logging.error(f"WebSocket Error: {e}, {user_id}")
@@ -48,6 +54,46 @@ def init_app(app: FastAPI, user_data: UserData, args: Args, pipeline):
             print(f"User disconnected: {user_id}")
             user_data.delete_user(user_id)
 
+    async def handle_websocket_data(user_id: uuid.UUID, websocket: WebSocket):
+        if not user_data.check_user(user_id):
+            return HTTPException(status_code=404, detail="User not found")
+        last_time = time.time()
+        try:
+            while True:
+                data = await websocket.receive_json()
+                if data["status"] != "next_frame":
+                    asyncio.sleep(1.0 / 24)
+                    continue
+
+                params = await websocket.receive_json()
+                params = pipeline.InputParams(**params)
+                info = pipeline.Info()
+                params = SimpleNamespace(**params.dict())
+                if info.input_mode == "image":
+                    image_data = await websocket.receive_bytes()
+                    params.image = bytes_to_pil(image_data)
+                await user_data.update_data(user_id, params)
+                await websocket.send_json(
+                    {
+                        "status": "wait",
+                    }
+                )
+                if args.timeout > 0 and time.time() - last_time > args.timeout:
+                    await websocket.send_json(
+                        {
+                            "status": "timeout",
+                            "message": "Your session has ended",
+                            "userId": user_id,
+                        }
+                    )
+                    await websocket.close()
+                    return
+                await asyncio.sleep(1.0 / 24)
+
+        except Exception as e:
+            logging.error(f"Error: {e}")
+            traceback.print_exc()
+
     @app.get("/queue_size")
     async def get_queue_size():
         queue_size = user_data.get_user_count()
@@ -59,10 +105,20 @@ def init_app(app: FastAPI, user_data: UserData, args: Args, pipeline):
             print(f"New stream request: {user_id}")
 
             async def generate():
+                websocket = user_data.get_websocket(user_id)
+                last_params = SimpleNamespace()
                 while True:
                     params = await user_data.get_latest_data(user_id)
-                    if not params:
+                    if not vars(params) or params.__dict__ == last_params.__dict__:
+                        await websocket.send_json(
+                            {
+                                "status": "send_frame",
+                            }
+                        )
+                        await asyncio.sleep(0.1)
                         continue
+
+                    last_params = params
                     image = pipeline.predict(params)
                     if image is None:
                         continue
@@ -71,6 +127,11 @@ def init_app(app: FastAPI, user_data: UserData, args: Args, pipeline):
                     # https://bugs.chromium.org/p/chromium/issues/detail?id=1250396
                     if not is_firefox(request.headers["user-agent"]):
                         yield frame
+                    await websocket.send_json(
+                        {
+                            "status": "send_frame",
+                        }
+                    )
 
             return StreamingResponse(
                 generate(),
@@ -82,37 +143,6 @@ def init_app(app: FastAPI, user_data: UserData, args: Args, pipeline):
             traceback.print_exc()
             return HTTPException(status_code=404, detail="User not found")
 
-    async def handle_websocket_data(user_id: uuid.UUID, websocket: WebSocket):
-        if not user_data.check_user(user_id):
-            return HTTPException(status_code=404, detail="User not found")
-        last_time = time.time()
-        try:
-            while True:
-                params = await websocket.receive_json()
-                params = pipeline.InputParams(**params)
-                info = pipeline.Info()
-                params = SimpleNamespace(**params.dict())
-                if info.input_mode == "image":
-                    image_data = await websocket.receive_bytes()
-                    params.image = bytes_to_pil(image_data)
-
-                await user_data.update_data(user_id, params)
-                if args.timeout > 0 and time.time() - last_time > args.timeout:
-                    await websocket.send_json(
-                        {
-                            "status": "timeout",
-                            "message": "Your session has ended",
-                            "userId": user_id,
-                        }
-                    )
-                    await websocket.close()
-                    return
-                await asyncio.sleep(1.0 / 24)
-
-        except Exception as e:
-            logging.error(f"Error: {e}")
-            traceback.print_exc()
-
     # route to setup frontend
     @app.get("/settings")
     async def settings():

frontend/src/lib/components/VideoInput.svelte

@@ -12,7 +12,6 @@
   let videoFrameCallbackId: number;
   const WIDTH = 512;
   const HEIGHT = 512;
-  const THROTTLE_FPS = 6;
 
   onDestroy(() => {
     if (videoFrameCallbackId) videoEl.cancelVideoFrameCallback(videoFrameCallbackId);
@@ -22,13 +21,9 @@
     videoEl.srcObject = $mediaStream;
   }
 
-  let last_millis = 0;
   async function onFrameChange(now: DOMHighResTimeStamp, metadata: VideoFrameCallbackMetadata) {
-    if (now - last_millis > 1000 / THROTTLE_FPS) {
-      const blob = await grapBlobImg();
-      onFrameChangeStore.set({ now, metadata, blob });
-      last_millis = now;
-    }
+    const blob = await grapBlobImg();
+    onFrameChangeStore.set({ blob });
     videoFrameCallbackId = videoEl.requestVideoFrameCallback(onFrameChange);
   }
 

frontend/src/lib/lcmLive.ts

@@ -6,6 +6,7 @@ export enum LCMLiveStatus {
     CONNECTED = "connected",
     DISCONNECTED = "disconnected",
     WAIT = "wait",
+    SEND_FRAME = "send_frame",
 }
 
 const initStatus: LCMLiveStatus = LCMLiveStatus.DISCONNECTED;
@@ -15,7 +16,7 @@ export const streamId = writable<string | null>(null);
 
 let websocket: WebSocket | null = null;
 export const lcmLiveActions = {
-    async start() {
+    async start(getSreamdata: () => any[]) {
         return new Promise((resolve, reject) => {
 
             try {
@@ -43,6 +44,17 @@ export const lcmLiveActions = {
                             streamId.set(userId);
                             resolve(userId);
                             break;
+                        case "send_frame":
+                            lcmLiveStatus.set(LCMLiveStatus.SEND_FRAME);
+                            const streamData = getSreamdata();
+                            websocket?.send(JSON.stringify({ status: "next_frame" }));
+                            for (const d of streamData) {
+                                this.send(d);
+                            }
+                            break;
+                        case "wait":
+                            lcmLiveStatus.set(LCMLiveStatus.WAIT);
+                            break;
                         case "timeout":
                             console.log("timeout");
                             lcmLiveStatus.set(LCMLiveStatus.DISCONNECTED);
@@ -60,7 +72,6 @@ export const lcmLiveActions = {
                 console.error(err);
                 lcmLiveStatus.set(LCMLiveStatus.DISCONNECTED);
                 streamId.set(null);
-
                 reject(err);
             }
         });

frontend/src/lib/mediaStream.ts

@@ -5,7 +5,7 @@ export enum MediaStreamStatusEnum {
     CONNECTED = "connected",
     DISCONNECTED = "disconnected",
 }
-export const onFrameChangeStore: Writable<{ now: Number, metadata: VideoFrameCallbackMetadata, blob: Blob }> = writable();
+export const onFrameChangeStore: Writable<{ blob: Blob }> = writable({ blob: new Blob() });
 
 export const mediaDevices = writable<MediaDeviceInfo[]>([]);
 export const mediaStreamStatus = writable(MediaStreamStatusEnum.INIT);

frontend/src/routes/+page.svelte

@@ -35,25 +35,18 @@
     pipelineParams = pipelineParams.filter((e) => e?.disabled !== true);
   }
 
-  // send Webcam stream to LCM if image mode
-  $: {
-    if (
-      isImageMode &&
-      $lcmLiveStatus === LCMLiveStatus.CONNECTED &&
-      $mediaStreamStatus === MediaStreamStatusEnum.CONNECTED
-    ) {
-      lcmLiveActions.send(getPipelineValues());
-      lcmLiveActions.send($onFrameChangeStore.blob);
-    }
-  }
-  $: {
-    if (!isImageMode && $lcmLiveStatus === LCMLiveStatus.CONNECTED) {
-      lcmLiveActions.send($deboucedPipelineValues);
+  function getSreamdata() {
+    if (isImageMode) {
+      return [getPipelineValues(), $onFrameChangeStore?.blob];
+    } else {
+      return [$deboucedPipelineValues];
     }
   }
 
   $: isLCMRunning = $lcmLiveStatus !== LCMLiveStatus.DISCONNECTED;
-
+  $: {
+    console.log('lcmLiveStatus', $lcmLiveStatus);
+  }
   let disabled = false;
   async function toggleLcmLive() {
     if (!isLCMRunning) {
@@ -62,7 +55,7 @@
         await mediaStreamActions.start();
       }
       disabled = true;
-      await lcmLiveActions.start();
+      await lcmLiveActions.start(getSreamdata);
       disabled = false;
     } else {
       if (isImageMode) {

pipelines/img2img.py

@@ -0,0 +1,136 @@
+from diffusers import (
+    AutoPipelineForImage2Image,
+    AutoencoderTiny,
+)
+from compel import Compel
+import torch
+
+try:
+    import intel_extension_for_pytorch as ipex  # type: ignore
+except:
+    pass
+
+import psutil
+from config import Args
+from pydantic import BaseModel, Field
+from PIL import Image
+
+base_model = "SimianLuo/LCM_Dreamshaper_v7"
+taesd_model = "madebyollin/taesd"
+
+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 = "img2img"
+        title: str = "Image-to-Image LCM"
+        description: str = "Generates an image from a text prompt"
+        input_mode: str = "image"
+
+    class InputParams(BaseModel):
+        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=20,
+            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",
+        )
+
+    def __init__(self, args: Args, device: torch.device, torch_dtype: torch.dtype):
+        if args.safety_checker:
+            self.pipe = AutoPipelineForImage2Image.from_pretrained(base_model)
+        else:
+            self.pipe = AutoPipelineForImage2Image.from_pretrained(
+                base_model,
+                safety_checker=None,
+            )
+        if args.use_taesd:
+            self.pipe.vae = AutoencoderTiny.from_pretrained(
+                taesd_model, torch_dtype=torch_dtype, use_safetensors=True
+            )
+
+        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:
+            self.pipe.enable_attention_slicing()
+
+        if args.torch_compile:
+            print("Running torch compile")
+            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
+            )
+
+            self.pipe(
+                prompt="warmup",
+                image=[Image.new("RGB", (768, 768))],
+            )
+
+        self.compel_proc = Compel(
+            tokenizer=self.pipe.tokenizer,
+            text_encoder=self.pipe.text_encoder,
+            truncate_long_prompts=False,
+        )
+
+    def predict(self, params: "Pipeline.InputParams") -> Image.Image:
+        generator = torch.manual_seed(params.seed)
+        prompt_embeds = self.compel_proc(params.prompt)
+        results = self.pipe(
+            image=params.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",
+        )
+
+        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]
+
+        return result_image

user_queue.py

@@ -36,6 +36,7 @@ class UserData:
     async def get_latest_data(self, user_id: UUID) -> SimpleNamespace:
         user_session = self.data_content[user_id]
         queue = user_session["queue"]
+
         try:
             return await queue.get()
         except asyncio.QueueEmpty:
@@ -55,5 +56,8 @@ class UserData:
     def get_user_count(self) -> int:
         return len(self.data_content)
 
+    def get_websocket(self, user_id: UUID) -> WebSocket:
+        return self.data_content[user_id]["websocket"]
+
 
 user_data = UserData()