Update app.py

app.py

@@ -80,14 +80,19 @@ class CreedBrattonAI:
             
             # Load model with proper device handling
             if self.device == "cuda":
+                print("🤖 Loading model for GPU...")
                 self.model = AutoModelForCausalLM.from_pretrained(
                     model_name,
                     torch_dtype=torch.float16,  # Use float16 for GPU efficiency
-                    device_map="auto",  # Auto device mapping for GPU
+                    device_map=None,  # Don't use auto device mapping in ZeroGPU
                     trust_remote_code=True,
                     low_cpu_mem_usage=True
                 )
+                # Explicitly move to CUDA
+                print("🔧 Explicitly moving model to CUDA...")
+                self.model = self.model.to(self.device)
             else:
+                print("🤖 Loading model for CPU...")
                 self.model = AutoModelForCausalLM.from_pretrained(
                     model_name,
                     torch_dtype=torch.float32,  # Use float32 for CPU
@@ -104,9 +109,13 @@ class CreedBrattonAI:
             
             self.model.eval()
             
+            # Verify final device placement
+            final_device = next(self.model.parameters()).device
+            print(f"🎯 Model final device: {final_device}")
+            
             self.model_loaded = True
             self.loading = False
-            print(f"✅ Creed's consciousness loaded on {self.device}!")
+            print(f"✅ Creed's consciousness loaded on {final_device}!")
             
             # GPU memory info
             if self.device == "cuda" and torch.cuda.is_available():
@@ -128,8 +137,9 @@ class CreedBrattonAI:
                     self.model = AutoModelForCausalLM.from_pretrained(
                         base_model,
                         torch_dtype=torch.float16,
-                        device_map="auto"
+                        device_map=None  # Don't use auto in ZeroGPU
                     )
+                    self.model = self.model.to(self.device)
                 else:
                     self.model = AutoModelForCausalLM.from_pretrained(
                         base_model,
@@ -140,23 +150,38 @@ class CreedBrattonAI:
                 
                 self.model.eval()
                 self.model_loaded = True
-                print(f"✅ Fallback model loaded on {self.device}")
+                fallback_device = next(self.model.parameters()).device
+                print(f"✅ Fallback model loaded on {fallback_device}")
             except Exception as fallback_error:
                 print(f"❌ Fallback also failed: {fallback_error}")
             self.loading = False
     
     @spaces.GPU if SPACES_AVAILABLE else lambda func: func
     def generate_response_gpu(self, conversation: str) -> str:
-        """Generate response using the loaded model - back to working version"""
+        """Generate response using the loaded model with proper device handling"""
         
         if not self.model_loaded:
             return "❌ Model not loaded"
         
         try:
+            # Always ensure model is on the correct device in ZeroGPU
+            current_model_device = next(self.model.parameters()).device
+            print(f"🔍 Current model device: {current_model_device}")
+            
+            if self.device == "cuda" and current_model_device.type != "cuda":
+                print(f"🔄 Moving model from {current_model_device} to {self.device}")
+                self.model = self.model.to(self.device)
+            
+            # Verify model device after potential move
+            actual_device = next(self.model.parameters()).device
+            print(f"🎯 Model now on: {actual_device}")
+            
             # Simple tokenization that was working before
             inputs = self.tokenizer.encode(conversation, return_tensors="pt")
-            if self.device == "cuda":
-                inputs = inputs.to(self.device)
+            
+            # Put inputs on same device as model
+            inputs = inputs.to(actual_device)
+            print(f"🔍 Inputs device: {inputs.device}")
             
             # Generate response with original settings that worked
             with torch.no_grad():
@@ -180,6 +205,7 @@ class CreedBrattonAI:
             return self._clean_response(response)
             
         except Exception as e:
+            print(f"❌ Generation error: {e}")
             return f"🎸 *Creed scratches his head* Something weird happened... {str(e)[:100]}"
     
     def generate_response(self, message: str, history: List[List[str]]) -> Iterator[str]:
@@ -342,6 +368,8 @@ def main():
     
     # Memory status if GPU available
     if torch.cuda.is_available() and creed_ai.model_loaded:
+        actual_model_device = next(creed_ai.model.parameters()).device
+        print(f"🎯 Model actually on: {actual_model_device}")
         print(f"🔥 Final GPU Memory: {torch.cuda.memory_allocated() // 1024**2} MB allocated")
         print(f"📊 GPU Memory Reserved: {torch.cuda.memory_reserved() // 1024**2} MB reserved")
     
@@ -678,10 +706,11 @@ def main():
     ) as demo:
         
         # Modern header
+        actual_device = next(creed_ai.model.parameters()).device if creed_ai.model_loaded else creed_ai.device
         gr.HTML(f"""
         <div class="header">
             <h1>🎸 Creed Bratton AI</h1>
-            <p>Powered by phxdev/creed-qwen-0.5b-lora • Running on {'🚀 GPU' if creed_ai.device == 'cuda' else '🖥️ CPU'}</p>
+            <p>Powered by phxdev/creed-qwen-0.5b-lora • Running on {'🚀 GPU' if 'cuda' in str(actual_device) else '🖥️ CPU'} ({actual_device})</p>
         </div>
         """)