fix: return fps from non-cache extraction path

The fps output was only returned on cache hits. Fresh extractions
returned only features, leaving fps null.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

.gitignore

@@ -0,0 +1,9 @@
+__pycache__/
+*.pyc
+*.pyo
+*.egg-info/
+dist/
+build/
+.eggs/
+*.so
+.env

README.md

@@ -0,0 +1,156 @@
+# ComfyUI-PrismAudio
+
+Custom nodes for [PrismAudio](https://huggingface.co/FunAudioLLM/PrismAudio) (ICLR 2026) — video-to-audio and text-to-audio generation using decomposed Chain-of-Thought reasoning with a 518M parameter DiT diffusion model and Stable Audio 2.0 VAE.
+
+## Installation
+
+Clone into your ComfyUI custom nodes directory:
+
+```bash
+cd ComfyUI/custom_nodes
+git clone https://github.com/Ethanfel/ComfyUI-Prismaudio.git ComfyUI-PrismAudio
+pip install -r ComfyUI-PrismAudio/requirements.txt
+```
+
+**flash-attn** is optional — detected at runtime, falls back to PyTorch SDPA if unavailable.
+
+## Nodes
+
+### PrismAudio Model Loader
+
+Loads the DiT diffusion model and VAE. Auto-downloads weights from HuggingFace on first use.
+
+| Input | Options | Description |
+|-------|---------|-------------|
+| `precision` | auto / fp32 / fp16 / bf16 | DiT and conditioner dtype. VAE is always fp32. |
+| `offload_strategy` | auto / keep_in_vram / offload_to_cpu | Memory management. |
+
+---
+
+### PrismAudio Feature Extractor
+
+Extracts video features (VideoPrism LvT, Synchformer) and text features (T5-Gemma) from a video in a subprocess. Results are cached on disk.
+
+| Input | Description |
+|-------|-------------|
+| `video` | IMAGE tensor from any ComfyUI video loader |
+| `caption_cot` | Chain-of-thought description of the audio scene |
+| `video_info` | *(optional)* `VHS_VIDEOINFO` from VHS LoadVideo — sets fps automatically |
+| `fps` | Source fps — ignored if `video_info` is connected |
+| `python_env` | `managed_env` (auto-created isolated venv, recommended) or `comfyui_env` (current Python, see warning below) |
+| `cache_dir` | Directory for cached `.npz` files. Empty = system temp dir. |
+| `hf_token` | HuggingFace token for gated models. Prefer `HF_TOKEN` env var instead. |
+
+**Outputs:** `features` (PRISMAUDIO_FEATURES), `fps` (FLOAT)
+
+**`managed_env`** auto-creates a venv at `_extract_env/` inside the plugin directory on first use and installs JAX, TF, VideoPrism, and Synchformer. This takes several minutes the first time.
+
+**`comfyui_env`** uses the current ComfyUI Python — JAX/TF/videoprism must already be installed. Installing them into the ComfyUI environment may conflict with existing packages.
+
+---
+
+### PrismAudio Feature Loader
+
+Loads a pre-computed `.npz` feature file. Use this to re-use extracted features without re-running the extractor.
+
+| Input | Description |
+|-------|-------------|
+| `npz_path` | Path to a `.npz` file produced by the Feature Extractor |
+
+---
+
+### PrismAudio Sampler
+
+Video-to-audio generation. Takes model + features, produces AUDIO.
+
+| Input | Description |
+|-------|-------------|
+| `model` | From Model Loader |
+| `features` | From Feature Extractor or Feature Loader |
+| `duration` | Audio duration in seconds. Set to `0` to use the video duration from features automatically. |
+| `steps` | Sampling steps (default: 100) |
+| `cfg_scale` | Classifier-free guidance scale (default: 7.0) |
+| `seed` | RNG seed |
+
+---
+
+### PrismAudio Text Only
+
+Text-to-audio generation without video. Uses the T5-Gemma encoder.
+
+| Input | Description |
+|-------|-------------|
+| `model` | From Model Loader |
+| `text_prompt` | Chain-of-thought audio scene description. Longer, more detailed prompts produce better results. |
+| `duration` | Audio duration in seconds |
+| `steps` | Sampling steps (default: 100) |
+| `cfg_scale` | Classifier-free guidance scale (default: 7.0) |
+| `seed` | RNG seed |
+
+---
+
+## Workflows
+
+### Video-to-Audio
+
+```
+VHS LoadVideo ──► PrismAudio Feature Extractor ──► PrismAudio Sampler ──► Save Audio
+                         (video_info) ──────────────────► (fps auto)
+                         (features) ────────────────────► (features)
+                         duration=0 ─────────────────────► (auto from features)
+```
+
+### Pre-computed Features
+
+```
+PrismAudio Feature Loader (.npz) ──► PrismAudio Sampler ──► Save Audio
+```
+
+### Text-to-Audio
+
+```
+PrismAudio Text Only ──► Save Audio
+```
+
+## HuggingFace Authentication
+
+Required for T5-Gemma (gated model) and PrismAudio weights.
+
+1. Visit <https://huggingface.co/FunAudioLLM/PrismAudio> and accept the license.
+2. Authenticate via one of:
+   - **Environment variable:** `export HF_TOKEN=hf_...`
+   - **CLI login:** `huggingface-cli login`
+
+There is no `hf_token` widget on the main nodes by design — ComfyUI saves all STRING values to workflow JSON, which would expose your token. The Feature Extractor has an `hf_token` input as a convenience but using `HF_TOKEN` env var is preferred.
+
+## Model Files
+
+Weights are auto-downloaded to `ComfyUI/models/prismaudio/`:
+
+| File | Size | Description |
+|------|------|-------------|
+| `prismaudio.ckpt` | ~2.7 GB | Diffusion model (DiT) |
+| `vae.ckpt` | ~2.5 GB | Stable Audio 2.0 VAE |
+| `synchformer_state_dict.pth` | ~950 MB | Synchformer visual encoder |
+
+T5-Gemma and VideoPrism LvT are cached in `~/.cache/huggingface/`.
+
+## VRAM Requirements
+
+| VRAM | Recommended settings |
+|------|----------------------|
+| 24 GB+ | `keep_in_vram`, any precision |
+| 12–24 GB | `offload_to_cpu`, bf16/fp16 |
+| 8–12 GB | `offload_to_cpu`, fp16 |
+| < 8 GB | May work with `offload_to_cpu` + fp16 |
+
+## Troubleshooting
+
+- **Gated model errors** — Accept the license at <https://huggingface.co/FunAudioLLM/PrismAudio> and set `HF_TOKEN`.
+- **VRAM errors** — Switch `offload_strategy` to `offload_to_cpu` and/or use `fp16` precision.
+- **Feature extraction fails** — Ensure `synchformer_state_dict.pth` is in `models/prismaudio/`. On first run with `managed_env`, installation takes several minutes.
+- **flash-attn** — Optional. Auto-detected at runtime; falls back to PyTorch SDPA.
+
+## Credits
+
+PrismAudio by [FunAudioLLM](https://github.com/FunAudioLLM) (ICLR 2026). [Model & weights](https://huggingface.co/FunAudioLLM/PrismAudio).

__init__.py

@@ -1,6 +1,10 @@
 """
 ComfyUI-PrismAudio: Video-to-Audio and Text-to-Audio generation using PrismAudio (ICLR 2026).
 """
+import sys
+import os
+sys.path.insert(0, os.path.dirname(__file__))
+
 from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
 
 __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

data_utils/v2a_utils/feature_utils_288.py

@@ -0,0 +1,337 @@
+"""
+PrismAudio feature extraction utilities.
+
+Implements FeaturesUtils used by scripts/extract_features.py to extract:
+  - Text features via T5-Gemma (transformers)
+  - Video features via VideoPrism (JAX/Flax, google-deepmind/videoprism)
+  - Sync features via Synchformer visual encoder (PyTorch)
+"""
+
+import os
+import torch
+import torch.nn as nn
+import numpy as np
+
+
+class FeaturesUtils:
+    def __init__(self, vae_config_path=None, synchformer_ckpt=None, device=None):
+        self.device = device or torch.device("cpu")
+        self._t5_tokenizer = None
+        self._t5_encoder = None
+        self._vp_model = None
+        self._vp_state = None
+        self._vp_text_tokenizer = None
+        self._sync_model = None
+
+        self._synchformer_ckpt = synchformer_ckpt
+        self._load_synchformer()
+
+    # ------------------------------------------------------------------
+    # T5-Gemma text encoding
+    # ------------------------------------------------------------------
+
+    def _ensure_t5(self):
+        if self._t5_encoder is not None:
+            return
+        from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+        model_id = "google/t5gemma-l-l-ul2-it"
+        print(f"[FeaturesUtils] Loading T5-Gemma: {model_id}")
+        self._t5_tokenizer = AutoTokenizer.from_pretrained(model_id)
+        self._t5_encoder = (
+            AutoModelForSeq2SeqLM.from_pretrained(model_id)
+            .get_encoder()
+            .to(self.device)
+            .eval()
+        )
+
+    def encode_t5_text(self, texts):
+        """
+        Args:
+            texts: list of str
+        Returns:
+            Tensor [seq_len, 1024]
+        """
+        self._ensure_t5()
+        tokens = self._t5_tokenizer(
+            texts, return_tensors="pt", padding=True
+        ).to(self.device)
+        with torch.no_grad():
+            out = self._t5_encoder(**tokens)
+        # Move encoder off GPU to save VRAM
+        self._t5_encoder.to("cpu")
+        torch.cuda.empty_cache()
+        return out.last_hidden_state.squeeze(0)  # [seq_len, 1024]
+
+    # ------------------------------------------------------------------
+    # VideoPrism video + text encoding (JAX)
+    # ------------------------------------------------------------------
+
+    def _ensure_videoprism(self):
+        if self._vp_model is not None:
+            return
+        from videoprism import models as vp
+        import jax
+        model_name = "videoprism_lvt_public_v1_large"
+        print(f"[FeaturesUtils] Loading VideoPrism LvT large (1024-dim joint video-text)...")
+        self._vp_model = vp.get_model(model_name)
+        self._vp_state = vp.load_pretrained_weights(model_name)
+        self._vp_text_tokenizer = vp.load_text_tokenizer("c4_en")
+        jax_dev = jax.devices()[0]
+        self._jax_forward = jax.jit(
+            lambda x, y, z: self._vp_model.apply(
+                self._vp_state, x, y, z, train=False, return_intermediate=True
+            ),
+            device=jax_dev,
+        )
+
+    def encode_video_and_text_with_videoprism(self, clip_input, texts):
+        """
+        Args:
+            clip_input: Tensor [1, T, C, H, W] float32, values in [-1, 1]
+            texts: list of str — CoT captions, passed to VideoPrism LvT text tower
+        Returns:
+            global_video_features: Tensor [1, D]
+            video_features:        Tensor [T, D]  — per-frame L2-normalized embeddings
+            global_text_features:  Tensor [1, D]
+        """
+        self._ensure_videoprism()
+        import jax.numpy as jnp
+        from videoprism import models as vp
+
+        # Normalise from [-1,1] to [0,1] and convert to [B, T, H, W, C] JAX array
+        frames = clip_input.squeeze(0)          # [T, C, H, W]
+        frames = (frames + 1.0) / 2.0          # [-1,1] → [0,1]
+        frames = frames.permute(0, 2, 3, 1)    # [T, H, W, C]
+        frames_np = frames.cpu().numpy().astype(np.float32)
+        frames_jax = jnp.array(frames_np)[None]  # [1, T, H, W, C]
+
+        # Tokenize text (padding value 1.0 = pad, 0.0 = real token)
+        text_ids, text_paddings = vp.tokenize_texts(self._vp_text_tokenizer, texts)
+
+        # Joint video+text forward with intermediate outputs
+        video_embeddings, text_embeddings, outputs = self._jax_forward(
+            frames_jax, text_ids, text_paddings
+        )
+
+        # Per-frame features: [B, T, 1024] L2-normalized
+        frame_embed_np = np.array(outputs["frame_embeddings"])   # [1, T, 1024]
+        per_frame = torch.from_numpy(frame_embed_np[0]).to(self.device)  # [T, 1024]
+
+        # Global video embedding: [1024] → [1, 1024]
+        global_video = torch.from_numpy(
+            np.array(video_embeddings[0])
+        ).unsqueeze(0).to(self.device)  # [1, 1024]
+
+        # Global text embedding: [1024] → [1, 1024]
+        global_text = torch.from_numpy(
+            np.array(text_embeddings[0])
+        ).unsqueeze(0).to(self.device)  # [1, 1024]
+
+        return global_video, per_frame, global_text
+
+    # ------------------------------------------------------------------
+    # Synchformer sync feature encoding
+    # ------------------------------------------------------------------
+
+    def _load_synchformer(self):
+        if not self._synchformer_ckpt or not os.path.exists(self._synchformer_ckpt):
+            return
+
+        print(f"[FeaturesUtils] Loading Synchformer from: {self._synchformer_ckpt}")
+        state = torch.load(self._synchformer_ckpt, map_location="cpu", weights_only=False)
+
+        # Checkpoint may be raw state_dict or wrapped in {"model": ...}
+        if isinstance(state, dict) and "model" in state:
+            state_dict = state["model"]
+        else:
+            state_dict = state
+
+        self._sync_model = _SynchformerVisualEncoder(state_dict, self.device)
+        self._sync_model.eval()
+
+    def encode_video_with_sync(self, sync_input):
+        """
+        Args:
+            sync_input: Tensor [1, T, C, H, W] float32, values in [-1, 1]
+        Returns:
+            sync_features: Tensor [num_segments, 768]
+        """
+        if self._sync_model is None:
+            raise RuntimeError(
+                "[FeaturesUtils] Synchformer checkpoint not loaded. "
+                "Pass synchformer_ckpt to FeaturesUtils or set --synchformer_ckpt."
+            )
+        frames = sync_input.squeeze(0).to(self.device)  # [T, C, H, W]
+        with torch.no_grad():
+            return self._sync_model(frames)
+
+
+# ------------------------------------------------------------------
+# Synchformer visual encoder — TimeSformer-style ViT-B/16
+# Architecture reverse-engineered from synchformer_state_dict.pth
+# ------------------------------------------------------------------
+
+import torch.nn.functional as F
+
+
+class _PatchEmbed(nn.Module):
+    """2D patch embedding: [B, 3, 224, 224] → [B, 196, 768]."""
+    def __init__(self):
+        super().__init__()
+        self.proj = nn.Conv2d(3, 768, kernel_size=16, stride=16)
+
+    def forward(self, x):
+        return self.proj(x).flatten(2).transpose(1, 2)
+
+
+class _ViTAttn(nn.Module):
+    """ViT-style QKV attention (timm convention: qkv as single Linear)."""
+    def __init__(self, dim=768, num_heads=12):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = self.head_dim ** -0.5
+        self.qkv = nn.Linear(dim, dim * 3)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(self, x):
+        B, N, D = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv.unbind(0)
+        attn = F.softmax((q @ k.transpose(-2, -1)) * self.scale, dim=-1)
+        return self.proj((attn @ v).transpose(1, 2).reshape(B, N, D))
+
+
+class _BlockMLP(nn.Module):
+    """Two-layer MLP with GELU, keys fc1/fc2 to match checkpoint."""
+    def __init__(self, dim=768, mlp_dim=3072):
+        super().__init__()
+        self.fc1 = nn.Linear(dim, mlp_dim)
+        self.fc2 = nn.Linear(mlp_dim, dim)
+
+    def forward(self, x):
+        return self.fc2(F.gelu(self.fc1(x)))
+
+
+class _TimeSformerBlock(nn.Module):
+    """
+    Factorized space-time attention block.
+    norm1 → spatial attn → norm3 → temporal attn → norm2 → MLP
+    """
+    def __init__(self, dim=768, num_heads=12):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim)
+        self.attn = _ViTAttn(dim, num_heads)
+        self.norm3 = nn.LayerNorm(dim)
+        self.timeattn = _ViTAttn(dim, num_heads)
+        self.norm2 = nn.LayerNorm(dim)
+        self.mlp = _BlockMLP(dim)
+
+    def forward(self, x, T):
+        # x: [T, N, D]  (T frames treated as batch, N=197 spatial tokens)
+        x = x + self.attn(self.norm1(x))
+        # Temporal attention: for each spatial position, attend across T frames
+        # [T, N, D] → [N, T, D] → attend → [N, T, D] → [T, N, D]
+        xt = x.permute(1, 0, 2)
+        xt = xt + self.timeattn(self.norm3(xt))
+        x = xt.permute(1, 0, 2)
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class _SpatialAttnAgg(nn.Module):
+    """
+    Aggregates 196 spatial patches → 1 feature per frame using a
+    TransformerEncoderLayer with a learnable CLS token.
+    Key names match nn.TransformerEncoderLayer: self_attn, linear1, linear2, norm1, norm2.
+    """
+    def __init__(self, dim=768, num_heads=12):
+        super().__init__()
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, dim))
+        self.self_attn = nn.MultiheadAttention(dim, num_heads, batch_first=True)
+        self.linear1 = nn.Linear(dim, dim * 4)
+        self.linear2 = nn.Linear(dim * 4, dim)
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+
+    def forward(self, x):
+        # x: [T, 196, 768] — spatial patches (CLS stripped)
+        T = x.shape[0]
+        cls = self.cls_token.expand(T, -1, -1)
+        x = torch.cat([cls, x], dim=1)              # [T, 197, 768]
+        xn = self.norm1(x)
+        x = x + self.self_attn(xn, xn, xn)[0]
+        x = x + self.linear2(F.gelu(self.linear1(self.norm2(x))))
+        return x[:, 0, :]                            # [T, 768] — CLS per frame
+
+
+class _SynchformerVisualEncoder(nn.Module):
+    """
+    TimeSformer-style ViT-B/16 visual encoder for the PrismAudio Synchformer checkpoint.
+    Processes video in segments of 8 frames → [T_aligned, 768] per-frame features.
+    """
+
+    def __init__(self, state_dict, device):
+        super().__init__()
+        self.device = device
+        self.segment_frames = 8
+
+        self.patch_embed = _PatchEmbed()
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, 768))
+        self.pos_embed = nn.Parameter(torch.zeros(1, 197, 768))
+        self.temp_embed = nn.Parameter(torch.zeros(1, 8, 768))
+        self.blocks = nn.ModuleList([_TimeSformerBlock() for _ in range(12)])
+        self.norm = nn.LayerNorm(768)
+        self.spatial_attn_agg = _SpatialAttnAgg()
+
+        # Load weights from vfeat_extractor.* prefix
+        prefix = "vfeat_extractor."
+        sub = {k[len(prefix):]: v for k, v in state_dict.items() if k.startswith(prefix)}
+        # Exclude 3D patch embed (we use 2D only)
+        sub = {k: v for k, v in sub.items() if not k.startswith("patch_embed_3d")}
+        missing, unexpected = self.load_state_dict(sub, strict=False)
+        print(f"[FeaturesUtils] Synchformer loaded — missing={len(missing)}, unexpected={len(unexpected)}")
+        if missing:
+            print(f"[FeaturesUtils]   missing keys (first 5): {missing[:5]}")
+
+        self.to(device)
+
+    def forward(self, frames):
+        """
+        Args:
+            frames: [T, C, H, W] float32 in [-1, 1], at 25fps
+        Returns:
+            [T_aligned, 768] — per-frame features (T_aligned = floor(T/8)*8)
+        """
+        T = frames.shape[0]
+        seg = self.segment_frames
+        num_seg = max(1, T // seg)
+        T_aligned = num_seg * seg
+
+        results = []
+        for i in range(num_seg):
+            chunk = frames[i * seg:(i + 1) * seg]   # [8, C, H, W]
+            results.append(self._forward_segment(chunk))
+        return torch.cat(results, dim=0)             # [T_aligned, 768]
+
+    def _forward_segment(self, x):
+        # x: [8, 3, 224, 224]
+        T = x.shape[0]  # 8
+
+        # Patch embedding + CLS token
+        x = self.patch_embed(x)                                      # [8, 196, 768]
+        cls = self.cls_token.expand(T, -1, -1)
+        x = torch.cat([cls, x], dim=1)                               # [8, 197, 768]
+
+        # Positional + temporal embeddings
+        x = x + self.pos_embed                                        # broadcast (1,197,768)
+        x = x + self.temp_embed.squeeze(0).unsqueeze(1)              # (8,1,768) broadcast
+
+        # Transformer blocks (factorized space-time)
+        for block in self.blocks:
+            x = block(x, T)
+
+        x = self.norm(x)
+
+        # Aggregate spatial patches → 1 feature per frame
+        return self.spatial_attn_agg(x[:, 1:, :])                    # [8, 768]

nodes/feature_extractor.py

@@ -0,0 +1,207 @@
+import os
+import sys
+import hashlib
+import subprocess
+import tempfile
+import torch
+
+from .utils import PRISMAUDIO_CATEGORY
+from .feature_loader import PrismAudioFeatureLoader
+
+# Managed venv created automatically when python_env is left as default
+_PLUGIN_DIR = os.path.dirname(os.path.dirname(__file__))
+_MANAGED_VENV = os.path.join(_PLUGIN_DIR, "_extract_env")
+_MANAGED_PYTHON = os.path.join(_MANAGED_VENV, "bin", "python")
+
+_EXTRACT_PACKAGES = [
+    "torch", "torchaudio", "torchvision",
+    # TF 2.15 only supports Python <=3.11; use >=2.16 for Python 3.12+
+    "tensorflow-cpu>=2.16.0",
+    # jax[cuda13] includes jaxlib; pip-managed CUDA libs (no local toolkit needed)
+    "jax[cuda13]", "flax",
+    "transformers", "decord", "einops", "numpy", "mediapy",
+    "git+https://github.com/google-deepmind/videoprism.git",
+]
+
+
+def _pip_install(pip, *packages, label=None):
+    """Install one or more packages with visible output; raise on failure."""
+    tag = label or packages[0]
+    print(f"[PrismAudio]   installing {tag} ...", flush=True)
+    result = subprocess.run(
+        [pip, "install", "--progress-bar", "on"] + list(packages),
+        capture_output=False,
+    )
+    if result.returncode != 0:
+        raise RuntimeError(
+            f"[PrismAudio] Failed to install {tag} (exit {result.returncode}). "
+            "See pip output above for details."
+        )
+    print(f"[PrismAudio]   {tag} OK", flush=True)
+
+
+def _ensure_extract_env():
+    """Create and populate the managed venv on first use."""
+    if os.path.exists(_MANAGED_PYTHON):
+        return _MANAGED_PYTHON
+
+    import shutil
+    if os.path.exists(_MANAGED_VENV):
+        print("[PrismAudio] Removing incomplete venv and retrying...", flush=True)
+        shutil.rmtree(_MANAGED_VENV)
+
+    print(f"[PrismAudio] Creating feature-extraction venv at: {_MANAGED_VENV}", flush=True)
+    subprocess.run([sys.executable, "-m", "venv", _MANAGED_VENV], check=True)
+
+    pip = os.path.join(_MANAGED_VENV, "bin", "pip")
+
+    print("[PrismAudio] Upgrading pip...", flush=True)
+    subprocess.run([pip, "install", "--upgrade", "pip"], check=True)
+
+    total = len(_EXTRACT_PACKAGES)
+    print(f"[PrismAudio] Installing {total} package groups — this may take several minutes...", flush=True)
+
+    for i, pkg in enumerate(_EXTRACT_PACKAGES, 1):
+        label = pkg.split("/")[-1] if pkg.startswith("git+") else pkg.split(">=")[0].split("==")[0].split("[")[0]
+        print(f"[PrismAudio] [{i}/{total}] {label}", flush=True)
+        _pip_install(pip, pkg, label=label)
+
+    print("[PrismAudio] Feature-extraction env ready.", flush=True)
+    return _MANAGED_PYTHON
+
+
+def _hash_inputs(video_tensor, cot_text):
+    """Create a hash of the inputs for caching."""
+    h = hashlib.sha256()
+    h.update(video_tensor.cpu().numpy().tobytes()[:1024 * 1024])  # First 1MB for speed
+    h.update(cot_text.encode())
+    return h.hexdigest()[:16]
+
+
+def _save_frames_to_npy(video_tensor, output_path):
+    """Save ComfyUI IMAGE tensor [T,H,W,C] float32 [0,1] to .npy as uint8.
+
+    Lossless — avoids H.264 encode/decode roundtrip.
+    """
+    import numpy as np
+    frames_np = (video_tensor.cpu().numpy() * 255).astype("uint8")
+    np.save(output_path, frames_np)
+
+
+class PrismAudioFeatureExtractor:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "video": ("IMAGE",),
+                "caption_cot": ("STRING", {"default": "", "multiline": True, "tooltip": "Chain-of-thought description"}),
+            },
+            "optional": {
+                "video_info": ("VHS_VIDEOINFO", {"tooltip": "Connect VHS LoadVideo info output to auto-set fps."}),
+                "fps": ("FLOAT", {"default": 30.0, "min": 1.0, "max": 120.0, "step": 0.001, "tooltip": "Frame rate of the input video. Ignored if video_info is connected."}),
+                "python_env": (["managed_env", "comfyui_env"], {"tooltip": "managed_env: auto-created isolated venv with JAX/TF (recommended). comfyui_env: current ComfyUI Python — WARNING: may conflict with existing packages and destabilize ComfyUI."}),
+                "cache_dir": ("STRING", {"default": "", "tooltip": "Directory to cache extracted features. Empty = temp dir"}),
+                "hf_token": ("STRING", {"default": "", "tooltip": "HuggingFace token for gated models (e.g. google/t5gemma). Get yours at huggingface.co/settings/tokens"}),
+            },
+        }
+
+    RETURN_TYPES = ("PRISMAUDIO_FEATURES", "FLOAT")
+    RETURN_NAMES = ("features", "fps")
+    FUNCTION = "extract_features"
+    CATEGORY = PRISMAUDIO_CATEGORY
+
+    def extract_features(self, video, caption_cot, video_info=None, fps=30.0, python_env="managed_env", cache_dir="", hf_token=""):
+        # Resolve fps from VHS video_info if connected
+        if video_info is not None:
+            fps = video_info["loaded_fps"]
+
+        # Resolve python binary
+        if python_env == "comfyui_env":
+            print("[PrismAudio] WARNING: using ComfyUI Python env — JAX/TF/videoprism must already be installed. "
+                  "Installing them here may conflict with existing packages and destabilize ComfyUI.", flush=True)
+            python_bin = sys.executable
+        else:
+            python_bin = _ensure_extract_env()
+
+        # Determine cache directory
+        if not cache_dir:
+            cache_dir = os.path.join(tempfile.gettempdir(), "prismaudio_features")
+        os.makedirs(cache_dir, exist_ok=True)
+
+        # Check cache
+        cache_hash = _hash_inputs(video, caption_cot)
+        cached_path = os.path.join(cache_dir, f"{cache_hash}.npz")
+        if os.path.exists(cached_path):
+            print(f"[PrismAudio] Using cached features: {cached_path}")
+            loader = PrismAudioFeatureLoader()
+            features, = loader.load_features(cached_path)
+            return (features, float(fps))
+
+        # Save frames to temp file (lossless .npy, no codec roundtrip)
+        import time
+        t0 = time.perf_counter()
+        frames = video.shape[0]
+        print(f"[PrismAudio] Saving {frames} frames to .npy (fps={fps})...", flush=True)
+        with tempfile.NamedTemporaryFile(suffix=".npy", delete=False) as tmp:
+            tmp_video = tmp.name
+        _save_frames_to_npy(video, tmp_video)
+        print(f"[PrismAudio] Frames saved in {time.perf_counter() - t0:.1f}s", flush=True)
+
+        # Build subprocess command
+        script_path = os.path.join(
+            os.path.dirname(os.path.dirname(__file__)),
+            "scripts", "extract_features.py"
+        )
+
+        import folder_paths
+        synchformer_ckpt = os.path.join(folder_paths.models_dir, "prismaudio", "synchformer_state_dict.pth")
+        if not os.path.exists(synchformer_ckpt):
+            raise RuntimeError(
+                f"[PrismAudio] Synchformer checkpoint not found: {synchformer_ckpt}\n"
+                "Download synchformer_state_dict.pth from FunAudioLLM/PrismAudio and place it in models/prismaudio/."
+            )
+
+        cmd = [
+            python_bin,
+            script_path,
+            "--video", tmp_video,
+            "--cot_text", caption_cot,
+            "--output", cached_path,
+            "--source_fps", str(fps),
+            "--synchformer_ckpt", synchformer_ckpt,
+        ]
+
+        # Build env: inherit current env, inject HF token if provided
+        import copy
+        env = copy.copy(os.environ)
+        token = hf_token.strip() if hf_token else os.environ.get("HF_TOKEN", "")
+        if token:
+            env["HF_TOKEN"] = token
+            env["HUGGING_FACE_HUB_TOKEN"] = token
+        else:
+            print("[PrismAudio] Warning: no HF_TOKEN set — gated models (e.g. t5gemma) will fail. "
+                  "Add your token in the hf_token input or set HF_TOKEN env var.", flush=True)
+
+        print(f"[PrismAudio] Extracting features via subprocess (output streams live)...")
+        try:
+            # capture_output=False: let stdout/stderr stream directly to ComfyUI logs
+            result = subprocess.run(
+                cmd,
+                capture_output=False,
+                timeout=600,  # 10 minute timeout
+                env=env,
+            )
+            if result.returncode != 0:
+                raise RuntimeError(
+                    f"[PrismAudio] Feature extraction subprocess exited with code {result.returncode}. "
+                    "See output above for details."
+                )
+            print("[PrismAudio] Feature extraction subprocess finished successfully.")
+        finally:
+            if os.path.exists(tmp_video):
+                os.unlink(tmp_video)
+
+        # Load the extracted features
+        loader = PrismAudioFeatureLoader()
+        features, = loader.load_features(cached_path)
+        return (features, float(fps))

nodes/feature_loader.py

@@ -0,0 +1,53 @@
+import os
+import numpy as np
+import torch
+from .utils import PRISMAUDIO_CATEGORY
+
+# Keys consumed by the conditioners (video_features, text_features, sync_features)
+# global_video_features and global_text_features are NOT consumed by any conditioner
+# in the prismaudio.json config — they are unused.
+REQUIRED_KEYS = [
+    "video_features",
+    "text_features",
+    "sync_features",
+]
+
+
+class PrismAudioFeatureLoader:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "npz_path": ("STRING", {"default": "", "tooltip": "Path to pre-computed .npz feature file"}),
+            },
+        }
+
+    RETURN_TYPES = ("PRISMAUDIO_FEATURES",)
+    RETURN_NAMES = ("features",)
+    FUNCTION = "load_features"
+    CATEGORY = PRISMAUDIO_CATEGORY
+
+    def load_features(self, npz_path):
+        if not os.path.exists(npz_path):
+            raise FileNotFoundError(f"[PrismAudio] Feature file not found: {npz_path}")
+
+        data = np.load(npz_path, allow_pickle=True)
+
+        features = {}
+        for key in REQUIRED_KEYS:
+            if key in data:
+                features[key] = torch.from_numpy(data[key]).float()
+            else:
+                print(f"[PrismAudio] Warning: key '{key}' not found in {npz_path}, using zeros")
+                # Provide zero tensor rather than None — Cond_MLP/Sync_MLP crash on None
+                # Sync_MLP requires length divisible by 8 (segments of 8 frames)
+                if key == "sync_features":
+                    features[key] = torch.zeros(8, 768)
+                else:
+                    features[key] = torch.zeros(1, 1024)
+
+        # Load duration if present
+        if "duration" in data:
+            features["duration"] = float(data["duration"])
+
+        return (features,)

nodes/model_loader.py

@@ -0,0 +1,154 @@
+import os
+import json
+import torch
+import folder_paths
+import comfy.model_management as mm
+import comfy.utils
+
+from .utils import (
+    PRISMAUDIO_CATEGORY, get_prismaudio_model_dir, register_model_folder,
+    get_device, get_offload_device, determine_precision, determine_offload_strategy,
+    soft_empty_cache, resolve_hf_token,
+)
+
+# HuggingFace repo for auto-download
+HF_REPO_ID = "FunAudioLLM/PrismAudio"
+REQUIRED_FILES = {
+    "diffusion": "prismaudio.ckpt",
+    "vae": "vae.ckpt",
+    "synchformer": "synchformer_state_dict.pth",
+}
+
+
+def _download_if_missing(filename, model_dir, hf_token=None):
+    """Download a model file from HuggingFace if not present locally."""
+    filepath = os.path.join(model_dir, filename)
+    if os.path.exists(filepath):
+        return filepath
+
+    from huggingface_hub import hf_hub_download
+    print(f"[PrismAudio] Downloading {filename} from {HF_REPO_ID}...")
+    try:
+        downloaded = hf_hub_download(
+            repo_id=HF_REPO_ID,
+            filename=filename,
+            local_dir=model_dir,
+            token=hf_token or None,
+        )
+        return downloaded
+    except Exception as e:
+        if "401" in str(e) or "403" in str(e) or "gated" in str(e).lower():
+            raise RuntimeError(
+                f"[PrismAudio] Model '{filename}' requires license acceptance. "
+                f"Visit https://huggingface.co/{HF_REPO_ID} to accept the license, "
+                f"then set HF_TOKEN env var or run: huggingface-cli login"
+            ) from e
+        raise
+
+
+class PrismAudioModelLoader:
+    @classmethod
+    def INPUT_TYPES(cls):
+        register_model_folder()
+        return {
+            "required": {
+                "precision": (["auto", "fp32", "fp16", "bf16"],),
+                "offload_strategy": (["auto", "keep_in_vram", "offload_to_cpu"],),
+            },
+        }
+
+    RETURN_TYPES = ("PRISMAUDIO_MODEL",)
+    RETURN_NAMES = ("model",)
+    FUNCTION = "load_model"
+    CATEGORY = PRISMAUDIO_CATEGORY
+
+    def load_model(self, precision, offload_strategy):
+        device = get_device()
+        dtype = determine_precision(precision, device)
+        strategy = determine_offload_strategy(offload_strategy)
+        token = resolve_hf_token()
+        model_dir = get_prismaudio_model_dir()
+
+        # Auto-download missing files
+        for key, filename in REQUIRED_FILES.items():
+            _download_if_missing(filename, model_dir, hf_token=token)
+
+        # Load config
+        config_path = os.path.join(
+            os.path.dirname(os.path.dirname(__file__)),
+            "prismaudio_core", "configs", "prismaudio.json"
+        )
+        with open(config_path) as f:
+            model_config = json.load(f)
+
+        # Create model from config
+        from prismaudio_core.factory import create_model_from_config
+        model = create_model_from_config(model_config)
+
+        # Load diffusion weights
+        diffusion_path = os.path.join(model_dir, REQUIRED_FILES["diffusion"])
+        diffusion_state = comfy.utils.load_torch_file(diffusion_path)
+        # Handle wrapped state dicts: some ckpts wrap in {"state_dict": ...}
+        if "state_dict" in diffusion_state:
+            diffusion_state = diffusion_state["state_dict"]
+        diff_result = model.load_state_dict(diffusion_state, strict=False)
+        print(f"[PrismAudio] Diffusion ckpt: {len(diffusion_state)} keys in file", flush=True)
+        print(f"[PrismAudio] Diffusion load: missing={len(diff_result.missing_keys)}, unexpected={len(diff_result.unexpected_keys)}", flush=True)
+        if diff_result.missing_keys:
+            print(f"[PrismAudio]   missing (first 10): {diff_result.missing_keys[:10]}", flush=True)
+        if diff_result.unexpected_keys:
+            print(f"[PrismAudio]   unexpected (first 5): {diff_result.unexpected_keys[:5]}", flush=True)
+        # Sample a few ckpt keys to verify prefix alignment
+        sample_keys = list(diffusion_state.keys())[:5]
+        print(f"[PrismAudio]   ckpt key samples: {sample_keys}", flush=True)
+
+        # Load VAE weights separately
+        # Use comfy.utils.load_torch_file for consistency and PyTorch 2.6+ compat
+        vae_path = os.path.join(model_dir, REQUIRED_FILES["vae"])
+        vae_full_state = comfy.utils.load_torch_file(vae_path)
+        print(f"[PrismAudio] VAE ckpt: {len(vae_full_state)} keys in file", flush=True)
+        # Sample raw keys to see actual prefix
+        vae_sample_keys = list(vae_full_state.keys())[:8]
+        print(f"[PrismAudio]   VAE raw key samples: {vae_sample_keys}", flush=True)
+        # Strip "autoencoder." prefix from keys
+        vae_state = {}
+        prefix = "autoencoder."
+        for k, v in vae_full_state.items():
+            if k.startswith(prefix):
+                vae_state[k[len(prefix):]] = v
+            else:
+                vae_state[k] = v
+        print(f"[PrismAudio]   VAE after strip: {len(vae_state)} keys", flush=True)
+        # Sample model keys to compare
+        model_vae_keys = list(model.pretransform.state_dict().keys())[:5]
+        print(f"[PrismAudio]   pretransform model key samples: {model_vae_keys}", flush=True)
+        # strict=False: vae.ckpt is a training checkpoint that also contains
+        # discriminator, loss modules, and EMA wrappers not present in the
+        # inference AudioAutoencoder — ignore those extra keys.
+        # Load directly into the inner AudioAutoencoder to get IncompatibleKeys back
+        # (AutoencoderPretransform.load_state_dict doesn't return the result)
+        vae_result = model.pretransform.model.load_state_dict(vae_state, strict=False)
+        print(f"[PrismAudio] VAE load: missing={len(vae_result.missing_keys)}, unexpected={len(vae_result.unexpected_keys)}", flush=True)
+        if vae_result.missing_keys:
+            print(f"[PrismAudio]   VAE missing (first 10): {vae_result.missing_keys[:10]}", flush=True)
+
+        # Apply precision: DiT + conditioners in user-selected dtype,
+        # but keep VAE (pretransform) in fp32 to avoid NaN from snake activations in fp16
+        model.model.to(dtype)        # DiTWrapper
+        model.conditioner.to(dtype)  # MultiConditioner
+        # model.pretransform stays in fp32
+
+        if strategy == "keep_in_vram":
+            model = model.to(device)
+        else:
+            model = model.to(get_offload_device())
+
+        model.eval()
+
+        return ({
+            "model": model,
+            "dtype": dtype,
+            "strategy": strategy,
+            "config": model_config,
+            "model_dir": model_dir,
+        },)

nodes/sampler.py

@@ -0,0 +1,165 @@
+import torch
+import comfy.model_management as mm
+import comfy.utils
+
+from .utils import (
+    PRISMAUDIO_CATEGORY, SAMPLE_RATE, DOWNSAMPLING_RATIO, IO_CHANNELS,
+    get_device, get_offload_device, soft_empty_cache,
+)
+
+
+class PrismAudioSampler:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("PRISMAUDIO_MODEL",),
+                "features": ("PRISMAUDIO_FEATURES",),
+                "duration": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 30.0, "step": 0.1, "tooltip": "Audio duration in seconds. Set to 0 to use the video duration from features automatically."}),
+                "steps": ("INT", {"default": 100, "min": 1, "max": 100, "tooltip": "Number of sampling steps"}),
+                "cfg_scale": ("FLOAT", {"default": 7.0, "min": 1.0, "max": 20.0, "step": 0.1, "tooltip": "Classifier-free guidance scale"}),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFF}),
+            },
+        }
+
+    RETURN_TYPES = ("AUDIO",)
+    RETURN_NAMES = ("audio",)
+    FUNCTION = "generate"
+    CATEGORY = PRISMAUDIO_CATEGORY
+
+    def generate(self, model, features, duration, steps, cfg_scale, seed):
+        device = get_device()
+        dtype = model["dtype"]
+        strategy = model["strategy"]
+        diffusion = model["model"]
+
+        # Resolve duration: 0 means use video duration from features
+        if duration <= 0:
+            if "duration" not in features:
+                raise ValueError("[PrismAudio] duration=0 but features contain no duration. Set duration manually or use PrismAudioFeatureExtractor.")
+            duration = features["duration"]
+            print(f"[PrismAudio] Using video duration from features: {duration:.2f}s", flush=True)
+
+        # Compute latent dimensions
+        latent_length = round(SAMPLE_RATE * duration / DOWNSAMPLING_RATIO)
+
+        # Note: no seq length config needed — the model adapts to input tensor shapes
+        # dynamically via its transformer architecture.
+
+        # Determine if video features are present (not all zeros)
+        has_video = features.get("video_features") is not None and features["video_features"].abs().sum() > 0
+
+        video_feat = features["video_features"].to(device, dtype=dtype)
+        sync_feat  = features["sync_features"].to(device, dtype=dtype)
+
+        # Build metadata as a TUPLE of dicts (one per batch sample)
+        # MultiConditioner.forward(batch_metadata: List[Dict]) iterates over this
+        sample_meta = {
+            "video_features": video_feat,
+            "text_features":  features["text_features"].to(device, dtype=dtype),
+            "sync_features":  sync_feat,
+            "video_exist": torch.tensor(has_video),
+        }
+        metadata = (sample_meta,)
+
+        # Move model to device if offloaded
+        if strategy == "offload_to_cpu":
+            diffusion.model.to(device)
+            diffusion.conditioner.to(device)
+            soft_empty_cache()
+
+        with torch.no_grad(), torch.amp.autocast(device_type=device.type, dtype=dtype):
+            # Run conditioning
+            conditioning = diffusion.conditioner(metadata, device)
+
+            # Handle missing video: substitute learned empty embeddings
+            if not has_video:
+                _substitute_empty_features(diffusion, conditioning, device, dtype)
+
+            # Assemble conditioning inputs for the DiT
+            cond_inputs = diffusion.get_conditioning_inputs(conditioning)
+
+            # Generate noise from seed (MPS doesn't support torch.Generator)
+            gen_device = "cpu" if device.type == "mps" else device
+            generator = torch.Generator(device=gen_device).manual_seed(seed)
+            noise = torch.randn(
+                [1, IO_CHANNELS, latent_length],
+                generator=generator,
+                device=gen_device,
+            ).to(device=device, dtype=dtype)
+
+            # Sample with progress bar
+            pbar = comfy.utils.ProgressBar(steps)
+
+            from prismaudio_core.inference.sampling import sample_discrete_euler
+
+            def on_step(info):
+                pbar.update(1)
+
+            fakes = sample_discrete_euler(
+                diffusion.model,
+                noise,
+                steps,
+                callback=on_step,
+                **cond_inputs,
+                cfg_scale=cfg_scale,
+                batch_cfg=True,
+            )
+
+            fakes_f = fakes.float()
+            print(f"[PrismAudio] latent stats: shape={tuple(fakes_f.shape)} mean={fakes_f.mean():.4f} std={fakes_f.std():.4f} min={fakes_f.min():.4f} max={fakes_f.max():.4f}", flush=True)
+
+            # Offload diffusion model and conditioner before VAE decode
+            if strategy == "offload_to_cpu":
+                diffusion.model.to(get_offload_device())
+                diffusion.conditioner.to(get_offload_device())
+                soft_empty_cache()
+                diffusion.pretransform.to(device)
+
+            # VAE decode in fp32 (snake activations overflow in fp16)
+            with torch.amp.autocast(device_type=device.type, enabled=False):
+                audio = diffusion.pretransform.decode(fakes_f)
+
+            # Offload VAE
+            if strategy == "offload_to_cpu":
+                diffusion.pretransform.to(get_offload_device())
+                soft_empty_cache()
+
+        # Peak normalize then clamp (matching reference: div by max abs before clamp)
+        audio = audio.float()
+        pre_norm_std = audio.std().item()
+        pre_norm_peak = audio.abs().max().item()
+        peak = audio.abs().max().clamp(min=1e-8)
+        audio = (audio / peak).clamp(-1, 1)
+        print(f"[PrismAudio] audio stats (pre-norm): std={pre_norm_std:.4f} peak={pre_norm_peak:.4f}", flush=True)
+
+        # Return as ComfyUI AUDIO: {"waveform": [B, channels, samples], "sample_rate": int}
+        return ({"waveform": audio.cpu(), "sample_rate": SAMPLE_RATE},)
+
+
+def _substitute_empty_features(diffusion, conditioning, device, dtype):
+    """Replace video/sync conditioning with learned empty embeddings when video is absent.
+
+    empty_clip_feat and empty_sync_feat are learned null embeddings in the conditioner
+    output space (1024-dim). Passing zero features through bias-free Cond_MLP produces
+    near-zero activations, NOT the learned null signal the model was trained with.
+
+    The conditioner returns {key: [tensor, mask]} where tensor is [B, seq, dim].
+    """
+    dit = diffusion.model.model if hasattr(diffusion.model, 'model') else diffusion.model
+
+    # Substitute video_features with learned empty_clip_feat
+    if hasattr(dit, 'empty_clip_feat') and 'video_features' in conditioning:
+        empty = dit.empty_clip_feat.to(device, dtype=dtype)           # [1, 1024]
+        batch_size = conditioning['video_features'][0].shape[0]
+        empty_expanded = empty.unsqueeze(0).expand(batch_size, -1, -1)  # [B, 1, 1024]
+        conditioning['video_features'][0] = empty_expanded
+        conditioning['video_features'][1] = torch.ones(batch_size, 1, device=device)
+
+    # Substitute sync_features with learned empty_sync_feat
+    if hasattr(dit, 'empty_sync_feat') and 'sync_features' in conditioning:
+        empty = dit.empty_sync_feat.to(device, dtype=dtype)           # [1, 1024]
+        batch_size = conditioning['sync_features'][0].shape[0]
+        empty_expanded = empty.unsqueeze(0).expand(batch_size, -1, -1)  # [B, 1, 1024]
+        conditioning['sync_features'][0] = empty_expanded
+        conditioning['sync_features'][1] = torch.ones(batch_size, 1, device=device)

nodes/text_only.py

@@ -0,0 +1,160 @@
+import torch
+import comfy.model_management as mm
+import comfy.utils
+
+from .utils import (
+    PRISMAUDIO_CATEGORY, SAMPLE_RATE, DOWNSAMPLING_RATIO, IO_CHANNELS,
+    get_device, get_offload_device, soft_empty_cache, resolve_hf_token,
+)
+from .sampler import _substitute_empty_features
+
+
+class PrismAudioTextOnly:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("PRISMAUDIO_MODEL",),
+                "text_prompt": ("STRING", {"default": "", "multiline": True, "tooltip": "Detailed chain-of-thought description of the audio scene. Use long, descriptive text — e.g. 'A large dog barks sharply twice, with ambient outdoor background noise. The sound is clear and close.' Short prompts produce lower quality."}),
+                "duration": ("FLOAT", {"default": 10.0, "min": 1.0, "max": 30.0, "step": 0.1}),
+                "steps": ("INT", {"default": 100, "min": 1, "max": 100}),
+                "cfg_scale": ("FLOAT", {"default": 7.0, "min": 1.0, "max": 20.0, "step": 0.1}),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFF}),
+            },
+        }
+
+    RETURN_TYPES = ("AUDIO",)
+    RETURN_NAMES = ("audio",)
+    FUNCTION = "generate"
+    CATEGORY = PRISMAUDIO_CATEGORY
+
+    def generate(self, model, text_prompt, duration, steps, cfg_scale, seed):
+        device = get_device()
+        dtype = model["dtype"]
+        strategy = model["strategy"]
+        diffusion = model["model"]
+
+        latent_length = round(SAMPLE_RATE * duration / DOWNSAMPLING_RATIO)
+
+        # Encode text with T5-Gemma
+        text_features = _encode_text_t5(text_prompt, device, dtype)
+
+        # Build metadata: tuple of one dict per sample
+        # Use zero tensors for video/sync (not None — Cond_MLP crashes on None via pad_sequence)
+        # Sync_MLP requires length divisible by 8 (segments of 8 frames) — minimum [8, 768]
+        # These will be substituted with learned empty embeddings after conditioning
+        sample_meta = {
+            "video_features": torch.zeros(1, 1024, device=device, dtype=dtype),
+            "text_features": text_features.to(device, dtype=dtype),
+            "sync_features": torch.zeros(8, 768, device=device, dtype=dtype),
+            "video_exist": torch.tensor(False),
+        }
+        metadata = (sample_meta,)
+
+        if strategy == "offload_to_cpu":
+            diffusion.model.to(device)
+            diffusion.conditioner.to(device)
+            soft_empty_cache()
+
+        with torch.no_grad(), torch.amp.autocast(device_type=device.type, dtype=dtype):
+            conditioning = diffusion.conditioner(metadata, device)
+
+            # Substitute empty features for video/sync
+            _substitute_empty_features(diffusion, conditioning, device, dtype)
+
+            cond_inputs = diffusion.get_conditioning_inputs(conditioning)
+
+            # Generate noise from seed (MPS doesn't support torch.Generator)
+            gen_device = "cpu" if device.type == "mps" else device
+            generator = torch.Generator(device=gen_device).manual_seed(seed)
+            noise = torch.randn(
+                [1, IO_CHANNELS, latent_length],
+                generator=generator,
+                device=gen_device,
+            ).to(device=device, dtype=dtype)
+
+            pbar = comfy.utils.ProgressBar(steps)
+
+            from prismaudio_core.inference.sampling import sample_discrete_euler
+
+            def on_step(info):
+                pbar.update(1)
+
+            fakes = sample_discrete_euler(
+                diffusion.model,
+                noise,
+                steps,
+                callback=on_step,
+                **cond_inputs,
+                cfg_scale=cfg_scale,
+                batch_cfg=True,
+            )
+
+            fakes_f = fakes.float()
+            print(f"[PrismAudio] latent stats: shape={tuple(fakes_f.shape)} mean={fakes_f.mean():.4f} std={fakes_f.std():.4f} min={fakes_f.min():.4f} max={fakes_f.max():.4f}", flush=True)
+
+            if strategy == "offload_to_cpu":
+                diffusion.model.to(get_offload_device())
+                diffusion.conditioner.to(get_offload_device())
+                soft_empty_cache()
+                diffusion.pretransform.to(device)
+
+            # VAE decode in fp32 (snake activations overflow in fp16)
+            with torch.amp.autocast(device_type=device.type, enabled=False):
+                audio = diffusion.pretransform.decode(fakes_f)
+
+            if strategy == "offload_to_cpu":
+                diffusion.pretransform.to(get_offload_device())
+                soft_empty_cache()
+
+        # Peak normalize then clamp
+        audio = audio.float()
+        pre_norm_std = audio.std().item()
+        pre_norm_peak = audio.abs().max().item()
+        peak = audio.abs().max().clamp(min=1e-8)
+        audio = (audio / peak).clamp(-1, 1)
+        print(f"[PrismAudio] audio stats (pre-norm): std={pre_norm_std:.4f} peak={pre_norm_peak:.4f}", flush=True)
+        print(f"[PrismAudio] audio shape: {tuple(audio.shape)}", flush=True)
+
+        return ({"waveform": audio.cpu(), "sample_rate": SAMPLE_RATE},)
+
+
+# T5-Gemma encoder singleton
+_t5_model = None
+_t5_tokenizer = None
+
+
+def _encode_text_t5(text, device, dtype):
+    """Encode text using T5-Gemma.
+
+    Uses AutoModelForSeq2SeqLM.get_encoder() to match the reference
+    FeaturesUtils.encode_t5_text() implementation.
+    No truncation applied (matching reference behavior).
+    """
+    global _t5_model, _t5_tokenizer
+
+    if _t5_model is None:
+        from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+        model_id = "google/t5gemma-l-l-ul2-it"
+        token = resolve_hf_token()
+        print(f"[PrismAudio] Loading T5-Gemma text encoder: {model_id}")
+        _t5_tokenizer = AutoTokenizer.from_pretrained(model_id, token=token)
+        _t5_model = AutoModelForSeq2SeqLM.from_pretrained(model_id, token=token).get_encoder()
+        _t5_model.eval()
+
+    _t5_model.to(device, dtype=dtype)
+
+    tokens = _t5_tokenizer(
+        text,
+        return_tensors="pt",
+        padding=True,
+    ).to(device)
+
+    with torch.no_grad():
+        outputs = _t5_model(**tokens)
+
+    # Move T5 off GPU after encoding to save VRAM
+    _t5_model.to("cpu")
+    soft_empty_cache()
+
+    return outputs.last_hidden_state.squeeze(0)  # [seq_len, dim]

prismaudio_core/factory.py

@@ -51,14 +51,7 @@ def create_pretransform_from_config(pretransform_config, sample_rate):
 
         pretransform = AutoencoderPretransform(autoencoder, scale=scale, model_half=model_half, iterate_batch=iterate_batch, chunked=chunked)
     elif pretransform_type == 'wavelet':
-        from prismaudio_core.models.pretransforms import WaveletPretransform
+        raise NotImplementedError("wavelet pretransform type is not supported")
-
-        wavelet_config = pretransform_config["config"]
-        channels = wavelet_config["channels"]
-        levels = wavelet_config["levels"]
-        wavelet = wavelet_config["wavelet"]
-
-        pretransform = WaveletPretransform(channels, levels, wavelet)
     elif pretransform_type == 'pqmf':
         from prismaudio_core.models.pretransforms import PQMFPretransform
         pqmf_config = pretransform_config["config"]
@@ -327,7 +320,6 @@ def create_diffusion_cond_from_config(config: tp.Dict[str, tp.Any]):
         UNetCFG1DWrapper,
         UNet1DCondWrapper,
         DiTWrapper,
-        MMDiTWrapper,
     )
 
     model_config = config["model"]
@@ -350,7 +342,7 @@ def create_diffusion_cond_from_config(config: tp.Dict[str, tp.Any]):
     elif diffusion_model_type == 'dit':
         diffusion_model = DiTWrapper(**diffusion_model_config)
     elif diffusion_model_type == 'mmdit':
-        diffusion_model = MMDiTWrapper(**diffusion_model_config)
+        raise NotImplementedError("mmdit diffusion model type is not supported")
 
     io_channels = model_config.get('io_channels', None)
     assert io_channels is not None, "Must specify io_channels in model config"
@@ -401,12 +393,7 @@ def create_diffusion_cond_from_config(config: tp.Dict[str, tp.Any]):
         extra_kwargs["diffusion_objective"] = diffusion_objective
 
     elif model_type == "diffusion_prior":
-        prior_type = model_config.get("prior_type", None)
+        raise NotImplementedError("diffusion_prior model type is not supported")
-        assert prior_type is not None, "Must specify prior_type in diffusion prior model config"
-
-        if prior_type == "mono_stereo":
-            from prismaudio_core.models.diffusion_prior import MonoToStereoDiffusionPrior
-            wrapper_fn = MonoToStereoDiffusionPrior
 
     return wrapper_fn(
         diffusion_model,

prismaudio_core/inference/__init__.py

@@ -0,0 +1,4 @@
+from .sampling import sample_discrete_euler
+from .utils import set_audio_channels, prepare_audio
+
+__all__ = ["sample_discrete_euler", "set_audio_channels", "prepare_audio"]

prismaudio_core/inference/sampling.py

@@ -0,0 +1,29 @@
+import torch
+
+
+@torch.no_grad()
+def sample_discrete_euler(model, x, steps, sigma_max=1, callback=None, **extra_args):
+    """Discrete Euler sampler for rectified flow, with optional callback.
+
+    Modified from PrismAudio to add callback parameter for ComfyUI progress reporting.
+    Original uses tqdm internally.
+
+    Args:
+        model: The diffusion model (DiTWrapper)
+        x: Initial noise tensor [B, C, T]
+        steps: Number of sampling steps
+        sigma_max: Maximum sigma (default 1.0 for rectified flow)
+        callback: Optional callable({"i": step, "x": current_x}) for progress
+        **extra_args: Passed to model() — includes cross_attn_cond, add_cond,
+                      sync_cond, cfg_scale, batch_cfg, etc.
+    """
+    t = torch.linspace(sigma_max, 0, steps + 1, device=x.device, dtype=x.dtype)
+
+    for i, (t_curr, t_next) in enumerate(zip(t[:-1], t[1:])):
+        dt = t_next - t_curr
+        t_curr_tensor = t_curr * torch.ones(x.shape[0], dtype=x.dtype, device=x.device)
+        x = x + dt * model(x, t_curr_tensor, **extra_args)
+        if callback is not None:
+            callback({"i": i, "x": x})
+
+    return x

prismaudio_core/inference/utils.py

@@ -0,0 +1,62 @@
+import torch
+import torch.nn.functional as F
+from torchaudio import transforms as T
+
+
+def set_audio_channels(audio, target_channels):
+    """Convert audio tensor to target number of channels.
+
+    Args:
+        audio: Audio tensor of shape [B, C, T]
+        target_channels: Desired number of channels (1 for mono, 2 for stereo)
+
+    Returns:
+        Audio tensor with the target number of channels.
+    """
+    if target_channels == 1:
+        # Convert to mono
+        audio = audio.mean(1, keepdim=True)
+    elif target_channels == 2:
+        # Convert to stereo
+        if audio.shape[1] == 1:
+            audio = audio.repeat(1, 2, 1)
+        elif audio.shape[1] > 2:
+            audio = audio[:, :2, :]
+    return audio
+
+
+def prepare_audio(audio, in_sr, target_sr, target_length, target_channels, device):
+    """Resample, pad/trim, and convert channels of an audio tensor.
+
+    Args:
+        audio: Audio tensor (1D, 2D [C, T], or 3D [B, C, T])
+        in_sr: Input sample rate
+        target_sr: Target sample rate
+        target_length: Target length in samples (padded or cropped)
+        target_channels: Target number of channels
+        device: Torch device to place the audio on
+
+    Returns:
+        Audio tensor of shape [B, target_channels, target_length] on device.
+    """
+    audio = audio.to(device)
+
+    if in_sr != target_sr:
+        resample_tf = T.Resample(in_sr, target_sr).to(device)
+        audio = resample_tf(audio)
+
+    # Add batch dimension
+    if audio.dim() == 1:
+        audio = audio.unsqueeze(0).unsqueeze(0)
+    elif audio.dim() == 2:
+        audio = audio.unsqueeze(0)
+
+    # Pad or crop to target_length
+    if audio.shape[-1] < target_length:
+        audio = F.pad(audio, (0, target_length - audio.shape[-1]))
+    elif audio.shape[-1] > target_length:
+        audio = audio[:, :, :target_length]
+
+    audio = set_audio_channels(audio, target_channels)
+
+    return audio

prismaudio_core/models/transformer.py

@@ -919,6 +919,12 @@ class ContinuousTransformer(nn.Module):
                 x = self.fusion_mlp(x)
 
             if sync_cond is not None:
+                # Resample sync_cond to match audio sequence length if needed
+                if sync_cond.shape[1] != x.shape[1]:
+                    sync_cond = torch.nn.functional.interpolate(
+                        sync_cond.transpose(1, 2), size=x.shape[1],
+                        mode='linear', align_corners=False,
+                    ).transpose(1, 2)
                 if self.sync_film_generator is not None:
                     scale, shift = self.sync_film_generator(sync_cond).chunk(2, dim=-1)
                     x = x * (1 + scale) + shift

requirements.txt

@@ -1,8 +1,11 @@
 einops>=0.7.0
+einops-exts
 safetensors
 huggingface_hub
 transformers>=4.52.3
 k-diffusion>=0.1.1
 alias-free-torch
 descript-audio-codec
+vector-quantize-pytorch
+scipy
 tqdm

scripts/environment.yml

@@ -0,0 +1,21 @@
+name: prismaudio-extract
+channels:
+  - conda-forge
+  - defaults
+dependencies:
+  - python=3.10
+  - pip
+  - ffmpeg<7
+  - pip:
+    - torch>=2.6.0
+    - torchaudio>=2.6.0
+    - torchvision>=0.21.0
+    - tensorflow-cpu==2.15.0
+    - jax
+    - jaxlib
+    - transformers>=4.52.3
+    - decord
+    - einops>=0.7.0
+    - numpy
+    - mediapy
+    - git+https://github.com/google-deepmind/videoprism.git

scripts/extract_features.py

@@ -0,0 +1,168 @@
+#!/usr/bin/env python3
+"""
+Standalone PrismAudio feature extraction script.
+Runs in a separate Python env with JAX/TF installed (auto-created by PrismAudioFeatureExtractor).
+
+Usage:
+    python extract_features.py --video input.mp4 --cot_text "description..." --output features.npz
+"""
+
+import argparse
+import os
+import sys
+import time
+import numpy as np
+import torch
+
+# Add plugin root to sys.path so data_utils (and prismaudio_core) are importable
+_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
+_PLUGIN_DIR = os.path.dirname(_SCRIPT_DIR)
+if _PLUGIN_DIR not in sys.path:
+    sys.path.insert(0, _PLUGIN_DIR)
+
+
+def _step(n, total, label):
+    """Print step header and return start time."""
+    print(f"[extract] Step {n}/{total} — {label}...", flush=True)
+    return time.perf_counter()
+
+
+def _done(t0, extra=""):
+    elapsed = time.perf_counter() - t0
+    suffix = f"  {extra}" if extra else ""
+    print(f"[extract]   done in {elapsed:.1f}s{suffix}", flush=True)
+
+
+def main():
+    t_total = time.perf_counter()
+
+    parser = argparse.ArgumentParser(description="PrismAudio feature extraction")
+    parser.add_argument("--video", required=True, help="Path to input video")
+    parser.add_argument("--cot_text", required=True, help="Chain-of-thought description")
+    parser.add_argument("--output", required=True, help="Output .npz path")
+    parser.add_argument("--synchformer_ckpt", default=None, help="Path to synchformer checkpoint")
+    parser.add_argument("--vae_config", default=None, help="Path to VAE config JSON")
+    parser.add_argument("--source_fps", type=float, default=30.0, help="Original video fps (used when --video is a .npy file)")
+    parser.add_argument("--clip_fps", type=float, default=4.0)
+    parser.add_argument("--clip_size", type=int, default=288)
+    parser.add_argument("--sync_fps", type=float, default=25.0)
+    parser.add_argument("--sync_size", type=int, default=224)
+    args = parser.parse_args()
+
+    print(f"[extract] Python     : {sys.executable}", flush=True)
+    print(f"[extract] Video      : {args.video}", flush=True)
+    print(f"[extract] Output     : {args.output}", flush=True)
+    print(f"[extract] CoT text   : {args.cot_text[:80]}{'...' if len(args.cot_text) > 80 else ''}", flush=True)
+
+    if not os.path.exists(args.video):
+        print(f"[extract] ERROR: video not found: {args.video}", flush=True)
+        sys.exit(1)
+
+    print(f"[extract] Device     : {'cuda' if torch.cuda.is_available() else 'cpu'}", flush=True)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # ------------------------------------------------------------------
+    t0 = _step(1, 6, "importing dependencies")
+    from data_utils.v2a_utils.feature_utils_288 import FeaturesUtils
+    import torchvision.transforms as T
+    _done(t0)
+
+    # ------------------------------------------------------------------
+    t0 = _step(2, 6, "loading models (T5, VideoPrism, Synchformer)")
+    feat_utils = FeaturesUtils(
+        vae_config_path=args.vae_config,
+        synchformer_ckpt=args.synchformer_ckpt,
+        device=device,
+    )
+    _done(t0)
+
+    # ------------------------------------------------------------------
+    t0 = _step(3, 6, "reading and preprocessing video")
+    if args.video.endswith(".npy"):
+        all_frames = np.load(args.video)   # [T, H, W, C] uint8
+        fps = args.source_fps
+        total_frames = all_frames.shape[0]
+        duration = total_frames / fps
+        print(f"[extract]   fps={fps:.3f}  frames={total_frames}  duration={duration:.2f}s", flush=True)
+
+        clip_indices = [int(i * fps / args.clip_fps) for i in range(int(duration * args.clip_fps))]
+        clip_indices = [min(i, total_frames - 1) for i in clip_indices]
+        clip_frames = all_frames[clip_indices]
+        print(f"[extract]   CLIP frames  : {len(clip_indices)} @ {args.clip_fps}fps → {args.clip_size}×{args.clip_size}", flush=True)
+
+        sync_indices = [int(i * fps / args.sync_fps) for i in range(int(duration * args.sync_fps))]
+        sync_indices = [min(i, total_frames - 1) for i in sync_indices]
+        sync_frames = all_frames[sync_indices]
+        print(f"[extract]   Sync frames  : {len(sync_indices)} @ {args.sync_fps}fps → {args.sync_size}×{args.sync_size}", flush=True)
+    else:
+        from decord import VideoReader, cpu
+        vr = VideoReader(args.video, ctx=cpu(0))
+        fps = vr.get_avg_fps()
+        total_frames = len(vr)
+        duration = total_frames / fps
+        print(f"[extract]   fps={fps:.3f}  frames={total_frames}  duration={duration:.2f}s", flush=True)
+
+        clip_indices = [int(i * fps / args.clip_fps) for i in range(int(duration * args.clip_fps))]
+        clip_indices = [min(i, total_frames - 1) for i in clip_indices]
+        clip_frames = vr.get_batch(clip_indices).asnumpy()
+        print(f"[extract]   CLIP frames  : {len(clip_indices)} @ {args.clip_fps}fps → {args.clip_size}×{args.clip_size}", flush=True)
+
+        sync_indices = [int(i * fps / args.sync_fps) for i in range(int(duration * args.sync_fps))]
+        sync_indices = [min(i, total_frames - 1) for i in sync_indices]
+        sync_frames = vr.get_batch(sync_indices).asnumpy()
+        print(f"[extract]   Sync frames  : {len(sync_indices)} @ {args.sync_fps}fps → {args.sync_size}×{args.sync_size}", flush=True)
+
+    clip_transform = T.Compose([
+        T.ToPILImage(),
+        T.Resize(args.clip_size),
+        T.CenterCrop(args.clip_size),
+        T.ToTensor(),
+        T.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+    ])
+    clip_input = torch.stack([clip_transform(f) for f in clip_frames]).unsqueeze(0).to(device)
+
+    sync_transform = T.Compose([
+        T.ToPILImage(),
+        T.Resize(args.sync_size),
+        T.CenterCrop(args.sync_size),
+        T.ToTensor(),
+        T.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+    ])
+    sync_input = torch.stack([sync_transform(f) for f in sync_frames]).unsqueeze(0).to(device)
+    _done(t0)
+
+    # ------------------------------------------------------------------
+    t0 = _step(4, 6, "encoding text with T5-Gemma")
+    text_features = feat_utils.encode_t5_text([args.cot_text])
+    _done(t0, f"shape={tuple(text_features.shape)}")
+
+    # ------------------------------------------------------------------
+    t0 = _step(5, 6, "encoding video with VideoPrism")
+    global_video_features, video_features, global_text_features = \
+        feat_utils.encode_video_and_text_with_videoprism(clip_input, [args.cot_text])
+    _done(t0, f"video={tuple(video_features.shape)}  global={tuple(global_video_features.shape)}")
+
+    # ------------------------------------------------------------------
+    t0 = _step(6, 6, "encoding video with Synchformer")
+    sync_features = feat_utils.encode_video_with_sync(sync_input)
+    _done(t0, f"shape={tuple(sync_features.shape)}")
+
+    # ------------------------------------------------------------------
+    t0 = time.perf_counter()
+    print(f"[extract] Saving features to {args.output} ...", flush=True)
+    np.savez(
+        args.output,
+        video_features=video_features.cpu().float().numpy(),
+        global_video_features=global_video_features.cpu().float().numpy(),
+        text_features=text_features.cpu().float().numpy(),
+        global_text_features=global_text_features.cpu().float().numpy(),
+        sync_features=sync_features.cpu().float().numpy(),
+        caption_cot=args.cot_text,
+        duration=duration,
+    )
+    print(f"[extract] Saved in {time.perf_counter() - t0:.1f}s", flush=True)
+    print(f"[extract] Total time: {time.perf_counter() - t_total:.1f}s", flush=True)
+
+
+if __name__ == "__main__":
+    main()

scripts/install_extract_env.sh

@@ -0,0 +1,44 @@
+#!/usr/bin/env bash
+# Install the PrismAudio feature-extraction environment using pip venv.
+# Use this instead of environment.yml when conda is unavailable (e.g. NVIDIA Docker).
+#
+# Usage:
+#   bash scripts/install_extract_env.sh [/path/to/venv]
+#
+# Default venv path: /opt/prismaudio-extract
+# After installation, point the PrismAudioFeatureExtractor node's python_env to:
+#   <venv>/bin/python   (Linux/Mac)
+#   <venv>\Scripts\python.exe  (Windows)
+
+set -euo pipefail
+
+VENV_DIR="${1:-/opt/prismaudio-extract}"
+
+echo "[PrismAudio] Creating venv at: ${VENV_DIR}"
+python3 -m venv "${VENV_DIR}"
+
+PIP="${VENV_DIR}/bin/pip"
+
+echo "[PrismAudio] Upgrading pip..."
+"${PIP}" install --upgrade pip
+
+echo "[PrismAudio] Installing PyTorch stack..."
+"${PIP}" install torch torchaudio torchvision
+
+echo "[PrismAudio] Installing feature-extraction dependencies..."
+"${PIP}" install \
+    "tensorflow-cpu>=2.16.0" \
+    "jax[cpu]" \
+    "jaxlib" \
+    "transformers" \
+    "decord" \
+    "einops" \
+    "numpy" \
+    "mediapy"
+
+echo "[PrismAudio] Installing VideoPrism..."
+"${PIP}" install "git+https://github.com/google-deepmind/videoprism.git"
+
+echo ""
+echo "[PrismAudio] Done. Set python_env in PrismAudioFeatureExtractor to:"
+echo "  ${VENV_DIR}/bin/python"

workflows/text_to_audio.json

@@ -0,0 +1,158 @@
+{
+  "id": "a1c3e5f7-b2d4-4e6a-8c0f-1a3b5c7d9e2f",
+  "revision": 0,
+  "last_node_id": 3,
+  "last_link_id": 2,
+  "nodes": [
+    {
+      "id": 1,
+      "type": "PrismAudioModelLoader",
+      "pos": [
+        -160,
+        -224
+      ],
+      "size": [
+        288,
+        96
+      ],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "model",
+          "type": "PRISMAUDIO_MODEL",
+          "slot_index": 0,
+          "links": [
+            1
+          ]
+        }
+      ],
+      "properties": {
+        "aux_id": "ethanfel/ComfyUI-Prismaudio",
+        "ver": "62a3c5d",
+        "Node name for S&R": "PrismAudioModelLoader",
+        "ue_properties": {
+          "widget_ue_connectable": {},
+          "version": "7.8",
+          "input_ue_unconnectable": {}
+        }
+      },
+      "widgets_values": [
+        "auto",
+        "auto"
+      ]
+    },
+    {
+      "id": 2,
+      "type": "PrismAudioTextOnly",
+      "pos": [
+        192,
+        -224
+      ],
+      "size": [
+        480,
+        222
+      ],
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "model",
+          "type": "PRISMAUDIO_MODEL",
+          "link": 1
+        }
+      ],
+      "outputs": [
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "slot_index": 0,
+          "links": [
+            2
+          ]
+        }
+      ],
+      "properties": {
+        "aux_id": "ethanfel/ComfyUI-Prismaudio",
+        "ver": "62a3c5d",
+        "Node name for S&R": "PrismAudioTextOnly",
+        "ue_properties": {
+          "widget_ue_connectable": {},
+          "version": "7.8",
+          "input_ue_unconnectable": {}
+        }
+      },
+      "widgets_values": [
+        "A large dog barks sharply twice in an outdoor setting, with ambient background noise of rustling leaves and a gentle breeze. The sound is clear and close, recorded at ground level.",
+        10.0,
+        100,
+        7.0,
+        0,
+        "randomize"
+      ]
+    },
+    {
+      "id": 3,
+      "type": "PreviewAudio",
+      "pos": [
+        736,
+        -224
+      ],
+      "size": [
+        300,
+        76
+      ],
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "link": 2
+        }
+      ],
+      "outputs": [],
+      "properties": {
+        "Node name for S&R": "PreviewAudio"
+      },
+      "widgets_values": []
+    }
+  ],
+  "links": [
+    [
+      1,
+      1,
+      0,
+      2,
+      0,
+      "PRISMAUDIO_MODEL"
+    ],
+    [
+      2,
+      2,
+      0,
+      3,
+      0,
+      "AUDIO"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 1.1674071890328979,
+      "offset": [
+        1814.5534800416863,
+        500.0421331448515
+      ]
+    },
+    "ue_links": [],
+    "links_added_by_ue": [],
+    "frontendVersion": "1.42.8"
+  },
+  "version": 0.4
+}

workflows/video_to_audio.json

@@ -0,0 +1,421 @@
+{
+  "id": "2481bfbf-ce24-46c5-abdc-1d9163ff78ae",
+  "revision": 0,
+  "last_node_id": 12,
+  "last_link_id": 30,
+  "nodes": [
+    {
+      "id": 1,
+      "type": "VHS_LoadVideo",
+      "pos": [
+        -704,
+        -256
+      ],
+      "size": [
+        288,
+        474.7188081936685
+      ],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "meta_batch",
+          "shape": 7,
+          "type": "VHS_BatchManager",
+          "link": null
+        },
+        {
+          "name": "vae",
+          "shape": 7,
+          "type": "VAE",
+          "link": null
+        }
+      ],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "slot_index": 0,
+          "links": [
+            12,
+            20
+          ]
+        },
+        {
+          "name": "frame_count",
+          "type": "INT",
+          "slot_index": 1,
+          "links": []
+        },
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "slot_index": 2,
+          "links": []
+        },
+        {
+          "name": "video_info",
+          "type": "VHS_VIDEOINFO",
+          "slot_index": 3,
+          "links": [
+            21
+          ]
+        }
+      ],
+      "properties": {
+        "cnr_id": "comfyui-videohelpersuite",
+        "ver": "1.7.9",
+        "Node name for S&R": "VHS_LoadVideo",
+        "ue_properties": {
+          "widget_ue_connectable": {},
+          "version": "7.8",
+          "input_ue_unconnectable": {}
+        }
+      },
+      "widgets_values": {
+        "video": "Railtransport_3_479.mp4",
+        "force_rate": 0,
+        "custom_width": 0,
+        "custom_height": 0,
+        "frame_load_cap": 0,
+        "skip_first_frames": 0,
+        "select_every_nth": 1,
+        "format": "AnimateDiff",
+        "videopreview": {
+          "hidden": false,
+          "paused": false,
+          "params": {
+            "force_rate": 0,
+            "frame_load_cap": 0,
+            "skip_first_frames": 0,
+            "select_every_nth": 1,
+            "filename": "Railtransport_3_479.mp4",
+            "type": "input",
+            "format": "video/mp4"
+          }
+        }
+      }
+    },
+    {
+      "id": 2,
+      "type": "PrismAudioModelLoader",
+      "pos": [
+        -160,
+        -224
+      ],
+      "size": [
+        288,
+        96
+      ],
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "model",
+          "type": "PRISMAUDIO_MODEL",
+          "slot_index": 0,
+          "links": [
+            26
+          ]
+        }
+      ],
+      "properties": {
+        "aux_id": "ethanfel/ComfyUI-Prismaudio",
+        "ver": "3894fcc9b40a19d959614d514d5dff65cdfb6eab",
+        "Node name for S&R": "PrismAudioModelLoader",
+        "ue_properties": {
+          "widget_ue_connectable": {},
+          "version": "7.8",
+          "input_ue_unconnectable": {}
+        }
+      },
+      "widgets_values": [
+        "auto",
+        "auto"
+      ]
+    },
+    {
+      "id": 12,
+      "type": "PrismAudioSampler",
+      "pos": [
+        256,
+        -224
+      ],
+      "size": [
+        384,
+        224
+      ],
+      "flags": {},
+      "order": 3,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "model",
+          "type": "PRISMAUDIO_MODEL",
+          "link": 26
+        },
+        {
+          "name": "features",
+          "type": "PRISMAUDIO_FEATURES",
+          "link": 27
+        }
+      ],
+      "outputs": [
+        {
+          "name": "audio",
+          "type": "AUDIO",
+          "links": [
+            29
+          ]
+        }
+      ],
+      "properties": {
+        "aux_id": "ethanfel/ComfyUI-Prismaudio",
+        "ver": "30631c0cb4d97cc6aed69a52e3ee4d89df03926c",
+        "ue_properties": {
+          "widget_ue_connectable": {},
+          "input_ue_unconnectable": {}
+        },
+        "Node name for S&R": "PrismAudioSampler"
+      },
+      "widgets_values": [
+        0,
+        100,
+        7,
+        4096333446,
+        "randomize"
+      ]
+    },
+    {
+      "id": 11,
+      "type": "PrismAudioFeatureExtractor",
+      "pos": [
+        -384,
+        -64
+      ],
+      "size": [
+        544,
+        288
+      ],
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "video",
+          "type": "IMAGE",
+          "link": 20
+        },
+        {
+          "name": "video_info",
+          "shape": 7,
+          "type": "VHS_VIDEOINFO",
+          "link": 21
+        }
+      ],
+      "outputs": [
+        {
+          "name": "features",
+          "type": "PRISMAUDIO_FEATURES",
+          "links": [
+            27
+          ]
+        },
+        {
+          "name": "fps",
+          "type": "FLOAT",
+          "links": [
+            30
+          ]
+        }
+      ],
+      "properties": {
+        "aux_id": "ethanfel/ComfyUI-Prismaudio",
+        "ver": "5b62be04471bf118b2cd3cc71431a302f5730b01",
+        "Node name for S&R": "PrismAudioFeatureExtractor",
+        "ue_properties": {
+          "widget_ue_connectable": {},
+          "input_ue_unconnectable": {},
+          "version": "7.8"
+        }
+      },
+      "widgets_values": [
+        "Generate ambient countryside sounds with a gentle breeze rustling the leaves of a large tree. From the right, introduce a faint rumble of wheels on a track and a steam engine chugging. Allow the sounds to grow louder and pan from right to left as the steam train travels across the landscape. Include the powerful chugging and clattering of carriages in the soundscape, then gradually recede the sounds to the left. Ensure no additional background noise or music is present.\n",
+        30,
+        "managed_env",
+        "/media/unraid/comfyui/output/prismaudiocache/",
+        ""
+      ]
+    },
+    {
+      "id": 9,
+      "type": "VHS_VideoCombine",
+      "pos": [
+        704,
+        -256
+      ],
+      "size": [
+        384,
+        552.75
+      ],
+      "flags": {},
+      "order": 4,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 12
+        },
+        {
+          "name": "audio",
+          "shape": 7,
+          "type": "AUDIO",
+          "link": 29
+        },
+        {
+          "name": "meta_batch",
+          "shape": 7,
+          "type": "VHS_BatchManager",
+          "link": null
+        },
+        {
+          "name": "vae",
+          "shape": 7,
+          "type": "VAE",
+          "link": null
+        },
+        {
+          "name": "frame_rate",
+          "type": "FLOAT",
+          "widget": {
+            "name": "frame_rate"
+          },
+          "link": 30
+        }
+      ],
+      "outputs": [
+        {
+          "name": "Filenames",
+          "type": "VHS_FILENAMES",
+          "links": null
+        }
+      ],
+      "properties": {
+        "cnr_id": "comfyui-videohelpersuite",
+        "ver": "1.7.9",
+        "Node name for S&R": "VHS_VideoCombine",
+        "ue_properties": {
+          "widget_ue_connectable": {},
+          "input_ue_unconnectable": {},
+          "version": "7.8"
+        }
+      },
+      "widgets_values": {
+        "frame_rate": 30,
+        "loop_count": 0,
+        "filename_prefix": "AnimateDiff",
+        "format": "video/h264-mp4",
+        "pix_fmt": "yuv420p",
+        "crf": 19,
+        "save_metadata": true,
+        "trim_to_audio": false,
+        "pingpong": false,
+        "save_output": false,
+        "videopreview": {
+          "hidden": false,
+          "paused": false,
+          "params": {
+            "filename": "AnimateDiff_00001-audio.mp4",
+            "subfolder": "",
+            "type": "temp",
+            "format": "video/h264-mp4",
+            "frame_rate": 30,
+            "workflow": "AnimateDiff_00001.png",
+            "fullpath": "/basedir/temp/AnimateDiff_00001-audio.mp4"
+          }
+        }
+      }
+    }
+  ],
+  "links": [
+    [
+      12,
+      1,
+      0,
+      9,
+      0,
+      "IMAGE"
+    ],
+    [
+      20,
+      1,
+      0,
+      11,
+      0,
+      "IMAGE"
+    ],
+    [
+      21,
+      1,
+      3,
+      11,
+      1,
+      "VHS_VIDEOINFO"
+    ],
+    [
+      26,
+      2,
+      0,
+      12,
+      0,
+      "PRISMAUDIO_MODEL"
+    ],
+    [
+      27,
+      11,
+      0,
+      12,
+      1,
+      "PRISMAUDIO_FEATURES"
+    ],
+    [
+      29,
+      12,
+      0,
+      9,
+      1,
+      "AUDIO"
+    ],
+    [
+      30,
+      11,
+      1,
+      9,
+      4,
+      "FLOAT"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 1.1674071890328979,
+      "offset": [
+        1814.5534800416863,
+        500.0421331448515
+      ]
+    },
+    "ue_links": [],
+    "links_added_by_ue": [],
+    "frontendVersion": "1.42.8",
+    "VHS_latentpreview": true,
+    "VHS_latentpreviewrate": 0,
+    "VHS_MetadataImage": true,
+    "VHS_KeepIntermediate": true
+  },
+  "version": 0.4
+}