Create app.py

app.py

@@ -0,0 +1,164 @@
+import os, io, base64, time, yaml, requests
+from PIL import Image
+import gradio as gr
+from ratelimits import RateLimiter  
+
+BACKEND_URL = os.getenv("BACKEND_URL") 
+
+cfg = yaml.safe_load(open("configs/infer.yaml"))
+limiter = RateLimiter()
+
+def b64_to_img(s: str):
+    data = base64.b64decode(s)
+    return Image.open(io.BytesIO(data)).convert("RGB")
+
+def _infer(p, st, sc, h, w, sd, et, sess_state, request: gr.Request):
+    # same session/IP logic as yours
+    if request is not None and request.headers:
+        hdrs = {k.lower(): v for k, v in request.headers.items()}
+        xff = hdrs.get("x-forwarded-for")
+        if xff:
+            ip = xff.split(",")[0].strip()
+        elif request.client:
+            ip = request.client.host
+        else:
+            ip = "unknown"
+    else:
+        ip = "unknown"
+
+    now = time.time()
+    if "started_at" not in sess_state:
+        sess_state["started_at"] = now
+    if "count" not in sess_state:
+        sess_state["count"] = 0
+    if now - sess_state["started_at"] > limiter.per_session_max_age:
+        sess_state["started_at"] = now
+        sess_state["count"] = 0
+
+    ok, reason = limiter.pre_check(ip, sess_state)
+    if not ok:
+        blank = Image.new("RGB", (int(w), int(h)), (30, 30, 30))
+        return blank, blank, f"Rate limited: {reason}", sess_state
+
+    payload = {
+        "prompt": p,
+        "steps": int(st),
+        "scale": float(sc),
+        "height": int(h),
+        "width": int(w),
+        "seed": str(sd),
+        "eta": float(et),
+    }
+    resp = requests.post(BACKEND_URL, json=payload, timeout=120)
+    resp.raise_for_status()
+    out = resp.json()
+    base_img = b64_to_img(out["base_image"])
+    lora_img = b64_to_img(out["lora_image"])
+
+    limiter.post_consume(ip, out.get("duration", 0.0))
+
+    return base_img, lora_img, out.get("status", "ok"), sess_state
+
+
+with gr.Blocks(title="Astro-Diffusion: Base vs LoRA") as demo:
+    session_state = gr.State({"count": 0, "started_at": time.time()})
+
+    gr.HTML(
+        """
+        <style>
+        .astro-header {
+            background: linear-gradient(90deg, #0f172a 0%, #1d4ed8 50%, #0ea5e9 100%);
+            padding: 0.9rem 1rem 0.85rem 1rem;
+            border-radius: 0.6rem;
+            margin-bottom: 0.9rem;
+            display: flex;
+            justify-content: space-between;
+            align-items: center;
+            gap: 1rem;
+        }
+        .astro-title {
+            color: #ffffff !important;
+            margin: 0;
+            font-weight: 700;
+            letter-spacing: 0.01em;
+        }
+        .astro-sub {
+            color: #ffffff !important;
+            margin: 0.3rem 0 0 0;
+            font-style: italic;
+            font-size: 0.8rem;
+        }
+        .astro-badge {
+            background: #facc15;
+            color: #0f172a;
+            padding: 0.4rem 1.05rem;
+            border-radius: 9999px;
+            font-weight: 700;
+            white-space: nowrap;
+            font-size: 0.95rem;
+        }
+        .prompt-panel {
+            background: #e8fff4;
+            padding: 0.5rem 0.5rem 0.2rem 0.5rem;
+            border-radius: 0.5rem;
+            margin-bottom: 0.5rem;
+        }
+        .gradio-container label, 
+        label,
+        .gradio-container [class*="label"],
+        .gradio-container [class^="svelte-"][class*="label"],
+        .gradio-container .block p > label {
+            color: #000000 !important;
+            font-weight: 600;
+        }
+        .gradio-container [data-testid="block-label"],
+        .gradio-container [data-testid="block-label"] * {
+            color: #000000 !important;
+            font-weight: 600;
+        }
+        </style>
+        <div class="astro-header">
+            <div>
+                <h2 class="astro-title">Astro-Diffusion : Base SD vs custom LoRA</h2>
+                <p class="astro-sub">Video generation and more features coming up..!</p>
+            </div>
+            <div class="astro-badge">by Srivatsava Kasibhatla</div>
+        </div>
+        """
+    )
+
+    with gr.Group(elem_classes=["prompt-panel"]):
+        prompt = gr.Textbox(
+            value="a high-resolution spiral galaxy with blue star-forming arms and a bright yellow core",
+            label="Prompt",
+        )
+
+    with gr.Row():
+        steps = gr.Slider(10, 60, value=cfg.get("num_inference_steps", 30), step=1, label="Steps")
+        scale = gr.Slider(1.0, 12.0, value=cfg.get("guidance_scale", 7.5), step=0.5, label="Guidance")
+        height = gr.Number(value=min(int(cfg.get("height", 512)), 512), label="Height", minimum=32, maximum=512)
+        width = gr.Number(value=min(int(cfg.get("width", 512)), 512), label="Width", minimum=32, maximum=512)
+        seed = gr.Textbox(value=str(cfg.get("seed", 1234)), label="Seed")
+        eta = gr.Slider(0.0, 1.0, value=0.0, step=0.01, label="Eta")
+
+    btn = gr.Button("Generate")
+    out_base = gr.Image(label="Base Model Output")
+    out_lora = gr.Image(label="LoRA Model Output")
+    status = gr.Textbox(label="Status", interactive=False)
+
+    btn.click(
+        _infer,
+        [prompt, steps, scale, height, width, seed, eta, session_state],
+        [out_base, out_lora, status, session_state],
+    )
+
+
+port = int(os.getenv("GRADIO_SERVER_PORT", "7861"))
+share = os.getenv("GRADIO_PUBLIC_SHARE", "False") == "True"
+demo.launch(
+    server_name="0.0.0.0",
+    server_port=port,
+    show_error=True,
+    share=share,
+)
+