[debug] zeroGPU

app.py

@@ -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."""
-        # For ZeroGPU, move model to GPU on-demand
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        if ZEROGPU_AVAILABLE:
-            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(
@@ -252,23 +247,16 @@ class ModelManager:
         thread.start()
         
         # Stream tokens
+        generated_text = ""
         for token in streamer:
-            yield token
-        thread.join()
-        
-        if ZEROGPU_AVAILABLE:
-            self.single_model.to("cpu")
-
-
+            generated_text += token
+            yield generated_text
+    
     @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)."""
-        # For ZeroGPU, move model to GPU on-demand
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        if ZEROGPU_AVAILABLE:
-            self.t2t_model.context_model.to(device)
-            self.t2t_model.answer_model.to(device)
-
+        # @spaces.GPU decorator handles GPU allocation automatically
+        
         # Stage 1: Context generation
         context_streamer = TextIteratorStreamer(
             self.t2t_model.context_tokenizer,
@@ -283,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,
@@ -332,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,
@@ -349,25 +337,16 @@ class ModelManager:
         for token in answer_streamer:
             answer_text += token
             yield context_text, answer_text
-        thread.join()
-
-        if ZEROGPU_AVAILABLE:
-            self.t2t_model.context_model.to("cpu")
-            self.t2t_model.answer_model.to("cpu")
-
+    
     @spaces.GPU(duration=60)
     def generate_c2c(self, user_input: str) -> Generator[str, None, None]:
         """Generate response from C2C model with streaming."""
-        # For ZeroGPU, move model to GPU on-demand
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        if ZEROGPU_AVAILABLE:
-            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(
@@ -380,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 = {
@@ -402,12 +381,10 @@ class ModelManager:
         thread.start()
         
         # Stream tokens
+        generated_text = ""
         for token in streamer:
-            yield token
-        thread.join()
-
-        if ZEROGPU_AVAILABLE:
-            self.c2c_model.to("cpu")
+            generated_text += token
+            yield generated_text
 
 
 def create_demo(model_manager: ModelManager):