Add model comparison dashboard

src/model_demo.py

@@ -0,0 +1,78 @@
+import base64
+import os
+import random
+import time
+from io import BytesIO
+
+import gradio as gr
+import requests
+from PIL import Image
+
+SERVER_API = os.environ["SERVER_API"]
+
+
+def image_from_base64(image_data: str) -> Image:
+    image_buffer = BytesIO(base64.b64decode(image_data))
+    image = Image.open(image_buffer)
+    return image
+
+
+def submit(prompt: str, seed: int | str | None, baseline: bool) -> tuple:
+    if isinstance(seed, str):
+        seed = hash(seed) % 2 ** 32
+    if not seed:
+        random.seed(int(time.time()))
+        seed = random.randint(0, 2 ** 32 - 1)
+
+    print(f"Making request with prompt: {prompt}, seed: {seed}, baseline: {baseline}")
+    response = requests.post(f"{SERVER_API}/generate", params={"prompt": prompt, "baseline": baseline, "seed": seed}, verify=False)
+    response.raise_for_status()
+    result = response.json()
+    generation_time = float(result["generation_time"])
+    image = image_from_base64(result["image"])
+    print(f"Received image with generation time: {generation_time:.3f}s")
+
+    return None, gr.Image(
+        image,
+        label=f"{generation_time:.3f}s",
+        show_label=True
+    )
+
+
+def create_textbox() -> gr.Textbox:
+    response = requests.get(f"{SERVER_API}/model", verify=False)
+    response.raise_for_status()
+    model = response.json()
+    return gr.Textbox(f"{model['uid']} - {model['url']}", interactive=False, show_label=False)
+
+
+def create_demo():
+    with gr.Accordion(f"EdgeMaxxing Model Comparison"):
+        with gr.Group():
+            with gr.Row():
+                with gr.Column():
+                    gr.Textbox("Baseline", interactive=False, show_label=False)
+                    baseline_image_component = gr.Image(show_label=False)
+
+                with gr.Column():
+                    textbox = gr.Textbox()
+                    textbox.attach_load_event(lambda: create_textbox(), None)
+                    optimized_image_component = gr.Image(show_label=False)
+        with gr.Row():
+            prompt = gr.Textbox(
+                placeholder="Enter prompt...",
+                interactive=True,
+                submit_btn=True,
+                show_label=False,
+                autofocus=True,
+                scale=10,
+            )
+
+            seed_input = gr.Textbox(
+                placeholder="Enter seed...",
+                interactive=True,
+                show_label=False,
+            )
+
+    prompt.submit(lambda prompt, seed: submit(prompt, seed, True), inputs=[prompt, seed_input], outputs=[prompt, baseline_image_component])
+    prompt.submit(lambda prompt, seed: submit(prompt, seed, False), inputs=[prompt, seed_input], outputs=[prompt, optimized_image_component])