Update app.py

app.py

@@ -146,14 +146,14 @@ class Diffusion:
         self.alpha_bars = torch.cumprod(self.alphas, dim=0)
 
     @torch.no_grad()
-    def sample(self, model, text_tokens, image_size=64, steps=None):
+    def sample(self, model, text_tokens, image_size=64, steps=None, progress_callback=None):
         model.eval()
         if steps is None:
             steps = self.timesteps
         
         x = torch.randn(1, 3, image_size, image_size).to(self.device)
         
-        for t in reversed(range(steps)):
+        for i, t in enumerate(reversed(range(steps))):
             t_batch = torch.full((x.shape[0],), t, device=self.device, dtype=torch.long)
             predicted_noise = model(x, t_batch, text_tokens)
             
@@ -169,6 +169,11 @@ class Diffusion:
             x = (1 / torch.sqrt(alpha)) * (x - ((1 - alpha) / torch.sqrt(1 - alpha_bar)) * predicted_noise)
             x = x + torch.sqrt(beta) * noise
             
+            # Report progress
+            if progress_callback is not None:
+                progress = (i + 1) / steps
+                progress_callback(progress)
+            
         model.train()
         return x
 
@@ -232,18 +237,31 @@ def tokenize_text(text, max_len=20):
         indices.append(0)  # PAD token
     return torch.tensor(indices).unsqueeze(0).to(device)
 
-# Generate image
-def generate_image(prompt):
+# Generate image with progress
+def generate_image(prompt, progress=gr.Progress()):
     global model, device, vocab_data
     
     if model is None or vocab_data is None:
         return None
     
+    progress(0, desc="Starting generation...")
+    
     diffusion = Diffusion(timesteps=500, device=device)  # Use 500 timesteps like training
     
+    def update_progress(pct):
+        progress(pct, desc=f"Generating... {pct*100:.1f}%")
+    
     with torch.no_grad():
         text_tokens = tokenize_text(prompt)
-        generated = diffusion.sample(model, text_tokens, image_size=64, steps=500)
+        generated = diffusion.sample(
+            model, 
+            text_tokens, 
+            image_size=64, 
+            steps=500,
+            progress_callback=update_progress
+        )
+    
+    progress(1.0, desc="Converting to image...")
     
     # Convert to image
     image = generated.cpu().squeeze(0)