perf: reduce prompt relay and audio ui overhead

load_audio_ui.py

@@ -1,7 +1,15 @@
-import folder_paths
-import os
-import torch
-import av
+import folder_paths
+import numpy as np
+import os
+import torch
+import av
+
+DEFAULT_SAMPLE_RATE = 44100
+DEFAULT_SILENCE_SECONDS = 1
+
+
+def build_silence(sample_rate: int = DEFAULT_SAMPLE_RATE, seconds: int = DEFAULT_SILENCE_SECONDS, channels: int = 2) -> torch.Tensor:
+    return torch.zeros((channels, int(sample_rate * seconds)), dtype=torch.float32)
 
 def f32_pcm(wav: torch.Tensor) -> torch.Tensor:
     """Convert audio to float 32 bits PCM format."""
@@ -13,30 +21,29 @@ def f32_pcm(wav: torch.Tensor) -> torch.Tensor:
         return wav.float() / (2 ** 31)
     raise ValueError(f"Unsupported wav dtype: {wav.dtype}")
 
-def load_audio_file(filepath: str) -> tuple[torch.Tensor, int]:
-    """Uses the latest ComfyUI av-based decoding for maximum compatibility."""
-    with av.open(filepath) as af:
-        if not af.streams.audio:
-            raise ValueError("No audio stream found in the file.")
+def load_audio_file(filepath: str) -> tuple[torch.Tensor, int]:
+    """Uses the latest ComfyUI av-based decoding for maximum compatibility."""
+    with av.open(filepath) as af:
+        if not af.streams.audio:
+            raise ValueError("No audio stream found in the file.")
 
         stream = af.streams.audio[0]
         sr = stream.codec_context.sample_rate
         n_channels = stream.channels
 
-        frames = []
-        for frame in af.decode(streams=stream.index):
-            buf = torch.from_numpy(frame.to_ndarray())
-            if buf.shape[0] != n_channels:
-                buf = buf.view(-1, n_channels).t()
-
-            frames.append(buf)
-
-        if not frames:
-            raise ValueError("No audio frames decoded.")
-
-        wav = torch.cat(frames, dim=1)
-        wav = f32_pcm(wav)
-        return wav, sr
+        frame_arrays = []
+        for frame in af.decode(streams=stream.index):
+            array = frame.to_ndarray()
+            if array.shape[0] != n_channels:
+                array = array.reshape(-1, n_channels).T
+            frame_arrays.append(array)
+
+        if not frame_arrays:
+            raise ValueError("No audio frames decoded.")
+
+        wav = torch.from_numpy(np.concatenate(frame_arrays, axis=1))
+        wav = f32_pcm(wav)
+        return wav, sr
 
 
 class LoadAudioUI:
@@ -95,21 +102,20 @@ class LoadAudioUI:
         
         # --- FALLBACK LOGIC ---
         # If the file is 'none' or doesn't exist on disk, provide 1 second of silence
-        if audio == "none" or not audio_path or not os.path.exists(audio_path):
-            missing_info = audio if audio != "none" else "None selected"
-            print(f"!!! [LoadAudioUI] Warning: Audio file '{missing_info}' not found. Outputting 1 second of silence.")
-            
-            sample_rate = 44100
-            # 1 second of silence (stereo) -> shape [channels, time]
-            waveform = torch.zeros((2, 44100))
-        else:
-            try:
-                waveform, sample_rate = load_audio_file(audio_path)
-            except Exception as e:
-                # If decoding fails for any reason, fallback to silence rather than crashing the workflow
-                print(f"!!! [LoadAudioUI] Error decoding {audio}: {e}. Falling back to silence.")
-                sample_rate = 44100
-                waveform = torch.zeros((2, 44100))
+        if audio == "none" or not audio_path or not os.path.exists(audio_path):
+            missing_info = audio if audio != "none" else "None selected"
+            print(f"!!! [LoadAudioUI] Warning: Audio file '{missing_info}' not found. Outputting 1 second of silence.")
+
+            sample_rate = DEFAULT_SAMPLE_RATE
+            waveform = build_silence(sample_rate)
+        else:
+            try:
+                waveform, sample_rate = load_audio_file(audio_path)
+            except Exception as e:
+                # If decoding fails for any reason, fallback to silence rather than crashing the workflow
+                print(f"!!! [LoadAudioUI] Error decoding {audio}: {e}. Falling back to silence.")
+                sample_rate = DEFAULT_SAMPLE_RATE
+                waveform = build_silence(sample_rate)
 
         # Convert seconds to frames
         start_frame = int(start_time * sample_rate)
@@ -138,4 +144,4 @@ class LoadAudioUI:
         # Calculate the final trimmed duration in seconds as a float
         final_duration = float(trimmed_waveform.shape[1] / sample_rate)
         
-        return (audio_output, final_duration)
+        return (audio_output, final_duration)

patches.py

@@ -1,9 +1,12 @@
-import logging
-import types
-
-import comfy.ldm.modules.attention
-
-log = logging.getLogger(__name__)
+import logging
+import os
+import types
+
+import comfy.ldm.modules.attention
+
+log = logging.getLogger(__name__)
+_DEBUG_PROMPT_RELAY = os.environ.get("PROMPT_RELAY_DEBUG", "").strip().lower() in {"1", "true", "yes", "on"}
+_DEBUG_LOG_PATH = os.path.join(os.path.dirname(__file__), "debug_prompt_relay.log")
 
 
 def _masked_attention(q, k, v, heads, mask, transformer_options={}, **kwargs):
@@ -72,14 +75,14 @@ def _make_masked_override(prev_override):
     return override
 
 
-def debug_log(msg):
-    try:
-        import os
-        log_path = os.path.join(os.path.dirname(__file__), "debug_prompt_relay.log")
-        with open(log_path, "a", encoding="utf-8") as f:
-            f.write(msg + "\n")
-    except Exception:
-        pass
+def debug_log(msg):
+    if not _DEBUG_PROMPT_RELAY:
+        return
+    try:
+        with open(_DEBUG_LOG_PATH, "a", encoding="utf-8") as f:
+            f.write(msg + "\n")
+    except Exception:
+        pass
 
 
 def _make_ltx_mask_wrapper(underlying, mask_fn, attr):