[debug] ZeroGPU

app.py

@@ -7,8 +7,8 @@ This creates a web interface to compare three inference modes simultaneously:
 3. C2C: Rosetta model with projectors
 
 ZeroGPU Support:
-- Models are loaded to CPU at startup
-- @spaces.GPU decorator moves models to GPU on-demand for each inference
+- Models are loaded to CUDA at startup
+- @spaces.GPU decorator handles GPU allocation automatically for each inference
 - Works seamlessly on both ZeroGPU and regular GPU environments
 """
 
@@ -51,16 +51,15 @@ class ModelManager:
             c2c_checkpoint_path: Path to C2C checkpoint directory
             device: Device to use (cuda, cpu, or auto)
         """
-        # For ZeroGPU, load models to CPU and move to GPU in decorated functions
+        # For ZeroGPU, models should be loaded to CUDA directly
+        # The @spaces.GPU decorator handles GPU allocation automatically
         if device == "auto":
-            if ZEROGPU_AVAILABLE:
-                self.device = torch.device("cpu")
-                print("ZeroGPU detected: Loading models to CPU (will move to GPU on-demand)")
-            else:
-                self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         else:
             self.device = torch.device(device)
         print(f"Using device: {self.device}")
+        if ZEROGPU_AVAILABLE:
+            print("ZeroGPU detected: Models will be loaded to CUDA (decorator handles allocation)")
         
         # Model configurations
         self.single_model_name = single_model_name
@@ -221,16 +220,12 @@ class ModelManager:
     @spaces.GPU(duration=60)
     def generate_single(self, user_input: str) -> Generator[str, None, None]:
         """Generate response from single model with streaming."""
-        # Move model to GPU for ZeroGPU
-        device = torch.device("cuda" if ZEROGPU_AVAILABLE else self.device)
-        if ZEROGPU_AVAILABLE and self.single_model.device.type != "cuda":
-            self.single_model.to(device)
-        
+        # @spaces.GPU decorator handles GPU allocation automatically
         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(device)
+        inputs = self.single_tokenizer(text, return_tensors="pt").to(self.device)
         
         # Setup streamer
         streamer = TextIteratorStreamer(
@@ -260,13 +255,7 @@ 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)."""
-        # Move models to GPU for ZeroGPU
-        device = torch.device("cuda" if ZEROGPU_AVAILABLE else self.device)
-        if ZEROGPU_AVAILABLE:
-            if self.t2t_model.context_model.device.type != "cuda":
-                self.t2t_model.context_model.to(device)
-            if self.t2t_model.answer_model.device.type != "cuda":
-                self.t2t_model.answer_model.to(device)
+        # @spaces.GPU decorator handles GPU allocation automatically
         
         # Stage 1: Context generation
         context_streamer = TextIteratorStreamer(
@@ -282,7 +271,7 @@ class ModelManager:
             add_generation_prompt=True,
             return_tensors="pt",
             enable_thinking=False
-        ).to(device)
+        ).to(self.device)
         
         generation_kwargs = {
             'input_ids': inputs,
@@ -331,7 +320,7 @@ class ModelManager:
             add_generation_prompt=True,
             return_tensors="pt",
             enable_thinking=False
-        ).to(device)
+        ).to(self.device)
         
         generation_kwargs = {
             'input_ids': inputs,
@@ -352,16 +341,12 @@ class ModelManager:
     @spaces.GPU(duration=60)
     def generate_c2c(self, user_input: str) -> Generator[str, None, None]:
         """Generate response from C2C model with streaming."""
-        # Move model to GPU for ZeroGPU
-        device = torch.device("cuda" if ZEROGPU_AVAILABLE else self.device)
-        if ZEROGPU_AVAILABLE and self.c2c_model.device.type != "cuda":
-            self.c2c_model.to(device)
-        
+        # @spaces.GPU decorator handles GPU allocation automatically
         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(device)
+        inputs = self.c2c_tokenizer(text, return_tensors="pt").to(self.device)
         
         # Setup streamer
         streamer = TextIteratorStreamer(
@@ -374,12 +359,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(device)
+        ).unsqueeze(0).to(self.device)
         label_index = torch.tensor([-1, 0], dtype=torch.long).repeat(
             1, 1
-        ).unsqueeze(0).to(device)
+        ).unsqueeze(0).to(self.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(device)
+                      torch.arange(full_length, dtype=torch.long).unsqueeze(0).to(self.device)
         
         # Generation parameters
         generation_kwargs = {