live params!

app-img2img.py

@@ -49,7 +49,7 @@ else:
 pipe.set_progress_bar_config(disable=True)
 pipe.to(torch_device="cuda", torch_dtype=torch.float32)
 pipe.unet.to(memory_format=torch.channels_last)
-pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+# pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
 user_queue_map = {}
 
 # for torch.compile
@@ -58,7 +58,7 @@ pipe(prompt="warmup", image=[Image.new("RGB", (512, 512))])
 def predict(input_image, prompt, guidance_scale=8.0, strength=0.5, seed=2159232):
     generator = torch.manual_seed(seed)
     # Can be set to 1~50 steps. LCM support fast inference even <= 4 steps. Recommend: 1~8 steps.
-    num_inference_steps = 4
+    num_inference_steps = 3
     results = pipe(
         prompt=prompt,
         # generator=generator,
@@ -66,7 +66,7 @@ def predict(input_image, prompt, guidance_scale=8.0, strength=0.5, seed=2159232)
         strength=strength,
         num_inference_steps=num_inference_steps,
         guidance_scale=guidance_scale,
-        lcm_origin_steps=30,
+        lcm_origin_steps=20,
         output_type="pil",
     )
     nsfw_content_detected = (
@@ -111,11 +111,8 @@ async def websocket_endpoint(websocket: WebSocket):
         await websocket.send_json(
             {"status": "success", "message": "Connected", "userId": uid}
         )
-        params = await websocket.receive_json()
-        params = InputParams(**params)
         user_queue_map[uid] = {
-            "queue": asyncio.Queue(),
-            "params": params,
+            "queue": asyncio.Queue()
         }
         await websocket.send_json(
             {"status": "start", "message": "Start Streaming", "userId": uid}
@@ -148,19 +145,16 @@ async def stream(user_id: uuid.UUID):
     try:
         user_queue = user_queue_map[uid]
         queue = user_queue["queue"]
-        params = user_queue["params"]
-        seed = params.seed
-        prompt = params.prompt
-        strength = params.strength
-        guidance_scale = params.guidance_scale
-
+        
         async def generate():
             while True:
-                input_image = await queue.get()
+                data = await queue.get()
+                input_image = data["image"]
+                params = data["params"]
                 if input_image is None:
                     continue
 
-                image = predict(input_image, prompt, guidance_scale, strength, seed)
+                image = predict(input_image, params.prompt, params.guidance_scale, params.strength, params.seed)
                 if image is None:
                     continue
                 frame_data = io.BytesIO()
@@ -190,6 +184,8 @@ async def handle_websocket_data(websocket: WebSocket, user_id: uuid.UUID):
     try:
         while True:
             data = await websocket.receive_bytes()
+            params = await websocket.receive_json()
+            params = InputParams(**params)
             pil_image = Image.open(io.BytesIO(data))
 
             while not queue.empty():
@@ -197,7 +193,10 @@ async def handle_websocket_data(websocket: WebSocket, user_id: uuid.UUID):
                     queue.get_nowait()
                 except asyncio.QueueEmpty:
                     continue
-            await queue.put(pil_image)
+            await queue.put({
+                "image": pil_image,
+                "params": params
+            })
             if TIMEOUT > 0 and time.time() - last_time > TIMEOUT:
                 await websocket.send_json(
                     {

img2img/index.html

@@ -21,10 +21,10 @@
         const queueSizeEl = document.querySelector("#queue_size");
         const errorEl = document.querySelector("#error");
 
-        function LCMLive(webcamVideo, liveImage) {
+        function LCMLive(webcamVideo, liveImage, seedEl, promptEl, guidanceEl, strengthEl) {
             let websocket;
 
-            async function start(params) {
+            async function start() {
                 return new Promise((resolve, reject) => {
                     const websocketURL = `${window.location.protocol === "https:" ? "wss" : "ws"
                         }:${window.location.host}/ws`;
@@ -46,7 +46,6 @@
                         const data = JSON.parse(event.data);
                         switch (data.status) {
                             case "success":
-                                socket.send(JSON.stringify(params));
                                 break;
                             case "start":
                                 const userId = data.userId;
@@ -71,6 +70,12 @@
                 ctx.drawImage(webcamVideo, 0, 0, canvas.width, canvas.height);
                 const blob = await canvas.convertToBlob({ type: "image/jpeg", quality: 1 });
                 websocket.send(blob);
+                websocket.send(JSON.stringify({
+                    "seed": seedEl.value,
+                    "prompt": promptEl.value,
+                    "guidance_scale": guidanceEl.value,
+                    "strength": strengthEl.value
+                }));
             }
 
             function initVideoStream(userId) {
@@ -124,15 +129,11 @@
         }
 
 
-        const lcmLive = LCMLive(videoEl, imageEl);
+        const lcmLive = LCMLive(videoEl, imageEl, seedEl, promptEl, guidanceEl, strengthEl);
         startBtn.addEventListener("click", async () => {
             try {
-                const seed = seedEl.value;
-                const prompt = promptEl.value;
-                const guidance_scale = guidanceEl.value;
-                const strength = strengthEl.value;
                 startBtn.disabled = true;
-                const res = await lcmLive.start({ seed, prompt, guidance_scale, strength });
+                const res = await lcmLive.start();
                 startBtn.disabled = false;
                 if (res.status === "timeout")
                     toggleMessage("success")
@@ -176,9 +177,6 @@
                     target="_blank" class="text-blue-500 hover:underline">Diffusers</a> with a MJPEG
                 stream server.
             </p>
-            <p class="text-sm">
-                To change settings or prompt, stop the current stream and start a new one.
-            </p>
             <p class="text-sm">
                 There are <span id="queue_size" class="font-bold">0</span> user(s) sharing the same GPU, affecting
                 real-time performance. Maximum queue size is 4. <a
@@ -218,7 +216,7 @@
                     <input type="number" id="seed" name="seed" value="299792458"
                         class="font-light border border-gray-700 text-right rounded-md p-2">
                     <button
-                        onclick="document.querySelector('#seed').value = BigInt(Math.floor(Math.random() * 2**64-1))"
+                        onclick="document.querySelector('#seed').value = Math.floor(Math.random() * Number.MAX_SAFE_INTEGER)"
                         class="bg-gray-700 hover:bg-gray-800 text-white font-normal py-1 w-[50px] rounded disabled:bg-gray-300 disabled:cursor-not-allowed text-sm">
                         Rand
                     </button>