Update app.py

app.py

@@ -13,7 +13,7 @@ import random
 import numpy as np
 import re
 import time
-from typing import List, Tuple, Generator
+from typing import List, Tuple
 import os
 import gc
 import spaces
@@ -21,14 +21,8 @@ import spaces
 # Global model variables for memory efficiency
 tokenizer = None
 model = None
-current_generator = None
 device = None
 
-def get_noising_schedule(i, max_it, sharpness=5.0):
-    """Exponential noise schedule for denoising"""
-    x = i / max_it
-    return (np.exp(-sharpness * x) - np.exp(-sharpness)) / (1 - np.exp(-sharpness))
-
 class ARDiffusionGenerator:
     """Base AR-Diffusion generator with shared functionality"""
     
@@ -58,7 +52,7 @@ class ARDiffusionGenerator:
 """
 
 class QualityGenerator(ARDiffusionGenerator):
-    """Quality-focused AR-Diffusion generator (from first script)"""
+    """Quality-focused AR-Diffusion generator"""
     
     def filter_logits(self, logits: torch.Tensor, top_k: int = 0, top_p: float = 1.0, 
                      temperature: float = 1.0) -> torch.Tensor:
@@ -194,8 +188,6 @@ class QualityGenerator(ARDiffusionGenerator):
         start_time = time.time()
         
         for step in range(steps):
-            step_start = time.time()
-            
             if progress_callback:
                 progress = 0.2 + (step / steps) * 0.7
                 elapsed = time.time() - start_time
@@ -222,7 +214,6 @@ class QualityGenerator(ARDiffusionGenerator):
                     max_replacements = min(3, len(mask_positions))
                 
                 sorted_positions = sorted(mask_positions.tolist())
-                step_replacements = 0
                 
                 for pos in sorted_positions[:max_replacements]:
                     if pos < len(logits):
@@ -257,7 +248,6 @@ class QualityGenerator(ARDiffusionGenerator):
                                     break
                         
                         current_ids[pos] = new_token
-                        step_replacements += 1
                         total_replacements += 1
         
         if progress_callback:
@@ -307,7 +297,7 @@ class QualityGenerator(ARDiffusionGenerator):
         return response
 
 class SpeedGenerator(ARDiffusionGenerator):
-    """Speed-focused AR-Diffusion generator (from second script)"""
+    """Speed-focused AR-Diffusion generator"""
     
     def filter_logits(self, logits: torch.Tensor, top_k: int = 15, top_p: float = 0.8, 
                      temperature: float = 1.0) -> torch.Tensor:
@@ -425,8 +415,6 @@ class SpeedGenerator(ARDiffusionGenerator):
         # Use mixed precision for speed on GPU
         with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=self.device.type == 'cuda'):
             for step in range(steps):
-                step_start = time.time()
-                
                 if progress_callback:
                     progress = 0.2 + (step / steps) * 0.7
                     elapsed = time.time() - start_time
@@ -448,7 +436,6 @@ class SpeedGenerator(ARDiffusionGenerator):
                     max_replace = min(8, len(mask_pos))
                     positions = sorted(mask_pos.tolist())[:max_replace]
                     
-                    step_replacements = 0
                     for pos in positions:
                         if pos < len(logits):
                             token_logits = logits[pos].clone()
@@ -475,7 +462,6 @@ class SpeedGenerator(ARDiffusionGenerator):
                                 new_token = top_indices[1].item()
                             
                             current_ids[pos] = new_token
-                            step_replacements += 1
                             total_replacements += 1
         
         if progress_callback:
@@ -519,21 +505,61 @@ class SpeedGenerator(ARDiffusionGenerator):
         
         return response
 
- {device}...")
-    
-    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
-    if tokenizer.pad_token is None:
-        tokenizer.pad_token = tokenizer.eos_token
+@spaces.GPU
+def load_model():
+    """Load model with Zero GPU optimization using @spaces.GPU"""
+    global tokenizer, model, device
     
-    model = AutoModelForCausalLM.from_pretrained(
-        model_path,
-        torch_dtype=torch.float16 if device.type == "cuda" else torch.float32,
-        device_map="auto" if device.type == "cuda" else None,
-        trust_remote_code=True,
-        low_cpu_mem_usage=True
-    )
+    if tokenizer is not None and model is not None:
+        return tokenizer, model, device
     
-    return tokenizer, model, device
+    try:
+        # This appears to be a LoRA adapter
+        adapter_path = "rootxhacker/llama-3B-diffusion-exp-fixed"
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        
+        print(f"Loading AR-Diffusion model on {device}...")
+        
+        # Load tokenizer from adapter
+        tokenizer = AutoTokenizer.from_pretrained(adapter_path, trust_remote_code=True)
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+        
+        # Load the adapter model
+        print("Loading adapter model...")
+        model = AutoModelForCausalLM.from_pretrained(
+            adapter_path,
+            torch_dtype=torch.float16 if device.type == "cuda" else torch.float32,
+            device_map="auto" if device.type == "cuda" else None,
+            trust_remote_code=True,
+            low_cpu_mem_usage=True
+        )
+        
+        print("✅ AR-Diffusion model loaded successfully!")
+        return tokenizer, model, device
+        
+    except Exception as e:
+        print(f"❌ Error loading {adapter_path}: {e}")
+        
+        # Fallback to a working model for demonstration
+        print("🔄 Falling back to demonstration model...")
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        fallback_model = "gpt2-medium"
+        
+        tokenizer = AutoTokenizer.from_pretrained(fallback_model)
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+            
+        model = AutoModelForCausalLM.from_pretrained(
+            fallback_model,
+            torch_dtype=torch.float16 if device.type == "cuda" else torch.float32,
+            device_map="auto" if device.type == "cuda" else None,
+            low_cpu_mem_usage=True
+        )
+        
+        print(f"✅ Fallback model {fallback_model} loaded successfully!")
+        print("⚠️ Note: Using fallback model - AR-Diffusion features may not work as expected")
+        return tokenizer, model, device
 
 def cleanup_memory():
     """Clean up GPU memory"""
@@ -563,7 +589,6 @@ def chat_function(message, history, mode, progress=gr.Progress()):
         # Generate response with progress callback
         def progress_callback(pct, status_msg):
             progress(pct)
-            # We'll show status in the performance display instead
         
         response, stats = generator.generate(message, progress_callback)
         
@@ -711,11 +736,11 @@ if __name__ == "__main__":
         show_error=True
     )
 
-# Updated requirements.txt should include:
+# Requirements.txt should include:
 # torch>=2.0.0
 # transformers>=4.30.0
 # gradio
 # numpy
 # accelerate
 # spaces
-# peft  # For LoRA adapter support
+# peft