app.py

"""gradio.Server backend with ZeroGPU dynamic allocation support.

Relocated to root level for native Hugging Face Gradio SDK Space compatibility.
Handles model weight downloads and environment path setup directly in Python.
"""
from __future__ import annotations

import os
import glob
import logging
import time
import base64
from contextlib import asynccontextmanager
from fastapi import Request
from fastapi.responses import HTMLResponse
from gradio import Server
from huggingface_hub import snapshot_download

# Initialize logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
logger = logging.getLogger(__name__)

# Ensure MFLUX_STUDIO_GPU_TOKEN is set for backend compatibility
os.environ.setdefault("MFLUX_STUDIO_GPU_TOKEN", "local-demo-unused")

# ── Dynamic Model Download and Path Configuration ───────────────────────────
def init_models_and_env():
    # Detect persistent storage mount at /data
    has_data = os.path.exists("/data") and os.access("/data", os.W_OK)
    model_root = "/data/models" if has_data else "models"
    
    # 1. Download models using huggingface_hub snapshot_download
    ternary_dir = os.path.join(model_root, "bonsai-image-4B-ternary-gemlite")
    logger.info(f"Syncing ternary model to: {ternary_dir}...")
    snapshot_download(
        repo_id='prism-ml/bonsai-image-ternary-4B-gemlite-2bit',
        local_dir=ternary_dir,
        max_workers=8
    )
    
    binary_dir = os.path.join(model_root, "bonsai-image-4B-binary-gemlite")
    logger.info(f"Syncing binary model to: {binary_dir}...")
    snapshot_download(
        repo_id='prism-ml/bonsai-image-binary-4B-gemlite-1bit',
        local_dir=binary_dir,
        max_workers=8
    )
    
    # 2. Resolve dynamic paths and configure environment variables
    os.environ["MFLUX_STUDIO_GPU_DEFAULT_BACKEND"] = "bonsai-ternary-gemlite"
    
    ternary_transformer = glob.glob(os.path.join(ternary_dir, "transformer-gemlite-*"))
    if not ternary_transformer:
        raise FileNotFoundError(f"Ternary transformer subdirectory not found under {ternary_dir}")
    os.environ["MFLUX_STUDIO_GPU_TERNARY_TRANSFORMER_PATH"] = ternary_transformer[0]
    
    binary_transformer = glob.glob(os.path.join(binary_dir, "transformer-gemlite-*"))
    if not binary_transformer:
        raise FileNotFoundError(f"Binary transformer subdirectory not found under {binary_dir}")
    os.environ["MFLUX_STUDIO_GPU_BINARY_TRANSFORMER_PATH"] = binary_transformer[0]
    
    os.environ["MFLUX_STUDIO_GPU_TEXT_ENCODER_PATH"] = os.path.join(ternary_dir, "text_encoder-hqq-4bit")
    os.environ["MFLUX_STUDIO_GPU_VAE_PATH"] = os.path.join(ternary_dir, "vae")
    os.environ["MFLUX_STUDIO_GPU_TOKENIZER_PATH"] = os.path.join(ternary_dir, "text_encoder-hqq-4bit/tokenizer")
    
    # Setup JIT compilation caches
    if has_data:
        os.environ["TRITON_CACHE_DIR"] = "/data/cache/triton"
        logger.info("Persistent Storage Bucket attached. Model weights and Triton caches will persist.")
    else:
        logger.warning("/data not mounted. Caches will be ephemeral.")

# Run model setup and environment pinning
init_models_and_env()

# Apply low-memory and Triton/gemlite caching monkeypatches from scripts.local_backend
try:
    import scripts.local_backend
    logger.info("Successfully imported scripts.local_backend to apply performance monkeypatches.")
except ImportError:
    logger.warning("Could not import scripts.local_backend. Continuing without low-memory patches.")

from backend_gpu.pipeline_gpu import GpuPipeline, DEFAULT_GPU_BACKEND, _normalize_gpu_backend

# ── ZeroGPU spaces module loading ────────────────────────────────────────────
try:
    import spaces
    has_spaces = True
    logger.info("ZeroGPU spaces library loaded successfully.")
except ImportError:
    has_spaces = False
    logger.warning("ZeroGPU spaces library not found. Running in standard GPU/CPU fallback mode.")

# Conditional decorator for ZeroGPU support
def spaces_gpu(func):
    if has_spaces:
        return spaces.GPU(func)
    return func

# ── Global GpuPipeline Reference (ZeroGPU compatible) ────────────────────────
pipeline = None

# Initialize the gradio.Server app
app = Server()

@app.api(name="generate")
@spaces_gpu
def generate(
    prompt: str,
    seed: int = 0,
    steps: int = 4,
    guidance: float = 1.0,
    backend: str = "bonsai-ternary-gemlite",
    height: int = 512,
    width: int = 512,
    max_sequence_length: int = 256,
) -> dict:
    """Generate an image using the GpuPipeline and return base64 data and metrics.
    
    Decorated with @spaces.GPU to dynamically lease Blackwell GPUs on Hugging Face.
    """
    global pipeline
    if pipeline is None:
        backend_env = os.environ.get("MFLUX_STUDIO_GPU_DEFAULT_BACKEND", DEFAULT_GPU_BACKEND)
        try:
            backend_env = _normalize_gpu_backend(backend_env)
        except ValueError:
            backend_env = DEFAULT_GPU_BACKEND
            
        logger.info(f"Initializing and prewarming global GpuPipeline (backend={backend_env}) inside ZeroGPU lease...")
        pipeline = GpuPipeline(backend=backend_env)
        pipeline.prewarm()

        # Run warmup shapes inside ZeroGPU lease
        warmup_shapes = os.environ.get("BONSAI_WARMUP_SHAPES", "")
        skip_warmup = os.environ.get("BONSAI_SKIP_WARMUP", "").lower() in {"1", "true", "yes"}

        if warmup_shapes and not skip_warmup:
            logger.info(f"Running boot-time warmup for shapes: {warmup_shapes} inside ZeroGPU lease...")
            for spec in warmup_shapes.split(","):
                spec = spec.strip().lower()
                if not spec:
                    continue
                try:
                    w, h = (int(x) for x in spec.split("x", 1))
                    logger.info(f"Warming up shape {w}x{h}...")
                    pipeline.generate_png(prompt="warmup", seed=0, steps=1, height=h, width=w)
                except Exception as e:
                    logger.warning(f"Warmup failed for shape {spec}: {e}")
                    
    try:
        pipeline.ensure_backend(backend=backend)
    except ValueError as exc:
        return {"error": str(exc)}

    logger.info(f"Generating image. Prompt: {prompt[:50]}... Backend: {backend} ({width}x{height})")
    
    gen_start = time.perf_counter()
    try:
        image_bytes = pipeline.generate_png(
            prompt=prompt,
            seed=seed,
            steps=steps,
            height=height,
            width=width,
            guidance=guidance,
            max_sequence_length=max_sequence_length,
        )
    except Exception as e:
        logger.error(f"Image generation failed: {e}")
        return {"error": f"Generation failed: {str(e)}"}
        
    wall_seconds = time.perf_counter() - gen_start
    peak_mem = pipeline.last_peak_memory_mb or 0.0
    
    logger.info(f"Generation successful. Time: {wall_seconds:.3f}s. Peak RAM: {peak_mem:.1f} MB.")
    
    # Return base64 representation of PNG alongside metrics
    img_b64 = base64.b64encode(image_bytes).decode("ascii")
    return {
        "image_b64": img_b64,
        "wall_seconds": round(wall_seconds, 3),
        "peak_memory_mb": round(peak_mem, 1),
    }

@app.api(name="backends")
def backends() -> dict:
    """List supported model backends and default configurations."""
    arm = os.environ.get("MFLUX_STUDIO_GPU_DEFAULT_BACKEND", "bonsai-ternary-gemlite")
    if arm.endswith("-gemlite"):
        default_family = arm[: -len("-gemlite")]
    else:
        default_family = arm
        
    supported = [
        f.strip()
        for f in os.environ.get(
            "BONSAI_SUPPORTED_FAMILIES", "bonsai-ternary,bonsai-binary"
        ).split(",")
        if f.strip()
    ]
    
    return {
        "supported_families": supported,
        "default_family": default_family,
    }

@app.get("/", response_class=HTMLResponse)
async def homepage():
    """Serve the custom index.html frontend."""
    html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")
    try:
        with open(html_path, "r", encoding="utf-8") as f:
            return HTMLResponse(content=f.read())
    except FileNotFoundError:
        return HTMLResponse(content="<h1>Bonsai Image Demo</h1><p>index.html not found.</p>", status_code=404)

if __name__ == "__main__":
    app.launch(show_error=True)