handle ZeroGPU tensor packing crash on huggingface

README.md

@@ -51,7 +51,7 @@ python download_cifar10.py 42 --split train --output assets/cifar10-train-42.png
 
 Push this repository to a Gradio Space and select ZeroGPU in the Space hardware settings. The Dockerfile and legacy server are not used by the Gradio deployment. Account eligibility and daily GPU quotas apply.
 
-The app imports `spaces` before PyTorch and places both GPU models on CUDA at startup when `SPACES_ZERO_GPU` is enabled. Reverse sampling tries GPU first and automatically restarts on CPU if GPU allocation, quota, or supported CUDA availability errors prevent the run. Locally, the app prefers CUDA, then Apple Silicon MPS, then CPU. MPS availability failures also fall back to CPU. Separate pipelines are cached for each model and device; the CPU pipeline loads on first use. CPU sampling is substantially slower. GPU sampling runs inside a decorated function called by the Gradio handler. `GPU_DURATION_SECONDS` defaults to 120; benchmark a full 1,000-step run and adjust this Space variable if necessary. A timeout is reported by the Gradio queue; selecting fewer checkpoints does not reduce the sampling steps.
+The app imports `spaces` before PyTorch and copies model tensors into owned CPU storage before placing both GPU models on CUDA at startup when `SPACES_ZERO_GPU` is enabled. Reverse sampling tries GPU first and automatically restarts on CPU if GPU allocation, quota, or supported CUDA availability errors prevent the run. Locally, the app prefers CUDA, then Apple Silicon MPS, then CPU. MPS availability failures also fall back to CPU. Separate pipelines are cached for each model and device; the CPU pipeline loads on first use. CPU sampling is substantially slower. GPU sampling runs inside a decorated function called by the Gradio handler. `GPU_DURATION_SECONDS` defaults to 120; benchmark a full 1,000-step run and adjust this Space variable if necessary. A timeout is reported by the Gradio queue; selecting fewer checkpoints does not reduce the sampling steps.
 
 For persistent caching, attach a writable storage bucket and set the Space variable `MODEL_CACHE_DIR` to its mount path (for example `/data/models`). The old Docker path `/app/models` is also usable if explicitly configured with that variable.
 

app.py

@@ -1,4 +1,8 @@
 """Gradio entry point for local use and Hugging Face ZeroGPU."""
+import faulthandler
+
+faulthandler.enable()  # Include Python stacks if native startup code crashes.
+
 import spaces  # Must initialize ZeroGPU before importing torch/diffusers.
 
 import base64

generate.py

@@ -52,6 +52,12 @@ def load_model(device=None, model_id=MODEL_ID):
             print(f"Loading {model_id} on device: {device}")
             pipeline = DDPMPipeline.from_pretrained(
                 model_id, cache_dir=Path(os.environ.get("MODEL_CACHE_DIR", ROOT / "models")))
+            if device == "cuda" and os.environ.get("SPACES_ZERO_GPU", "").lower() in {"1", "true", "yes"}:
+                # ZeroGPU packs CUDA aliases at startup. Give it owned CPU
+                # storage rather than tensors backed by mapped checkpoint files.
+                for component in pipeline.components.values():
+                    if isinstance(component, torch.nn.Module):
+                        component._apply(lambda tensor: tensor.clone())
             pipeline.to(device)
             _PIPELINES[key] = pipeline
         _PIPELINE = _PIPELINES[key]

requirements.txt

@@ -1,3 +1,4 @@
+accelerate>=0.30.0
 diffusers>=0.30.0
 matplotlib>=3.8.0
 Pillow>=10.0.0