From 126fb3a63a71a19ef622604c30fbf985672c832e Mon Sep 17 00:00:00 2001 From: Varshith B Date: Thu, 30 Jan 2025 17:46:41 +0530 Subject: [PATCH] fix: cleanup --- nodes/audio_utils/save_audio_tensor.py | 4 ---- server/pipeline.py | 8 +++----- 2 files changed, 3 insertions(+), 9 deletions(-) diff --git a/nodes/audio_utils/save_audio_tensor.py b/nodes/audio_utils/save_audio_tensor.py index a3dd43ef..31ffc3b5 100644 --- a/nodes/audio_utils/save_audio_tensor.py +++ b/nodes/audio_utils/save_audio_tensor.py @@ -1,9 +1,5 @@ -import queue -import logging from comfystream import tensor_cache -logger = logging.getLogger(__name__) - class SaveAudioTensor: CATEGORY = "audio_utils" RETURN_TYPES = () diff --git a/server/pipeline.py b/server/pipeline.py index 7f79c4cd..c2e597d9 100644 --- a/server/pipeline.py +++ b/server/pipeline.py @@ -11,16 +11,14 @@ class Pipeline: def __init__(self, **kwargs): - self.client = ComfyStreamClient(**kwargs, max_workers=5) # hardcoded max workers + self.client = ComfyStreamClient(**kwargs, max_workers=5) # TODO: hardcoded max workers, should it be configurable? self.video_futures = asyncio.Queue() self.audio_futures = asyncio.Queue() - - self.resampler = av.audio.resampler.AudioResampler(format='s16', layout='mono', rate=48000) # find a better way to convert to mono async def warm(self): dummy_video_inp = torch.randn(1, 512, 512, 3) - dummy_audio_inp = np.random.randint(-32768, 32767, 48 * 20, dtype=np.int16) # has to be more than the buffer size in comfy workflow + 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) @@ -54,7 +52,7 @@ def video_preprocess(self, frame: av.VideoFrame) -> torch.Tensor: return torch.from_numpy(frame_np).unsqueeze(0) def audio_preprocess(self, frame: av.AudioFrame) -> torch.Tensor: - return self.resampler.resample(frame)[0].to_ndarray().flatten() + return frame.to_ndarray().ravel().reshape(-1, 2).mean(axis=1).astype(np.int16) def video_postprocess(self, output: torch.Tensor) -> av.VideoFrame: return av.VideoFrame.from_ndarray(