[debug] zeroGPU

app.py

@@ -106,6 +106,7 @@ class ModelManager:
             self.single_model_name, self.device
         )
         set_default_chat_template(self.single_tokenizer, self.single_model_name)
+        
         # Move to CUDA if available (following HuggingFace ZeroGPU pattern)
         if torch.cuda.is_available():
             self.single_model = self.single_model.to('cuda')
@@ -186,6 +187,7 @@ class ModelManager:
         self.c2c_model, self.c2c_tokenizer = load_rosetta_model(
             model_config, eval_config, self.device
         )
+        
         # Move to CUDA if available (following HuggingFace ZeroGPU pattern)
         if torch.cuda.is_available():
             self.c2c_model = self.c2c_model.to('cuda')
@@ -259,9 +261,15 @@ class ModelManager:
         
         # Stream tokens
         generated_text = ""
-        for token in streamer:
-            generated_text += token
-            yield generated_text
+        try:
+            for token in streamer:
+                generated_text += token
+                yield generated_text
+        except Exception as e:
+            print(f"[Single] Streaming error: {e}")
+            yield f"Error generating response: {e}"
+        finally:
+            thread.join()
     
     @spaces.GPU(duration=90)
     def generate_t2t(self, user_input: str) -> Generator[tuple[str, str], None, None]:
@@ -391,9 +399,15 @@ class ModelManager:
         
         # Stream tokens
         generated_text = ""
-        for token in streamer:
-            generated_text += token
-            yield generated_text
+        try:
+            for token in streamer:
+                generated_text += token
+                yield generated_text
+        except Exception as e:
+            print(f"[C2C] Streaming error: {e}")
+            yield f"Error generating response: {e}"
+        finally:
+            thread.join()
 
 
 def create_demo(model_manager: ModelManager):