app.py

"""Gradio UI for the DiffusionDB SD 1.5 LoRA.

Holds no model code. Generation happens in the inference Space
(whosouravsharma/diffusiondb-sd15-lora-inference), which is called here over
gradio_client. That keeps this Space light: no torch, no CUDA image, no
weights, so it builds in a minute and runs fine on free CPU while only the
backend needs a GPU.
"""

import json
import os
import time

import gradio as gr
from gradio_client import Client, handle_file  # noqa: F401
from huggingface_hub import HfApi, hf_hub_download

from src.tracing import configure_logging, logger, traced_predict

configure_logging()

# ZeroGPU refuses to start a Space in which no @spaces.GPU function is
# defined. This Space runs no model — every generation happens in the
# inference Space over HTTP — so there is no real GPU work to decorate and
# this probe exists solely to satisfy that startup check. It is never called.
# Moving this Space to CPU hardware would let it go.
try:
    import spaces

    @spaces.GPU(duration=1)
    def _zerogpu_probe():
        return "ok"

except ImportError:  # not installed off-Spaces; local runs do not need it
    pass


BACKEND = os.environ.get(
    "BACKEND_SPACE", "whosouravsharma/diffusiondb-sd15-lora-inference"
)
DATASET_REPO = "whosouravsharma/text-to-image-diffusiondb-2M"
MODEL_REPO = "whosouravsharma/diffusiondb-sd15-lora"

_client = None


# The backend sleeps after five idle minutes. Waking it means scheduling a
# T4, starting the container, importing torch and loading ~2 GB of weights.
# That is a wait, not a failure, so the UI waits it out instead of erroring.
WAKE_TIMEOUT = 600      # give up after ten minutes
POLL_INTERVAL = 5       # seconds between connection attempts

# Stages a Space cannot leave on its own. Waiting on these would just burn
# the full timeout and then report something unhelpful.
STUCK_STAGES = {"PAUSED", "RUNTIME_ERROR", "BUILD_ERROR", "CONFIG_ERROR"}


def _require_token() -> str:
    token = os.environ.get("HF_TOKEN")
    if not token:
        raise gr.Error(
            "HF_TOKEN is not set on this Space. The inference backend is "
            "private and cannot be reached without it."
        )
    return token


def _backend_stage(token: str):
    """Current runtime stage of the backend Space, or None if unknown."""
    try:
        return HfApi(token=token).get_space_runtime(BACKEND).stage
    except Exception:
        return None


def _waking_message(stage, elapsed: int) -> str:
    detail = {
        "SLEEPING": "waking it up",
        "BUILDING": "it is rebuilding",
        "APP_STARTING": "loading the model",
        "RUNNING": "loading the model",
    }.get(stage, "waking it up")
    return (
        f"⏳ **Starting the server…** the GPU backend sleeps when idle, so "
        f"{detail}. The first image takes about 40 seconds; later ones take "
        f"a few.\n\n`{elapsed}s elapsed`"
    )


def connect_backend():
    """Generator: yields status tuples while waiting, returns a live Client.

    A sleeping Space wakes on its first HTTP request, so simply attempting
    the connection starts it; the loop then keeps retrying until the app
    inside is actually serving. Consumed with `yield from`, so every status
    update reaches the browser as it happens.
    """
    global _client
    if _client is not None:
        return _client

    token = _require_token()
    started = time.monotonic()

    while True:
        elapsed = int(time.monotonic() - started)
        try:
            # Passed positionally on purpose: this argument is named hf_token
            # in gradio_client 1.x and token in 2.x, and the Space installs
            # whichever version sdk_version pulls in. It is the second
            # parameter in both. This call also wakes a sleeping Space.
            _client = Client(BACKEND, token)
            return _client
        except Exception as error:
            last_error = error

        stage = _backend_stage(token)
        if stage in STUCK_STAGES:
            raise gr.Error(
                f"The inference Space is {stage.lower().replace('_', ' ')} and "
                f"cannot start itself. It needs to be restarted from its "
                f"Settings page."
            )
        if time.monotonic() - started > WAKE_TIMEOUT:
            raise gr.Error(
                f"The inference Space did not come up within "
                f"{WAKE_TIMEOUT // 60} minutes ({last_error})."
            )

        yield _waking_message(stage, elapsed), None, None
        time.sleep(POLL_INTERVAL)


def load_examples() -> list[str]:
    """Real held-out prompts rather than invented ones.

    These are the prompts the model was evaluated on and never trained on.
    """
    try:
        path = hf_hub_download(
            DATASET_REPO, "eval_prompts.json",
            repo_type="dataset", revision="v2-clean",
        )
        return json.loads(open(path).read())[:6]
    except Exception:
        return [
            "a steampunk owl inside a glass jar, intricate detail",
            "a cyberpunk silkscreen pop art portrait",
        ]


# The backend sleeps after five idle minutes. Waking it means scheduling a
# T4, starting the container, importing torch, and loading ~2 GB of weights
# before the first step runs -- around 40 seconds during which a bare spinner
# is indistinguishable from a hang. These say what is actually happening.
STATUS_IDLE = ""
STATUS_STARTING = "⏳ **Starting the server…** contacting the GPU backend."
STATUS_GENERATING = "🎨 **Generating…**"


def generate(prompt, negative, steps, guidance, lora_scale, seed,
             request: gr.Request = None):
    """Streams status while the backend wakes, then returns the image.

    A generator rather than a plain function so the first yield paints
    immediately: the user sees that the server is starting instead of
    watching a spinner for forty seconds. A sleeping backend is a wait, not
    an error -- the only failures raised here are ones waiting cannot fix.

    `request` is filled in by Gradio from the type annotation; it is not a UI
    input and never appears in `inputs`.
    """
    if not prompt or not prompt.strip():
        raise gr.Error("Enter a prompt.")

    yield STATUS_STARTING, None, None

    # Waits out a cold backend, yielding progress the whole time.
    client = yield from connect_backend()

    yield STATUS_GENERATING, None, None

    try:
        image, used = traced_predict(
            client,
            prompt=prompt,
            negative=negative,
            steps=int(steps),
            guidance=float(guidance),
            lora_scale=float(lora_scale),
            seed=int(seed),
            request=request,
        )
    except Exception as error:
        yield STATUS_IDLE, None, None
        raise gr.Error(
            f"Generation failed ({error}). The inference Space may have gone "
            f"back to sleep mid-request — try again."
        )

    yield STATUS_IDLE, image, used


with gr.Blocks(title="DiffusionDB SD 1.5 LoRA") as demo:
    gr.Markdown(
        f"# DiffusionDB SD 1.5 LoRA\n"
        f"Stable Diffusion 1.5 with a rank-32 LoRA trained on 13,598 "
        f"prompt-image pairs from DiffusionDB "
        f"([model]({'https://huggingface.co/' + MODEL_REPO}) · "
        f"[dataset]({'https://huggingface.co/datasets/' + DATASET_REPO})).\n\n"
        f"The adapter shifts style toward the DiffusionDB aesthetic — the "
        f"keyword-heavy *artstation / intricate / octane render* idiom its "
        f"users prompted with. Set **LoRA strength to 0** to render vanilla "
        f"SD 1.5 at the same seed for a direct comparison."
    )

    with gr.Row():
        with gr.Column(scale=3):
            prompt = gr.Textbox(
                label="Prompt", lines=3,
                placeholder="a steampunk owl inside a glass jar, intricate detail",
            )
            gr.Examples(
                examples=[[p] for p in load_examples()],
                inputs=[prompt],
                label="Example prompts",
                examples_per_page=6,
            )
            negative = gr.Textbox(
                label="Negative prompt", lines=1,
                placeholder="blurry, watermark, text",
            )
            run = gr.Button("Generate", variant="primary")

            with gr.Accordion("Settings", open=False):
                lora_scale = gr.Slider(
                    0.0, 1.5, value=1.0, step=0.05, label="LoRA strength",
                    info="0 = base SD 1.5, 1 = as trained, >1 exaggerates",
                )
                steps = gr.Slider(10, 50, value=25, step=1, label="Steps")
                guidance = gr.Slider(
                    1.0, 15.0, value=7.5, step=0.5, label="Guidance scale",
                )
                seed = gr.Number(
                    value=-1, precision=0, label="Seed", info="-1 for random",
                )

        with gr.Column(scale=4):
            status = gr.Markdown(STATUS_IDLE)
            output = gr.Image(label="Output", height=512)
            used_seed = gr.Number(label="Seed used", interactive=False)

    gr.Markdown(
        "Trained on Stable Diffusion 1.x outputs, so it reproduces that "
        "model's artifacts along with its style. Images are 512×512, the "
        "resolution it was fine-tuned at."
    )

    for trigger in (run.click, prompt.submit):
        trigger(
            generate,
            inputs=[prompt, negative, steps, guidance, lora_scale, seed],
            outputs=[status, output, used_seed],
        )

if __name__ == "__main__":
    demo.queue(max_size=12).launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860)),
    )