[debug] zeroGPU

app.py

@@ -57,7 +57,11 @@ class ModelManager:
             self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         else:
             self.device = torch.device(device)
+        
+        # Debug information
         print(f"Using device: {self.device}")
+        print(f"CUDA available: {torch.cuda.is_available()}")
+        print(f"CUDA device count: {torch.cuda.device_count() if torch.cuda.is_available() else 0}")
         if ZEROGPU_AVAILABLE:
             print("ZeroGPU detected: Models will be loaded to CUDA (decorator handles allocation)")
         
@@ -104,8 +108,10 @@ class ModelManager:
         self.single_model, self.single_tokenizer = load_hf_model(
             self.single_model_name, self.device
         )
+        # Explicitly move model to device (required for ZeroGPU)
+        self.single_model = self.single_model.to(self.device)
         set_default_chat_template(self.single_tokenizer, self.single_model_name)
-        print("[Single] ✓ Model loaded")
+        print(f"[Single] ✓ Model loaded on {self.single_model.device}")
     
     def _load_t2t_model(self):
         """Load two-stage model."""
@@ -121,7 +127,10 @@ class ModelManager:
             device=str(self.device),
             background_prompt=self.t2t_background_prompt
         )
-        print("[T2T] ✓ Model loaded")
+        # Explicitly move models to device (required for ZeroGPU)
+        self.t2t_model.context_model = self.t2t_model.context_model.to(self.device)
+        self.t2t_model.answer_model = self.t2t_model.answer_model.to(self.device)
+        print(f"[T2T] ✓ Models loaded on {self.t2t_model.context_model.device} and {self.t2t_model.answer_model.device}")
     
     def _load_c2c_model(self):
         """Load Rosetta (C2C) model."""
@@ -178,7 +187,9 @@ class ModelManager:
         self.c2c_model, self.c2c_tokenizer = load_rosetta_model(
             model_config, eval_config, self.device
         )
-        print("[C2C] ✓ Model loaded")
+        # Explicitly move model to device (required for ZeroGPU)
+        self.c2c_model = self.c2c_model.to(self.device)
+        print(f"[C2C] ✓ Model loaded on {self.c2c_model.device}")
     
     def _load_all_models(self):
         """Load all models sequentially."""
@@ -221,6 +232,11 @@ class ModelManager:
     def generate_single(self, user_input: str) -> Generator[str, None, None]:
         """Generate response from single model with streaming."""
         # @spaces.GPU decorator handles GPU allocation automatically
+        # Ensure model is on correct device (ZeroGPU may move it)
+        if self.single_model.device != self.device:
+            print(f"[Single] Moving model from {self.single_model.device} to {self.device}")
+            self.single_model = self.single_model.to(self.device)
+        
         messages = [{"role": "system", "content": ""}, {"role": "user", "content": user_input}]
         text = self.single_tokenizer.apply_chat_template(
             messages, tokenize=False, add_generation_prompt=True, enable_thinking=False
@@ -256,6 +272,13 @@ class ModelManager:
     def generate_t2t(self, user_input: str) -> Generator[tuple[str, str], None, None]:
         """Generate response from T2T model with streaming (returns context, answer)."""
         # @spaces.GPU decorator handles GPU allocation automatically
+        # Ensure models are on correct device (ZeroGPU may move them)
+        if self.t2t_model.context_model.device != self.device:
+            print(f"[T2T] Moving context model from {self.t2t_model.context_model.device} to {self.device}")
+            self.t2t_model.context_model = self.t2t_model.context_model.to(self.device)
+        if self.t2t_model.answer_model.device != self.device:
+            print(f"[T2T] Moving answer model from {self.t2t_model.answer_model.device} to {self.device}")
+            self.t2t_model.answer_model = self.t2t_model.answer_model.to(self.device)
         
         # Stage 1: Context generation
         context_streamer = TextIteratorStreamer(
@@ -342,6 +365,11 @@ class ModelManager:
     def generate_c2c(self, user_input: str) -> Generator[str, None, None]:
         """Generate response from C2C model with streaming."""
         # @spaces.GPU decorator handles GPU allocation automatically
+        # Ensure model is on correct device (ZeroGPU may move it)
+        if self.c2c_model.device != self.device:
+            print(f"[C2C] Moving model from {self.c2c_model.device} to {self.device}")
+            self.c2c_model = self.c2c_model.to(self.device)
+        
         messages = [{"role": "system", "content": ""}, {"role": "user", "content": user_input}]
         text = self.c2c_tokenizer.apply_chat_template(
             messages, tokenize=False, add_generation_prompt=True, enable_thinking=False