changed to gradio

README.md

@@ -3,14 +3,15 @@ title: Diffusion
 emoji: 🌍
 colorFrom: indigo
 colorTo: purple
-sdk: docker
-app_port: 8000
+sdk: gradio
+app_file: app.py
+python_version: "3.10"
 pinned: false
 ---
 
 # Diffusion Lab
 
-The reverse process runs the pretrained `google/ddpm-celebahq-256` model through Hugging Face Diffusers. This is an unconditional CelebA-HQ face generator: reverse sampling starts with random Gaussian noise and does not denoise the uploaded image. The upload is used only by the browser-side forward-noising demonstration.
+The Gradio app supports ZeroGPU and streams checkpoint images. The reverse process runs the pretrained `google/ddpm-celebahq-256` model through Hugging Face Diffusers. This is an unconditional CelebA-HQ face generator: reverse sampling normally starts with random Gaussian noise. You can also select the final forward-noise image as its input; this is a demonstration, not a guaranteed reconstruction of the uploaded image.
 
 ## Run
 
@@ -20,10 +21,10 @@ Use Python 3.10 or newer. Install a PyTorch build appropriate for your machine f
 python -m venv .venv
 source .venv/bin/activate
 python -m pip install -r requirements.txt
-python server.py
+python app.py
 ```
 
-Open <http://127.0.0.1:8000>. The first generation downloads the model weights into `models/`; later runs use that cache. The backend prefers CUDA, then Apple Silicon MPS, then CPU. Model inference uses all 1,000 DDPM steps and sends the selected intermediate images to the browser as they are produced.
+Open <http://127.0.0.1:7860>. Startup downloads the model weights into `models/`; later runs use that cache. The backend prefers CUDA, then Apple Silicon MPS, then CPU. Model inference uses all 1,000 DDPM steps and sends the selected intermediate images to the browser as they are produced.
 
 To generate and display one image without starting the server, run:
 
@@ -38,3 +39,20 @@ To download a particular zero-based row from the public CelebA-HQ dataset:
 python download_celeba_hq.py 1374
 python download_celeba_hq.py 42 --split train --output face-42.png
 ```
+
+## ZeroGPU deployment
+
+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.
+
+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.
+
+For local use, install `gradio` and `spaces` alongside `requirements.txt` (these packages are provided by the Gradio Spaces runtime):
+
+```sh
+python -m pip install gradio spaces
+python app.py
+```
+
+The original browser interface remains available with `python server.py` for local or dedicated-GPU use.

app.py

@@ -0,0 +1,85 @@
+"""Gradio entry point for local use and Hugging Face ZeroGPU."""
+import spaces  # Must initialize ZeroGPU before importing torch/diffusers.
+
+import os
+import secrets
+
+import gradio as gr
+import numpy as np
+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 parse_settings(seed, checkpoint_text):
+    try:
+        seed = int(str(seed).strip())
+        checkpoints = sorted(set(int(part.strip()) for part in checkpoint_text.split(",")))
+    except (ValueError, TypeError):
+        raise gr.Error("Enter an integer seed and comma-separated integer checkpoints.")
+    if not 0 <= seed <= 2147483647:
+        raise gr.Error("Seed must be between 0 and 2147483647.")
+    if not 2 <= len(checkpoints) <= 12 or checkpoints[0] != 0 or checkpoints[-1] != 1000:
+        raise gr.Error("Use 2–12 checkpoints between 0 and 1000, including both endpoints.")
+    return seed, checkpoints
+
+
+def diffuse(image, seed, checkpoint_text):
+    seed, checkpoints = parse_settings(seed, checkpoint_text)
+    if image is None:
+        raise gr.Error("Upload an image first.")
+    pixels = np.asarray(image.convert("RGB").resize((256, 256)), dtype=np.float32) / 127.5 - 1
+    noise = np.random.default_rng(seed).standard_normal(pixels.shape)
+    alpha_bars = np.cumprod(1 - np.linspace(0.0001, 0.02, 1000))
+    frames = []
+    final_noise = None
+    for step in checkpoints:
+        alpha = 1 if step == 0 else alpha_bars[step - 1]
+        sample = np.sqrt(alpha) * pixels + np.sqrt(1 - alpha) * noise
+        frame = Image.fromarray(np.round((sample.clip(-1, 1) + 1) * 127.5).astype(np.uint8))
+        frames.append((frame, f"t = {step}"))
+        final_noise = frame
+    return frames, final_noise
+
+
+@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.")
+    frames = []
+    for step, frame in generate_images(seed, checkpoints, initial_image):
+        if frame is not None:
+            frames.append((frame, f"t = {step}"))
+        if frame is not None or step % 25 == 0:
+            yield list(frames), f"Denoising: {1000 - step}/1000 steps · seed {seed}"
+
+
+with gr.Blocks(title="Diffusion Lab") as demo:
+    gr.Markdown("# Diffusion Lab\nExplore forward noising and the 1,000-step DDPM reverse process.")
+    with gr.Row():
+        seed = gr.Textbox(value="42", label="Seed")
+        randomize = gr.Button("Randomize seed")
+        checkpoints = gr.Textbox(value=", ".join(map(str, CHECKPOINTS)), label="Checkpoint timesteps")
+    randomize.click(lambda: str(secrets.randbelow(2**31)), outputs=seed, queue=False)
+    default_image = ROOT / "assets/celeba_hq_validation_91_256.png"
+    source = gr.Image(value=str(default_image) if default_image.exists() else None,
+                      type="pil", label="Source image")
+    diffuse_button = gr.Button("Diffuse")
+    forward_gallery = gr.Gallery(label="Forward process: image to noise", columns=4)
+    forward_noise = gr.State(None)
+    diffuse_button.click(diffuse, [source, seed, checkpoints], [forward_gallery, forward_noise])
+    input_mode = gr.Radio(["Gaussian noise", "Forward noise"], value="Gaussian noise", label="Reverse input")
+    gr.Markdown("The face model is unconditional. Forward noise is an experiment, not a guaranteed reconstruction.")
+    denoise_button = gr.Button("Denoise", variant="primary")
+    reverse_gallery = gr.Gallery(label="Reverse process: noise to face", columns=4)
+    status = gr.Textbox(label="Progress", interactive=False)
+    denoise_button.click(denoise, [seed, checkpoints, input_mode, forward_noise],
+                         [reverse_gallery, status], concurrency_limit=1)
+
+if __name__ == "__main__":
+    demo.queue().launch()

generate.py

@@ -1,4 +1,5 @@
 import argparse
+import os
 import secrets
 import threading
 from pathlib import Path
@@ -16,7 +17,7 @@ _TORCH = None
 _DEVICE = None
 
 
-def load_model():
+def load_model(device=None):
     """Load and cache the diffusion pipeline, returning it and its device."""
     global _PIPELINE, _TORCH, _DEVICE
 
@@ -31,7 +32,9 @@ def load_model():
                 f"Import detail: {error}"
             ) from error
 
-        if torch.cuda.is_available():
+        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"
@@ -39,7 +42,7 @@ def load_model():
             device = "cpu"
 
         print(f"Using device: {device}")
-        pipeline = DDPMPipeline.from_pretrained(MODEL_ID, cache_dir=ROOT / "models")
+        pipeline = DDPMPipeline.from_pretrained(MODEL_ID, cache_dir=Path(os.environ.get("MODEL_CACHE_DIR", ROOT / "models")))
         pipeline.to(device)
         _TORCH = torch
         _DEVICE = device

requirements.txt

@@ -2,4 +2,4 @@ diffusers>=0.30.0
 matplotlib>=3.8.0
 Pillow>=10.0.0
 safetensors>=0.4.0
-torch>=2.1.0
+torch>=2.8.0