test

app-txt2imglora.py

@@ -0,0 +1,258 @@
+import asyncio
+import json
+import logging
+import traceback
+from pydantic import BaseModel
+
+from fastapi import FastAPI, WebSocket, HTTPException, WebSocketDisconnect
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import (
+    StreamingResponse,
+    JSONResponse,
+    HTMLResponse,
+    FileResponse,
+)
+
+from diffusers import DiffusionPipeline, LCMScheduler, AutoencoderTiny
+from compel import Compel, ReturnedEmbeddingsType
+import torch
+
+try:
+    import intel_extension_for_pytorch as ipex
+except:
+    pass
+from PIL import Image
+import numpy as np
+import gradio as gr
+import io
+import uuid
+import os
+import time
+import psutil
+
+
+MAX_QUEUE_SIZE = int(os.environ.get("MAX_QUEUE_SIZE", 0))
+TIMEOUT = float(os.environ.get("TIMEOUT", 0))
+SAFETY_CHECKER = os.environ.get("SAFETY_CHECKER", None)
+TORCH_COMPILE = os.environ.get("TORCH_COMPILE", None)
+
+WIDTH = 768
+HEIGHT = 768
+# disable tiny autoencoder for better quality speed tradeoff
+USE_TINY_AUTOENCODER = False
+
+# check if MPS is available OSX only M1/M2/M3 chips
+mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
+xpu_available = hasattr(torch, "xpu") and torch.xpu.is_available()
+device = torch.device(
+    "cuda" if torch.cuda.is_available() else "xpu" if xpu_available else "cpu"
+)
+torch_device = device
+# change to torch.float16 to save GPU memory
+torch_dtype = torch.float16
+
+print(f"TIMEOUT: {TIMEOUT}")
+print(f"SAFETY_CHECKER: {SAFETY_CHECKER}")
+print(f"MAX_QUEUE_SIZE: {MAX_QUEUE_SIZE}")
+print(f"device: {device}")
+
+if mps_available:
+    device = torch.device("mps")
+    torch_device = "cpu"
+    torch_dtype = torch.float32
+
+model_id = "stabilityai/stable-diffusion-xl-base-1.0"
+
+if SAFETY_CHECKER == "True":
+    pipe = DiffusionPipeline.from_pretrained(model_id)
+else:
+    pipe = DiffusionPipeline.from_pretrained(model_id, safety_checker=None)
+    
+if USE_TINY_AUTOENCODER:
+    pipe.vae = AutoencoderTiny.from_pretrained(
+        "madebyollin/taesd", torch_dtype=torch_dtype, use_safetensors=True
+    )
+pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+pipe.set_progress_bar_config(disable=True)
+pipe.to(device=torch_device, dtype=torch_dtype).to(device)
+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()
+
+if TORCH_COMPILE:
+    pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+    pipe.vae = torch.compile(pipe.vae, mode="reduce-overhead", fullgraph=True)
+
+    pipe(prompt="warmup", num_inference_steps=1, guidance_scale=8.0)
+
+# Load LCM LoRA
+pipe.load_lora_weights("lcm-sd/lcm-sdxl-lora", weight_name="lcm_sdxl_lora.safetensors", adapter_name="lcm")
+
+compel_proc = Compel(
+    tokenizer=[pipe.tokenizer, pipe.tokenizer_2],
+    text_encoder=[pipe.text_encoder, pipe.text_encoder_2],
+    returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,
+    requires_pooled=[False, True],
+)
+user_queue_map = {}
+
+
+class InputParams(BaseModel):
+    seed: int = 2159232
+    prompt: str
+    guidance_scale: float = 0.5
+    strength: float = 0.5
+    steps: int = 4
+    lcm_steps: int = 50
+    width: int = WIDTH
+    height: int = HEIGHT
+
+
+def predict(params: InputParams):
+    generator = torch.manual_seed(params.seed)
+    prompt_embeds, pooled_prompt_embeds = compel_proc(params.prompt)
+    results = pipe(
+        prompt_embeds=prompt_embeds,
+        pooled_prompt_embeds=pooled_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]
+
+
+app = FastAPI()
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+@app.websocket("/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    await websocket.accept()
+    if MAX_QUEUE_SIZE > 0 and len(user_queue_map) >= MAX_QUEUE_SIZE:
+        print("Server is full")
+        await websocket.send_json({"status": "error", "message": "Server is full"})
+        await websocket.close()
+        return
+
+    try:
+        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(),
+        }
+        await websocket.send_json(
+            {"status": "start", "message": "Start Streaming", "userId": uid}
+        )
+        await handle_websocket_data(websocket, uid)
+    except WebSocketDisconnect as e:
+        logging.error(f"WebSocket Error: {e}, {uid}")
+        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
+
+
+@app.get("/queue_size")
+async def get_queue_size():
+    queue_size = len(user_queue_map)
+    return JSONResponse({"queue_size": queue_size})
+
+
+@app.get("/stream/{user_id}")
+async def stream(user_id: uuid.UUID):
+    uid = str(user_id)
+    try:
+        user_queue = user_queue_map[uid]
+        queue = user_queue["queue"]
+
+        async def generate():
+            while True:
+                params = await queue.get()
+                if params is None:
+                    continue
+
+                image = predict(params)
+                if image is None:
+                    continue
+                frame_data = io.BytesIO()
+                image.save(frame_data, format="JPEG")
+                frame_data = frame_data.getvalue()
+                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)
+
+        return StreamingResponse(
+            generate(), media_type="multipart/x-mixed-replace;boundary=frame"
+        )
+    except Exception as e:
+        logging.error(f"Streaming Error: {e}, {user_queue_map}")
+        traceback.print_exc()
+        return HTTPException(status_code=404, detail="User not found")
+
+
+async def handle_websocket_data(websocket: WebSocket, user_id: uuid.UUID):
+    uid = str(user_id)
+    user_queue = user_queue_map[uid]
+    queue = user_queue["queue"]
+    if not queue:
+        return HTTPException(status_code=404, detail="User not found")
+    last_time = time.time()
+    try:
+        while True:
+            params = await websocket.receive_json()
+            params = InputParams(**params)
+            while not queue.empty():
+                try:
+                    queue.get_nowait()
+                except asyncio.QueueEmpty:
+                    continue
+            await queue.put(params)
+            if TIMEOUT > 0 and time.time() - last_time > TIMEOUT:
+                await websocket.send_json(
+                    {
+                        "status": "timeout",
+                        "message": "Your session has ended",
+                        "userId": uid,
+                    }
+                )
+                await websocket.close()
+                return
+
+    except Exception as e:
+        logging.error(f"Error: {e}")
+        traceback.print_exc()
+
+
+@app.get("/", response_class=HTMLResponse)
+async def root():
+    return FileResponse("./static/txt2imglora.html")

static/txt2imglora.html

@@ -0,0 +1,310 @@
+<!doctype html>
+<html>
+
+<head>
+    <meta charset="UTF-8">
+    <title>Real-Time Latent Consistency Model</title>
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <script
+        src="https://cdnjs.cloudflare.com/ajax/libs/iframe-resizer/4.3.1/iframeResizer.contentWindow.min.js"></script>
+    <script src="https://cdn.jsdelivr.net/npm/piexifjs@1.0.6/piexif.min.js"></script>
+    <script src="https://cdn.tailwindcss.com"></script>
+    <style type="text/tailwindcss">
+        .button {
+          @apply bg-gray-700 hover:bg-gray-800 text-white font-normal p-2 rounded disabled:bg-gray-300 dark:disabled:bg-gray-700 disabled:cursor-not-allowed dark:disabled:text-black
+        }
+    </style>
+    <script type="module">
+        const getValue = (id) => {
+            const el = document.querySelector(`${id}`)
+            if (el.type === "checkbox")
+                return el.checked;
+            return el.value;
+        }
+        const startBtn = document.querySelector("#start");
+        const stopBtn = document.querySelector("#stop");
+        const videoEl = document.querySelector("#webcam");
+        const imageEl = document.querySelector("#player");
+        const queueSizeEl = document.querySelector("#queue_size");
+        const errorEl = document.querySelector("#error");
+        const snapBtn = document.querySelector("#snap");
+        const paramsEl = document.querySelector("#params");
+        const promptEl = document.querySelector("#prompt");
+        paramsEl.addEventListener("submit", (e) => e.preventDefault());
+        function LCMLive(promptEl, paramsEl, liveImage) {
+            let 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;
+                                initPromptStream(userId);
+                                break;
+                            case "timeout":
+                                stop();
+                                resolve({ "status": "timeout" });
+                            case "error":
+                                stop();
+                                reject(data.message);
+                        }
+                    };
+                    websocket = socket;
+                })
+            }
+
+            async function promptUpdateStream(e) {
+                const dimension = getValue("input[name=dimension]:checked");
+                const [WIDTH, HEIGHT] = JSON.parse(dimension);
+                websocket.send(JSON.stringify({
+                    "seed": getValue("#seed"),
+                    "prompt": getValue("#prompt"),
+                    "guidance_scale": getValue("#guidance-scale"),
+                    "steps": getValue("#steps"),
+                    "lcm_steps": getValue("#lcm_steps"),
+                    "width": WIDTH,
+                    "height": HEIGHT,
+                }));
+            }
+            function debouceInput(fn, delay) {
+                let timer;
+                return function (...args) {
+                    clearTimeout(timer);
+                    timer = setTimeout(() => {
+                        fn(...args);
+                    }, delay);
+                }
+            }
+            const debouncedInput = debouceInput(promptUpdateStream, 200);
+            function initPromptStream(userId) {
+                liveImage.src = `/stream/${userId}`;
+                paramsEl.addEventListener("change", debouncedInput);
+                promptEl.addEventListener("input", debouncedInput);
+            }
+
+            async function stop() {
+                websocket.close();
+                paramsEl.removeEventListener("change", debouncedInput);
+                promptEl.removeEventListener("input", debouncedInput);
+            }
+            return {
+                start,
+                stop
+            }
+        }
+        function toggleMessage(type) {
+            errorEl.hidden = false;
+            errorEl.scrollIntoView();
+            switch (type) {
+                case "error":
+                    errorEl.innerText = "To many users are using the same GPU, please try again later.";
+                    errorEl.classList.toggle("bg-red-300", "text-red-900");
+                    break;
+                case "success":
+                    errorEl.innerText = "Your session has ended, please start a new one.";
+                    errorEl.classList.toggle("bg-green-300", "text-green-900");
+                    break;
+            }
+            setTimeout(() => {
+                errorEl.hidden = true;
+            }, 2000);
+        }
+        function snapImage() {
+            try {
+                const zeroth = {};
+                const exif = {};
+                const gps = {};
+                zeroth[piexif.ImageIFD.Make] = "LCM Text-to-Image";
+                zeroth[piexif.ImageIFD.ImageDescription] = `prompt: ${getValue("#prompt")} | seed: ${getValue("#seed")} | guidance_scale: ${getValue("#guidance-scale")} | lcm_steps: ${getValue("#lcm_steps")} | steps: ${getValue("#steps")}`;
+                zeroth[piexif.ImageIFD.Software] = "https://github.com/radames/Real-Time-Latent-Consistency-Model";
+
+                exif[piexif.ExifIFD.DateTimeOriginal] = new Date().toISOString();
+
+                const exifObj = { "0th": zeroth, "Exif": exif, "GPS": gps };
+                const exifBytes = piexif.dump(exifObj);
+
+                const canvas = document.createElement("canvas");
+                canvas.width = imageEl.naturalWidth;
+                canvas.height = imageEl.naturalHeight;
+                const ctx = canvas.getContext("2d");
+                ctx.drawImage(imageEl, 0, 0);
+                const dataURL = canvas.toDataURL("image/jpeg");
+                const withExif = piexif.insert(exifBytes, dataURL);
+
+                const a = document.createElement("a");
+                a.href = withExif;
+                a.download = `lcm_txt_2_img${Date.now()}.png`;
+                a.click();
+            } catch (err) {
+                console.log(err);
+            }
+        }
+
+
+        const lcmLive = LCMLive(promptEl, paramsEl, imageEl);
+        startBtn.addEventListener("click", async () => {
+            try {
+                startBtn.disabled = true;
+                snapBtn.disabled = false;
+                const res = await lcmLive.start();
+                startBtn.disabled = false;
+                if (res.status === "timeout")
+                    toggleMessage("success")
+            } catch (err) {
+                console.log(err);
+                toggleMessage("error")
+                startBtn.disabled = false;
+            }
+        });
+        stopBtn.addEventListener("click", () => {
+            lcmLive.stop();
+        });
+        window.addEventListener("beforeunload", () => {
+            lcmLive.stop();
+        });
+        snapBtn.addEventListener("click", snapImage);
+        setInterval(() =>
+            fetch("/queue_size")
+                .then((res) => res.json())
+                .then((data) => {
+                    queueSizeEl.innerText = data.queue_size;
+                })
+                .catch((err) => {
+                    console.log(err);
+                })
+            , 5000);
+    </script>
+</head>
+
+<body class="text-black dark:bg-gray-900 dark:text-white">
+    <div class="fixed right-2 top-2 p-4 font-bold text-sm rounded-lg max-w-xs text-center" id="error">
+    </div>  
+    <main class="container mx-auto px-4 py-4 max-w-4xl flex flex-col gap-4">
+        <article class="text-center max-w-xl mx-auto">
+            <h1 class="text-3xl font-bold">Real-Time Latent Consistency Model</h1>
+            <h2 class="text-2xl font-bold mb-4">Text to Image</h2>
+            <p class="text-sm">
+                This demo showcases
+                <a href="https://huggingface.co/SimianLuo/LCM_Dreamshaper_v7" target="_blank"
+                    class="text-blue-500 underline hover:no-underline">LCM</a> Text to Image model
+                using
+                <a href="https://github.com/huggingface/diffusers/tree/main/examples/community#latent-consistency-pipeline"
+                    target="_blank" class="text-blue-500 underline hover:no-underline">Diffusers</a> with a MJPEG
+                stream server.
+            </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 10. <a
+                    href="https://huggingface.co/spaces/radames/Real-Time-Latent-Consistency-Model?duplicate=true"
+                    target="_blank" class="text-blue-500 underline hover:no-underline">Duplicate</a> and run it on your
+                own GPU.
+            </p>
+        </article>
+        <div>
+            <h2 class="font-medium">Prompt</h2>
+            <p class="text-sm text-gray-500 dark:text-gray-400">
+                Start your session and type your prompt here, 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>
+            <div class="flex text-normal px-1 py-1 border border-gray-700 rounded-md items-center">
+                <textarea type="text" id="prompt" class="font-light w-full px-3 py-2 mx-1  outline-none dark:text-black"
+                    title=" Start your session and type your prompt here, you can see the result in real-time."
+                    placeholder="Add your prompt here...">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</textarea>
+            </div>
+
+        </div>
+        <div class="">
+            <details>
+                <summary class="font-medium cursor-pointer">Advanced Options</summary>
+                <form class="grid grid-cols-3 items-center gap-3 py-3" id="params" action="">
+                    <label class="text-sm font-medium" for="dimension">Image Dimensions</label>
+                    <div class="col-span-2 flex gap-2">
+                        <div class="flex gap-1">
+                            <input type="radio" id="dimension512" name="dimension" value="[512,512]"
+                                class="cursor-pointer">
+                            <label for="dimension512" class="text-sm cursor-pointer">512x512</label>
+                        </div>
+                        <div class="flex gap-1">
+                            <input type="radio" id="dimension768" name="dimension" value="[768,768]"
+                                lass="cursor-pointer">
+                            <label for="dimension768" class="text-sm cursor-pointer">768x768</label>
+                        </div>
+                        <div class="flex gap-1">
+                            <input type="radio" id="dimension1024" name="dimension" value="[1024,1024]" checked
+                                class="cursor-pointer">
+                            <label for="dimension1024" class="text-sm cursor-pointer">1024x1024</label>
+                        </div>
+                    </div>
+                    <!--  -->
+                    <label class="text-sm font-medium " for="steps">Inference Steps
+                    </label>
+                    <input type="range" id="steps" name="steps" min="1" max="20" value="4"
+                        oninput="this.nextElementSibling.value = Number(this.value)">
+                    <output class="text-xs w-[50px] text-center font-light px-1 py-1 border border-gray-700 rounded-md">
+                        4</output>
+                    <!--  -->
+                    <label class="text-sm font-medium" for="lcm_steps">LCM Inference Steps
+                    </label>
+                    <input type="range" id="lcm_steps" name="lcm_steps" min="2" max="60" value="50"
+                        oninput="this.nextElementSibling.value = Number(this.value)">
+                    <output class="text-xs w-[50px] text-center font-light px-1 py-1 border border-gray-700 rounded-md">
+                        50</output>
+                    <!--  -->
+                    <label class="text-sm font-medium" for="guidance-scale">Guidance Scale
+                    </label>
+                    <input type="range" id="guidance-scale" name="guidance-scale" min="0" max="5" step="0.0001"
+                        value="0.8" oninput="this.nextElementSibling.value = Number(this.value).toFixed(2)">
+                    <output class="text-xs w-[50px] text-center font-light px-1 py-1 border border-gray-700 rounded-md">
+                        8.0</output>
+                    <!--  -->
+                    <label class="text-sm font-medium" for="seed">Seed</label>
+                    <input type="number" id="seed" name="seed" value="299792458"
+                        class="font-light border border-gray-700 text-right rounded-md p-2 dark:text-black">
+                    <button class="button"
+                        onclick="document.querySelector('#seed').value = Math.floor(Math.random() * 1000000000); document.querySelector('#params').dispatchEvent(new Event('change'))">
+                        Rand
+                    </button>
+                    <!--  -->
+                </form>
+            </details>
+        </div>
+        <div class="flex gap-3">
+            <button id="start" class="button">
+                Start
+            </button>
+            <button id="stop" class="button">
+                Stop
+            </button>
+            <button id="snap" disabled class="button ml-auto">
+                Snapshot
+            </button>
+        </div>
+        <div class="relative rounded-lg border border-slate-300 overflow-hidden">
+            <img id="player" class="w-full aspect-square rounded-lg"
+                src="data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mNkYAAAAAYAAjCB0C8AAAAASUVORK5CYII=">
+        </div>
+    </main>
+</body>
+
+</html>