app.py

import time, base64, io, os, requests
from typing import Optional
from PIL import Image
import gradio as gr

# =========================
# Stable Horde config
# =========================
HORDE_URL = "https://stablehorde.net/api/v2/generate/async"
HORDE_STATUS = "https://stablehorde.net/api/v2/generate/status/{id}"

# Optional: set an API key in your Space secrets for better priority:
# Settings → Repository secrets → add HORDE_API_KEY
HORDE_API_KEY = os.getenv("HORDE_API_KEY", "")

DEFAULT_STEPS = 24
DEFAULT_W = 768
DEFAULT_H = 768
POLL_INTERVAL = 2.5        # seconds between polls
POLL_TIMEOUT = 180         # max seconds to wait for a job
MODEL = None               # let Horde choose; or set e.g. "SDXL 1.0"

# =========================
# Horde client
# =========================
def horde_txt2img(prompt: str,
                  steps: int = DEFAULT_STEPS,
                  width: int = DEFAULT_W,
                  height: int = DEFAULT_H,
                  model: Optional[str] = MODEL) -> Image.Image:
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt.")

    payload = {
        "prompt": prompt.strip(),
        "params": {
            "steps": int(steps),
            "width": int(width),
            "height": int(height),
            "n": 1
        },
        "nsfw": False,
        "censor_nsfw": True
    }
    if model:
        payload["models"] = [model]

    # Submit
    headers = {"apikey": HORDE_API_KEY} if HORDE_API_KEY else {}
    submit = requests.post(HORDE_URL, json=payload, headers=headers, timeout=30)
    submit.raise_for_status()
    job_id = submit.json()["id"]

    # Poll
    start = time.time()
    while True:
        status = requests.get(HORDE_STATUS.format(id=job_id), timeout=30).json()
        if status.get("faulted"):
            raise gr.Error(f"Horde job faulted: {status}")
        if status.get("done"):
            gens = status.get("generations") or []
            if not gens:
                raise gr.Error("Horde finished but returned no generations.")
            b64 = gens[0]["img"]
            img_bytes = base64.b64decode(b64)
            return Image.open(io.BytesIO(img_bytes)).convert("RGB")

        if time.time() - start > POLL_TIMEOUT:
            raise gr.Error("Timed out waiting for Horde. Try again or reduce steps/size.")
        time.sleep(POLL_INTERVAL)

# Wrapper for Gradio (adds optional negative prompt + long-shot toggle if you want later)
def generate_image(prompt, steps, size):
    # size like "768x768"
    try:
        w, h = [int(x.strip()) for x in size.lower().split("x")]
    except Exception:
        w, h = DEFAULT_W, DEFAULT_H
    img = horde_txt2img(prompt, steps=steps, width=w, height=h)
    return img

# =========================
# UI
# =========================
CUSTOM_CSS = """
.gradio-container { padding: 24px; }

/* Rounded prompt boxes */
.prompt-box textarea {
  border-radius: 18px !important;
  min-height: 90px;
  font-size: 16px;
  line-height: 1.4;
  padding: 14px 16px;
}

/* Pill buttons */
.pill button {
  border-radius: 999px !important;
  padding: 10px 18px;
  font-size: 15px;
}

/* Blue-ish image boxes feel (border radius) */
.image-out .wrap, .image-out .svelte-1ipelgc {
  border-radius: 22px !important;
}
"""

with gr.Blocks(css=CUSTOM_CSS, title="Stitch – Image Checkpoints (Stable Horde)") as demo:
    gr.Markdown("### Image Checkpoints (Stable Horde) — generate per-prompt frames")

    # Global controls (applies to all generators)
    with gr.Row():
        steps = gr.Slider(8, 50, value=DEFAULT_STEPS, step=1, label="Steps (quality/time)")
        size = gr.Dropdown(
            choices=["512x512", "768x768", "1024x576", "1024x768"],
            value=f"{DEFAULT_W}x{DEFAULT_H}",
            label="Resolution"
        )

    # 4 rows: [Prompt + Button]  |  [Image Output]
    prompt_boxes = []
    gen_buttons = []
    img_outputs = []

    for i in range(1, 5):
        with gr.Row():
            with gr.Column(scale=1, min_width=320):
                p = gr.Textbox(
                    placeholder=f"Prompt input (Image {i})",
                    lines=4,
                    label=None,
                    elem_classes=["prompt-box"]
                )
                b = gr.Button(f"Generate image {i}", elem_classes=["pill"])
            with gr.Column(scale=2, min_width=380):
                img = gr.Image(label=f"Image {i} output", type="pil", elem_classes=["image-out"])
        prompt_boxes.append(p)
        gen_buttons.append(b)
        img_outputs.append(img)

    # Wire callbacks
    for i in range(4):
        gen_buttons[i].click(
            fn=generate_image,
            inputs=[prompt_boxes[i], steps, size],
            outputs=[img_outputs[i]]
        )

    gr.Markdown(
        "> Tip: You can generate images independently per row. Steps ↑ = higher quality but slower. "
        "Stable Horde is free (crowd GPUs), so expect occasional queue time."
    )

if __name__ == "__main__":
    demo.queue().launch()