Update app.py

app.py

@@ -55,8 +55,8 @@ WARMUP_BASE_ASSISTANT_TEXT = (
 def _apply_seed(seed: int | None) -> None:
     if seed is None:
         return
+    # ZeroGPU: avoid touching torch.cuda outside the managed GPU call.
     torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
 
 
 def _load_audio(path: Path, target_sample_rate: int = SAMPLE_RATE) -> torch.Tensor:
@@ -553,14 +553,13 @@ def _load_backend(
     device_str: str,
     attn_impl: str,
 ):
-
+    # ZeroGPU: do not call torch.cuda.is_available() here; it may trigger low-level CUDA init.
     device = torch.device(device_str)
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
     processor = MossTTSRealtimeProcessor(tokenizer)
 
-    # dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+    # ZeroGPU: avoid torch.cuda.is_bf16_supported() before CUDA is fully managed.
     dtype = torch.float16
-
     if attn_impl and attn_impl.lower() not in {"none", ""}:
         model = MossTTSRealtime.from_pretrained(model_path, attn_implementation=attn_impl, torch_dtype=dtype).to(device)
         if (
@@ -798,10 +797,11 @@ class WarmupManager:
         self._lock = threading.Lock()
         self._thread: threading.Thread | None = None
         self._started = False
-        self._state = "pending"
-        self._progress = 0.0
-        self._message = "Waiting for startup warmup."
-        self._detail = "The app warms the streaming path before the first real request."
+        # ZeroGPU: startup warmup is disabled because it initializes CUDA outside @spaces.GPU.
+        self._state = "ready"
+        self._progress = 1.0
+        self._message = "Ready."
+        self._detail = "Startup warmup disabled for ZeroGPU; the first generation will load the model."
         self._error: str | None = None
 
     def start(self) -> None:
@@ -1288,8 +1288,7 @@ def _build_demo(
     tts_demo: StreamingTTSDemo,
     warmup_manager: WarmupManager,
 ):
-    # initial_warmup_snapshot = warmup_manager.snapshot()
-
+    initial_warmup_snapshot = warmup_manager.snapshot()
     with gr.Blocks(title="MossTTSRealtime") as demo:
         gr.Markdown("MossTTSRealtime demo")
         gr.Markdown("Note: The first run may take a while to load the model.")
@@ -1324,15 +1323,10 @@ def _build_demo(
                     chunk_duration = gr.Slider(0.01, 1.0, value=0.24, step=0.01, label="Codec Chunk Duration (s)")
                     stream_prebuffer_seconds = gr.Slider(0.0, 20.0, value=0.0, step=0.05, label="Initial Buffer (s)")
 
-                # run_btn = gr.Button(
-                #     "Generate" if initial_warmup_snapshot.ready else "Warming Up...",
-                #     elem_id="tts_generate",
-                #     interactive=initial_warmup_snapshot.ready,
-                # )
                 run_btn = gr.Button(
-                    "Generate",
+                    "Generate" if initial_warmup_snapshot.ready else "Warming Up...",
                     elem_id="tts_generate",
-                    interactive=True,
+                    interactive=initial_warmup_snapshot.ready,
                 )
 
             with gr.Column():
@@ -1341,7 +1335,7 @@ def _build_demo(
                 initial_status = _status_from_snapshot(initial_warmup_snapshot)
                 status = gr.Textbox(label="Status", lines=3, value=initial_status)
 
-        warmup_timer = gr.Timer(value=WARMUP_POLL_INTERVAL_SECONDS, active=True)
+        warmup_timer = gr.Timer(value=WARMUP_POLL_INTERVAL_SECONDS, active=not initial_warmup_snapshot.ready)
 
         def _poll_warmup_state():
             snapshot = warmup_manager.snapshot()
@@ -1350,7 +1344,7 @@ def _build_demo(
                 _warmup_status_update(snapshot),
                 _warmup_timer_update(snapshot),
             )
-            
+
         @spaces.GPU
         def _on_generate(
             user_text_value,
@@ -1374,10 +1368,6 @@ def _build_demo(
             chunk_duration_value,
             stream_prebuffer_seconds_value,
         ):
-            # warmup_snapshot = warmup_manager.snapshot()
-            # if not warmup_snapshot.ready:
-            #     yield json.dumps({"reset": True}), gr.update(value=None), _warmup_gate_message(warmup_snapshot)
-            #     return
             try:
                 started_at = time.monotonic()
                 full_chunks: list[np.ndarray] = []
@@ -1530,6 +1520,8 @@ def main():
             attn_impl=args.attn_implementation,
         ),
     )
+    # ZeroGPU: do not run startup warmup, because it would initialize CUDA
+    # in a background thread outside @spaces.GPU.
     # warmup_manager.start()
     demo = _build_demo(args, tts_demo, warmup_manager)
     demo.queue(max_size=10, default_concurrency_limit=1).launch(
@@ -1540,4 +1532,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()