Upload app.py

app.py

@@ -0,0 +1,233 @@
+"""Boomer FLA — Hugging Face Space demo (ZeroGPU / RTX Pro 6000 Blackwell + Gradio)."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+# ZeroGPU Spaces: ~/.cache is often read-only — use /tmp for HF + diffusers caches.
+_hf_root = Path(os.environ.get("HF_HOME", "/tmp/huggingface"))
+for _sub in ("hub", "modules", "transformers", "diffusers"):
+    (_hf_root / _sub).mkdir(parents=True, exist_ok=True)
+os.environ.setdefault("HF_HOME", str(_hf_root))
+os.environ.setdefault("HUGGINGFACE_HUB_CACHE", str(_hf_root / "hub"))
+os.environ.setdefault("HF_MODULES_CACHE", str(_hf_root / "modules"))
+os.environ.setdefault("TRANSFORMERS_CACHE", str(_hf_root / "transformers"))
+os.environ.setdefault("DIFFUSERS_CACHE", str(_hf_root / "diffusers"))
+
+import gc
+import sys
+
+import gradio as gr
+import spaces
+import torch
+from huggingface_hub import snapshot_download
+
+MODEL_ID = os.environ.get("BOOMER_MODEL_ID", "akrao9/Boomer-T2I")
+DEFAULT_STEPS = 32
+DEFAULT_CFG = 4.5
+DEFAULT_CFG_RESCALE = 0.5
+
+EXAMPLE_PROMPTS: list[tuple[str, str]] = [
+    (
+        "Cyberpunk Highland",
+        "A sweeping cinematic view of a futuristic Scottish highland at twilight, glowing neon purple and blue moss creeping over ancient rock formations, distant futuristic spires reflecting off a dark loch, dramatic low-angle photography.",
+    ),
+    (
+        "Volcanic Caldera",
+        "A high-contrast shot looking down into an active volcanic caldera, vibrant flowing rivers of orange lava cutting through pitch-black obsidian stone, thick plumes of dark smoke catching the crimson glow, hyper-detailed geological textures.",
+    ),
+    (
+        "Nordic Winter Fjords",
+        "A breathtaking winter view of deep Norwegian fjords during the blue hour, snow-laden mountain slopes plunging into dark mirror-like ocean water, vibrant emerald green Northern Lights stretching across the sky, crisp and ultra-clear atmospheric rendering.",
+    ),
+    (
+        "Ancient Sunken Ruins",
+        "A wide landscape shot of ancient stone ruins submerged in a shallow, crystal-clear tropical ocean, vibrant coral reefs and schools of exotic fish visible beneath the surface, sun rays slicing through the water, warm cinematic lighting.",
+    ),
+]
+
+
+def _hf_token() -> str | None:
+    """Return HF token from Space secrets; None if unset (do not pass token=True)."""
+    for key in ("HF_TOKEN", "HUGGING_FACE_HUB_TOKEN", "HUGGING_HUB_TOKEN"):
+        value = os.environ.get(key, "").strip()
+        if value:
+            return value
+    return None
+
+
+_hf = _hf_token()
+if _hf is None:
+    print(
+        "WARNING: No HF_TOKEN secret found. Add one in Space Settings → Secrets "
+        "(required for gated Gemma text encoder).",
+        flush=True,
+    )
+
+print(f"Loading Boomer pipeline from {MODEL_ID} ...", flush=True)
+_model_dir = Path(
+    snapshot_download(
+        MODEL_ID,
+        token=_hf,
+        ignore_patterns=["*.png", "*.jpg", "*.jpeg"],
+    )
+)
+if str(_model_dir) not in sys.path:
+    sys.path.insert(0, str(_model_dir))
+
+from pipeline_boomer import BoomerPipeline  # noqa: E402
+
+pipe = BoomerPipeline.from_pretrained(str(_model_dir), torch_dtype=torch.bfloat16, token=_hf)
+pipe.to("cuda")
+pipe._hf_token = _hf
+
+print("Pre-loading VAE on cuda ...", flush=True)
+pipe._ensure_vae()
+if _hf is not None:
+    print("Pre-loading text encoder on cuda ...", flush=True)
+    pipe._ensure_text_encoder()
+else:
+    print("Skipping text encoder preload (HF_TOKEN not set).", flush=True)
+print("Model ready.", flush=True)
+
+
+def _resolve_seed(seed: float | int | None) -> int:
+    if seed is None or int(seed) < 0:
+        return torch.randint(0, 2**32 - 1, (1,)).item()
+    return int(seed)
+
+
+# ZeroGPU "large" = half NVIDIA RTX Pro 6000 Blackwell MIG slice (48 GB VRAM).
+@spaces.GPU(size="large", duration=150)
+def generate_image(
+    prompt: str,
+    seed: float,
+    steps: float,
+    cfg_scale: float,
+) -> tuple[object, str]:
+    prompt = (prompt or "").strip()
+    if not prompt:
+        raise gr.Error("Please enter a prompt before generating.")
+    if pipe._hf_token is None:
+        raise gr.Error(
+            "HF_TOKEN Space secret is required. In Space Settings → Secrets, add "
+            "HF_TOKEN with a Hugging Face token that has accepted the "
+            "Gemma Terms of Use (https://ai.google.dev/gemma/terms)."
+        )
+
+    resolved_seed = _resolve_seed(seed)
+    step_count = max(1, int(steps))
+    cfg = float(cfg_scale)
+
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    result = pipe(
+        prompt,
+        steps=step_count,
+        seed=resolved_seed,
+        cfg_scale=cfg,
+        cfg_rescale=DEFAULT_CFG_RESCALE,
+        offload_text_encoder=True,
+    )
+    image = result[0]
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    status = (
+        f"Done — seed {resolved_seed}, steps {step_count}, CFG {cfg:.1f}. "
+        "1024×1024px via STORK-2."
+    )
+    return image, status
+
+
+def reset_form() -> tuple[str, None, str, float]:
+    return "", None, "Reset.", -1.0
+
+
+def build_ui() -> gr.Blocks:
+    example_rows = [[text, -1, DEFAULT_STEPS, DEFAULT_CFG] for _, text in EXAMPLE_PROMPTS]
+
+    with gr.Blocks(title="Boomer FLA") as demo:
+        gr.Markdown(
+            """
+            # Boomer FLA — Text to Image
+            Generate **1024×1024** images from text. **Best for landscapes and scenic
+            environments** — not reliable for humans or portraits (JourneyDB + FineT2I training).
+            Runs on **ZeroGPU** (NVIDIA RTX Pro 6000 Blackwell, 48 GB). First run may
+            take a minute while a GPU is allocated.
+            """
+        )
+
+        with gr.Row(equal_height=False):
+            with gr.Column(scale=3):
+                prompt = gr.Textbox(
+                    label="Prompt",
+                    placeholder="Describe the image you want to generate ...",
+                    lines=5,
+                )
+                with gr.Accordion("Advanced settings", open=False):
+                    seed = gr.Number(
+                        label="Seed (-1 = random)",
+                        value=-1,
+                        precision=0,
+                    )
+                    steps = gr.Slider(
+                        label="Denoising steps",
+                        minimum=16,
+                        maximum=64,
+                        step=1,
+                        value=DEFAULT_STEPS,
+                    )
+                    cfg_scale = gr.Slider(
+                        label="CFG scale",
+                        minimum=1.0,
+                        maximum=8.0,
+                        step=0.1,
+                        value=DEFAULT_CFG,
+                    )
+                with gr.Row():
+                    generate_btn = gr.Button("Generate", variant="primary")
+                    reset_btn = gr.Button("Reset")
+                output = gr.Image(label="Generated image", type="pil", height=640)
+                status = gr.Textbox(label="Status", interactive=False)
+
+            with gr.Column(scale=2):
+                gr.Markdown("### Example prompts")
+                gr.Markdown("Click an example to **generate immediately**.")
+                gr.Examples(
+                    examples=example_rows,
+                    inputs=[prompt, seed, steps, cfg_scale],
+                    outputs=[output, status],
+                    fn=generate_image,
+                    run_on_click=True,
+                    cache_examples=False,
+                    label="",
+                    examples_per_page=4,
+                )
+
+        generate_btn.click(
+            fn=generate_image,
+            inputs=[prompt, seed, steps, cfg_scale],
+            outputs=[output, status],
+        )
+        reset_btn.click(
+            fn=reset_form,
+            outputs=[prompt, output, status, seed],
+        )
+
+        prompt.submit(
+            fn=generate_image,
+            inputs=[prompt, seed, steps, cfg_scale],
+            outputs=[output, status],
+        )
+
+    return demo
+
+
+demo = build_ui()
+
+if __name__ == "__main__":
+    demo.launch()