Initial commit

.gitignore

@@ -0,0 +1,4 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class

README.md

@@ -1,13 +1,12 @@
 ---
-title: Realtime Stable Diffusion
-emoji: ⚡
-colorFrom: green
-colorTo: green
+title: Real Time Latent Consistency Models
+emoji: 👀
+colorFrom: pink
+colorTo: indigo
 sdk: gradio
-sdk_version: 4.3.0
+sdk_version: 3.50.2
 app_file: app.py
 pinned: false
-license: other
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,205 @@
+import json
+from collections import deque
+from dataclasses import dataclass
+from typing import Optional
+
+import gradio as gr
+import websockets
+from gradio.processing_utils import decode_base64_to_image, encode_pil_to_base64
+from PIL import Image
+from websockets.sync.client import connect
+
+from constants import DESCRIPTION, WS_ADDRESS, LOGO
+from utils import replace_background
+
+from gradio_examples import EXAMPLES
+
+MAX_QUEUE_SIZE = 4
+
+
+@dataclass
+class GenerationState:
+    prompts: deque
+    responses: deque
+
+
+def get_initial_state() -> GenerationState:
+    return GenerationState(
+        prompts=deque(maxlen=MAX_QUEUE_SIZE),
+        responses=deque(maxlen=MAX_QUEUE_SIZE),
+    )
+
+
+def load_initial_state(request: gr.Request) -> GenerationState:
+    print("Loading initial state for", request.client.host)
+
+    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 send_inference_request(state: GenerationState) -> Image.Image:
+    prompts_queue = state.prompts
+    response_queue = state.responses
+
+    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)))
+
+    arguments = {
+        "prompt": prompt,
+        "image_url": encode_pil_to_base64(image),
+        "strength": strength,
+        "negative_prompt": "cartoon, illustration, animation. face. male, female",
+        "seed": seed,
+        "guidance_scale": 1,
+        "num_inference_steps": 4,
+        "sync_mode": 1,
+        "num_images": 1,
+    }
+
+    connection = connect(WS_ADDRESS)
+    connection.send(json.dumps(arguments))
+
+    try:
+        response = json.loads(connection.recv())
+    except websockets.exceptions.ConnectionClosedOK:
+        print("Connection closed, reconnecting...")
+        # TODO: This is a hacky way to reconnect, but it works for now
+        #      Ideally, we should be able to reconnect to the same connection
+        #      and not have to create a new one
+        connection = connect(WS_ADDRESS)
+        try:
+            response = json.loads(connection.recv())
+        except websockets.exceptions.ConnectionClosedOK:
+            print("Connection closed again, aborting...")
+            return state
+
+    # TODO: If a new connection is created, the response do not contain the images.
+    if "images" in response:
+        response_queue.append((response, original_image_size))
+
+    return state
+
+
+def update_output_image(state: GenerationState):
+    image_update = gr.update()
+    inference_time_update = gr.update()
+
+    response_queue = state.responses
+
+    if len(response_queue) > 0:
+        response, original_image_size = response_queue.popleft()
+        generated_image = decode_base64_to_image(response["images"][0]["url"])
+        inference_time = response["timings"]["inference"]
+
+        generated_image.resize(original_image_size).save("generated.png")
+
+        image_update = gr.update(value=generated_image.resize(original_image_size))
+        inference_time_update = gr.update(value=round(inference_time, 4))
+
+    return image_update, inference_time_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(scale=23):
+            prompt_box = gr.Textbox(label="Prompt")
+        with gr.Column(scale=1):
+            inference_time_box = gr.Number(
+                label="Inference Time (s)", interactive=False
+            )
+
+    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(
+        send_inference_request,
+        inputs=[generation_state],
+        outputs=[generation_state],
+        every=0.1,
+    )
+    demo.load(
+        update_output_image,
+        inputs=[generation_state],
+        outputs=[output_image, inference_time_box, 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=4, api_open=False).launch()

constants.py

@@ -0,0 +1,209 @@
+import os
+
+DESCRIPTION = """
+# Real Time Latent Consistency Model Hosted on [fal.ai](https://fal.ai)
+"""
+
+WS_ADDRESS = os.environ["WS_ADDRESS"]
+
+
+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",
+]

style.css

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

utils.py

@@ -0,0 +1,14 @@
+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