fix: warmup issue fix

Author: varshith15

server/app.py

@@ -107,7 +107,7 @@ async def offer(request):
     params = await request.json()
 
     await pipeline.set_prompts(params["prompts"])
-    await pipeline.warm()
+    # await pipeline.warm()
 
     offer_params = params["offer"]
     offer = RTCSessionDescription(sdp=offer_params["sdp"], type=offer_params["type"])
@@ -198,6 +198,11 @@ async def on_connectionstatechange():
 
     await pc.setRemoteDescription(offer)
 
+    if "m=audio" in pc.remoteDescription.sdp:
+        await pipeline.warm_audio()
+    if "m=video" in pc.remoteDescription.sdp:
+        await pipeline.warm_video()
+
     answer = await pc.createAnswer()
     await pc.setLocalDescription(answer)
 

server/pipeline.py

@@ -16,14 +16,16 @@ def __init__(self, **kwargs):
         self.video_futures = asyncio.Queue()
         self.audio_futures = asyncio.Queue()
 
-    async def warm(self):
+    async def warm_video(self):
         dummy_video_inp = torch.randn(1, 512, 512, 3)
-        dummy_audio_inp = np.random.randint(-32768, 32767, 48 * 20, dtype=np.int16)  # TODO: might affect the workflow, due to buffering
 
         for _ in range(WARMUP_RUNS):
             image_out_fut = self.client.put_video_input(dummy_video_inp)
             await image_out_fut
 
+    async def warm_audio(self):
+        dummy_audio_inp = np.random.randint(-32768, 32767, 48 * 20, dtype=np.int16)  # TODO: might affect the workflow, due to buffering
+
         futs = []
         for _ in range(WARMUP_RUNS):
             audio_out_fut = self.client.put_audio_input(dummy_audio_inp)