spaceui.py

"""Builds a Space UI from a declarative spec.

Shared verbatim by every Space in the wavespeed org. Generated from
_shared/spaceapp/ — edit there and re-run _shared/build_apps.py.

Handling of the user's API key
------------------------------
The key is typed into the browser and travels to this server on each request.
Keeping it from leaking takes more than masking the textbox, because Gradio has
several features that will happily persist or republish an input:

  * `type="password"`          - not echoed back into the DOM.
  * `api_visibility="private"` - Gradio 6 otherwise documents this event on the
                                 app's public API page, generating client
                                 snippets that include every input, the key
                                 among them.
  * `analytics_enabled=False`  - on Blocks, plus the env vars set in app.py
                                 before gradio is imported.
  * no `gr.State`/`gr.Examples` ever holds the key, so it is not serialised
    into the page or cached to disk.
  * every error is passed through `wavespeed.redact` before display, because a
    requests exception can stringify the Authorization header.

Gradio's flagging feature (which writes raw inputs to a CSV) belongs to
gr.Interface; this UI is built from gr.Blocks, which has no flagging, so there
is nothing to switch off there.

The key is a plain function argument: it lives for the duration of one request
and is not retained between them.
"""

from __future__ import annotations

import gradio as gr

import wavespeed as ws

SITE = "https://wavespeed.ai"


def link(path: str, campaign: str) -> str:
    """Build an outbound wavespeed.ai URL carrying UTM attribution.

    Without these, traffic from Hugging Face lands in analytics as plain
    referral with no way to tell which Space produced it.
    """
    sep = "&" if "?" in path else "?"
    return (
        f"{SITE}{path}{sep}utm_source=huggingface&utm_medium=space"
        f"&utm_campaign={campaign}"
    )


def _collect(spec, key, values, progress):
    """Turn UI values into an API payload, uploading any local files first."""
    payload = dict(spec.get("extra", {}))
    for field, value in zip(spec["fields"], values):
        kind, api_key_name = field["kind"], field["key"]

        if kind in ("image", "audio", "video"):
            if not value:
                if field.get("required", True):
                    raise ws.WaveSpeedError(f"{field['label']} is required.")
                continue
            progress(0.1, desc=f"Uploading {field['label'].lower()}…")
            payload[api_key_name] = ws.upload(key, value)

        elif kind == "images":
            if not value:
                if field.get("required", True):
                    raise ws.WaveSpeedError(f"{field['label']} is required.")
                continue
            progress(0.1, desc=f"Uploading {field['label'].lower()}…")
            payload[api_key_name] = [ws.upload(key, value)]

        elif kind == "prompt":
            text = (value or "").strip()
            if not text and field.get("required", True):
                raise ws.WaveSpeedError("Enter a prompt.")
            if text:
                payload[api_key_name] = text

        elif kind == "seed":
            # -1 means "let the service choose"; sending it would pin the seed.
            if value is not None and int(value) >= 0:
                payload[api_key_name] = int(value)

        elif value is not None and value != "":
            payload[api_key_name] = value

    return payload


def build(spec):
    """Return a configured gr.Blocks for this Space."""
    css = spec["css"]
    camp = spec["campaign"]
    outputs_are_video = spec["output"] == "video"

    with gr.Blocks(
        title=f"{spec['title']} - WaveSpeed AI",
        analytics_enabled=False,
    ) as demo:
        gr.HTML(
            f"""
            <div class="hero-container">
              <a class="hero-badge" href="{link('/', camp)}"
                 target="_blank" rel="noopener">WAVESPEED AI</a>
              <h1 class="hero-title">{spec['title']}</h1>
              <p class="hero-desc">{spec['tagline']}</p>
            </div>
            """
        )

        with gr.Row(elem_classes="api-key-row"):
            api_key = gr.Textbox(
                label="WaveSpeed API key",
                placeholder="Paste your API key — it is used for this request only",
                type="password",          # never echoed back to the page
                show_label=False,
                container=False,
                scale=4,
            )
            gr.HTML(
                f'<a class="get-key-btn" href="{link("/dashboard", camp)}" '
                'target="_blank" rel="noopener">Get a key</a>'
            )
        gr.Markdown(
            "Your key is sent only to `api.wavespeed.ai` to run this model. "
            "It is not stored, logged, or shared, and generations are billed to "
            "your own account.",
            elem_classes="key-note",
        )

        controls = []
        with gr.Row():
            with gr.Column(scale=1):
                for f in spec["fields"]:
                    controls.append(_make_control(f))
                run_btn = gr.Button(
                    spec.get("button", "Generate"),
                    variant="primary",
                    elem_classes="primary-btn",
                )
            with gr.Column(scale=1):
                if spec["output"] == "compare":
                    outs = [
                        gr.Image(label=m["label"], type="filepath")
                        for m in spec["compare"]
                    ]
                elif outputs_are_video:
                    outs = [gr.Video(label="Result")]
                else:
                    outs = [gr.Image(label="Result", type="filepath")]

        gr.HTML(
            f"""
            <div class="cta-container">
              <p class="cta-desc">Runs
                <a href="{link('/models/' + spec['model'], camp)}"
                   target="_blank" rel="noopener"><code>{spec['model']}</code></a>
                on WaveSpeed &middot;
                <a href="{link('/models', camp)}" target="_blank"
                   rel="noopener">Browse all models</a> &middot;
                <a href="{link('/docs', camp)}" target="_blank"
                   rel="noopener">API docs</a>
              </p>
            </div>
            """
        )

        def _run(key, *values, progress=gr.Progress()):
            blank = [None] * len(outs)
            if not key or not key.strip():
                gr.Warning("Enter your WaveSpeed API key first.")
                return blank[0] if len(blank) == 1 else tuple(blank)
            try:
                payload = _collect(spec, key, values, progress)
                progress(0.3, desc="Submitting…")
                if spec["output"] == "compare":
                    models = spec["compare"]
                    results = []
                    for i, m in enumerate(models):
                        progress(
                            0.3 + 0.6 * i / len(models),
                            desc=f"Running {m['label']}…",
                        )
                        merged = dict(payload, **m.get("extra", {}))
                        results.append(ws.run(key, m["model"], merged)[0])
                    return tuple(results)
                outputs = ws.run(
                    key, spec["model"], payload,
                    on_tick=lambda s: progress(0.6, desc=f"Generating ({s})…"),
                )
                return outputs[0]
            except ws.WaveSpeedError as e:
                # Message is already redacted by the client.
                gr.Warning(str(e))
            except Exception as e:  # noqa: BLE001 - never surface a raw trace
                gr.Warning(ws.redact(f"Unexpected error: {e}", key))
            return blank[0] if len(blank) == 1 else tuple(blank)

        run_btn.click(
            _run,
            inputs=[api_key, *controls],
            outputs=outs,
            # Keep this event off the public API page — its generated snippets
            # would include the api_key input.
            api_visibility="private",
        )

    return demo


def _make_control(f):
    # seed fields carry no explicit label; they get the default below.
    kind, label = f["kind"], f.get("label", "")
    if kind == "prompt":
        return gr.Textbox(
            label=label, placeholder=f.get("placeholder", ""),
            lines=f.get("lines", 3),
        )
    if kind == "image":
        return gr.Image(label=label, type="filepath")
    if kind == "images":
        return gr.Image(label=label, type="filepath")
    if kind == "audio":
        return gr.Audio(label=label, type="filepath")
    if kind == "video":
        return gr.Video(label=label)
    if kind == "choice":
        return gr.Dropdown(
            label=label, choices=f["choices"], value=f.get("default", f["choices"][0])
        )
    if kind == "bool":
        return gr.Checkbox(label=label, value=f.get("default", False))
    if kind == "seed":
        return gr.Number(label=f.get("label", "Seed (-1 = random)"), value=-1, precision=0)
    if kind == "slider":
        return gr.Slider(
            label=label, minimum=f["min"], maximum=f["max"],
            step=f.get("step", 1), value=f["default"],
        )
    raise ValueError(f"unknown field kind: {kind}")


def launch(demo, css):
    """Launch the app. Gradio 6 takes css here rather than on Blocks."""
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        css=css,
        quiet=True,
    )