Migrate to native Hugging Face Gradio SDK Space with ZeroGPU support

Dockerfile

@@ -1,41 +0,0 @@
-# syntax=docker/dockerfile:1.6
-# CUDA 12.8 runtime — gemlite/Triton kernels JIT against the runtime ptxas
-# that comes with this image; no need for the larger -devel variant.
-FROM nvidia/cuda:12.8.0-runtime-ubuntu22.04
-
-# ── system deps ──────────────────────────────────────────────────────────────
-# ca-certificates, curl, git: basic utilities.
-# build-essential: needed for JIT kernel compile and some sdist setups.
-# python3-venv: python baseline for uv package manager.
-# procps: provides process management.
-RUN apt-get update && apt-get install -y --no-install-recommends \
-        ca-certificates curl git build-essential python3 python3-venv \
-        libgomp1 libssl3 procps \
-    && rm -rf /var/lib/apt/lists/*
-
-# Non-root user: HF Spaces convention is uid 1000.
-RUN useradd -m -u 1000 user
-USER user
-ENV HOME=/home/user PATH="/home/user/.local/bin:$PATH"
-
-# uv (Python venv + package manager). The demo's setup.sh assumes it's on PATH.
-RUN curl -fsSL https://astral.sh/uv/install.sh | sh
-
-WORKDIR /home/user/app
-
-# ── clone + run setup.sh ─────────────────────────────────────────────────────
-# SKIP_DOWNLOAD=1 keeps setup.sh from pulling the 3.5 GB model at build time
-# — entrypoint.sh handles that at boot so a Space restart doesn't have to
-# rebuild the image.
-RUN git clone https://github.com/PrismML-Eng/Bonsai-image-demo.git . \
-    && SKIP_DOWNLOAD=1 BONSAI_PACKAGE_MIN_AGE_DAYS=0 ./setup.sh
-
-# ── Space-local files ────────────────────────────────────────────────────────
-# All Space-specific code lives under space/ (Python gradio.Server app + HTML frontend).
-# The demo's own code stays at the repo root (cloned earlier) so they don't collide.
-COPY --chown=user space/ /home/user/app/space/
-RUN chmod +x /home/user/app/space/entrypoint.sh
-
-EXPOSE 7860
-
-CMD ["/home/user/app/space/entrypoint.sh"]

README.md

@@ -3,8 +3,7 @@ title: Bonsai Image GPU
 emoji: 🎨
 colorFrom: green
 colorTo: blue
-sdk: docker
-app_port: 7860
+sdk: gradio
 suggested_hardware: l40sx1
 pinned: true
 short_description: Run Bonsai-Image-4B models on GPU

space/app.py → app.py

@@ -1,20 +1,20 @@
-"""gradio.Server backend for Bonsai-Image-Demo.
+"""gradio.Server backend with ZeroGPU dynamic allocation support.
 
-Replaces the complex custom metrics middleware and N-replica queues
-with Gradio's native queue engine and Server class.
+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
-import uuid
-import tempfile
 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")
@@ -23,6 +23,56 @@ 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
@@ -32,6 +82,21 @@ except ImportError:
 
 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
+
 @asynccontextmanager
 async def lifespan(fastapi_app: Server):
     logger.info("Initializing GpuPipeline inside gradio.Server lifespan...")
@@ -46,7 +111,7 @@ async def lifespan(fastapi_app: Server):
     pipeline.prewarm()
     fastapi_app.state.pipeline = pipeline
     
-    # Run warmup shapes if configured
+    # Run warmup shapes if configured (ZeroGPU AOT caching benefit)
     warmup_shapes = os.environ.get("BONSAI_WARMUP_SHAPES", "")
     skip_warmup = os.environ.get("BONSAI_SKIP_WARMUP", "").lower() in {"1", "true", "yes"}
     
@@ -70,6 +135,7 @@ async def lifespan(fastapi_app: Server):
 app = Server(lifespan=lifespan)
 
 @app.api(name="generate")
+@spaces_gpu
 def generate(
     prompt: str,
     seed: int = 0,
@@ -82,7 +148,7 @@ def generate(
 ) -> dict:
     """Generate an image using the GpuPipeline and return base64 data and metrics.
     
-    Runs inside Gradio's native queue engine with serial concurrency.
+    Decorated with @spaces.GPU to dynamically lease Blackwell GPUs on Hugging Face.
     """
     pipeline = app.state.pipeline
     try:

requirements.txt

@@ -0,0 +1,13 @@
+torch
+torchvision
+transformers
+diffusers
+accelerate
+huggingface_hub
+gradio>=5.0.0
+hqq
+uvicorn
+spaces
+gemlite
+git+https://github.com/PrismML-Eng/mflux-prism.git
+git+https://github.com/PrismML-Eng/image-studio.git#subdirectory=backend_gpu

space/entrypoint.sh

@@ -1,108 +0,0 @@
-#!/bin/bash
-# Streamlined entrypoint for Bonsai Image Studio (gradio.Server)
-#
-# Boot order:
-#   1. Detect platform and active GPU.
-#   2. Set up persistent storage caches (/data) if available.
-#   3. Sync model weights (ternary and binary) — public, no tokens needed.
-#   4. Start the single gradio.Server app on :7860 via uvicorn.
-#
-set -euo pipefail
-
-APP_DIR="${HOME:-/home/user}/app"
-cd "$APP_DIR"
-
-export PATH="$APP_DIR/.venv/bin:$PATH"
-export HF_HUB_ENABLE_HF_TRANSFER=1
-
-echo "==>  Starting Bonsai Image Studio initialization ..."
-
-# ── GPU detection ───────────────────────────────────────────────────────────
-GPU_NAME=$(nvidia-smi --query-gpu=name --format=csv,noheader 2>/dev/null | awk 'NR==1' | xargs)
-GPU_CAP=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null | awk 'NR==1' | tr -d '.')
-[ -z "$GPU_NAME" ] && GPU_NAME="unknown"
-[ -z "$GPU_CAP" ]  && GPU_CAP="00"
-echo "[OK]  GPU: $GPU_NAME (sm_${GPU_CAP})"
-
-# ── persistent storage detection ─────────────────────────────────────────────
-_setup_persistent() {
-    [ -d /data ] && [ -w /data ] || return 1
-
-    # Namespaced caches
-    _gemlite_dir="/data/cache/gemlite-sm${GPU_CAP}"
-    _triton_dir="/data/cache/triton-sm${GPU_CAP}"
-
-    # Migrations
-    if [ -d /data/cache/gemlite ] && [ ! -e "$_gemlite_dir" ]; then
-        echo "[INFO] migrating /data/cache/gemlite → gemlite-sm${GPU_CAP}"
-        mv /data/cache/gemlite "$_gemlite_dir" 2>/dev/null || true
-    fi
-    if [ -d /data/cache/triton ] && [ ! -e "$_triton_dir" ]; then
-        echo "[INFO] migrating /data/cache/triton → triton-sm${GPU_CAP}"
-        mv /data/cache/triton "$_triton_dir" 2>/dev/null || true
-    fi
-
-    mkdir -p /data/models "$_gemlite_dir" "$_triton_dir" /data/state /data/state/daily 2>/dev/null || return 1
-    rm -rf "$APP_DIR/models" 2>/dev/null || return 1
-    ln -s /data/models "$APP_DIR/models" 2>/dev/null || return 1
-    mkdir -p "$APP_DIR/outputs" 2>/dev/null || return 1
-    rm -rf "$APP_DIR/outputs/.gemlite_cache" "$APP_DIR/outputs/.triton_cache" 2>/dev/null || true
-    ln -s "$_gemlite_dir" "$APP_DIR/outputs/.gemlite_cache" 2>/dev/null || return 1
-    ln -s "$_triton_dir"  "$APP_DIR/outputs/.triton_cache"  2>/dev/null || return 1
-    return 0
-}
-
-if _setup_persistent; then
-    echo "[OK]  /data Storage Bucket attached — model + caches will persist"
-    export BONSAI_STATE_DIR=/data/state
-    export BONSAI_PERSISTENT_STORAGE=1
-else
-    echo "[WARN] /data not mounted — caches will reset on every Space restart."
-    export BONSAI_STATE_DIR="$APP_DIR/outputs/.state"
-    export BONSAI_PERSISTENT_STORAGE=0
-    mkdir -p "$BONSAI_STATE_DIR" 2>/dev/null || true
-fi
-
-# ── model download / sync ────────────────────────────────────────────────────
-MODEL_DIR="$APP_DIR/models/bonsai-image-4B-ternary-gemlite"
-BINARY_MODEL_DIR="$APP_DIR/models/bonsai-image-4B-binary-gemlite"
-
-# Set token if provided, but download works token-free for public models
-if [ -n "${HF_TOKEN:-}" ]; then
-    export BONSAI_TOKEN="$HF_TOKEN"
-fi
-
-echo "==>  syncing bonsai-image-ternary-4B-gemlite-2bit ..."
-./scripts/download_model.sh --model ternary-gemlite || echo "[WARN] Ternary model sync encountered errors."
-echo "==>  syncing bonsai-image-binary-4B-gemlite-1bit ..."
-./scripts/download_model.sh --model binary-gemlite || echo "[WARN] Binary model sync encountered errors."
-
-# ── pin model paths once; shared across all workers ──────────────────────────
-export MFLUX_STUDIO_GPU_DEFAULT_BACKEND="bonsai-ternary-gemlite"
-
-_ternary_transformer_dir=$(ls -d "$MODEL_DIR"/transformer-gemlite-* 2>/dev/null | awk 'NR==1')
-if [ -z "$_ternary_transformer_dir" ]; then
-    echo "[ERR] no transformer-gemlite-* subdir under $MODEL_DIR" >&2
-    exit 1
-fi
-export MFLUX_STUDIO_GPU_TERNARY_TRANSFORMER_PATH="$_ternary_transformer_dir"
-
-_binary_transformer_dir=$(ls -d "$BINARY_MODEL_DIR"/transformer-gemlite-* 2>/dev/null | awk 'NR==1')
-if [ -z "$_binary_transformer_dir" ]; then
-    echo "[ERR] no transformer-gemlite-* subdir under $BINARY_MODEL_DIR" >&2
-    exit 1
-fi
-export MFLUX_STUDIO_GPU_BINARY_TRANSFORMER_PATH="$_binary_transformer_dir"
-export MFLUX_STUDIO_GPU_TEXT_ENCODER_PATH="$MODEL_DIR/text_encoder-hqq-4bit"
-export MFLUX_STUDIO_GPU_VAE_PATH="$MODEL_DIR/vae"
-export MFLUX_STUDIO_GPU_TOKENIZER_PATH="$MODEL_DIR/text_encoder-hqq-4bit/tokenizer"
-
-# Warmup configurations
-: "${BONSAI_WARMUP_SHAPES:=512x512,1024x1024}"
-export BONSAI_WARMUP_SHAPES
-
-# ── Start the gradio.Server app directly on port 7860 ────────────────────────
-echo "==>  Starting gradio.Server on public port 7860 ..."
-exec uvicorn space.app:app \
-    --host 0.0.0.0 --port 7860 \
-    --no-access-log