feat: add local version of lcm model

README.md

@@ -4,7 +4,7 @@ emoji: 📈
 colorFrom: gray
 colorTo: purple
 sdk: gradio
-sdk_version: 4.4.1
+sdk_version: 3.50.2
 app_file: app.py
 pinned: false
 ---

app.py

@@ -0,0 +1,175 @@
+import threading
+from collections import deque
+from dataclasses import dataclass
+from typing import Optional
+
+import gradio as gr
+from PIL import Image
+
+from constants import DESCRIPTION, LOGO
+from gradio_examples import EXAMPLES
+from model import get_pipeline
+from utils import replace_background
+
+MAX_QUEUE_SIZE = 4
+
+pipeline = get_pipeline()
+
+
+@dataclass
+class GenerationState:
+    prompts: deque
+    generations: deque
+
+
+def get_initial_state() -> GenerationState:
+    return GenerationState(
+        prompts=deque(maxlen=MAX_QUEUE_SIZE),
+        generations=deque(maxlen=MAX_QUEUE_SIZE),
+    )
+
+
+def load_initial_state(request: gr.Request) -> GenerationState:
+    print("Loading initial state for", request.client.host)
+    print("Total number of active threads", threading.active_count())
+
+    return get_initial_state()
+
+
+async def put_to_queue(
+    image: Optional[Image.Image],
+    prompt: str,
+    seed: int,
+    strength: float,
+    state: GenerationState,
+):
+    prompts_queue = state.prompts
+
+    if prompt and image is not None:
+        prompts_queue.append((image, prompt, seed, strength))
+
+    return state
+
+
+def inference(state: GenerationState) -> Image.Image:
+    prompts_queue = state.prompts
+    generations_queue = state.generations
+
+    if len(prompts_queue) == 0:
+        return state
+
+    image, prompt, seed, strength = prompts_queue.popleft()
+
+    original_image_size = image.size
+    image = replace_background(image.resize((512, 512)))
+
+    result = pipeline(
+        prompt=prompt,
+        image=image,
+        strength=strength,
+        seed=seed,
+        guidance_scale=1,
+        num_inference_steps=4,
+    )
+
+    output_image = result.images[0].resize(original_image_size)
+
+    generations_queue.append(output_image)
+
+    return state
+
+
+def update_output_image(state: GenerationState):
+    image_update = gr.update()
+
+    generations_queue = state.generations
+
+    if len(generations_queue) > 0:
+        generated_image = generations_queue.popleft()
+        image_update = gr.update(value=generated_image)
+
+    return image_update, state
+
+
+with gr.Blocks(css="style.css", title=f"Realtime Latent Consistency Model") as demo:
+    generation_state = gr.State(get_initial_state())
+
+    gr.HTML(f'<div style="width: 70px;">{LOGO}</div>')
+    gr.Markdown(DESCRIPTION)
+    with gr.Row(variant="default"):
+        input_image = gr.Image(
+            tool="color-sketch",
+            source="canvas",
+            label="Initial Image",
+            type="pil",
+            height=512,
+            width=512,
+            brush_radius=40.0,
+        )
+
+        output_image = gr.Image(
+            label="Generated Image",
+            type="pil",
+            interactive=False,
+            elem_id="output_image",
+        )
+    with gr.Row():
+        with gr.Column():
+            prompt_box = gr.Textbox(label="Prompt", value=EXAMPLES[0])
+
+    with gr.Accordion(label="Advanced Options", open=False):
+        with gr.Row():
+            with gr.Column():
+                strength = gr.Slider(
+                    label="Strength",
+                    minimum=0.1,
+                    maximum=1.0,
+                    step=0.05,
+                    value=0.8,
+                    info="""
+                    Strength of the initial image that will be applied during inference.
+                    """,
+                )
+            with gr.Column():
+                seed = gr.Slider(
+                    label="Seed",
+                    minimum=0,
+                    maximum=2**31 - 1,
+                    step=1,
+                    randomize=True,
+                    info="""
+                    Seed for the random number generator.
+                    """,
+                )
+
+    demo.load(
+        load_initial_state,
+        outputs=[generation_state],
+    )
+    demo.load(
+        inference,
+        inputs=[generation_state],
+        outputs=[generation_state],
+        every=0.1,
+    )
+    demo.load(
+        update_output_image,
+        inputs=[generation_state],
+        outputs=[output_image, generation_state],
+        every=0.1,
+    )
+    for event in [input_image.change, prompt_box.change, strength.change, seed.change]:
+        event(
+            put_to_queue,
+            [input_image, prompt_box, seed, strength, generation_state],
+            [generation_state],
+            show_progress=False,
+            queue=True,
+        )
+
+    gr.Markdown("## Example Prompts")
+    gr.Examples(examples=EXAMPLES, inputs=[prompt_box], label="Examples")
+
+
+if __name__ == "__main__":
+    demo.queue(concurrency_count=20, api_open=False).launch(max_threads=1024)

constants.py

@@ -0,0 +1,205 @@
+DESCRIPTION = """
+# Real Time Latent Consistency Model
+"""
+
+
+LOGO = """
+<svg
+    width="100%"
+    height="100%"
+    viewBox="0 0 89 32"
+    fill="none"
+    xmlns="http://www.w3.org/2000/svg"
+>
+    <path
+    d="M52.308 3.07812H57.8465V4.92428H56.0003V6.77043H54.1541V10.4627H57.8465V12.3089H54.1541V25.232H52.308V27.0781H46.7695V25.232H48.6157V12.3089H46.7695V10.4627H48.6157V6.77043H50.4618V4.92428H52.308V3.07812Z"
+    fill="currentColor"
+    ></path>
+    <path
+    d="M79.3849 23.3858H81.2311V25.232H83.0772V27.0781H88.6157V25.232H86.7695V23.3858H84.9234V4.92428H79.3849V23.3858Z"
+    fill="currentColor"
+    ></path>
+    <path
+    d="M57.8465 14.155H59.6926V12.3089H61.5388V10.4627H70.7695V12.3089H74.4618V23.3858H76.308V25.232H78.1541V27.0781H72.6157V25.232H70.7695V23.3858H68.9234V14.155H67.0772V12.3089H65.2311V14.155H63.3849V23.3858H65.2311V25.232H67.0772V27.0781H61.5388V25.232H59.6926V23.3858H57.8465V14.155Z"
+    fill="currentColor"
+    ></path>
+    <path
+    d="M67.0772 25.232V23.3858H68.9234V25.232H67.0772Z"
+    fill="currentColor"
+    ></path>
+    <rect
+    opacity="0.22"
+    x="7.38477"
+    y="29.5391"
+    width="2.46154"
+    height="2.46154"
+    fill="#5F4CD9"
+    ></rect>
+    <rect
+    opacity="0.85"
+    x="2.46094"
+    y="19.6914"
+    width="12.3077"
+    height="2.46154"
+    fill="#5F4CD9"
+    ></rect>
+    <rect
+    x="4.92383"
+    y="17.2305"
+    width="9.84615"
+    height="2.46154"
+    fill="#5F4CD9"
+    ></rect>
+    <rect
+    opacity="0.4"
+    x="7.38477"
+    y="27.0781"
+    width="4.92308"
+    height="2.46154"
+    fill="#5F4CD9"
+    ></rect>
+    <rect
+    opacity="0.7"
+    y="22.1562"
+    width="14.7692"
+    height="2.46154"
+    fill="#5F4CD9"
+    ></rect>
+    <rect
+    opacity="0.5"
+    x="7.38477"
+    y="24.6133"
+    width="7.38462"
+    height="2.46154"
+    fill="#5F4CD9"
+    ></rect>
+    <rect
+    opacity="0.22"
+    x="7.38477"
+    y="12.3086"
+    width="2.46154"
+    height="2.46154"
+    fill="#5F4CD9"
+    ></rect>
+    <rect
+    opacity="0.85"
+    x="2.46094"
+    y="2.46094"
+    width="12.3077"
+    height="2.46154"
+    fill="#5F4CD9"
+    ></rect>
+    <rect x="4.92383" width="9.84615" height="2.46154" fill="#5F4CD9"></rect>
+    <rect
+    opacity="0.4"
+    x="7.38477"
+    y="9.84375"
+    width="4.92308"
+    height="2.46154"
+    fill="#5F4CD9"
+    ></rect>
+    <rect
+    opacity="0.7"
+    y="4.92188"
+    width="14.7692"
+    height="2.46154"
+    fill="#5F4CD9"
+    ></rect>
+    <rect
+    opacity="0.5"
+    x="7.38477"
+    y="7.38281"
+    width="7.38462"
+    height="2.46154"
+    fill="#5F4CD9"
+    ></rect>
+    <rect
+    opacity="0.22"
+    x="24.6152"
+    y="29.5391"
+    width="2.46154"
+    height="2.46154"
+    fill="#5F4CD9"
+    ></rect>
+    <rect
+    opacity="0.85"
+    x="19.6914"
+    y="19.6914"
+    width="12.3077"
+    height="2.46154"
+    fill="#5F4CD9"
+    ></rect>
+    <rect
+    x="22.1543"
+    y="17.2305"
+    width="9.84615"
+    height="2.46154"
+    fill="#5F4CD9"
+    ></rect>
+    <rect
+    opacity="0.4"
+    x="24.6152"
+    y="27.0781"
+    width="4.92308"
+    height="2.46154"
+    fill="#5F4CD9"
+    ></rect>
+    <rect
+    opacity="0.7"
+    x="17.2305"
+    y="22.1562"
+    width="14.7692"
+    height="2.46154"
+    fill="#5F4CD9"
+    ></rect>
+    <rect
+    opacity="0.5"
+    x="24.6152"
+    y="24.6133"
+    width="7.38462"
+    height="2.46154"
+    fill="#5F4CD9"
+    ></rect>
+    <rect
+    opacity="0.22"
+    x="24.6152"
+    y="12.3086"
+    width="2.46154"
+    height="2.46154"
+    fill="#5F4CD9"
+    ></rect>
+    <rect
+    opacity="0.85"
+    x="19.6914"
+    y="2.46094"
+    width="12.3077"
+    height="2.46154"
+    fill="#5F4CD9"
+    ></rect>
+    <rect x="22.1543" width="9.84615" height="2.46154" fill="#5F4CD9"></rect>
+    <rect
+    opacity="0.4"
+    x="24.6152"
+    y="9.84375"
+    width="4.92308"
+    height="2.46154"
+    fill="#5F4CD9"
+    ></rect>
+    <rect
+    opacity="0.7"
+    x="17.2305"
+    y="4.92188"
+    width="14.7692"
+    height="2.46154"
+    fill="#5F4CD9"
+    ></rect>
+    <rect
+    opacity="0.5"
+    x="24.6152"
+    y="7.38281"
+    width="7.38462"
+    height="2.46154"
+    fill="#5F4CD9"
+    ></rect>
+</svg>
+"""

gradio_examples.py

@@ -0,0 +1,4 @@
+EXAMPLES = [
+    "a house on the water, oil painting",
+    "a sunset at a tropical beach with palm trees",
+]

model.py

@@ -0,0 +1,43 @@
+from typing import Any
+
+
+def get_pipeline():
+    import torch
+    from diffusers import AutoencoderTiny, AutoPipelineForImage2Image
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    pipe = AutoPipelineForImage2Image.from_pretrained(
+        "SimianLuo/LCM_Dreamshaper_v7",
+        use_safetensors=True,
+    )
+    pipe.vae = AutoencoderTiny.from_pretrained(
+        "madebyollin/taesd",
+        torch_dtype=torch.float16,
+        use_safetensors=True,
+    )
+    pipe = pipe.to(device, dtype=torch.float16)
+    pipe.unet.to(memory_format=torch.channels_last)
+    return pipe
+
+
+def get_test_pipeline():
+    from PIL import Image
+    from dataclasses import dataclass
+    import random
+    import time
+
+    @dataclass
+    class Images:
+        images: list[Image.Image]
+
+    class Pipeline:
+        def __call__(self, *args: Any, **kwds: Any) -> Any:
+            time.sleep(0.5)
+            r = random.randint(0, 255)
+            g = random.randint(0, 255)
+            b = random.randint(0, 255)
+
+            return Images(images=[Image.new("RGB", (512, 512), color=(r, g, b))])
+
+    return Pipeline()

style.css

@@ -0,0 +1,3 @@
+h1 {
+    text-align: center;
+}

utils.py

@@ -0,0 +1,15 @@
+from PIL import Image
+import numpy as np
+
+
+def replace_background(image: Image.Image, new_background_color=(0, 255, 255)):
+    image_np = np.array(image)
+
+    white_threshold = 255 * 3
+    white_pixels = np.sum(image_np, axis=-1) == white_threshold
+
+    image_np[white_pixels] = new_background_color
+
+    result = Image.fromarray(image_np)
+
+    return result