try device in the order of gpu/mps/cpu

README.md

@@ -44,7 +44,7 @@ python download_celeba_hq.py 42 --split train --output face-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, loads the model at startup, and places it on CUDA when `SPACES_ZERO_GPU` is enabled. GPU sampling runs inside a decorated 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 places the GPU model 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 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

@@ -16,7 +16,32 @@ from PIL import Image
 from generate import CHECKPOINTS, ROOT, generate_images, load_model
 
 ZERO_GPU = os.environ.get("SPACES_ZERO_GPU", "").lower() in {"1", "true", "yes"}
-load_model(device="cuda" if ZERO_GPU else None)
+def gpu_unavailable(error):
+    """Recognize allocation/quota/CUDA failures without masking model errors."""
+    message = str(error).lower()
+    return any(text in message for text in (
+        "quota exceeded", "exceeded your gpu quota", "exceeded your zerogpu quota",
+        "no gpu was available", "no gpu available", "cuda out of memory",
+        "cuda is not available", "no nvidia driver", "not compiled with cuda",
+        "cuda driver", "cuda-capable device", "gpu task aborted",
+    ))
+
+
+# ZeroGPU requires CUDA model placement before its worker is allocated.
+GPU_READY = False
+MPS_READY = False
+if ZERO_GPU:
+    try:
+        load_model(device="cuda")
+        GPU_READY = True
+    except Exception as error:
+        if not gpu_unavailable(error):
+            raise
+        print(f"GPU unavailable at startup; using CPU: {error}")
+else:
+    import torch
+    GPU_READY = torch.cuda.is_available()
+    MPS_READY = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
 
 
 def parse_settings(seed, checkpoint_text):
@@ -81,23 +106,51 @@ def diffuse(image, seed, checkpoint_text):
         yield checkpoint_slots(checkpoints, frames), final_noise if step == 1000 else None, progress_html(step, started)
 
 
-@spaces.GPU(duration=int(os.environ.get("GPU_DURATION_SECONDS", "120")))
-def denoise(seed, checkpoint_text, input_mode, forward_noise):
-    seed, checkpoints = parse_settings(seed, checkpoint_text)
-    initial_image = forward_noise if input_mode == "Forward noise" else None
-    if input_mode == "Forward noise" and initial_image is None:
-        raise gr.Error("Run Diffuse first to create forward noise.")
+def stream_denoising(seed, checkpoints, initial_image, device):
     started = time.monotonic()
     frames = {1000: initial_image if initial_image is not None else gaussian_noise(seed)}
     yield checkpoint_slots(checkpoints, frames), progress_html(0, started)
-    for step, frame in generate_images(seed, checkpoints, initial_image):
+    print('using', device)
+    for step, frame in generate_images(seed, checkpoints, initial_image, device=device):
         if frame is not None:
             frames[step] = frame
-        # Stream every pipeline status; only resend gallery images at checkpoints.
         yield (checkpoint_slots(checkpoints, frames) if frame is not None else gr.skip(),
                progress_html(1000 - step, started))
 
 
+@spaces.GPU(duration=int(os.environ.get("GPU_DURATION_SECONDS", "120")))
+def denoise_gpu(seed, checkpoints, initial_image):
+    yield from stream_denoising(seed, checkpoints, initial_image, "cuda")
+
+
+def denoise(seed, checkpoint_text, input_mode, forward_noise):
+    seed, checkpoints = parse_settings(seed, checkpoint_text)
+    initial_image = forward_noise if input_mode == "Forward noise" else None
+    if input_mode == "Forward noise" and initial_image is None:
+        raise gr.Error("Run Diffuse first to create forward noise.")
+    if GPU_READY:
+        try:
+            yield from denoise_gpu(seed, checkpoints, initial_image)
+            return
+        except Exception as error:
+            if not gpu_unavailable(error):
+                raise
+            fallback = "MPS" if MPS_READY else "CPU"
+            print(f"CUDA unavailable; restarting on {fallback}: {error}")
+            gr.Info(f"CUDA unavailable or quota exhausted. Running on {fallback}.")
+    if MPS_READY:
+        try:
+            yield from stream_denoising(seed, checkpoints, initial_image, "mps")
+            return
+        except (RuntimeError, NotImplementedError) as error:
+            message = str(error).lower()
+            if "mps" not in message and "metal" not in message:
+                raise
+            print(f"MPS unavailable; restarting on CPU: {error}")
+            gr.Info("MPS unavailable. Running on CPU.")
+    yield from stream_denoising(seed, checkpoints, initial_image, "cpu")
+
+
 # Reuse the original page's palette, typography, header, and panel styling.
 PAGE_CSS = (ROOT / "styles.css").read_text() + """
 .gradio-container { max-width:1440px !important; padding:0 6% !important;

generate.py

@@ -11,7 +11,11 @@ ROOT = Path(__file__).resolve().parent
 MODEL_ID = "google/ddpm-celebahq-256"
 CHECKPOINTS = (0, 50, 100, 200, 350, 500, 700, 1000)
 
+# Gradio may resume a generator on a different worker thread between yields.
+# A primitive Lock can be released there; RLock requires the acquiring thread.
 _MODEL_LOCK = threading.Lock()
+_CACHE_LOCK = threading.Lock()
+_PIPELINES = {}
 _PIPELINE = None
 _TORCH = None
 _DEVICE = None
@@ -21,7 +25,7 @@ def load_model(device=None):
     """Load and cache the diffusion pipeline, returning it and its device."""
     global _PIPELINE, _TORCH, _DEVICE
 
-    if _PIPELINE is None:
+    with _CACHE_LOCK:
         try:
             import torch
             from diffusers import DDPMPipeline
@@ -32,23 +36,24 @@ def load_model(device=None):
                 f"Import detail: {error}"
             ) from error
 
-        if device is not None:
-            pass
-        elif torch.cuda.is_available():
-            device = "cuda"
-        elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
-            device = "mps"
-        else:
-            device = "cpu"
-
-        print(f"Using device: {device}")
-        pipeline = DDPMPipeline.from_pretrained(MODEL_ID, cache_dir=Path(os.environ.get("MODEL_CACHE_DIR", ROOT / "models")))
-        pipeline.to(device)
+        if device is None:
+            if torch.cuda.is_available():
+                device = "cuda"
+            elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+                device = "mps"
+            else:
+                device = "cpu"
+        device = str(device)
+        if device not in _PIPELINES:
+            print(f"Loading model on device: {device}")
+            pipeline = DDPMPipeline.from_pretrained(
+                MODEL_ID, cache_dir=Path(os.environ.get("MODEL_CACHE_DIR", ROOT / "models")))
+            pipeline.to(device)
+            _PIPELINES[device] = pipeline
+        _PIPELINE = _PIPELINES[device]
         _TORCH = torch
         _DEVICE = device
-        _PIPELINE = pipeline
-
-    return _PIPELINE, _TORCH, _DEVICE
+        return _PIPELINE, _TORCH, _DEVICE
 
 
 def is_model_loaded():
@@ -62,14 +67,17 @@ def sample_to_image(sample):
     return Image.fromarray(image, mode="RGB")
 
 
-def generate_images(seed, checkpoints=CHECKPOINTS, initial_image=None):
+def generate_images(seed, checkpoints=CHECKPOINTS, initial_image=None, device=None):
     """Yield every remaining-step count and an image at requested checkpoints."""
     with _MODEL_LOCK:
-        pipeline, torch, device = load_model()
+        pipeline, torch, device = load_model(device=device)
         scheduler = pipeline.scheduler
         training_steps = scheduler.config.num_train_timesteps
         scheduler.set_timesteps(training_steps, device=device)
-        generator = torch.Generator(device=device).manual_seed(seed)
+        # MPS has no torch.Generator implementation. Keep its random stream
+        # on CPU; Diffusers also accepts this generator for scheduler noise.
+        random_device = "cpu" if device == "mps" else device
+        generator = torch.Generator(device=random_device).manual_seed(seed)
         unet = pipeline.unet
         sample_size = unet.config.sample_size
         if isinstance(sample_size, (tuple, list)):
@@ -80,9 +88,9 @@ def generate_images(seed, checkpoints=CHECKPOINTS, initial_image=None):
             sample = torch.randn(
                 (1, unet.config.in_channels, height, width),
                 generator=generator,
-                device=device,
+                device=random_device,
                 dtype=unet.dtype,
-            ) * scheduler.init_noise_sigma
+            ).to(device) * scheduler.init_noise_sigma
         else:
             import numpy as np