diff --git a/nodes/preview_override_node.py b/nodes/preview_override_node.py index f4527cbe..990c7971 100644 --- a/nodes/preview_override_node.py +++ b/nodes/preview_override_node.py @@ -10,6 +10,7 @@ import comfy.model_management import comfy.patcher_extension +import comfy.utils import latent_preview from comfy_api.latest import io from PIL import Image, ImageOps @@ -296,6 +297,49 @@ def _ltx_full_vae_decode_to_pil(vae, x0_5d, max_frames=None, stride=1): return [Image.fromarray(u8[i]) for i in range(u8.shape[0])] +def _decode_audio_waveform(audio_first_stage, ax, target_device): + # Decode on first_stage_model directly (VAE pinned to GPU in __call__) to skip comfy + # VAE.decode()'s per-step load_models_gpu. first_stage_model.decode gives (B, C, samples). + if ax is None or ax.numel() == 0: + return None + try: + ax = ax.to(device=target_device, dtype=torch.float32) + waveform = audio_first_stage.decode(ax) + except Exception as e: + logging.warning(f"[KJ PreviewOverride] audio VAE decode failed: {e}") + return None + if waveform.ndim == 2: + waveform = waveform.unsqueeze(1) # (B, samples) → (B, 1, samples) + if waveform.ndim != 3: + return None + # Mirror comfy_extras.nodes_audio.vae_decode_audio loudness normalization. + std = torch.std(waveform, dim=[1, 2], keepdim=True) * 5.0 + std[std < 1.0] = 1.0 + waveform = waveform / std + return waveform[0].clamp(-1.0, 1.0).detach().float().cpu() + + +def _encode_wav_b64(waveform_cpu, sample_rate): + # 16-bit PCM WAV — browser-native via