docs: add SelVA integration implementation plan

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]

docs/plans/2026-04-04-selva-integration-design.md

@@ -0,0 +1,167 @@
+# SelVA Integration Design
+
+**Date:** 2026-04-04
+**Branch:** feature/selva-integration (new from master)
+**Status:** Approved, ready for implementation
+
+---
+
+## Problem
+
+PrismAudio's sync conditioning is text-agnostic: Synchformer extracts features from
+all visual motion equally. In multi-source videos (person walking near a car), the DiT
+receives unfocused sync guidance and struggles to match audio events to the correct
+visual source.
+
+SelVA (CVPR 2026, arXiv:2512.02650) solves this with TextSynchformer — text conditioning
+is injected inside the Synchformer encoder via cross-attention, so sync features only
+encode motion relevant to the requested sound. This is the core architectural improvement
+needed for reliable V2A sync.
+
+---
+
+## Architecture
+
+### New directory layout
+
+```
+selva_core/          ← vendored SelVA source (model + ext + utils)
+nodes/
+  selva_model_loader.py
+  selva_feature_extractor.py
+  selva_sampler.py
+```
+
+### New custom types
+
+- `SELVA_MODEL` — `{generator, video_enc, feature_utils, variant, strategy, dtype}`
+- `SELVA_FEATURES` — `{clip_features, sync_features, duration}`
+
+### No subprocess
+
+SelVA is pure PyTorch. Feature extraction runs inline in ComfyUI — no managed venv,
+no JAX/TF, no pip install on first run.
+
+### Dependencies
+
+Zero new pip packages. ComfyUI already ships:
+- `open_clip_torch` (CLIP ViT-H-14-384, auto-downloads via `hf-hub:` on first use)
+- `transformers` (flan-t5-base, auto-downloads from HuggingFace on first use)
+- `torch`, `torchaudio`, `einops`
+
+---
+
+## Nodes
+
+### `SelvaModelLoader` → `SELVA_MODEL`
+
+| Input | Type | Default | Notes |
+|---|---|---|---|
+| variant | dropdown | medium_44k | small_16k / small_44k / medium_44k / large_44k |
+| precision | dropdown | bf16 | bf16 / fp16 / fp32 |
+| offload_strategy | dropdown | auto | auto / keep_in_vram / offload_to_cpu |
+
+Resolves weights from `models/selva/`. Raises descriptive errors with download
+instructions if files are missing.
+
+### `SelvaFeatureExtractor` → `SELVA_FEATURES`, `FLOAT` (fps)
+
+| Input | Type | Default | Notes |
+|---|---|---|---|
+| video | IMAGE | — | ComfyUI video tensor [T,H,W,C] |
+| prompt | STRING | — | Used by TextSynchformer to select relevant motion |
+| video_info | VHS_VIDEOINFO | opt | Auto-sets fps when connected |
+| fps | FLOAT | 30.0 | Fallback fps if video_info not connected |
+| cache_dir | STRING | "" | Empty = system temp dir |
+
+Feature extraction steps (all inline, no subprocess):
+1. Resize frames to 384×384 → CLIP video features `[B, T, 1024]`
+2. Resize frames to 224×224 + encode prompt with flan-T5 → TextSynchformer → text-conditioned sync features `[B, T, 768]`
+3. Save to `.npz` cache keyed by hash(frames[:1MB] + prompt + fps)
+
+### `SelvaSampler` → `AUDIO`
+
+| Input | Type | Default | Notes |
+|---|---|---|---|
+| model | SELVA_MODEL | — | |
+| features | SELVA_FEATURES | — | |
+| prompt | STRING | — | Should match extractor prompt; drives CLIP text guidance |
+| negative_prompt | STRING | "" | Steers away from unwanted sounds |
+| duration | FLOAT | 0.0 | 0 = auto from features duration |
+| steps | INT | 25 | Euler steps (25 is SelVA default, fast) |
+| cfg_strength | FLOAT | 4.5 | CFG scale (SelVA default) |
+| seed | INT | 0 | |
+
+Generation steps:
+1. Encode prompt → CLIP text features (for MMAudio)
+2. Encode negative prompt → empty conditions for CFG
+3. `net_generator.preprocess_conditions(clip_f, sync_f, text_clip)`
+4. Flow matching Euler ODE (`num_steps` iterations) with CFG
+5. `feature_utils.decode(latent)` → mel spectrogram
+6. `feature_utils.vocode(spec)` → waveform (BigVGAN for 16k, direct for 44k)
+
+**Note on dual prompt:** The extractor prompt is baked into sync_features via
+TextSynchformer at extraction time. The sampler prompt drives CLIP text conditioning
+at generation time. They should match — a tooltip explains this.
+
+---
+
+## Data Flow
+
+```
+[VHS LoadVideo] ──► [SelvaFeatureExtractor]
+                         │  prompt: "dog barking"
+                         │  video_info: (fps auto)
+                         ▼
+                    SELVA_FEATURES
+                    {clip_features [B,T,1024],
+                     sync_features [B,T,768],  ← text-conditioned
+                     duration: 8.2s}
+                         │
+[SelvaModelLoader] ──► [SelvaSampler]
+  variant: medium_44k    │  prompt: "dog barking"
+  precision: bf16        │  negative: "wind noise"
+                         │  cfg_strength: 4.5, steps: 25
+                         ▼
+                       AUDIO (44.1kHz or 16kHz)
+```
+
+---
+
+## Model Weights
+
+Location: `models/selva/`
+
+```
+video_enc_sup_5.pth                  ← TextSynch, shared across all variants
+generator_small_16k_sup_5.pth
+generator_small_44k_sup_5.pth
+generator_medium_44k_sup_5.pth
+generator_large_44k_sup_5.pth
+ext/
+  v1-16.pth                          ← VAE for 16k variants
+  v1-44.pth                          ← VAE for 44k variants
+  best_netG.pt                       ← BigVGAN vocoder (16k only)
+```
+
+`synchformer_state_dict.pth` is reused from `models/prismaudio/` — no duplicate.
+
+---
+
+## selva_core vendoring
+
+Copy from `jnwnlee/selva` (pinned to a specific commit for stability):
+- `selva_core/model/` — MMAudio, TextSynch, transformer layers, embeddings, flow matching
+- `selva_core/ext/` — autoencoder, BigVGAN, synchformer, rotary embeddings, mel converters
+- `selva_core/utils/` — transforms, generate() helper
+
+Rename all internal imports from `selva.*` → `selva_core.*`.
+
+---
+
+## What stays the same
+
+- All PrismAudio nodes unchanged
+- `models/prismaudio/` unchanged
+- Synchformer checkpoint shared (not duplicated)
+- Branch: new `feature/selva-integration` off master (LoRA work stays separate)

docs/plans/2026-04-04-selva-integration.md

@@ -0,0 +1,738 @@
+# SelVA Integration Implementation Plan
+
+> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.
+
+**Goal:** Add three new ComfyUI nodes (SelvaModelLoader, SelvaFeatureExtractor, SelvaSampler) that run SelVA's text-conditioned V2A pipeline inline — no subprocess, no JAX, pure PyTorch.
+
+**Architecture:** Vendor SelVA source into `selva_core/`, implement three nodes that mirror the PrismAudio pattern. `SelvaFeatureExtractor` takes `SELVA_MODEL` (needs TextSynchformer + CLIP/T5 from FeaturesUtils). `SelvaSampler` runs flow matching ODE with CFG and negative prompts.
+
+**Tech Stack:** PyTorch, open_clip (already in ComfyUI), transformers (already in ComfyUI), torchaudio, einops, torchvision
+
+---
+
+## Design reference
+
+`docs/plans/2026-04-04-selva-integration-design.md`
+
+**Key facts from SelVA source:**
+- CLIP input: `[B, T, C, 384, 384]` float32 `[0,1]` — normalization applied inside FeaturesUtils
+- Sync input: `[B, T, C, 224, 224]` float32 `[-1,1]` — normalize with `mean=std=[0.5,0.5,0.5]` before passing
+- CLIP frame rate: 8fps, Sync frame rate: 25fps
+- CONFIG_16K: latent=250, clip=64, sync=192 at 8s
+- CONFIG_44K: latent=345, clip=64, sync=192 at 8s
+- Sync segments: 16-frame windows, 8-frame stride (overlapping, unlike PrismAudio's 8-frame non-overlapping)
+- `net_generator.update_seq_lengths(latent_seq_len, clip_seq_len, sync_seq_len)` must be called before each generation when duration ≠ 8s
+
+---
+
+## Task 1: Create branch and vendor selva_core
+
+**Files:**
+- Create: `selva_core/` (full directory tree)
+
+**Step 1: Create new branch off master (not off feature/lora-trainer)**
+
+```bash
+git checkout master
+git checkout -b feature/selva-integration
+```
+
+**Step 2: Clone SelVA and copy source**
+
+```bash
+git clone https://github.com/jnwnlee/selva.git /tmp/selva_src
+cp -r /tmp/selva_src/selva /media/p5/Comfyui-Prismaudio/selva_core
+```
+
+**Step 3: Rename all internal imports**
+
+```bash
+cd /media/p5/Comfyui-Prismaudio/selva_core
+find . -name "*.py" -exec sed -i \
+  's/from selva\./from selva_core./g;
+   s/import selva\./import selva_core./g' {} \;
+```
+
+**Step 4: Record the pinned commit**
+
+```bash
+cd /tmp/selva_src && git rev-parse HEAD
+# Paste the hash into a comment at the top of selva_core/__init__.py
+```
+
+Edit `selva_core/__init__.py` to add at the top:
+```python
+# Vendored from https://github.com/jnwnlee/selva
+# Pinned commit: <PASTE_HASH_HERE>
+# Imports rewritten from selva.* → selva_core.*
+```
+
+**Step 5: Verify imports work**
+
+```bash
+cd /media/p5/Comfyui-Prismaudio
+python -c "
+from selva_core.model.networks_generator import MMAudio, get_my_mmaudio
+from selva_core.model.networks_video_enc import TextSynch, get_my_textsynch
+from selva_core.model.utils.features_utils import FeaturesUtils
+from selva_core.model.flow_matching import FlowMatching
+from selva_core.model.sequence_config import CONFIG_16K, CONFIG_44K, SequenceConfig
+print('selva_core imports OK')
+print(f'CONFIG_16K: latent={CONFIG_16K.latent_seq_len} clip={CONFIG_16K.clip_seq_len} sync={CONFIG_16K.sync_seq_len}')
+print(f'CONFIG_44K: latent={CONFIG_44K.latent_seq_len} clip={CONFIG_44K.clip_seq_len} sync={CONFIG_44K.sync_seq_len}')
+"
+```
+
+Expected:
+```
+selva_core imports OK
+CONFIG_16K: latent=250 clip=64 sync=192
+CONFIG_44K: latent=345 clip=64 sync=192
+```
+
+**Step 6: Commit**
+
+```bash
+git add selva_core/
+git commit -m "chore: vendor selva_core from jnwnlee/selva@<HASH>
+
+Pure PyTorch SelVA source for SelvaModelLoader/FeatureExtractor/Sampler nodes.
+Imports rewritten from selva.* to selva_core.*. No training code included."
+```
+
+---
+
+## Task 2: Implement SelvaModelLoader
+
+**Files:**
+- Create: `nodes/selva_model_loader.py`
+- Modify: `nodes/__init__.py`
+
+**Step 1: Create `nodes/selva_model_loader.py`**
+
+```python
+import os
+import torch
+import folder_paths
+
+from .utils import PRISMAUDIO_CATEGORY, get_offload_device, determine_offload_strategy
+
+# Variant → (generator filename, mode, has_bigvgan)
+_VARIANTS = {
+    "small_16k":   ("generator_small_16k_sup_5.pth",   "16k", True),
+    "small_44k":   ("generator_small_44k_sup_5.pth",   "44k", False),
+    "medium_44k":  ("generator_medium_44k_sup_5.pth",  "44k", False),
+    "large_44k":   ("generator_large_44k_sup_5.pth",   "44k", False),
+}
+
+_SELVA_DIR = os.path.join(folder_paths.models_dir, "selva")
+
+
+def _selva_path(*parts):
+    return os.path.join(_SELVA_DIR, *parts)
+
+
+def _require(path, hint):
+    if not os.path.exists(path):
+        raise RuntimeError(
+            f"[SelVA] Missing: {path}\n{hint}"
+        )
+    return path
+
+
+class SelvaModelLoader:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "variant": (list(_VARIANTS.keys()),),
+                "precision": (["bf16", "fp16", "fp32"],),
+                "offload_strategy": (["auto", "keep_in_vram", "offload_to_cpu"],),
+            }
+        }
+
+    RETURN_TYPES = ("SELVA_MODEL",)
+    RETURN_NAMES = ("model",)
+    FUNCTION = "load_model"
+    CATEGORY = PRISMAUDIO_CATEGORY
+
+    def load_model(self, variant, precision, offload_strategy):
+        from selva_core.model.networks_generator import get_my_mmaudio
+        from selva_core.model.networks_video_enc import get_my_textsynch
+        from selva_core.model.utils.features_utils import FeaturesUtils
+        from selva_core.model.sequence_config import CONFIG_16K, CONFIG_44K
+
+        gen_filename, mode, has_bigvgan = _VARIANTS[variant]
+
+        dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[precision]
+        strategy = determine_offload_strategy(offload_strategy)
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        # Resolve weight paths
+        video_enc_path = _require(
+            _selva_path("video_enc_sup_5.pth"),
+            "Download from https://huggingface.co/jnwnlee/selva and place in models/selva/"
+        )
+        gen_path = _require(
+            _selva_path(gen_filename),
+            f"Download {gen_filename} from https://huggingface.co/jnwnlee/selva and place in models/selva/"
+        )
+        vae_path = _require(
+            _selva_path("ext", f"v1-{mode}.pth"),
+            f"Download v1-{mode}.pth from MMAudio/SelVA release and place in models/selva/ext/"
+        )
+        synch_path = _require(
+            os.path.join(folder_paths.models_dir, "prismaudio", "synchformer_state_dict.pth"),
+            "Synchformer checkpoint missing from models/prismaudio/ — download from FunAudioLLM/PrismAudio"
+        )
+        bigvgan_path = None
+        if has_bigvgan:
+            bigvgan_path = _require(
+                _selva_path("ext", "best_netG.pt"),
+                "Download best_netG.pt (BigVGAN 16k vocoder) from MMAudio release and place in models/selva/ext/"
+            )
+
+        print(f"[SelVA] Loading TextSynch from {video_enc_path}", flush=True)
+        net_video_enc = get_my_textsynch("depth1").to(device, dtype).eval()
+        net_video_enc.load_weights(
+            torch.load(video_enc_path, map_location="cpu", weights_only=True)
+        )
+
+        print(f"[SelVA] Loading MMAudio ({variant}) from {gen_path}", flush=True)
+        seq_cfg = CONFIG_16K if mode == "16k" else CONFIG_44K
+        net_generator = get_my_mmaudio(variant).to(device, dtype).eval()
+        net_generator.load_weights(
+            torch.load(gen_path, map_location="cpu", weights_only=True)
+        )
+
+        print(f"[SelVA] Loading FeaturesUtils (CLIP + T5 + Synchformer + VAE)...", flush=True)
+        feature_utils = FeaturesUtils(
+            tod_vae_ckpt=vae_path,
+            synchformer_ckpt=synch_path,
+            enable_conditions=True,
+            mode=mode,
+            bigvgan_vocoder_ckpt=bigvgan_path,
+        ).to(device, dtype).eval()
+
+        if strategy == "offload_to_cpu":
+            net_generator.to(get_offload_device())
+            net_video_enc.to(get_offload_device())
+            feature_utils.to(get_offload_device())
+
+        print(f"[SelVA] Model ready: variant={variant} dtype={dtype} strategy={strategy}", flush=True)
+
+        return ({
+            "generator": net_generator,
+            "video_enc": net_video_enc,
+            "feature_utils": feature_utils,
+            "variant": variant,
+            "mode": mode,
+            "strategy": strategy,
+            "dtype": dtype,
+            "seq_cfg": seq_cfg,
+        },)
+```
+
+**Step 2: Register in `nodes/__init__.py`**
+
+In the `NODE_CLASS_MAPPINGS` dict, add:
+```python
+"SelvaModelLoader": (".selva_model_loader", "SelvaModelLoader", "SelVA Model Loader"),
+```
+
+**Step 3: Verify node registers**
+
+```bash
+cd /media/p5/Comfyui-Prismaudio
+python -c "
+import sys; sys.path.insert(0, '.')
+from nodes.selva_model_loader import SelvaModelLoader
+print('inputs:', list(SelvaModelLoader.INPUT_TYPES()['required'].keys()))
+print('outputs:', SelvaModelLoader.RETURN_TYPES)
+"
+```
+
+Expected: `inputs: ['variant', 'precision', 'offload_strategy']`
+
+**Step 4: Commit**
+
+```bash
+git add nodes/selva_model_loader.py nodes/__init__.py
+git commit -m "feat: SelvaModelLoader node — loads TextSynch + MMAudio + FeaturesUtils"
+```
+
+---
+
+## Task 3: Implement SelvaFeatureExtractor
+
+**Files:**
+- Create: `nodes/selva_feature_extractor.py`
+- Modify: `nodes/__init__.py`
+
+**Step 1: Create `nodes/selva_feature_extractor.py`**
+
+```python
+import os
+import hashlib
+import tempfile
+
+import torch
+import torch.nn.functional as F
+import numpy as np
+
+from .utils import PRISMAUDIO_CATEGORY, get_device, get_offload_device, soft_empty_cache
+
+# SelVA video preprocessing constants (from selva/utils/eval_utils.py)
+_CLIP_SIZE = 384
+_SYNC_SIZE = 224
+_CLIP_FPS  = 8
+_SYNC_FPS  = 25
+
+# Sync normalization: [-1, 1] (from selva/utils/eval_utils.py load_video)
+_SYNC_MEAN = torch.tensor([0.5, 0.5, 0.5]).view(1, 3, 1, 1)
+_SYNC_STD  = torch.tensor([0.5, 0.5, 0.5]).view(1, 3, 1, 1)
+
+
+def _sample_frames(video, source_fps, target_fps, duration):
+    """Sample frames from [T,H,W,C] float32 [0,1] at target_fps."""
+    T = video.shape[0]
+    n_out = max(1, int(duration * target_fps))
+    indices = [min(int(i / target_fps * source_fps), T - 1) for i in range(n_out)]
+    return video[indices]  # [N, H, W, C]
+
+
+def _resize_frames(frames, size):
+    """Resize [N,H,W,C] float32 [0,1] → [N,C,H,W] at target size."""
+    x = frames.permute(0, 3, 1, 2)  # [N, C, H, W]
+    x = F.interpolate(x, size=(size, size), mode="bicubic", align_corners=False)
+    return x.clamp(0, 1)  # [N, C, H, W] float32
+
+
+def _hash_inputs(video_tensor, prompt, fps, variant):
+    h = hashlib.sha256()
+    h.update(video_tensor.cpu().numpy().tobytes()[:1024 * 1024])
+    h.update(prompt.encode())
+    h.update(str(fps).encode())
+    h.update(variant.encode())
+    return h.hexdigest()[:16]
+
+
+class SelvaFeatureExtractor:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model":   ("SELVA_MODEL",),
+                "video":   ("IMAGE",),
+                "prompt":  ("STRING", {"default": "", "multiline": True,
+                                       "tooltip": "Text prompt used by TextSynchformer to focus sync features on the relevant sound source. Should match the prompt used in SelvaSampler."}),
+            },
+            "optional": {
+                "video_info": ("VHS_VIDEOINFO", {"tooltip": "Connect VHS LoadVideo info to auto-set fps."}),
+                "fps":   ("FLOAT", {"default": 30.0, "min": 1.0, "max": 120.0, "step": 0.001}),
+                "duration": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 30.0, "step": 0.1,
+                                       "tooltip": "Override duration in seconds. 0 = infer from video length and fps."}),
+                "cache_dir": ("STRING", {"default": "", "tooltip": "Directory for cached .npz features. Empty = temp dir."}),
+            },
+        }
+
+    RETURN_TYPES = ("SELVA_FEATURES", "FLOAT")
+    RETURN_NAMES = ("features", "fps")
+    FUNCTION = "extract_features"
+    CATEGORY = PRISMAUDIO_CATEGORY
+
+    def extract_features(self, model, video, prompt, video_info=None, fps=30.0,
+                         duration=0.0, cache_dir=""):
+        if video_info is not None:
+            fps = video_info["loaded_fps"]
+
+        T = video.shape[0]
+        if duration <= 0:
+            duration = T / fps
+        duration = min(duration, T / fps)  # clamp to actual video length
+
+        if not prompt.strip():
+            print("[SelVA] Warning: empty prompt — TextSynchformer sync features will be unfocused.", flush=True)
+
+        # Cache
+        if not cache_dir:
+            cache_dir = os.path.join(tempfile.gettempdir(), "selva_features")
+        os.makedirs(cache_dir, exist_ok=True)
+        cache_key = _hash_inputs(video, prompt, fps, model["variant"])
+        cached_path = os.path.join(cache_dir, f"{cache_key}.npz")
+
+        if os.path.exists(cached_path):
+            print(f"[SelVA] Using cached features: {cached_path}", flush=True)
+            return (_load_cached(cached_path), float(fps))
+
+        device = get_device()
+        dtype  = model["dtype"]
+        strategy = model["strategy"]
+        feature_utils = model["feature_utils"]
+        net_video_enc = model["video_enc"]
+
+        # Move feature models to device
+        if strategy == "offload_to_cpu":
+            feature_utils.to(device)
+            net_video_enc.to(device)
+            soft_empty_cache()
+
+        print(f"[SelVA] Extracting features: duration={duration:.2f}s fps={fps:.3f} prompt='{prompt[:60]}'", flush=True)
+
+        with torch.no_grad():
+            # --- CLIP frames: 384×384, [0,1], 8fps ---
+            clip_frames = _sample_frames(video, fps, _CLIP_FPS, duration)  # [N, H, W, C]
+            clip_frames = _resize_frames(clip_frames, _CLIP_SIZE)          # [N, C, 384, 384]
+            clip_input  = clip_frames.unsqueeze(0).to(device, dtype)       # [1, N, C, 384, 384]
+            print(f"[SelVA]   CLIP frames: {clip_frames.shape[0]} @ {_CLIP_FPS}fps", flush=True)
+
+            clip_features = feature_utils.encode_video_with_clip(clip_input)  # [1, N, 1024]
+
+            # --- Sync frames: 224×224, [-1,1], 25fps ---
+            n_sync = max(16, int(duration * _SYNC_FPS))  # minimum 16 for segmentation
+            sync_frames = _sample_frames(video, fps, _SYNC_FPS, duration)
+            if sync_frames.shape[0] < 16:
+                # Pad by repeating last frame to reach minimum 16
+                pad = 16 - sync_frames.shape[0]
+                sync_frames = torch.cat([sync_frames, sync_frames[-1:].expand(pad, -1, -1, -1)], dim=0)
+            sync_frames = _resize_frames(sync_frames, _SYNC_SIZE)          # [N, C, 224, 224]
+            # Normalize to [-1, 1]
+            mean = _SYNC_MEAN.to(sync_frames.device)
+            std  = _SYNC_STD.to(sync_frames.device)
+            sync_frames = (sync_frames - mean) / std
+            sync_input  = sync_frames.unsqueeze(0).to(device, dtype)       # [1, N, C, 224, 224]
+            print(f"[SelVA]   Sync frames: {sync_frames.shape[0]} @ {_SYNC_FPS}fps", flush=True)
+
+            # Encode T5 text + prepend supplementary tokens → text-conditioned sync features
+            text_f_t5, text_mask = feature_utils.encode_text_t5([prompt])  # [1, L, 768], [1, L]
+            text_f_t5, text_mask = net_video_enc.prepend_sup_text_tokens(text_f_t5, text_mask)
+            sync_features = net_video_enc.encode_video_with_sync(
+                sync_input, text_f=text_f_t5, text_mask=text_mask
+            )  # [1, T_sync, 768]
+
+        print(f"[SelVA]   clip_features: {tuple(clip_features.shape)}", flush=True)
+        print(f"[SelVA]   sync_features: {tuple(sync_features.shape)}", flush=True)
+
+        # Offload back if needed
+        if strategy == "offload_to_cpu":
+            feature_utils.to(get_offload_device())
+            net_video_enc.to(get_offload_device())
+            soft_empty_cache()
+
+        # Save cache
+        np.savez(
+            cached_path,
+            clip_features=clip_features.cpu().float().numpy(),
+            sync_features=sync_features.cpu().float().numpy(),
+            duration=duration,
+        )
+        print(f"[SelVA] Features cached: {cached_path}", flush=True)
+
+        features = {
+            "clip_features": clip_features.cpu(),
+            "sync_features": sync_features.cpu(),
+            "duration": duration,
+        }
+        return (features, float(fps))
+
+
+def _load_cached(path):
+    data = np.load(path, allow_pickle=False)
+    return {
+        "clip_features": torch.from_numpy(data["clip_features"]),
+        "sync_features": torch.from_numpy(data["sync_features"]),
+        "duration": float(data["duration"]),
+    }
+```
+
+**Step 2: Register in `nodes/__init__.py`**
+
+```python
+"SelvaFeatureExtractor": (".selva_feature_extractor", "SelvaFeatureExtractor", "SelVA Feature Extractor"),
+```
+
+**Step 3: Verify node registers**
+
+```bash
+python -c "
+import sys; sys.path.insert(0, '.')
+from nodes.selva_feature_extractor import SelvaFeatureExtractor
+inputs = SelvaFeatureExtractor.INPUT_TYPES()
+print('required:', list(inputs['required'].keys()))
+print('optional:', list(inputs['optional'].keys()))
+print('outputs:', SelvaFeatureExtractor.RETURN_TYPES)
+"
+```
+
+Expected: `required: ['model', 'video', 'prompt']`
+
+**Step 4: Commit**
+
+```bash
+git add nodes/selva_feature_extractor.py nodes/__init__.py
+git commit -m "feat: SelvaFeatureExtractor — inline CLIP + TextSynchformer feature extraction"
+```
+
+---
+
+## Task 4: Implement SelvaSampler
+
+**Files:**
+- Create: `nodes/selva_sampler.py`
+- Modify: `nodes/__init__.py`
+
+**Step 1: Create `nodes/selva_sampler.py`**
+
+```python
+import math
+import torch
+import comfy.utils
+
+from .utils import (
+    PRISMAUDIO_CATEGORY,
+    get_device, get_offload_device, soft_empty_cache,
+)
+
+
+def _make_seq_cfg(duration, mode):
+    """Compute sequence lengths for a given duration and mode."""
+    from selva_core.model.sequence_config import SequenceConfig
+    if mode == "16k":
+        return SequenceConfig(duration=duration, sampling_rate=16000, spectrogram_frame_rate=256)
+    else:
+        return SequenceConfig(duration=duration, sampling_rate=44100, spectrogram_frame_rate=512)
+
+
+class SelvaSampler:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model":    ("SELVA_MODEL",),
+                "features": ("SELVA_FEATURES",),
+                "prompt":   ("STRING", {"default": "", "multiline": True,
+                                        "tooltip": "Should match the prompt used in SelvaFeatureExtractor."}),
+                "negative_prompt": ("STRING", {"default": "", "multiline": True,
+                                               "tooltip": "Sounds to steer away from, e.g. 'wind noise, background music'."}),
+                "duration": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 30.0, "step": 0.1,
+                                       "tooltip": "Audio duration in seconds. 0 = use duration from features."}),
+                "steps":    ("INT",   {"default": 25,  "min": 1,   "max": 200}),
+                "cfg_strength": ("FLOAT", {"default": 4.5, "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, features, prompt, negative_prompt, duration, steps, cfg_strength, seed):
+        from selva_core.model.flow_matching import FlowMatching
+
+        device   = get_device()
+        dtype    = model["dtype"]
+        strategy = model["strategy"]
+        net_generator = model["generator"]
+        feature_utils = model["feature_utils"]
+        mode          = model["mode"]
+
+        # Resolve duration
+        if duration <= 0:
+            if "duration" not in features:
+                raise ValueError("[SelVA] duration=0 but features contain no duration field.")
+            duration = features["duration"]
+            print(f"[SelVA] Using video duration from features: {duration:.2f}s", flush=True)
+
+        seq_cfg = _make_seq_cfg(duration, mode)
+        sample_rate = seq_cfg.sampling_rate
+
+        # Move models to device
+        if strategy == "offload_to_cpu":
+            net_generator.to(device)
+            feature_utils.to(device)
+            soft_empty_cache()
+
+        clip_f = features["clip_features"].to(device, dtype)  # [1, T_clip, 1024]
+        sync_f = features["sync_features"].to(device, dtype)  # [1, T_sync, 768]
+
+        print(f"[SelVA] clip_f={tuple(clip_f.shape)} sync_f={tuple(sync_f.shape)}", flush=True)
+        print(f"[SelVA] seq_cfg: latent={seq_cfg.latent_seq_len} clip={seq_cfg.clip_seq_len} sync={seq_cfg.sync_seq_len}", flush=True)
+
+        # Update model sequence lengths for this duration
+        net_generator.update_seq_lengths(
+            latent_seq_len=seq_cfg.latent_seq_len,
+            clip_seq_len=seq_cfg.clip_seq_len,
+            sync_seq_len=seq_cfg.sync_seq_len,
+        )
+
+        with torch.no_grad():
+            # Encode text
+            text_clip = feature_utils.encode_text_clip([prompt])    # [1, 77, D]
+
+            # Build empty (negative) conditions for CFG
+            neg_text_clip = feature_utils.encode_text_clip([negative_prompt]) \
+                if negative_prompt.strip() else None
+
+            conditions = net_generator.preprocess_conditions(clip_f, sync_f, text_clip)
+            empty_conditions = net_generator.get_empty_conditions(
+                bs=1, negative_text_features=neg_text_clip
+            )
+
+            # Sample initial noise
+            rng = torch.Generator(device=device).manual_seed(seed)
+            x0 = torch.randn(
+                1, seq_cfg.latent_seq_len, net_generator.latent_dim,
+                device=device, dtype=dtype, generator=rng
+            )
+
+            # Flow matching ODE (Euler)
+            fm = FlowMatching(min_sigma=0, inference_mode="euler", num_steps=steps)
+            pbar = comfy.utils.ProgressBar(steps)
+
+            _step_count = [0]
+            orig_to_data = fm.to_data
+
+            def tracked_to_data(fn, x0_):
+                # ProgressBar update via step counting in ode_wrapper
+                return orig_to_data(fn, x0_)
+
+            # Wrap ODE to update progress bar
+            def ode_wrapper_tracked(t, x):
+                _step_count[0] += 1
+                pbar.update(1)
+                return net_generator.ode_wrapper(t, x, conditions, empty_conditions, cfg_strength)
+
+            x1 = fm.to_data(ode_wrapper_tracked, x0)
+
+        print(f"[SelVA] latent stats: mean={x1.float().mean():.4f} std={x1.float().std():.4f}", flush=True)
+
+        # Decode: latent → mel → audio
+        if strategy == "offload_to_cpu":
+            feature_utils.to(device)
+            soft_empty_cache()
+
+        with torch.no_grad():
+            x1_unnorm = net_generator.unnormalize(x1)
+            spec  = feature_utils.decode(x1_unnorm)
+            audio = feature_utils.vocode(spec)   # [1, samples] or [1, 1, samples]
+
+        if strategy == "offload_to_cpu":
+            net_generator.to(get_offload_device())
+            feature_utils.to(get_offload_device())
+            soft_empty_cache()
+
+        # Normalise to [-1, 1]
+        audio = audio.float()
+        if audio.dim() == 2:
+            audio = audio.unsqueeze(1)   # [1, 1, samples]
+        elif audio.dim() == 3 and audio.shape[1] != 1:
+            audio = audio.mean(dim=1, keepdim=True)  # stereo → mono
+
+        peak = audio.abs().max().clamp(min=1e-8)
+        audio = (audio / peak).clamp(-1, 1)
+        print(f"[SelVA] audio: shape={tuple(audio.shape)} sr={sample_rate}", flush=True)
+
+        return ({"waveform": audio.cpu(), "sample_rate": sample_rate},)
+```
+
+**Step 2: Register in `nodes/__init__.py`**
+
+```python
+"SelvaSampler": (".selva_sampler", "SelvaSampler", "SelVA Sampler"),
+```
+
+**Step 3: Verify node registers**
+
+```bash
+python -c "
+import sys; sys.path.insert(0, '.')
+from nodes.selva_sampler import SelvaSampler
+inputs = SelvaSampler.INPUT_TYPES()
+print('inputs:', list(inputs['required'].keys()))
+print('outputs:', SelvaSampler.RETURN_TYPES)
+"
+```
+
+Expected: `inputs: ['model', 'features', 'prompt', 'negative_prompt', 'duration', 'steps', 'cfg_strength', 'seed']`
+
+**Step 4: Commit**
+
+```bash
+git add nodes/selva_sampler.py nodes/__init__.py
+git commit -m "feat: SelvaSampler — flow matching ODE with CFG + negative prompts"
+```
+
+---
+
+## Task 5: Create example workflow and push
+
+**Files:**
+- Create: `workflows/selva_video_to_audio.json`
+
+**Step 1: Create workflow JSON**
+
+Create `workflows/selva_video_to_audio.json` with this node graph:
+- LoadVideo (VHS) → IMAGE + VHS_VIDEOINFO
+- SelvaModelLoader → SELVA_MODEL
+- SelvaFeatureExtractor (takes IMAGE + VHS_VIDEOINFO + SELVA_MODEL, prompt) → SELVA_FEATURES
+- SelvaSampler (takes SELVA_MODEL + SELVA_FEATURES, prompt, negative_prompt) → AUDIO
+- PreviewAudio (takes AUDIO)
+
+Set defaults: variant=medium_44k, precision=bf16, steps=25, cfg_strength=4.5, duration=0.
+
+**Step 2: Push branch**
+
+```bash
+git push -u origin feature/selva-integration
+```
+
+---
+
+## Task 6: Smoke test
+
+**Step 1: Check all three nodes are importable from ComfyUI's perspective**
+
+```bash
+cd /media/p5/Comfyui-Prismaudio
+python -c "
+import sys; sys.path.insert(0, '.')
+import nodes
+m = nodes.NODE_CLASS_MAPPINGS
+print('SelVA nodes:', [k for k in m if 'Selva' in k])
+assert 'SelvaModelLoader' in m
+assert 'SelvaFeatureExtractor' in m
+assert 'SelvaSampler' in m
+print('All SelVA nodes registered OK')
+"
+```
+
+**Step 2: Verify no import errors in full node load**
+
+```bash
+python -c "
+import sys; sys.path.insert(0, '.')
+from nodes.selva_model_loader import SelvaModelLoader
+from nodes.selva_feature_extractor import SelvaFeatureExtractor
+from nodes.selva_sampler import SelvaSampler
+print('All imports clean')
+"
+```
+
+**Step 3: Final commit with any fixes**
+
+```bash
+git add -A
+git commit -m "fix: selva integration smoke test fixes (if any)"
+git push
+```
+
+---
+
+## Notes
+
+- The `FeaturesUtils.train()` is overridden to always call `super().train(False)` — SelVA models are always in eval mode
+- `net_generator.update_seq_lengths` recalculates rotary position embeddings; call it before every generation when duration may vary
+- ProgressBar tracking: `FlowMatching.to_data` calls `fn(t, x)` for each Euler step; wrapping `ode_wrapper` with a counter gives accurate progress
+- The `feature_utils.vocode` returns audio at 16kHz for small_16k (uses BigVGAN) and 44.1kHz for 44k variants (uses VAE mel decoder directly)
+- If `encode_text_t5` or `encode_text_clip` fail with missing model errors on first run, it's HuggingFace downloading `flan-t5-base` and `apple/DFN5B-CLIP-ViT-H-14-384` — this is expected and takes a few minutes once

nodes/__init__.py

@@ -7,6 +7,8 @@ _NODES = {
     "PrismAudioFeatureExtractor": (".feature_extractor", "PrismAudioFeatureExtractor", "PrismAudio Feature Extractor"),
     "PrismAudioSampler": (".sampler", "PrismAudioSampler", "PrismAudio Sampler"),
     "PrismAudioTextOnly": (".text_only", "PrismAudioTextOnly", "PrismAudio Text Only"),
+    "PrismAudioLoRATrainer": (".lora_trainer", "PrismAudioLoRATrainer", "PrismAudio LoRA Trainer"),
+    "PrismAudioLoRALoader":  (".lora_loader",  "PrismAudioLoRALoader",  "PrismAudio LoRA Loader"),
 }
 
 for key, (module_path, class_name, display_name) in _NODES.items():

nodes/feature_extractor.py

@@ -0,0 +1,228 @@
+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")
+
+def _jax_package():
+    """Return the correct jax extra for the current CUDA version."""
+    try:
+        import torch
+        if torch.cuda.is_available():
+            cuda_ver = torch.version.cuda or ""
+            major = int(cuda_ver.split(".")[0]) if cuda_ver else 0
+            if major >= 13:
+                return "jax[cuda13]"
+            elif major >= 12:
+                return "jax[cuda12]"
+    except Exception:
+        pass
+    return "jax"  # CPU fallback
+
+
+_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 CUDA extra is resolved at install time based on detected CUDA version
+    _jax_package(), "flax",
+    "transformers", "decord", "einops", "numpy",
+    "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, fps):
+    """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())
+    h.update(str(fps).encode())  # fps affects frame sampling — must be part of the key
+    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"]
+
+        if not caption_cot.strip():
+            print("[PrismAudio] Warning: caption_cot is empty — text features will be degenerate. "
+                  "Provide a descriptive chain-of-thought caption for best results.", flush=True)
+
+        # 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, fps)
+        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/lora_loader.py

@@ -0,0 +1,106 @@
+import os
+import json
+import torch
+import torch.nn as nn
+
+from .utils import PRISMAUDIO_CATEGORY
+
+
+def _merge_lora_weights(dit: nn.Module, lora_state: dict, rank: int, alpha: float, strength: float):
+    """Add LoRA delta weights directly into the base model's nn.Linear tensors.
+
+    delta_W = lora_B @ lora_A * scale * strength
+    applied as: linear.weight += delta_W
+
+    This is equivalent to LoRALinear at inference but requires no wrapper,
+    no extra memory, and no change to the model's forward call graph.
+    """
+    scale = (alpha / rank) * strength
+
+    # Group saved keys by module path
+    a_map = {
+        k.replace(".lora_A.weight", ""): v
+        for k, v in lora_state.items() if k.endswith("lora_A.weight")
+    }
+    b_map = {
+        k.replace(".lora_B.weight", ""): v
+        for k, v in lora_state.items() if k.endswith("lora_B.weight")
+    }
+
+    merged = 0
+    for path, lora_A in a_map.items():
+        if path not in b_map:
+            print(f"[PrismAudio] LoRA merge: missing lora_B for {path}, skipping", flush=True)
+            continue
+        lora_B = b_map[path]  # [out_features, rank]
+        # delta_W: [out_features, in_features]
+        delta_W = (lora_B.float() @ lora_A.float()) * scale
+
+        # Navigate to the parent module using PyTorch's get_submodule
+        *parent_parts, child_name = path.split(".")
+        try:
+            parent = dit.get_submodule(".".join(parent_parts)) if parent_parts else dit
+        except AttributeError as e:
+            print(f"[PrismAudio] LoRA merge: could not find module '{path}': {e}", flush=True)
+            continue
+
+        linear = getattr(parent, child_name, None)
+        if not isinstance(linear, nn.Linear):
+            print(f"[PrismAudio] LoRA merge: expected nn.Linear at '{path}', got {type(linear)}", flush=True)
+            continue
+
+        linear.weight.data.add_(delta_W.to(linear.weight.dtype))
+        merged += 1
+
+    print(f"[PrismAudio] LoRA merged {merged} layer(s) (strength={strength:.3f})", flush=True)
+
+
+class PrismAudioLoRALoader:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model":     ("PRISMAUDIO_MODEL",),
+                "lora_path": ("STRING", {"default": "", "tooltip": "Path to .safetensors LoRA file produced by PrismAudio LoRA Trainer"}),
+                "strength":  ("FLOAT",  {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05, "tooltip": "LoRA influence scale. 1.0 = full strength, 0.0 = base model only"}),
+            },
+        }
+
+    RETURN_TYPES = ("PRISMAUDIO_MODEL",)
+    RETURN_NAMES = ("model",)
+    FUNCTION = "load_lora"
+    CATEGORY = PRISMAUDIO_CATEGORY
+
+    def load_lora(self, model, lora_path, strength):
+        from safetensors.torch import load_file
+
+        if not os.path.exists(lora_path):
+            raise FileNotFoundError(f"[PrismAudio] LoRA file not found: {lora_path}")
+
+        config_path = lora_path.replace(".safetensors", "_config.json")
+        if not os.path.exists(config_path):
+            raise FileNotFoundError(
+                f"[PrismAudio] LoRA config not found: {config_path}\n"
+                "Expected a _config.json alongside the .safetensors file."
+            )
+
+        with open(config_path) as f:
+            config = json.load(f)
+
+        rank  = config["rank"]
+        alpha = config["alpha"]
+
+        lora_state = load_file(lora_path)
+
+        # Merge LoRA weights in-place into the DiT's base linear layers.
+        # ComfyUI re-executes the upstream ModelLoader on the next queue run
+        # when inputs change, providing a fresh base model as needed.
+        dit = model["model"].model  # DiTWrapper
+
+        if strength == 0.0:
+            print("[PrismAudio] LoRA strength=0.0 — skipping merge, base model unchanged.", flush=True)
+            return (model,)
+
+        _merge_lora_weights(dit, lora_state, rank, alpha, strength)
+
+        return (model,)

nodes/lora_trainer.py

@@ -0,0 +1,284 @@
+import os
+import math
+import json
+import random
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import comfy.utils
+
+from .utils import (
+    PRISMAUDIO_CATEGORY, SAMPLE_RATE,
+    get_device, get_offload_device, soft_empty_cache,
+)
+
+
+# ---------------------------------------------------------------------------
+# LoRA primitives
+# ---------------------------------------------------------------------------
+
+class LoRALinear(nn.Module):
+    """Low-rank adapter wrapping a frozen nn.Linear."""
+
+    def __init__(self, linear: nn.Linear, rank: int, alpha: float):
+        super().__init__()
+        self.linear = linear
+        self.scale = alpha / rank
+        in_f, out_f = linear.in_features, linear.out_features
+        self.lora_A = nn.Linear(in_f, rank, bias=False)
+        self.lora_B = nn.Linear(rank, out_f, bias=False)
+        nn.init.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5))
+        nn.init.zeros_(self.lora_B.weight)
+
+    def forward(self, x):
+        return self.linear(x) + self.lora_B(self.lora_A(x)) * self.scale
+
+
+_TARGET_MODULE_PRESETS = {
+    "attn_only": {"to_q", "to_kv", "to_qkv", "to_out"},
+    "attn_ffn":  {"to_q", "to_kv", "to_qkv", "to_out", "proj"},
+    "full":      {"to_q", "to_kv", "to_qkv", "to_out", "proj", "project_in", "project_out"},
+}
+
+
+def _apply_lora(module: nn.Module, target_attrs: set, rank: int, alpha: float):
+    """Recursively replace matching nn.Linear layers with LoRALinear."""
+    for name, child in list(module.named_children()):
+        if isinstance(child, nn.Linear) and name in target_attrs:
+            setattr(module, name, LoRALinear(child, rank, alpha))
+        else:
+            _apply_lora(child, target_attrs, rank, alpha)
+
+
+def _unapply_lora(module: nn.Module):
+    """Replace LoRALinear back with the original frozen Linear (no weight merge)."""
+    for name, child in list(module.named_children()):
+        if isinstance(child, LoRALinear):
+            child.linear.weight.requires_grad_(False)
+            setattr(module, name, child.linear)
+        else:
+            _unapply_lora(child)
+
+
+def _get_lora_state_dict(module: nn.Module) -> dict:
+    """Return only LoRA parameter tensors from a module's state dict."""
+    return {k: v for k, v in module.state_dict().items()
+            if "lora_A" in k or "lora_B" in k}
+
+
+# ---------------------------------------------------------------------------
+# Dataset helpers
+# ---------------------------------------------------------------------------
+
+_AUDIO_EXTS = (".wav", ".flac", ".mp3")
+
+
+def _scan_dataset(dataset_dir: str):
+    """Return list of (npz_path, audio_path) pairs matched by stem."""
+    pairs = []
+    for fname in os.listdir(dataset_dir):
+        if not fname.endswith(".npz"):
+            continue
+        stem = os.path.join(dataset_dir, fname[:-4])
+        for ext in _AUDIO_EXTS:
+            audio_path = stem + ext
+            if os.path.exists(audio_path):
+                pairs.append((stem + ".npz", audio_path))
+                break
+    return sorted(pairs)
+
+
+def _load_audio(audio_path: str, device: torch.device) -> torch.Tensor:
+    """Load audio to [1, 2, samples] float32 tensor at SAMPLE_RATE."""
+    import torchaudio
+    waveform, sr = torchaudio.load(audio_path)
+    if sr != SAMPLE_RATE:
+        waveform = torchaudio.functional.resample(waveform, sr, SAMPLE_RATE)
+    if waveform.shape[0] == 1:
+        waveform = waveform.expand(2, -1)
+    elif waveform.shape[0] > 2:
+        waveform = waveform[:2]
+    return waveform.unsqueeze(0).to(device)  # [1, 2, samples]
+
+
+def _load_metadata(npz_path: str, device: torch.device, dtype: torch.dtype) -> dict:
+    """Load .npz features into a conditioner metadata dict."""
+    import numpy as np
+    data = np.load(npz_path, allow_pickle=True)
+    video_feat = torch.from_numpy(data["video_features"]).float().to(device, dtype=dtype)
+    text_feat  = torch.from_numpy(data["text_features"]).float().to(device, dtype=dtype)
+    sync_feat  = torch.from_numpy(data["sync_features"]).float().to(device, dtype=dtype)
+    has_video = bool(video_feat.abs().sum() > 0)
+    return {
+        "video_features": video_feat,
+        "text_features":  text_feat,
+        "sync_features":  sync_feat,
+        "video_exist":    torch.tensor(has_video),
+    }
+
+
+# ---------------------------------------------------------------------------
+# Trainer node
+# ---------------------------------------------------------------------------
+
+class PrismAudioLoRATrainer:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model":            ("PRISMAUDIO_MODEL",),
+                "dataset_dir":      ("STRING", {"default": "", "tooltip": "Directory containing paired .npz feature files and .wav/.flac audio files (matched by filename stem)"}),
+                "output_path":      ("STRING", {"default": "", "tooltip": "Save path for .safetensors weights. Empty = models/prismaudio/lora/"}),
+                "lora_rank":        ("INT",   {"default": 64,   "min": 1,     "max": 512}),
+                "lora_alpha":       ("FLOAT", {"default": 64.0, "min": 1.0,   "max": 1024.0}),
+                "target_modules":   (["attn_ffn", "attn_only", "full"], {"tooltip": "attn_only: Q/K/V/out only. attn_ffn: + FFN input (recommended). full: + transformer I/O projections"}),
+                "learning_rate":    ("FLOAT", {"default": 1e-4, "min": 1e-7,  "max": 1e-2,  "step": 1e-6}),
+                "train_steps":      ("INT",   {"default": 1000, "min": 1,     "max": 100000}),
+                "cfg_dropout_prob": ("FLOAT", {"default": 0.1,  "min": 0.0,   "max": 0.5,   "step": 0.01, "tooltip": "Probability of dropping conditioning per step — preserves CFG ability at inference"}),
+                "save_every":       ("INT",   {"default": 500,  "min": 1,     "max": 100000, "tooltip": "Save a checkpoint every N steps (in addition to final save)"}),
+                "seed":             ("INT",   {"default": 0,    "min": 0,     "max": 0xFFFFFFFF}),
+            },
+        }
+
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("lora_path",)
+    FUNCTION = "train"
+    CATEGORY = PRISMAUDIO_CATEGORY
+
+    def train(self, model, dataset_dir, output_path, lora_rank, lora_alpha,
+              target_modules, learning_rate, train_steps, cfg_dropout_prob, save_every, seed):
+        from safetensors.torch import save_file
+
+        device    = get_device()
+        dtype     = model["dtype"]
+        diffusion = model["model"]
+        strategy  = model["strategy"]
+
+        torch.manual_seed(seed)
+        random.seed(seed)
+
+        # Scan dataset
+        pairs = _scan_dataset(dataset_dir)
+        if not pairs:
+            raise RuntimeError(f"[PrismAudio] No (.npz + audio) pairs found in: {dataset_dir}")
+        print(f"[PrismAudio] LoRA training — {len(pairs)} sample(s), {train_steps} steps", flush=True)
+
+        # Resolve output path
+        if not output_path:
+            import folder_paths
+            out_dir = os.path.join(folder_paths.models_dir, "prismaudio", "lora")
+            os.makedirs(out_dir, exist_ok=True)
+            output_path = os.path.join(out_dir, f"prismaudio_lora_r{lora_rank}.safetensors")
+
+        # Move model to device
+        diffusion.model.to(device)
+        diffusion.conditioner.to(device)
+        diffusion.pretransform.to(device)
+
+        # Freeze all DiT params, then apply LoRA (adds trainable lora_A/lora_B)
+        dit = diffusion.model  # DiTWrapper
+        for p in dit.parameters():
+            p.requires_grad_(False)
+
+        target_attrs = _TARGET_MODULE_PRESETS[target_modules]
+        _apply_lora(dit, target_attrs, lora_rank, lora_alpha)
+
+        # Cast LoRA params to model dtype and move to device
+        for m in dit.modules():
+            if isinstance(m, LoRALinear):
+                m.lora_A.to(device=device, dtype=dtype)
+                m.lora_B.to(device=device, dtype=dtype)
+
+        trainable = [p for p in dit.parameters() if p.requires_grad]
+        n_params = sum(p.numel() for p in trainable)
+        print(f"[PrismAudio] LoRA trainable params: {n_params:,}  ({n_params/1e6:.2f}M)", flush=True)
+
+        diffusion.conditioner.eval()
+        diffusion.pretransform.eval()
+        dit.train()
+
+        optimizer = torch.optim.AdamW(trainable, lr=learning_rate)
+
+        # GradScaler for fp16 to prevent underflow
+        use_scaler = (dtype == torch.float16)
+        scaler = torch.cuda.amp.GradScaler() if use_scaler else None
+
+        pbar = comfy.utils.ProgressBar(train_steps)
+
+        try:
+            for step in range(1, train_steps + 1):
+                npz_path, audio_path = random.choice(pairs)
+
+                with torch.no_grad():
+                    # Encode audio to latent space
+                    audio = _load_audio(audio_path, device)
+                    x0 = diffusion.pretransform.encode(audio.float()).to(dtype)  # [1, 64, L]
+
+                    # Build conditioning from features
+                    metadata = (_load_metadata(npz_path, device, dtype),)
+                    conditioning = diffusion.conditioner(metadata, device)
+                    cond_inputs = diffusion.get_conditioning_inputs(conditioning)
+
+                # Rectified flow: interpolate between data and noise
+                t      = torch.rand(x0.shape[0], device=device, dtype=dtype)   # [1]
+                noise  = torch.randn_like(x0)
+                # t expanded for broadcast: [1] -> [1, 1, 1]
+                t_bcast = t[:, None, None]
+                x_t    = (1.0 - t_bcast) * x0 + t_bcast * noise
+                v_target = noise - x0
+
+                with torch.amp.autocast(device_type=device.type, dtype=dtype):
+                    v_pred = dit(x_t, t,
+                                 cfg_scale=1.0,
+                                 cfg_dropout_prob=cfg_dropout_prob,
+                                 **cond_inputs)
+
+                loss = F.mse_loss(v_pred.float(), v_target.float())
+
+                if use_scaler:
+                    scaler.scale(loss).backward()
+                    scaler.step(optimizer)
+                    scaler.update()
+                else:
+                    loss.backward()
+                    optimizer.step()
+                optimizer.zero_grad()
+
+                if step % 50 == 0:
+                    print(f"[PrismAudio] step {step}/{train_steps}  loss={loss.item():.6f}", flush=True)
+
+                if step % save_every == 0:
+                    ckpt_path = output_path.replace(".safetensors", f"_step{step}.safetensors")
+                    save_file(_get_lora_state_dict(dit), ckpt_path)
+                    print(f"[PrismAudio] Checkpoint: {ckpt_path}", flush=True)
+
+                pbar.update(1)
+
+            # Save final weights
+            save_file(_get_lora_state_dict(dit), output_path)
+
+            # Save config alongside weights so the loader knows the structure
+            config_path = output_path.replace(".safetensors", "_config.json")
+            with open(config_path, "w") as f:
+                json.dump({
+                    "rank": lora_rank,
+                    "alpha": lora_alpha,
+                    "target_modules": sorted(target_attrs),
+                }, f, indent=2)
+
+            print(f"[PrismAudio] LoRA saved: {output_path}", flush=True)
+
+        finally:
+            # Always restore model to base state — even on exception.
+            # Without this, LoRA wrappers would persist in the cached model and
+            # subsequent training runs would apply LoRA on top of existing LoRA.
+            dit.eval()
+            _unapply_lora(dit)
+
+            if strategy == "offload_to_cpu":
+                diffusion.model.to(get_offload_device())
+                diffusion.conditioner.to(get_offload_device())
+                diffusion.pretransform.to(get_offload_device())
+            soft_empty_cache()
+
+        return (output_path,)

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,183 @@
+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"}),
+                "sync_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.05, "tooltip": "Scale factor for sync conditioning. Higher values tighten audio-visual sync at the cost of audio naturalness; 0.0 disables sync guidance entirely."}),
+                "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, sync_strength, 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)
+
+        # Sync temporal coverage diagnostic
+        sync_frames = features["sync_features"].shape[0]
+        sync_duration_covered = sync_frames / 25.0  # Synchformer always extracts at 25fps
+        print(f"[PrismAudio] sync: {sync_frames} frames @ 25fps = {sync_duration_covered:.2f}s | "
+              f"audio target: {latent_length} latent frames = {duration:.2f}s", flush=True)
+        if abs(sync_duration_covered - duration) > 0.5:
+            print(f"[PrismAudio] Warning: sync coverage ({sync_duration_covered:.2f}s) differs from "
+                  f"audio duration ({duration:.2f}s) by more than 0.5s — consider re-extracting features "
+                  f"with the correct video duration.", flush=True)
+
+        # 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)
+
+            # Scale sync conditioning after the conditioner MLP (clean linear scale,
+            # avoids SiLU nonlinearity in Sync_MLP). The CFG null path always uses zeros,
+            # so this directly scales the sync guidance magnitude: cfg_scale * (strength*cond - 0).
+            # Only applied when video is present — T2A uses learned empty_sync_feat, not raw sync.
+            if has_video and sync_strength != 1.0 and 'sync_features' in conditioning:
+                conditioning['sync_features'][0] = conditioning['sync_features'][0] * sync_strength
+
+            # 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,12 +919,18 @@ 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
                 elif self.sync_gate is not None:
                     gate_value = torch.sigmoid(self.sync_gate)
-                    x = x + gate_value * sync_cond 
+                    x = x + gate_value * sync_cond
                 # else:
                 #     x = x + sync_cond
 

requirements.txt

@@ -1,8 +1,12 @@
 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
+torchaudio

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,170 @@
+#!/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(max(1, 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)
+
+        # Synchformer processes in segments of 8; ensure at least 8 frames
+        sync_indices = [int(i * fps / args.sync_fps) for i in range(max(8, 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(max(1, 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)
+
+        # Synchformer processes in segments of 8; ensure at least 8 frames
+        sync_indices = [int(i * fps / args.sync_fps) for i in range(max(8, 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
+}