Warm PersonaPlex under live inference mode

streaming_app.py

@@ -128,16 +128,20 @@ def prepare_assets(token):
 
 
 def warmup(mimi, other_mimi, lm_gen, frame_size):
-    for _ in range(2):
-        chunk = torch.zeros(1, 1, frame_size, dtype=torch.float32, device=DEVICE)
-        codes = mimi.encode(chunk)
-        _ = other_mimi.encode(chunk)
-        for index in range(codes.shape[-1]):
-            tokens = lm_gen.step(codes[:, :, index:index + 1])
-            if tokens is not None:
-                _ = mimi.decode(tokens[:, 1:9])
-                _ = other_mimi.decode(tokens[:, 1:9])
-    torch.cuda.synchronize()
+    # The live loop uses inference_mode. PyTorch specializes compiled graphs on
+    # that dispatch state, so warming under no_grad would force a recompile when
+    # CUDAGraphed starts capturing the first live depformer call.
+    with torch.inference_mode():
+        for _ in range(2):
+            chunk = torch.zeros(1, 1, frame_size, dtype=torch.float32, device=DEVICE)
+            codes = mimi.encode(chunk)
+            _ = other_mimi.encode(chunk)
+            for index in range(codes.shape[-1]):
+                tokens = lm_gen.step(codes[:, :, index:index + 1])
+                if tokens is not None:
+                    _ = mimi.decode(tokens[:, 1:9])
+                    _ = other_mimi.decode(tokens[:, 1:9])
+        torch.cuda.synchronize()
     mimi.reset_streaming()
     other_mimi.reset_streaming()
     lm_gen.reset_streaming()