new frontend

.gitignore

@@ -0,0 +1 @@
+.DS_Store

app.py

@@ -17,8 +17,8 @@ from vq_vae import VQVAE
 
 app = Flask(__name__, template_folder="templates")
 
-
-device = "mps" if torch.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu"  # don't know why but mps runs slower than cpu on m2 mac
+# Detect device — adjust if you want to force 'mps' or 'cuda'
+device = 'cpu'  # don't know why but mps runs slower than cpu
 
 # ------------------------------------------------------------------------------
 # Model Status Tracking
@@ -28,14 +28,13 @@ available_models = {
     "moe": False,       # MoEPixelTransformer (mixture of experts)
     "conv": False,      # ConvGenerator (direct generation)
     "vq": False,        # VQTransformer (autoregressive by token)
-    "vq-vae": False     # VQ-VAE only (encode/decode)
+    "vq-vae": False,    # VQ-VAE only (encode/decode)
+    "diffusion": False  # Diffusion model (DDPM)
 }
 
 # ------------------------------------------------------------------------------
 # Load models
 # ------------------------------------------------------------------------------
-
-
 # 1. Load ConvGenerator
 try:
     conv_config = ConvConfig.from_pretrained("my_conv")
@@ -93,6 +92,22 @@ try:
 except Exception as e:
     print(f"✗ Error loading VQ models: {str(e)}")
 
+# 5. Load Diffusion pipeline if available
+diffusion_pipe = None
+try:
+    from diffusers import DDPMPipeline
+    diffusion_model_dir = "my_diffusion_model"
+    if os.path.exists(diffusion_model_dir):
+        diffusion_pipe = DDPMPipeline.from_pretrained(
+            diffusion_model_dir, torch_dtype=torch.float32
+        ).to(device)
+        available_models["diffusion"] = True
+        print("✓ Diffusion pipeline loaded successfully")
+    else:
+        print(f"✗ Diffusion model directory '{diffusion_model_dir}' not found, skipping diffusion.")
+except Exception as e:
+    diffusion_pipe = None
+    print(f"✗ Error loading Diffusion pipeline: {str(e)}")
 # Select default model (use the first available one)
 for model_name, is_available in available_models.items():
     if is_available:
@@ -254,6 +269,52 @@ def generate_vq_vae_digit():
         print(f"Error in VQ-VAE reconstruction: {str(e)}")
         return Response(str(e), status=500)
 
+
+# ------------------------------------------------------------------------------
+# DIFFUSION GENERATION (using DDPM pipeline)
+# ------------------------------------------------------------------------------
+@app.route("/generate_diffusion_digit", methods=["GET"])
+def generate_diffusion_digit():
+    """Generate image using diffusion model (DDPM)."""
+    if diffusion_pipe is None:
+        return Response("Diffusion model not loaded", status=500)
+    try:
+        digit = int(request.args.get("digit", 0))
+        steps = int(request.args.get("steps", 50))
+        print(f"Generating diffusion image for digit {digit} with {steps} steps...")
+        num_steps = steps
+        scheduler = diffusion_pipe.scheduler
+        scheduler.set_timesteps(num_steps)
+
+        img = torch.randn(
+            (
+                1,
+                diffusion_pipe.unet.config.in_channels,
+                diffusion_pipe.unet.config.sample_size,
+                diffusion_pipe.unet.config.sample_size,
+            ),
+            device=device,
+            dtype=torch.float32,
+        )
+        labels = torch.tensor([digit], device=device)
+
+        for t in scheduler.timesteps:
+            with torch.no_grad():
+                model_output = diffusion_pipe.unet(img, t, class_labels=labels).sample
+            img = scheduler.step(model_output, t, img).prev_sample
+
+        img = (img / 2 + 0.5).clamp(0, 1)
+        array = img.cpu().permute(0, 2, 3, 1).numpy()[0]
+        array = (array * 255).astype(np.uint8)
+        image = Image.fromarray(array.squeeze(), mode="L").resize((28, 28))
+        buf = BytesIO()
+        image.save(buf, format="PNG")
+        buf.seek(0)
+        return Response(buf.getvalue(), mimetype="image/png")
+    except Exception as e:
+        print(f"Error generating diffusion image: {str(e)}")
+        return Response(str(e), status=500)
+
 # ------------------------------------------------------------------------------
 # STREAM DIGIT (Pixel-by-pixel generation or token-by-token for VQ)
 # ------------------------------------------------------------------------------
@@ -292,28 +353,27 @@ def stream_digit():
                             time.sleep(0.005)
                             
                     elif model_name == "vq" and vq_transformer_model and vq_model:
-                        # VQ-Transformer (token by token, then decode)
+                        # VQ-Transformer (token by token, with streaming decode)
                         generator = vq_transformer_model.generate_token_stream(digit, device)
                         tokens = []
-                        
-                        # Stream token generation progress
+
+                        # Stream token generation progress and partial image patches
                         for i, token in enumerate(generator):
                             tokens.append(token)
-                            progress = int((i+1) * 100 / 49)  # 49 tokens total
+                            progress = int((i + 1) * 100 / 49)
                             yield f"data: token:{i+1}:{progress}\n\n"
                             time.sleep(0.01)
-                        
-                        # Then decode tokens to image
-                        if len(tokens) == 49:  # Make sure we have all tokens
-                            token_tensor = torch.tensor(tokens, dtype=torch.long, device=device).reshape(1, 7, 7)
+
+                            # Partial decode: pad remaining tokens with zero index
+                            pad_tokens = tokens + [0] * (49 - len(tokens))
+                            token_tensor = torch.tensor(pad_tokens, dtype=torch.long, device=device).reshape(1, 7, 7)
                             decoded_img = vq_model.decode(token_tensor)
                             img_array = (decoded_img.cpu().squeeze().numpy() * 255).astype(np.uint8)
-                            
-                            # Stream the final image pixels
-                            flattened_pixels = img_array.flatten()
-                            for pixel in flattened_pixels:
-                                yield f"data: {int(pixel)}\n\n"
-                                time.sleep(0.001)
+
+                            # Stream full frame as CSV
+                            flat_pixels = img_array.flatten().tolist()
+                            yield f"data: frame:{','.join(str(int(p)) for p in flat_pixels)}\n\n"
+                            time.sleep(0.001)
                     else:
                         yield "data: Error: Invalid model selected or model not available.\n\n"
 

dataset.py

@@ -12,6 +12,7 @@ class ConditionalMNISTDataset(Dataset):
         super().__init__()
         self.label_offset = label_offset
         
+        # Load MNIST from torchvision
         transform = transforms.ToTensor()
         self.data = datasets.MNIST(root="./data", train=(split=="train"), 
                                  download=True, transform=transform)

templates/index.html

@@ -1,104 +1,68 @@
 <!DOCTYPE html>
-<html>
+<html lang="en">
 <head>
-  <meta charset="utf-8" />
-  <title>Conditional MNIST Generation (Pixel-by-Pixel)</title>
-  <style>
-    body {
-      font-family: sans-serif;
-      margin: 20px;
-    }
-    #canvas {
-      width: 280px;   /* 10x zoom for 28px images */
-      height: 280px;
-      border: 1px solid #ccc;
-      image-rendering: pixelated; /* keep blocky pixels */
-      display: block;
-      margin-top: 10px;
-      background: #fff;
-    }
-    #log {
-      margin-top: 10px;
-      white-space: pre-wrap;
-      font-size: 14px;
-      color: #666;
-    }
-    .error {
-      color: #ff0000;
-    }
-    button:disabled {
-      opacity: 0.5;
-      cursor: not-allowed;
-    }
-    .progress-bar {
-      height: 20px;
-      background-color: #f0f0f0;
-      border-radius: 5px;
-      margin-top: 10px;
-      display: none;
-    }
-    .progress-fill {
-      height: 100%;
-      background-color: #4CAF50;
-      border-radius: 5px;
-      width: 0%;
-      transition: width 0.1s;
-    }
-  </style>
+  <meta charset="UTF-8">
+  <meta name="viewport" content="width=device-width, initial-scale=1.0">
+  <title>Image Generator</title>
+  <script src="https://cdn.tailwindcss.com"></script>
 </head>
-<body>
-  <h2>Conditional MNIST Generator</h2>
-  <p>Enter a digit (0-9) to generate:</p>
-  <input type="number" id="digitInput" value="7" min="0" max="9" style="width: 60px;">
-  <button id="generateBtn" onclick="generateDigit()">Generate</button>
-
-  <select id="modelSelector" onchange="selectModel()">
-    {% if available_models.get('pixel', False) %}
-    <option value="pixel" {% if selected_model == 'pixel' %}selected{% endif %}>PixelTransformer</option>
-    {% endif %}
-    {% if available_models.get('moe', False) %}
-    <option value="moe" {% if selected_model == 'moe' %}selected{% endif %}>MoEPixelTransformer</option>
-    {% endif %}
-    {% if available_models.get('conv', False) %}
-    <option value="conv" {% if selected_model == 'conv' %}selected{% endif %}>ConvGenerator</option>
-    {% endif %}
-    {% if available_models.get('vq', False) %}
-    <option value="vq" {% if selected_model == 'vq' %}selected{% endif %}>VQ-Transformer</option>
-    {% endif %}
-    {% if available_models.get('vq-vae', False) %}
-    <option value="vq-vae" {% if selected_model == 'vq-vae' %}selected{% endif %}>VQ-VAE Only</option>
-    {% endif %}
-  </select>
-
-  <canvas id="canvas" width="28" height="28"></canvas>
-  <div id="progress-container" class="progress-bar">
-    <div id="progress-fill" class="progress-fill"></div>
+<body class="bg-black text-white">
+  <div class="flex flex-col items-center justify-center min-h-screen space-y-6 p-4">
+    <h1 class="text-3xl font-semibold">Image Generator</h1>
+    <div class="flex flex-wrap items-center justify-center space-x-2">
+      <input id="digitInput" type="number" min="0" max="9" value="7"
+             class="w-16 px-2 py-1 bg-gray-800 border border-gray-600 rounded text-white focus:outline-none"/>
+      <input id="stepsInput" type="number" min="1" max="1000" value="50" placeholder="steps"
+             class="hidden w-20 px-2 py-1 bg-gray-800 border border-gray-600 rounded text-white focus:outline-none"/>
+      <button id="generateBtn" onclick="generateDigit()"
+              class="px-4 py-1 bg-gray-700 hover:bg-gray-600 rounded disabled:opacity-50">
+        Generate
+      </button>
+      <select id="modelSelector" onchange="selectModel()"
+              class="px-2 py-1 bg-gray-800 border border-gray-600 rounded text-white focus:outline-none">
+        {% for name, available in available_models.items() %}
+          {% if available %}
+            <option value="{{ name }}" {% if selected_model == name %}selected{% endif %}>
+              {{ name|capitalize }}
+            </option>
+          {% endif %}
+        {% endfor %}
+      </select>
+    </div>
+    <canvas id="canvas" width="28" height="28"
+            class="w-[280px] h-[280px] border border-gray-600 bg-black"
+            style="image-rendering: pixelated;"></canvas>
+    <div id="progress-container" class="w-[280px] bg-gray-800 rounded overflow-hidden"
+         style="display: none;">
+      <div id="progress-fill" class="bg-white h-1 w-0 transition-all"></div>
+    </div>
+    <div id="log" class="text-sm text-gray-400"></div>
   </div>
-  <div id="log"></div>
-
   <script>
     let currentModel = '{{ selected_model }}';
     let currentEventSource = null;
     let isGenerating = false;
-    let pixelCounter = 0;  // Track pixels for VQ model rendering
+    let pixelCounter = 0;
 
     function selectModel() {
       const modelSelector = document.getElementById('modelSelector');
       currentModel = modelSelector.value;
-      
-      // Update model selection on server
       fetch('/select_model', {
         method: 'POST',
-        headers: {
-          'Content-Type': 'application/json',
-        },
-        body: JSON.stringify({ model_type: currentModel })
+        headers: {'Content-Type': 'application/json'},
+        body: JSON.stringify({model_type: currentModel})
       });
-      
-      // Show/hide progress bar if VQ model
-      document.getElementById('progress-container').style.display = 
+      document.getElementById('progress-container').style.display =
         (currentModel === 'vq' || currentModel === 'vq-vae') ? 'block' : 'none';
+      const stepsInput = document.getElementById('stepsInput');
+      if (currentModel === 'diffusion') {
+        stepsInput.classList.remove('hidden');
+      } else {
+        stepsInput.classList.add('hidden');
+      }
     }
+    // initialize UI based on default selected model
+    selectModel();
 
     function setGenerating(generating) {
       isGenerating = generating;
@@ -109,35 +73,27 @@
     function generateDigit() {
       if (isGenerating) return;
       setGenerating(true);
-
-      // Clean up any existing EventSource
       if (currentEventSource) {
         currentEventSource.close();
         currentEventSource = null;
       }
-
       const digit = document.getElementById('digitInput').value;
       const canvas = document.getElementById('canvas');
       const ctx = canvas.getContext('2d');
       const log = document.getElementById('log');
       const progressBar = document.getElementById('progress-fill');
-      pixelCounter = 0;  // Reset pixel counter
-
-      // Clear previous content
-      ctx.fillStyle = 'white';
+      pixelCounter = 0;
+      ctx.fillStyle = 'black';
       ctx.fillRect(0, 0, canvas.width, canvas.height);
       log.textContent = 'Generating...';
       log.className = '';
       progressBar.style.width = '0%';
 
       if (currentModel === 'conv') {
-        // For ConvGenerator (instant generation)
         fetch(`/generate_conv_digit?digit=${digit}`)
           .then(response => {
             if (!response.ok) {
-              return response.text().then(text => {
-                throw new Error(text || `HTTP error! status: ${response.status}`);
-              });
+              return response.text().then(text => { throw new Error(text || `HTTP error! status: ${response.status}`); });
             }
             return response.blob();
           })
@@ -148,34 +104,51 @@
               log.textContent = 'Generated!';
               setGenerating(false);
             };
-            img.onerror = () => {
-              throw new Error('Failed to load generated image');
+            img.onerror = () => { throw new Error('Failed to load generated image'); };
+            img.src = URL.createObjectURL(blob);
+          })
+          .catch(error => {
+            console.error('Error:', error);
+            log.textContent = `Error generating image: ${error.message}`;
+            log.className = 'text-red-500';
+            setGenerating(false);
+          });
+      } else if (currentModel === 'diffusion') {
+        const steps = document.getElementById('stepsInput').value;
+        fetch(`/generate_diffusion_digit?digit=${digit}&steps=${steps}`)
+          .then(response => {
+            if (!response.ok) {
+              return response.text().then(text => { throw new Error(text || `HTTP error! status: ${response.status}`); });
+            }
+            return response.blob();
+          })
+          .then(blob => {
+            const img = new Image();
+            img.onload = () => {
+              ctx.drawImage(img, 0, 0);
+              log.textContent = 'Generated!';
+              setGenerating(false);
             };
+            img.onerror = () => { throw new Error('Failed to load generated image'); };
             img.src = URL.createObjectURL(blob);
           })
           .catch(error => {
             console.error('Error:', error);
             log.textContent = `Error generating image: ${error.message}`;
-            log.className = 'error';
+            log.className = 'text-red-500';
             setGenerating(false);
           });
       } else if (currentModel === 'vq' || currentModel === 'vq-vae') {
-        // Special handling for VQ models
         const imageData = ctx.createImageData(28, 28);
-        
-        // Use a specific endpoint for vq-vae direct reconstruction
-        const endpoint = currentModel === 'vq-vae' ? 
-          `/generate_vq_vae_digit?digit=${digit}` : 
-          `/stream_digit?digit=${digit}`;
-        
+        const endpoint = currentModel === 'vq-vae'
+          ? `/generate_vq_vae_digit?digit=${digit}`
+          : `/stream_digit?digit=${digit}`;
+
         if (currentModel === 'vq-vae') {
-          // For VQ-VAE direct reconstruction (non-streamed)
           fetch(endpoint)
             .then(response => {
               if (!response.ok) {
-                return response.text().then(text => {
-                  throw new Error(text || `HTTP error! status: ${response.status}`);
-                });
+                return response.text().then(text => { throw new Error(text || `HTTP error! status: ${response.status}`); });
               }
               return response.blob();
             })
@@ -186,135 +159,102 @@
                 log.textContent = 'Generated!';
                 setGenerating(false);
               };
-              img.onerror = () => {
-                throw new Error('Failed to load generated image');
-              };
+              img.onerror = () => { throw new Error('Failed to load generated image'); };
               img.src = URL.createObjectURL(blob);
             })
             .catch(error => {
               console.error('Error:', error);
               log.textContent = `Error generating image: ${error.message}`;
-              log.className = 'error';
+              log.className = 'text-red-500';
               setGenerating(false);
             });
         } else {
-          // For VQ-Transformer (streamed)
           currentEventSource = new EventSource(endpoint);
-          
           currentEventSource.onmessage = function(event) {
             const data = event.data;
-            
             if (data.startsWith('Error:')) {
               log.textContent = data;
-              log.className = 'error';
+              log.className = 'text-red-500';
               currentEventSource.close();
               setGenerating(false);
               return;
             }
-            
-            // Check if it's a token progress update
             if (data.startsWith('token:')) {
-              const parts = data.split(':');
-              const tokenNum = parseInt(parts[1]);
-              const progress = parseInt(parts[2]);
-              
-              // Update progress bar
+              const [, tokenNum, progress] = data.split(':');
               progressBar.style.width = `${progress}%`;
               log.textContent = `Generating tokens: ${tokenNum}/49 (${progress}%)`;
               return;
             }
-            
-            // Otherwise it's a pixel value
-            const pixelValue = parseInt(data);
-            if (isNaN(pixelValue)) {
-              console.error('Invalid pixel value:', data);
+            if (data.startsWith('frame:')) {
+              const pixels = data.slice(6).split(',').map(Number);
+              for (let idx = 0; idx < pixels.length; idx++) {
+                const x = idx % 28;
+                const y = Math.floor(idx / 28);
+                const i = (y * 28 + x) * 4;
+                imageData.data[i] = pixels[idx];
+                imageData.data[i + 1] = pixels[idx];
+                imageData.data[i + 2] = pixels[idx];
+                imageData.data[i + 3] = 255;
+              }
+              ctx.putImageData(imageData, 0, 0);
               return;
             }
-            
-            // Calculate pixel position 
+            const pixelValue = parseInt(data);
+            if (isNaN(pixelValue)) return;
             const x = pixelCounter % 28;
             const y = Math.floor(pixelCounter / 28);
-            
-            // Set RGB values for this pixel
             const idx = (y * 28 + x) * 4;
-            imageData.data[idx] = pixelValue;     // R
-            imageData.data[idx + 1] = pixelValue; // G
-            imageData.data[idx + 2] = pixelValue; // B
-            imageData.data[idx + 3] = 255;        // A (opacity)
-            
+            imageData.data[idx] = pixelValue;
+            imageData.data[idx + 1] = pixelValue;
+            imageData.data[idx + 2] = pixelValue;
+            imageData.data[idx + 3] = 255;
             pixelCounter++;
-            
-            // Update canvas every 28 pixels (full row)
-            if (x === 27 || pixelCounter === 28*28) {
+            if (x === 27 || pixelCounter === 28 * 28) {
               ctx.putImageData(imageData, 0, 0);
-              
-              if (pixelCounter >= 28*28) {
+              if (pixelCounter >= 28 * 28) {
                 currentEventSource.close();
                 log.textContent = 'Generation complete!';
                 setGenerating(false);
               }
             }
           };
-          
-          currentEventSource.onerror = function(error) {
-            console.error('EventSource error:', error);
+          currentEventSource.onerror = function(e) {
             currentEventSource.close();
-            log.textContent = 'Error in streaming!';
-            log.className = 'error';
             setGenerating(false);
           };
         }
       } else {
-        // For PixelTransformer & MoEPixelTransformer (pixel streaming)
         const imageData = ctx.createImageData(28, 28);
         let index = 0;
-
         currentEventSource = new EventSource(`/stream_digit?digit=${digit}`);
-        
         currentEventSource.onmessage = function(event) {
           const data = event.data;
           if (data.startsWith('Error:')) {
             log.textContent = data;
-            log.className = 'error';
+            log.className = 'text-red-500';
             currentEventSource.close();
-            currentEventSource = null;
             setGenerating(false);
             return;
           }
-
           const pixelValue = parseInt(data);
-          if (isNaN(pixelValue)) {
-            console.error('Invalid pixel value:', data);
-            return;
-          }
-          
-          // Set RGB values to the same value for grayscale
-          imageData.data[index] = pixelValue;     // R
-          imageData.data[index + 1] = pixelValue; // G
-          imageData.data[index + 2] = pixelValue; // B
-          imageData.data[index + 3] = 255;        // A (opacity)
+          if (isNaN(pixelValue)) return;
+          imageData.data[index] = pixelValue;
+          imageData.data[index + 1] = pixelValue;
+          imageData.data[index + 2] = pixelValue;
+          imageData.data[index + 3] = 255;
           index += 4;
-
-          // Update canvas every row (28 pixels)
           if (index % (28 * 4) === 0) {
             ctx.putImageData(imageData, 0, 0);
           }
-          
           if (index >= 28 * 28 * 4) {
             currentEventSource.close();
-            currentEventSource = null;
             log.textContent = 'Generation complete!';
             setGenerating(false);
           }
         };
-
-        currentEventSource.onerror = function(error) {
-          console.error('EventSource error:', error);
-          currentEventSource.close();
-          currentEventSource = null;
-          log.textContent = 'Error in streaming!';
-          log.className = 'error';
-          setGenerating(false);
+        currentEventSource.onerror = function() {
+        currentEventSource.close();
+        setGenerating(false);
         };
       }
     }

train_conv.py

@@ -5,12 +5,18 @@ from torch.utils.data import DataLoader
 from transformers import PreTrainedModel, PretrainedConfig
 from dataset import ConditionalMNISTDataset
 
+############################
+#      Config Class        #
+############################
 class ConvConfig(PretrainedConfig):
     model_type = "conv_generator"
     def __init__(self, latent_dim=100, **kwargs):
         super().__init__(**kwargs)
         self.latent_dim = latent_dim
 
+############################
+#      Model Class         #
+############################
 class ConvGeneratorModel(PreTrainedModel):
     config_class = ConvConfig
     def __init__(self, config):
@@ -51,11 +57,16 @@ class ConvGeneratorModel(PreTrainedModel):
         return out
 
 def main():
-    device = torch.device("mps" if torch.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu")
+    # Ensure MPS is available
+    if not torch.backends.mps.is_available():
+        print("MPS not available. Falling back to CPU.")
+        device = "cpu"
+    else:
+        device = "mps"
     print(f"Using device: {device}")
 
     dataset = ConditionalMNISTDataset("train")
-    loader = DataLoader(dataset, batch_size=64, shuffle=True, num_workers=0)
+    loader = DataLoader(dataset, batch_size=64, shuffle=True, num_workers=0)  # Reduced workers for MPS
 
     config = ConvConfig(latent_dim=100)
     model = ConvGeneratorModel(config).to(device)
@@ -67,6 +78,7 @@ def main():
             model.train()
             total_loss = 0
             for step, (x, y) in enumerate(loader, 1):
+                # Move both inputs to device
                 x = x.to(device)
                 y = y.to(device)
                 
@@ -76,7 +88,7 @@ def main():
                 optimizer.zero_grad()
                 generated_images = model(labels)                      # (bsz, 1, 28, 28)
 
-                # Mean Squared Error loss is typical for image generation
+                # Use Mean Squared Error loss
                 loss = F.mse_loss(generated_images, real_images)
                 loss.backward()
                 optimizer.step()

train_diff.py

@@ -0,0 +1,67 @@
+import os
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from torchvision import transforms, datasets
+from diffusers import UNet2DModel, DDPMScheduler, DDPMPipeline
+from tqdm import tqdm
+
+def train_diffusion():
+    # Train and save a DDPM diffusion model on MNIST.
+    device = "mps" if torch.backends.mps.is_available() else "cpu"
+    print(f"Using device: {device}")
+
+    transform = transforms.Compose([transforms.ToTensor()])
+    train_ds = datasets.MNIST(root="./data", train=True, download=True, transform=transform)
+    loader = DataLoader(train_ds, batch_size=128, shuffle=True)
+
+    # Conditional DDPM UNet for MNIST digits
+    unet = UNet2DModel(
+        sample_size=28,
+        in_channels=1,
+        out_channels=1,
+        block_out_channels=(32, 64, 128),
+        down_block_types=("DownBlock2D", "AttnDownBlock2D", "DownBlock2D"),
+        up_block_types=("UpBlock2D", "AttnUpBlock2D", "UpBlock2D"),
+        num_class_embeds=10,
+    ).to(device)
+    scheduler = DDPMScheduler(num_train_timesteps=1000)
+    pipeline = DDPMPipeline(unet=unet, scheduler=scheduler).to(device)
+    optimizer = torch.optim.AdamW(unet.parameters(), lr=1e-4, weight_decay=1e-4) # changed from Adam
+
+    epochs = 5
+    print(f"Training DDPM for {epochs} epochs...")
+    try:
+        for epoch in range(1, epochs + 1):
+            pbar = tqdm(loader, desc=f"Epoch {epoch}/{epochs}")
+            for images, labels in pbar:
+                images = images.to(device)
+                labels = labels.to(device)
+                noise = torch.randn_like(images)
+                timesteps = torch.randint(
+                    0, scheduler.num_train_timesteps, (images.shape[0],), device=device
+                ).long()
+                noisy = scheduler.add_noise(images, noise, timesteps)
+
+                # Conditional noise prediction
+                model_pred = unet(noisy, timesteps, class_labels=labels, return_dict=False)[0]
+                loss = F.mse_loss(model_pred, noise)
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+                pbar.set_postfix(loss=f"{loss.item():.4f}")
+    except KeyboardInterrupt:
+        print("\nKeyboard interrupt, saving model...")
+        output_dir = "my_diffusion_model"
+        pipeline.save_pretrained(output_dir)
+        print(f"Model saved to {output_dir}/")
+        return pipeline
+
+    output_dir = "my_diffusion_model"
+    pipeline.save_pretrained(output_dir)
+    print(f"Training complete. Model saved to {output_dir}/")
+    return pipeline
+
+if __name__ == "__main__":
+    train_diffusion()

vq_transformer.py

@@ -19,7 +19,7 @@ class VQTransformerConfig:
     epochs: int = 10
     warmup_steps: int = 500
     label_offset: int = 512  # Labels are tokens 512-521 (for digits 0-9)
-    device: str = field(default="mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu")
+    device: str = field(default="mps" if torch.backends.mps.is_available() else "cpu")
     
     @classmethod
     def from_pretrained(cls, path: str):

vq_vae.py

@@ -64,6 +64,7 @@ class VQVAE(nn.Module):
             nn.Conv2d(32, embedding_dim, 1, stride=1)  # 7x7 -> 7x7xembedding_dim
         )
         
+        # Vector Quantizer
         self.vq = VectorQuantizer(num_embeddings, embedding_dim)
         
         # Decoder for MNIST (7x7 -> 28x28)
@@ -94,6 +95,7 @@ class VQVAE(nn.Module):
         quantized = torch.matmul(one_hot.permute(0, 2, 3, 1), self.vq.embedding.weight)
         quantized = quantized.permute(0, 3, 1, 2)
         
+        # Decode
         reconstructed = self.decoder(quantized)
         return reconstructed
         
@@ -137,14 +139,18 @@ class VQVAE(nn.Module):
     
     @staticmethod
     def train_and_save(output_path="vq_vae_model.pt", device='cpu', batch_size=128, epochs=10):
+        # Setup data
         transform = transforms.Compose([transforms.ToTensor()])
         train_dataset = datasets.MNIST(root="./data", train=True, download=True, transform=transform)
         train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
-
+        
+        # Create model
         model = VQVAE().to(device)
-
+        
+        # Train
         model.train_model(train_loader, epochs=epochs, device=device)
         
+        # Save model
         torch.save(model.state_dict(), output_path)
         print(f"Model saved to {output_path}")