Interactive Demo

1import gradio as gr
2import torch
3import numpy as np
4from PIL import Image
5from typing import Optional, Tuple
6import time
7
8class DiffusionDemo:
9    """Interactive demo for diffusion image generation."""
10
11    def __init__(
12        self,
13        model_path: str,
14        device: str = "cuda",
15    ):
16        self.device = device
17        self.model = self._load_model(model_path)
18        self.sampler = None  # Initialize sampler
19
20    def _load_model(self, path: str):
21        """Load trained model from checkpoint."""
22        checkpoint = torch.load(path, map_location=self.device)
23        # Initialize model architecture
24        model = UNet(**checkpoint["config"])
25        model.load_state_dict(checkpoint["model_state_dict"])
26        model = model.to(self.device).eval()
27
28        # Initialize sampler
29        self.sampler = DDIMSampler(
30            model,
31            checkpoint["alphas_cumprod"].to(self.device),
32            device=self.device,
33        )
34
35        return model
36
37    @torch.no_grad()
38    def generate(
39        self,
40        num_images: int = 1,
41        num_steps: int = 50,
42        seed: Optional[int] = None,
43        eta: float = 0.0,
44        progress=gr.Progress(),
45    ) -> list:
46        """Generate images with progress tracking."""
47        if seed is not None:
48            torch.manual_seed(seed)
49
50        shape = (num_images, 3, 64, 64)
51        images = []
52
53        # Generate with progress callback
54        for step in range(num_steps):
55            progress(step / num_steps, desc=f"Step {step}/{num_steps}")
56            time.sleep(0.01)  # Allow UI update
57
58        samples = self.sampler.sample(
59            shape,
60            num_steps=num_steps,
61            eta=eta,
62        )
63
64        # Convert to PIL images
65        for sample in samples:
66            img = (sample.permute(1, 2, 0).cpu().numpy() + 1) / 2
67            img = (img.clip(0, 1) * 255).astype(np.uint8)
68            images.append(Image.fromarray(img))
69
70        return images
71
72    @torch.no_grad()
73    def generate_with_steps(
74        self,
75        seed: int = 42,
76        num_steps: int = 50,
77    ) -> Tuple[Image.Image, list]:
78        """Generate with intermediate steps for visualization."""
79        torch.manual_seed(seed)
80
81        shape = (1, 3, 64, 64)
82        samples, intermediates = self.sampler.sample(
83            shape,
84            num_steps=num_steps,
85            return_intermediates=True,
86        )
87
88        # Convert all to PIL
89        final = self._tensor_to_pil(samples[0])
90        steps = [self._tensor_to_pil(x[0]) for x in intermediates]
91
92        return final, steps
93
94    def _tensor_to_pil(self, tensor: torch.Tensor) -> Image.Image:
95        """Convert tensor to PIL Image."""
96        img = (tensor.permute(1, 2, 0).cpu().numpy() + 1) / 2
97        img = (img.clip(0, 1) * 255).astype(np.uint8)
98        return Image.fromarray(img)
99
100
101def create_gradio_interface(demo: DiffusionDemo) -> gr.Blocks:
102    """Create a full-featured Gradio interface."""
103
104    with gr.Blocks(title="Diffusion Image Generator", theme=gr.themes.Soft()) as interface:
105        gr.Markdown(
106            """
107            # Diffusion Image Generator
108
109            Generate images using a trained denoising diffusion model.
110            Adjust parameters below to control the generation process.
111            """
112        )
113
114        with gr.Tabs():
115            # Tab 1: Basic Generation
116            with gr.Tab("Generate"):
117                with gr.Row():
118                    with gr.Column(scale=1):
119                        num_images = gr.Slider(
120                            minimum=1, maximum=16, value=4, step=1,
121                            label="Number of Images"
122                        )
123                        num_steps = gr.Slider(
124                            minimum=10, maximum=200, value=50, step=10,
125                            label="Sampling Steps"
126                        )
127                        seed = gr.Number(
128                            value=42, precision=0,
129                            label="Random Seed (leave empty for random)"
130                        )
131                        eta = gr.Slider(
132                            minimum=0.0, maximum=1.0, value=0.0, step=0.1,
133                            label="Eta (0=deterministic, 1=stochastic)"
134                        )
135                        generate_btn = gr.Button("Generate", variant="primary")
136
137                    with gr.Column(scale=2):
138                        gallery = gr.Gallery(
139                            label="Generated Images",
140                            columns=4,
141                            height="auto",
142                        )
143
144                generate_btn.click(
145                    fn=demo.generate,
146                    inputs=[num_images, num_steps, seed, eta],
147                    outputs=gallery,
148                )
149
150            # Tab 2: Step-by-Step Visualization
151            with gr.Tab("Step-by-Step"):
152                with gr.Row():
153                    with gr.Column(scale=1):
154                        step_seed = gr.Number(value=42, precision=0, label="Seed")
155                        step_num_steps = gr.Slider(
156                            minimum=10, maximum=100, value=20, step=5,
157                            label="Number of Steps"
158                        )
159                        visualize_btn = gr.Button("Visualize Denoising")
160
161                    with gr.Column(scale=2):
162                        final_image = gr.Image(label="Final Result")
163                        steps_gallery = gr.Gallery(
164                            label="Denoising Steps",
165                            columns=5,
166                        )
167
168                visualize_btn.click(
169                    fn=demo.generate_with_steps,
170                    inputs=[step_seed, step_num_steps],
171                    outputs=[final_image, steps_gallery],
172                )
173
174            # Tab 3: Interpolation
175            with gr.Tab("Interpolation"):
176                with gr.Row():
177                    seed1 = gr.Number(value=42, precision=0, label="Start Seed")
178                    seed2 = gr.Number(value=123, precision=0, label="End Seed")
179                    num_interp = gr.Slider(
180                        minimum=5, maximum=20, value=10, step=1,
181                        label="Interpolation Steps"
182                    )
183
184                interpolate_btn = gr.Button("Generate Interpolation")
185                interp_gallery = gr.Gallery(label="Interpolation", columns=10)
186
187                def interpolate(seed1, seed2, num_interp):
188                    # Implementation would go here
189                    pass
190
191                interpolate_btn.click(
192                    fn=interpolate,
193                    inputs=[seed1, seed2, num_interp],
194                    outputs=interp_gallery,
195                )
196
197        gr.Markdown(
198            """
199            ### About This Demo
200
201            This demo uses a diffusion model trained on CIFAR-10.
202            The model learns to denoise images by predicting the noise
203            added at each timestep.
204
205            **Tips:**
206            - More steps = better quality but slower
207            - Eta=0 gives deterministic results
208            - Same seed = same output (when eta=0)
209            """
210        )
211
212    return interface
213
214
215# Launch the demo
216if __name__ == "__main__":
217    demo = DiffusionDemo("checkpoints/best_model.pt")
218    interface = create_gradio_interface(demo)
219    interface.launch(share=True)

1import streamlit as st
2import torch
3import numpy as np
4from PIL import Image
5import plotly.graph_objects as go
6import plotly.express as px
7from pathlib import Path
8import json
9
10# Configure page
11st.set_page_config(
12    page_title="Diffusion Model Explorer",
13    page_icon="🎨",
14    layout="wide",
15)
16
17@st.cache_resource
18def load_model(model_path: str):
19    """Load model with caching for efficiency."""
20    checkpoint = torch.load(model_path, map_location="cuda")
21    model = UNet(**checkpoint["config"])
22    model.load_state_dict(checkpoint["model_state_dict"])
23    model = model.cuda().eval()
24
25    sampler = DDIMSampler(model, checkpoint["alphas_cumprod"].cuda())
26    return model, sampler, checkpoint
27
28
29def main():
30    st.title("🎨 Diffusion Model Explorer")
31
32    # Sidebar for model selection
33    st.sidebar.header("Model Configuration")
34
35    model_path = st.sidebar.selectbox(
36        "Select Model",
37        options=list(Path("checkpoints").glob("*.pt")),
38        format_func=lambda x: x.stem,
39    )
40
41    if model_path:
42        model, sampler, checkpoint = load_model(str(model_path))
43
44        # Display model info
45        st.sidebar.subheader("Model Info")
46        config = checkpoint.get("config", {})
47        st.sidebar.json(config)
48
49    # Main content tabs
50    tab1, tab2, tab3, tab4 = st.tabs([
51        "Generate", "Noise Schedule", "Training History", "Comparison"
52    ])
53
54    with tab1:
55        st.header("Image Generation")
56
57        col1, col2 = st.columns([1, 2])
58
59        with col1:
60            num_images = st.slider("Number of images", 1, 16, 4)
61            num_steps = st.slider("Sampling steps", 10, 200, 50)
62            seed = st.number_input("Random seed", value=42, step=1)
63            eta = st.slider("Eta (stochasticity)", 0.0, 1.0, 0.0)
64
65            if st.button("Generate", type="primary"):
66                with st.spinner("Generating..."):
67                    torch.manual_seed(int(seed))
68                    shape = (num_images, 3, 64, 64)
69
70                    # Progress bar
71                    progress_bar = st.progress(0)
72                    status_text = st.empty()
73
74                    samples = sampler.sample(shape, num_steps=num_steps, eta=eta)
75
76                    progress_bar.progress(100)
77                    status_text.text("Done!")
78
79                    # Store in session state
80                    st.session_state["generated"] = samples
81
82        with col2:
83            if "generated" in st.session_state:
84                samples = st.session_state["generated"]
85
86                # Display as grid
87                cols = st.columns(4)
88                for i, sample in enumerate(samples):
89                    img = (sample.permute(1, 2, 0).cpu().numpy() + 1) / 2
90                    img = (img.clip(0, 1) * 255).astype(np.uint8)
91                    cols[i % 4].image(Image.fromarray(img), use_container_width=True)
92
93    with tab2:
94        st.header("Noise Schedule Analysis")
95
96        alphas_cumprod = checkpoint["alphas_cumprod"].cpu().numpy()
97        betas = checkpoint.get("betas", 1 - checkpoint["alphas"]).cpu().numpy()
98
99        col1, col2 = st.columns(2)
100
101        with col1:
102            fig = go.Figure()
103            fig.add_trace(go.Scatter(
104                y=alphas_cumprod,
105                mode="lines",
106                name="Alpha Bar (Signal)",
107            ))
108            fig.add_trace(go.Scatter(
109                y=1 - alphas_cumprod,
110                mode="lines",
111                name="1 - Alpha Bar (Noise)",
112            ))
113            fig.update_layout(
114                title="Noise Schedule",
115                xaxis_title="Timestep",
116                yaxis_title="Value",
117            )
118            st.plotly_chart(fig, use_container_width=True)
119
120        with col2:
121            fig = go.Figure()
122            fig.add_trace(go.Scatter(
123                y=betas,
124                mode="lines",
125                name="Beta (Noise Rate)",
126            ))
127            fig.update_layout(
128                title="Beta Schedule",
129                xaxis_title="Timestep",
130                yaxis_title="Beta",
131            )
132            st.plotly_chart(fig, use_container_width=True)
133
134        # SNR visualization
135        st.subheader("Signal-to-Noise Ratio")
136        snr = alphas_cumprod / (1 - alphas_cumprod + 1e-8)
137        log_snr = np.log(snr + 1e-8)
138
139        fig = go.Figure()
140        fig.add_trace(go.Scatter(y=log_snr, mode="lines", name="log(SNR)"))
141        fig.update_layout(
142            title="Log Signal-to-Noise Ratio",
143            xaxis_title="Timestep",
144            yaxis_title="log(SNR)",
145        )
146        st.plotly_chart(fig, use_container_width=True)
147
148    with tab3:
149        st.header("Training History")
150
151        # Load training logs if available
152        log_path = Path(model_path).parent / "training_log.json"
153
154        if log_path.exists():
155            with open(log_path) as f:
156                logs = json.load(f)
157
158            col1, col2 = st.columns(2)
159
160            with col1:
161                if "loss" in logs:
162                    fig = px.line(
163                        x=range(len(logs["loss"])),
164                        y=logs["loss"],
165                        title="Training Loss",
166                        labels={"x": "Step", "y": "Loss"},
167                    )
168                    st.plotly_chart(fig, use_container_width=True)
169
170            with col2:
171                if "fid" in logs:
172                    fig = px.line(
173                        x=range(len(logs["fid"])),
174                        y=logs["fid"],
175                        title="FID Over Training",
176                        labels={"x": "Epoch", "y": "FID"},
177                    )
178                    st.plotly_chart(fig, use_container_width=True)
179        else:
180            st.info("No training logs found for this model.")
181
182    with tab4:
183        st.header("Step Comparison")
184
185        st.write("Compare quality across different sampling steps:")
186
187        seed = st.number_input("Seed for comparison", value=42, key="comp_seed")
188        steps_to_compare = st.multiselect(
189            "Steps to compare",
190            options=[5, 10, 20, 50, 100, 200],
191            default=[10, 50, 100],
192        )
193
194        if st.button("Compare"):
195            cols = st.columns(len(steps_to_compare))
196
197            for i, steps in enumerate(steps_to_compare):
198                torch.manual_seed(int(seed))
199                sample = sampler.sample((1, 3, 64, 64), num_steps=steps)[0]
200                img = (sample.permute(1, 2, 0).cpu().numpy() + 1) / 2
201                img = (img.clip(0, 1) * 255).astype(np.uint8)
202
203                cols[i].image(Image.fromarray(img), caption=f"{steps} steps")
204
205
206if __name__ == "__main__":
207    main()

1from fastapi import FastAPI, HTTPException, BackgroundTasks
2from fastapi.responses import StreamingResponse
3from pydantic import BaseModel, Field
4from typing import Optional, List
5import torch
6import io
7import base64
8from PIL import Image
9import asyncio
10from concurrent.futures import ThreadPoolExecutor
11import uuid
12
13app = FastAPI(
14    title="Diffusion Image Generation API",
15    description="Generate images using a trained diffusion model",
16    version="1.0.0",
17)
18
19# Global model (loaded once at startup)
20model = None
21sampler = None
22executor = ThreadPoolExecutor(max_workers=4)
23
24
25class GenerationRequest(BaseModel):
26    """Request schema for image generation."""
27    num_images: int = Field(1, ge=1, le=16, description="Number of images to generate")
28    num_steps: int = Field(50, ge=10, le=200, description="Number of sampling steps")
29    seed: Optional[int] = Field(None, description="Random seed for reproducibility")
30    eta: float = Field(0.0, ge=0.0, le=1.0, description="Stochasticity parameter")
31    return_format: str = Field("base64", pattern="^(base64|url)$")
32
33
34class GenerationResponse(BaseModel):
35    """Response schema for generated images."""
36    request_id: str
37    images: List[str]
38    generation_time: float
39    parameters: dict
40
41
42@app.on_event("startup")
43async def load_model():
44    """Load model on server startup."""
45    global model, sampler
46
47    checkpoint = torch.load("checkpoints/best_model.pt", map_location="cuda")
48    model = UNet(**checkpoint["config"])
49    model.load_state_dict(checkpoint["model_state_dict"])
50    model = model.cuda().eval()
51
52    sampler = DDIMSampler(model, checkpoint["alphas_cumprod"].cuda())
53    print("Model loaded successfully")
54
55
56def generate_images_sync(
57    num_images: int,
58    num_steps: int,
59    seed: Optional[int],
60    eta: float,
61) -> List[Image.Image]:
62    """Synchronous generation (runs in thread pool)."""
63    if seed is not None:
64        torch.manual_seed(seed)
65
66    shape = (num_images, 3, 64, 64)
67
68    with torch.no_grad():
69        samples = sampler.sample(shape, num_steps=num_steps, eta=eta)
70
71    images = []
72    for sample in samples:
73        img = (sample.permute(1, 2, 0).cpu().numpy() + 1) / 2
74        img = (img.clip(0, 1) * 255).astype("uint8")
75        images.append(Image.fromarray(img))
76
77    return images
78
79
80def image_to_base64(image: Image.Image) -> str:
81    """Convert PIL image to base64 string."""
82    buffer = io.BytesIO()
83    image.save(buffer, format="PNG")
84    return base64.b64encode(buffer.getvalue()).decode()
85
86
87@app.post("/generate", response_model=GenerationResponse)
88async def generate_images(request: GenerationRequest):
89    """Generate images using the diffusion model."""
90    import time
91    start_time = time.time()
92
93    # Run generation in thread pool to avoid blocking
94    loop = asyncio.get_event_loop()
95    images = await loop.run_in_executor(
96        executor,
97        generate_images_sync,
98        request.num_images,
99        request.num_steps,
100        request.seed,
101        request.eta,
102    )
103
104    # Convert to base64
105    image_data = [image_to_base64(img) for img in images]
106
107    generation_time = time.time() - start_time
108
109    return GenerationResponse(
110        request_id=str(uuid.uuid4()),
111        images=image_data,
112        generation_time=generation_time,
113        parameters={
114            "num_images": request.num_images,
115            "num_steps": request.num_steps,
116            "seed": request.seed,
117            "eta": request.eta,
118        },
119    )
120
121
122@app.get("/generate/stream")
123async def generate_stream(
124    seed: int = 42,
125    num_steps: int = 50,
126):
127    """Stream generation progress and final image."""
128    async def generate():
129        torch.manual_seed(seed)
130        shape = (1, 3, 64, 64)
131
132        # Start generation
133        yield f"data: {{\"status\": \"starting\", \"step\": 0}}\n\n"
134
135        with torch.no_grad():
136            x = torch.randn(shape, device="cuda")
137
138            for step in range(num_steps):
139                # Denoise step
140                t = torch.full((1,), num_steps - 1 - step, device="cuda", dtype=torch.long)
141                x = sampler.denoise_step(x, t)
142
143                # Send progress
144                progress = (step + 1) / num_steps * 100
145                yield f"data: {{\"status\": \"generating\", \"step\": {step + 1}, \"progress\": {progress:.1f}}}\n\n"
146
147                await asyncio.sleep(0.01)  # Allow other requests
148
149            # Convert final image
150            img = (x[0].permute(1, 2, 0).cpu().numpy() + 1) / 2
151            img = (img.clip(0, 1) * 255).astype("uint8")
152            pil_img = Image.fromarray(img)
153            b64 = image_to_base64(pil_img)
154
155            yield f"data: {{\"status\": \"complete\", \"image\": \"{b64}\"}}\n\n"
156
157    return StreamingResponse(
158        generate(),
159        media_type="text/event-stream",
160    )
161
162
163@app.get("/health")
164async def health_check():
165    """Health check endpoint."""
166    return {
167        "status": "healthy",
168        "model_loaded": model is not None,
169        "device": str(next(model.parameters()).device) if model else None,
170    }
171
172
173# Docker deployment
174# Dockerfile:
175# FROM python:3.10-slim
176# WORKDIR /app
177# COPY requirements.txt .
178# RUN pip install -r requirements.txt
179# COPY . .
180# EXPOSE 8000
181# CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]

1diffusion-project/
2├── src/
3│   ├── models/
4│   │   ├── __init__.py
5│   │   ├── unet.py              # U-Net architecture
6│   │   ├── attention.py         # Attention blocks
7│   │   ├── embedding.py         # Time embeddings
8│   │   └── blocks.py            # ResBlocks, Up/Downsample
9│   │
10│   ├── diffusion/
11│   │   ├── __init__.py
12│   │   ├── gaussian.py          # GaussianDiffusion class
13│   │   ├── schedules.py         # Noise schedules
14│   │   ├── ddpm.py              # DDPM sampler
15│   │   └── ddim.py              # DDIM sampler
16│   │
17│   ├── data/
18│   │   ├── __init__.py
19│   │   ├── datasets.py          # Dataset classes
20│   │   ├── transforms.py        # Augmentations
21│   │   └── dataloader.py        # DataLoader setup
22│   │
23│   ├── training/
24│   │   ├── __init__.py
25│   │   ├── trainer.py           # Main training loop
26│   │   ├── optimizer.py         # Optimizer config
27│   │   ├── scheduler.py         # LR scheduling
28│   │   └── ema.py               # EMA implementation
29│   │
30│   ├── evaluation/
31│   │   ├── __init__.py
32│   │   ├── fid.py               # FID calculation
33│   │   ├── inception_score.py   # IS calculation
34│   │   ├── precision_recall.py  # PR metrics
35│   │   └── qualitative.py       # Visual analysis
36│   │
37│   └── utils/
38│       ├── __init__.py
39│       ├── logging.py           # W&B, TensorBoard
40│       ├── checkpoints.py       # Save/load models
41│       └── visualization.py     # Plotting utilities
42│
43├── scripts/
44│   ├── train.py                 # Main training script
45│   ├── sample.py                # Generation script
46│   ├── evaluate.py              # Evaluation script
47│   └── demo.py                  # Gradio demo
48│
49├── configs/
50│   ├── cifar10.yaml             # CIFAR-10 config
51│   ├── celeba.yaml              # CelebA config
52│   └── imagenet64.yaml          # ImageNet 64x64 config
53│
54├── checkpoints/                 # Saved models
55├── logs/                        # Training logs
56├── outputs/                     # Generated samples
57│
58├── requirements.txt
59├── setup.py
60├── Dockerfile
61└── README.md

1# 1. Install dependencies
2pip install -r requirements.txt
3
4# 2. Train the model
5python scripts/train.py --config configs/cifar10.yaml
6
7# 3. Monitor training
8tensorboard --logdir logs/
9
10# 4. Generate samples
11python scripts/sample.py \
12    --checkpoint checkpoints/best_model.pt \
13    --num_samples 1000 \
14    --output outputs/samples/
15
16# 5. Evaluate the model
17python scripts/evaluate.py \
18    --checkpoint checkpoints/best_model.pt \
19    --real_stats data/cifar10_stats.npz
20
21# 6. Launch interactive demo
22python scripts/demo.py --share