[debug] zeroGPU

app.py

@@ -7,8 +7,9 @@ This creates a web interface to compare three inference modes simultaneously:
 3. C2C: Rosetta model with projectors
 
 ZeroGPU Support:
-- Models are loaded to CUDA at startup
-- @spaces.GPU decorator handles GPU allocation automatically for each inference
+- Models are loaded to CUDA if available
+- @spaces.GPU decorator handles device allocation automatically
+- Inputs are moved to match the model's actual device
 - Works seamlessly on both ZeroGPU and regular GPU environments
 """
 
@@ -51,19 +52,15 @@ class ModelManager:
             c2c_checkpoint_path: Path to C2C checkpoint directory
             device: Device to use (cuda, cpu, or auto)
         """
-        # For ZeroGPU, models should be loaded to CUDA directly
-        # The @spaces.GPU decorator handles GPU allocation automatically
+        # Always use CUDA if available, ZeroGPU handles the rest
         if device == "auto":
             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)")
+            print("ZeroGPU environment detected")
         
         # Model configurations
         self.single_model_name = single_model_name
@@ -108,10 +105,11 @@ 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(f"[Single] ✓ Model loaded on {self.single_model.device}")
+        # Move to CUDA if available (following HuggingFace ZeroGPU pattern)
+        if torch.cuda.is_available():
+            self.single_model = self.single_model.to('cuda')
+        print(f"[Single] ✓ Model loaded")
     
     def _load_t2t_model(self):
         """Load two-stage model."""
@@ -127,10 +125,11 @@ class ModelManager:
             device=str(self.device),
             background_prompt=self.t2t_background_prompt
         )
-        # 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}")
+        # Move to CUDA if available (following HuggingFace ZeroGPU pattern)
+        if torch.cuda.is_available():
+            self.t2t_model.context_model = self.t2t_model.context_model.to('cuda')
+            self.t2t_model.answer_model = self.t2t_model.answer_model.to('cuda')
+        print("[T2T] ✓ Models loaded")
     
     def _load_c2c_model(self):
         """Load Rosetta (C2C) model."""
@@ -187,9 +186,10 @@ class ModelManager:
         self.c2c_model, self.c2c_tokenizer = load_rosetta_model(
             model_config, eval_config, self.device
         )
-        # 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}")
+        # Move to CUDA if available (following HuggingFace ZeroGPU pattern)
+        if torch.cuda.is_available():
+            self.c2c_model = self.c2c_model.to('cuda')
+        print("[C2C] ✓ Model loaded")
     
     def _load_all_models(self):
         """Load all models sequentially."""
@@ -231,17 +231,12 @@ class ModelManager:
     @spaces.GPU(duration=30)
     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
         )
-        inputs = self.single_tokenizer(text, return_tensors="pt").to(self.device)
+        # Use the model's actual device (ZeroGPU handles device placement)
+        inputs = self.single_tokenizer(text, return_tensors="pt").to(self.single_model.device)
         
         # Setup streamer
         streamer = TextIteratorStreamer(
@@ -271,15 +266,6 @@ class ModelManager:
     @spaces.GPU(duration=90)
     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(
             self.t2t_model.context_tokenizer,
@@ -294,7 +280,7 @@ class ModelManager:
             add_generation_prompt=True,
             return_tensors="pt",
             enable_thinking=False
-        ).to(self.device)
+        ).to(self.t2t_model.context_model.device)
         
         generation_kwargs = {
             'input_ids': inputs,
@@ -343,7 +329,7 @@ class ModelManager:
             add_generation_prompt=True,
             return_tensors="pt",
             enable_thinking=False
-        ).to(self.device)
+        ).to(self.t2t_model.answer_model.device)
         
         generation_kwargs = {
             'input_ids': inputs,
@@ -364,17 +350,12 @@ class ModelManager:
     @spaces.GPU(duration=30)
     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
         )
-        inputs = self.c2c_tokenizer(text, return_tensors="pt").to(self.device)
+        # Use the model's actual device (ZeroGPU handles device placement)
+        inputs = self.c2c_tokenizer(text, return_tensors="pt").to(self.c2c_model.device)
         
         # Setup streamer
         streamer = TextIteratorStreamer(
@@ -387,12 +368,12 @@ class ModelManager:
         full_length = inputs.input_ids.shape[1]
         instruction_index = torch.tensor([1, 0], dtype=torch.long).repeat(
             full_length - 1, 1
-        ).unsqueeze(0).to(self.device)
+        ).unsqueeze(0).to(self.c2c_model.device)
         label_index = torch.tensor([-1, 0], dtype=torch.long).repeat(
             1, 1
-        ).unsqueeze(0).to(self.device)
+        ).unsqueeze(0).to(self.c2c_model.device)
         position_ids = inputs.attention_mask.long().cumsum(-1) - 1 if inputs.attention_mask is not None else \
-                      torch.arange(full_length, dtype=torch.long).unsqueeze(0).to(self.device)
+                      torch.arange(full_length, dtype=torch.long).unsqueeze(0).to(self.c2c_model.device)
         
         # Generation parameters
         generation_kwargs = {