Optimization fix for 500 error

.gitignore

@@ -1,3 +1,5 @@
 venv
 diffusion_model_best2.pth
 diffusion_model.pth
+__pycache__/
+*.pyc

app.py

@@ -63,8 +63,7 @@ def generate_step_by_step(history):
         
     prompt = history[-1][0]
     
-    steps = 40
-    # Lower temperature restricts the model from making massive unpredicted jumps
+    steps = 20
     temp = 0.2
     top_k = 10
 
@@ -101,6 +100,10 @@ def generate_step_by_step(history):
 
     current_seq = seq_tensor.squeeze(0).clone()
 
+    # Pre-calculate penalty indices once
+    vocab_size = model.fc_out.out_features
+    penalty_indices = torch.tensor([i for i in range(vocab_size) if i not in special_ids], device=device)
+
     for step in range(1, steps + 1):
         t_val = max(1.0 - step / steps, 0.05)
         t = torch.tensor([t_val], device=device)
@@ -111,19 +114,19 @@ def generate_step_by_step(history):
             
         response_logits = logits[0, mask_indices]
         
-        # Enhanced Subword Spatial Repetition Penalty
+        # Optimized Vectorized Repetition Penalty
         if step > 1:
             prev_pred = current_seq[mask_indices]
-            counts = torch.bincount(prev_pred, minlength=model.fc_out.out_features)
-            rep_mask = counts > 1
-            for tok_id in range(len(counts)):
-                if rep_mask[tok_id] and tok_id not in special_ids:
-                    # Devastating penalty for tokens that loop more than once globally 
-                    penalty = 5.0 ** (counts[tok_id].item() - 1)
-                    response_logits[:, tok_id] = torch.where(
-                        response_logits[:, tok_id] < 0,
-                        response_logits[:, tok_id] * penalty,
-                        response_logits[:, tok_id] / penalty
+            # Get counts only for tokens present in context to avoid giant loops
+            unique_tokens, counts = torch.unique(prev_pred, return_counts=True)
+            for i, tok_id in enumerate(unique_tokens):
+                t_id = tok_id.item()
+                if t_id not in special_ids and counts[i] > 1:
+                    penalty = 5.0 ** (counts[i].item() - 1)
+                    response_logits[:, t_id] = torch.where(
+                        response_logits[:, t_id] < 0,
+                        response_logits[:, t_id] * penalty,
+                        response_logits[:, t_id] / penalty
                     )
         
         if top_k > 0: