import os, io, base64, time, yaml, requests
from PIL import Image
import gradio as gr
from requests.exceptions import ConnectionError, Timeout, HTTPError
# frontend-only: call your backend (RunPod/pod/etc.)
BACKEND_URL = os.getenv("BACKEND_URL", "http://localhost:7861").rstrip("/")
print(f"[HF] BACKEND_URL resolved to: {BACKEND_URL}")
# sample prompts
SAMPLE_PROMPTS = [
"a high-resolution spiral galaxy with blue star-forming arms and a bright yellow core",
"a crimson emission nebula with dark dust lanes and scattered newborn stars",
"a ringed gas giant with visible storm bands and subtle shadow on rings",
"an accretion disk around a black hole with relativistic jets, high contrast",
]
# default UI values if no YAML
cfg = {
"height": 512,
"width": 512,
"num_inference_steps": 30,
"guidance_scale": 7.5,
"seed": 1234,
"eta": 0,
}
# ---- health check ----
def check_backend():
try:
r = requests.get(f"{BACKEND_URL}/health", timeout=5)
r.raise_for_status()
data = r.json()
if data.get("status") == "ok":
return "backend=READY"
except Exception:
pass
return "backend=DOWN"
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, session_id):
# make sure we always have ints for blank images
h = int(h)
w = int(w)
payload = {
"prompt": p,
"steps": int(st),
"scale": float(sc),
"height": h,
"width": w,
"seed": str(sd),
"eta": float(et),
}
# send session_id if we have one
if session_id:
payload["session_id"] = session_id
try:
r = requests.post(f"{BACKEND_URL}/infer", json=payload, timeout=120)
if r.status_code == 429:
blank = Image.new("RGB", (w, h), (30, 30, 30))
out = r.json()
# backend also returns session_id on 429
new_sid = out.get("session_id", session_id)
msg = out.get("error", "rate limited by backend")
return blank, blank, msg, new_sid
r.raise_for_status()
out = r.json()
base_img = b64_to_img(out["base_image"])
lora_img = b64_to_img(out["lora_image"])
new_sid = out.get("session_id", session_id)
return base_img, lora_img, out.get("status", "ok"), new_sid
except ConnectionError:
blank = Image.new("RGB", (w, h), (120, 50, 50))
return blank, blank, "Backend not reachable (connection refused). Start the backend and retry.", session_id
except Timeout:
blank = Image.new("RGB", (w, h), (120, 50, 50))
return blank, blank, "Backend took too long. Please try again later.", session_id
except HTTPError as e:
blank = Image.new("RGB", (w, h), (120, 50, 50))
return blank, blank, f"Backend returned HTTP Error: {e.response.status_code}", session_id
except Exception as e:
blank = Image.new("RGB", (w, h), (120, 50, 50))
return blank, blank, f"Unknown client error: {e}", session_id
def build_ui():
with gr.Blocks(title="Astro-Diffusion: Base vs LoRA") as demo:
# session state lives in the browser/tab
session_state = gr.State(value="")
# header + status
status_lbl = gr.Markdown("checking backend...")
gr.HTML(
"""
"""
)
with gr.Group(elem_classes=["prompt-panel"]):
sample_dropdown = gr.Dropdown(
choices=SAMPLE_PROMPTS,
value=SAMPLE_PROMPTS[0],
label="Sample prompts",
)
prompt = gr.Textbox(
value=SAMPLE_PROMPTS[0],
label="Prompt",
)
# when user picks a sample, copy it into the textbox
sample_dropdown.change(fn=lambda x: x, inputs=sample_dropdown, outputs=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)
# send session_state, receive updated session_state
btn.click(
_infer,
[prompt, steps, scale, height, width, seed, eta, session_state],
[out_base, out_lora, status, session_state],
)
# ping once when UI loads
demo.load(fn=check_backend, inputs=None, outputs=status_lbl)
return demo
if __name__ == "__main__":
interface = build_ui()
port = int(os.getenv("PORT", os.getenv("GRADIO_SERVER_PORT", "8080")))
share = os.getenv("GRADIO_PUBLIC_SHARE", "True").lower() == "true"
interface.launch(server_name="0.0.0.0", ssr_mode = False, server_port=port, share=share)