docs: add SelVA integration implementation plan

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

README.md

@@ -1,6 +1,6 @@
 # ComfyUI-PrismAudio
 
-Custom nodes for [PrismAudio](https://github.com/FunAudioLLM/ThinkSound) (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.
+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
 
@@ -8,56 +8,120 @@ Clone into your ComfyUI custom nodes directory:
 
 ```bash
 cd ComfyUI/custom_nodes
-git clone -b prismaudio https://github.com/FunAudioLLM/ThinkSound ComfyUI-PrismAudio
+git clone https://github.com/Ethanfel/ComfyUI-Prismaudio.git ComfyUI-PrismAudio
 pip install -r ComfyUI-PrismAudio/requirements.txt
 ```
 
-**flash-attn** is optional. It is detected at runtime and falls back to PyTorch SDPA if unavailable.
+**flash-attn** is optional — detected at runtime, falls back to PyTorch SDPA if unavailable.
-
-For the **Feature Extractor** node (video feature extraction), a separate conda environment is required — see [Feature Extraction Environment](#feature-extraction-environment) below.
 
 ## Nodes
 
-| Node | Description |
+### PrismAudio Model Loader
-|------|-------------|
+
-| **PrismAudio Model Loader** | Loads the diffusion model and VAE. Auto-downloads weights from HuggingFace. Inputs: `precision` (auto/fp32/fp16/bf16), `offload_strategy` (auto/keep_in_vram/offload_to_cpu). |
+Loads the DiT diffusion model and VAE. Auto-downloads weights from HuggingFace on first use.
-| **PrismAudio Feature Loader** | Loads pre-computed `.npz` feature files for use with the sampler. |
+
-| **PrismAudio Feature Extractor** | Subprocess bridge that extracts features from video. Requires a separate conda env with JAX/TF. |
+| Input | Options | Description |
-| **PrismAudio Sampler** | Main generation node. Takes model + features, produces AUDIO. Inputs: `duration`, `steps`, `cfg_scale`, `seed`. |
+|-------|---------|-------------|
-| **PrismAudio Text Only** | Text-to-audio generation without video. Uses the T5-Gemma text encoder. Inputs: `text_prompt`, `duration`, `steps`, `cfg_scale`, `seed`. |
+| `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
 
-### Quality Path (Video-to-Audio)
+### Video-to-Audio
 
 ```
-Video → PrismAudio Feature Extractor → PrismAudio Sampler → Save Audio
+VHS LoadVideo ──► PrismAudio Feature Extractor ──► PrismAudio Sampler ──► Save Audio
+                         (video_info) ──────────────────► (fps auto)
+                         (features) ────────────────────► (features)
+                         duration=0 ─────────────────────► (auto from features)
 ```
 
-### Pre-computed Path
+### Pre-computed Features
 
 ```
-PrismAudio Feature Loader (.npz) → PrismAudio Sampler → Save Audio
+PrismAudio Feature Loader (.npz) ──► PrismAudio Sampler ──► Save Audio
 ```
 
-### Text-Only
+### Text-to-Audio
 
 ```
-PrismAudio Text Only → Save Audio
+PrismAudio Text Only ──► Save Audio
 ```
 
-> **Note:** CoT text is a STRING input on the sampler. You can use any existing ComfyUI LLM nodes to generate it.
-
 ## HuggingFace Authentication
 
-Required for gated models (T5-Gemma, and possibly Stable Audio VAE).
+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 nodes by design — ComfyUI saves all STRING values to workflow JSON, which would expose your token.
+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
 
@@ -65,42 +129,28 @@ Weights are auto-downloaded to `ComfyUI/models/prismaudio/`:
 
 | File | Size | Description |
 |------|------|-------------|
-| `prismaudio.ckpt` | ~2.7 GB | Diffusion model |
+| `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 |
+| `synchformer_state_dict.pth` | ~950 MB | Synchformer visual encoder |
 
-T5-Gemma is cached in the standard HuggingFace cache directory (`~/.cache/huggingface/`).
+T5-Gemma and VideoPrism LvT are cached in `~/.cache/huggingface/`.
 
 ## VRAM Requirements
 
-| VRAM | Strategy |
+| VRAM | Recommended settings |
-|------|----------|
+|------|----------------------|
-| 24 GB+ | Keep all models in VRAM |
+| 24 GB+ | `keep_in_vram`, any precision |
-| 12–24 GB | Sequential offload |
+| 12–24 GB | `offload_to_cpu`, bf16/fp16 |
-| 8–12 GB | Aggressive offload + fp16 |
+| 8–12 GB | `offload_to_cpu`, fp16 |
-| < 8 GB | May work with aggressive offload |
+| < 8 GB | May work with `offload_to_cpu` + fp16 |
-
-## Feature Extraction Environment
-
-The **PrismAudio Feature Extractor** node runs extraction in a subprocess using a separate Python environment (JAX/TF dependencies).
-
-```bash
-conda env create -f scripts/environment.yml
-conda activate prismaudio-extract
-```
-
-Then set the `python_env` input on the Feature Extractor node to:
-
-```
-/path/to/conda/envs/prismaudio-extract/bin/python
-```
 
 ## 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`, or use `fp16` precision.
+- **VRAM errors** — Switch `offload_strategy` to `offload_to_cpu` and/or use `fp16` precision.
-- **flash-attn** — Purely optional. Auto-detected at runtime; falls back to PyTorch SDPA.
+- **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). [Paper & code](https://github.com/FunAudioLLM/ThinkSound).
+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

@@ -13,48 +13,96 @@ _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",
-    "tensorflow-cpu==2.15.0",
+    # TF 2.15 only supports Python <=3.11; use >=2.16 for Python 3.12+
-    "jax[cpu]", "jaxlib",
+    "tensorflow-cpu>=2.16.0",
-    "transformers", "decord", "einops", "numpy", "mediapy",
+    # 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
 
-    print("[PrismAudio] Feature-extraction env not found — creating venv at:", _MANAGED_VENV)
+    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)
 
-    print("[PrismAudio] Installing feature-extraction dependencies (this takes a few minutes)...")
+    total = len(_EXTRACT_PACKAGES)
-    subprocess.run([pip, "install"] + _EXTRACT_PACKAGES, check=True)
+    print(f"[PrismAudio] Installing {total} package groups — this may take several minutes...", flush=True)
 
-    print("[PrismAudio] Feature-extraction env ready.")
+    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):
+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_video_tensor_to_mp4(video_tensor, output_path, fps=30):
+def _save_frames_to_npy(video_tensor, output_path):
-    """Save ComfyUI IMAGE tensor [T,H,W,C] to MP4."""
+    """Save ComfyUI IMAGE tensor [T,H,W,C] float32 [0,1] to .npy as uint8.
-    import torchvision.io as tvio
+
-    # ComfyUI IMAGE is [T,H,W,C] float32 [0,1]
+    Lossless — avoids H.264 encode/decode roundtrip.
-    frames = (video_tensor * 255).to(torch.uint8)
+    """
-    # torchvision write_video expects [T,H,W,C] uint8
+    import numpy as np
-    tvio.write_video(output_path, frames, fps=fps)
+    frames_np = (video_tensor.cpu().numpy() * 255).astype("uint8")
+    np.save(output_path, frames_np)
 
 
 class PrismAudioFeatureExtractor:
@@ -66,21 +114,35 @@ class PrismAudioFeatureExtractor:
                 "caption_cot": ("STRING", {"default": "", "multiline": True, "tooltip": "Chain-of-thought description"}),
             },
             "optional": {
-                "python_env": ("STRING", {"default": "python", "tooltip": "Path to python binary with JAX/TF. Leave as 'python' to auto-install a managed venv on first use."}),
+                "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"}),
-                "synchformer_ckpt": ("STRING", {"default": "", "tooltip": "Path to synchformer checkpoint (auto-resolved if empty)"}),
+                "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",)
+    RETURN_TYPES = ("PRISMAUDIO_FEATURES", "FLOAT")
-    RETURN_NAMES = ("features",)
+    RETURN_NAMES = ("features", "fps")
     FUNCTION = "extract_features"
     CATEGORY = PRISMAUDIO_CATEGORY
 
-    def extract_features(self, video, caption_cot, python_env="python", cache_dir="", synchformer_ckpt=""):
+    def extract_features(self, video, caption_cot, video_info=None, fps=30.0, python_env="managed_env", cache_dir="", hf_token=""):
-        # Resolve python binary — auto-install managed venv if using default
+        # Resolve fps from VHS video_info if connected
-        if python_env == "python":
+        if video_info is not None:
-            python_env = _ensure_extract_env()
+            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:
@@ -88,17 +150,23 @@ class PrismAudioFeatureExtractor:
         os.makedirs(cache_dir, exist_ok=True)
 
         # Check cache
-        cache_hash = _hash_inputs(video, caption_cot)
+        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()
-            return loader.load_features(cached_path)
+            features, = loader.load_features(cached_path)
+            return (features, float(fps))
 
-        # Save video to temp file
+        # Save frames to temp file (lossless .npy, no codec roundtrip)
-        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp:
+        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_video_tensor_to_mp4(video, tmp_video)
+        _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(
@@ -106,33 +174,55 @@ class PrismAudioFeatureExtractor:
             "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_env,
+            python_bin,
             script_path,
             "--video", tmp_video,
             "--cot_text", caption_cot,
             "--output", cached_path,
+            "--source_fps", str(fps),
+            "--synchformer_ckpt", synchformer_ckpt,
         ]
-        if synchformer_ckpt:
-            cmd.extend(["--synchformer_ckpt", synchformer_ckpt])
 
-        print(f"[PrismAudio] Extracting features via subprocess...")
+        # 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=True,
+                capture_output=False,
-                text=True,
                 timeout=600,  # 10 minute timeout
+                env=env,
             )
             if result.returncode != 0:
                 raise RuntimeError(
-                    f"[PrismAudio] Feature extraction failed:\n{result.stderr}"
+                    f"[PrismAudio] Feature extraction subprocess exited with code {result.returncode}. "
+                    "See output above for details."
                 )
-            print(result.stdout)
+            print("[PrismAudio] Feature extraction subprocess finished successfully.")
         finally:
             if os.path.exists(tmp_video):
                 os.unlink(tmp_video)
 
         # Load the extracted features
         loader = PrismAudioFeatureLoader()
-        return loader.load_features(cached_path)
+        features, = loader.load_features(cached_path)
+        return (features, float(fps))

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

@@ -91,12 +91,25 @@ class PrismAudioModelLoader:
         # Handle wrapped state dicts: some ckpts wrap in {"state_dict": ...}
         if "state_dict" in diffusion_state:
             diffusion_state = diffusion_state["state_dict"]
-        model.load_state_dict(diffusion_state, strict=False)
+        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."
@@ -105,7 +118,19 @@ class PrismAudioModelLoader:
                 vae_state[k[len(prefix):]] = v
             else:
                 vae_state[k] = v
-        model.pretransform.load_state_dict(vae_state)
+        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

nodes/sampler.py

@@ -15,9 +15,10 @@ class PrismAudioSampler:
             "required": {
                 "model": ("PRISMAUDIO_MODEL",),
                 "features": ("PRISMAUDIO_FEATURES",),
-                "duration": ("FLOAT", {"default": 10.0, "min": 1.0, "max": 30.0, "step": 0.1, "tooltip": "Audio duration in seconds"}),
+                "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": 24, "min": 1, "max": 100, "tooltip": "Number of sampling steps"}),
+                "steps": ("INT", {"default": 100, "min": 1, "max": 100, "tooltip": "Number of sampling steps"}),
-                "cfg_scale": ("FLOAT", {"default": 5.0, "min": 1.0, "max": 20.0, "step": 0.1, "tooltip": "Classifier-free guidance scale"}),
+                "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}),
             },
         }
@@ -27,27 +28,47 @@ class PrismAudioSampler:
     FUNCTION = "generate"
     CATEGORY = PRISMAUDIO_CATEGORY
 
-    def generate(self, model, features, duration, steps, cfg_scale, seed):
+    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": features["video_features"].to(device, dtype=dtype),
+            "video_features": video_feat,
-            "text_features": features["text_features"].to(device, dtype=dtype),
+            "text_features":  features["text_features"].to(device, dtype=dtype),
-            "sync_features": features["sync_features"].to(device, dtype=dtype),
+            "sync_features":  sync_feat,
             "video_exist": torch.tensor(has_video),
         }
         metadata = (sample_meta,)
@@ -66,6 +87,13 @@ class PrismAudioSampler:
             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)
 
@@ -96,6 +124,9 @@ class PrismAudioSampler:
                 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())
@@ -105,7 +136,7 @@ class PrismAudioSampler:
 
             # VAE decode in fp32 (snake activations overflow in fp16)
             with torch.amp.autocast(device_type=device.type, enabled=False):
-                audio = diffusion.pretransform.decode(fakes.float())
+                audio = diffusion.pretransform.decode(fakes_f)
 
             # Offload VAE
             if strategy == "offload_to_cpu":
@@ -114,30 +145,39 @@ class PrismAudioSampler:
 
         # 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 sync conditioning with learned empty embedding when video is absent.
+    """Replace video/sync conditioning with learned empty embeddings when video is absent.
 
-    Only substitutes sync_features — NOT video_features. The reference code
+    empty_clip_feat and empty_sync_feat are learned null embeddings in the conditioner
-    (predict.py/app.py) checks for 'metaclip_features' which doesn't exist in the
+    output space (1024-dim). Passing zero features through bias-free Cond_MLP produces
-    prismaudio.json config, so video substitution never runs. Cond_MLP with zero
+    near-zero activations, NOT the learned null signal the model was trained with.
-    input + bias-free linear layers naturally produces near-zero output.
 
     The conditioner returns {key: [tensor, mask]} where tensor is [B, seq, dim].
     """
     dit = diffusion.model.model if hasattr(diffusion.model, 'model') else diffusion.model
 
-    # Only substitute sync_features (matching reference behavior for prismaudio config)
+    # 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)
+        empty = dit.empty_sync_feat.to(device, dtype=dtype)           # [1, 1024]
-        cond_tensor = conditioning['sync_features'][0]
+        batch_size = conditioning['sync_features'][0].shape[0]
-        batch_size = cond_tensor.shape[0]
+        empty_expanded = empty.unsqueeze(0).expand(batch_size, -1, -1)  # [B, 1, 1024]
-        empty_expanded = empty.unsqueeze(0).expand(batch_size, -1, -1)
         conditioning['sync_features'][0] = empty_expanded
         conditioning['sync_features'][1] = torch.ones(batch_size, 1, device=device)

nodes/text_only.py

@@ -15,10 +15,10 @@ class PrismAudioTextOnly:
         return {
             "required": {
                 "model": ("PRISMAUDIO_MODEL",),
-                "text_prompt": ("STRING", {"default": "", "multiline": True, "tooltip": "Text description for audio generation"}),
+                "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": 24, "min": 1, "max": 100}),
+                "steps": ("INT", {"default": 100, "min": 1, "max": 100}),
-                "cfg_scale": ("FLOAT", {"default": 5.0, "min": 1.0, "max": 20.0, "step": 0.1}),
+                "cfg_scale": ("FLOAT", {"default": 7.0, "min": 1.0, "max": 20.0, "step": 0.1}),
                 "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFF}),
             },
         }
@@ -90,6 +90,9 @@ class PrismAudioTextOnly:
                 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())
@@ -98,7 +101,7 @@ class PrismAudioTextOnly:
 
             # VAE decode in fp32 (snake activations overflow in fp16)
             with torch.amp.autocast(device_type=device.type, enabled=False):
-                audio = diffusion.pretransform.decode(fakes.float())
+                audio = diffusion.pretransform.decode(fakes_f)
 
             if strategy == "offload_to_cpu":
                 diffusion.pretransform.to(get_offload_device())
@@ -106,8 +109,12 @@ class PrismAudioTextOnly:
 
         # 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},)
 

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/extract_features.py

@@ -1,64 +1,118 @@
 #!/usr/bin/env python3
 """
 Standalone PrismAudio feature extraction script.
-Run in a separate conda env with JAX/TF installed.
+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
-
-Setup:
-    conda env create -f environment.yml
-    conda activate prismaudio-extract
 """
 
 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"Error: Video not found: {args.video}")
+        print(f"[extract] ERROR: video not found: {args.video}", flush=True)
         sys.exit(1)
 
-    # Import feature extraction utils (requires JAX/TF)
+    print(f"[extract] Device     : {'cuda' if torch.cuda.is_available() else 'cpu'}", flush=True)
-    from data_utils.v2a_utils.feature_utils_288 import FeaturesUtils
-    import torchvision.transforms as T
-    from decord import VideoReader, cpu
-
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-    # Initialize feature extractor
+    # ------------------------------------------------------------------
+    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)
 
-    # Load and preprocess video
+    # ------------------------------------------------------------------
-    vr = VideoReader(args.video, ctx=cpu(0))
+    t0 = _step(3, 6, "reading and preprocessing video")
-    fps = vr.get_avg_fps()
+    if args.video.endswith(".npy"):
-    total_frames = len(vr)
+        all_frames = np.load(args.video)   # [T, H, W, C] uint8
-    duration = total_frames / fps
+        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)
 
-    # Extract CLIP frames (4fps, 288x288)
+        clip_indices = [int(i * fps / args.clip_fps) for i in range(max(1, int(duration * args.clip_fps)))]
-    clip_indices = [int(i * fps / args.clip_fps) for i in range(int(duration * args.clip_fps))]
+        clip_indices = [min(i, total_frames - 1) for i in clip_indices]
-    clip_indices = [min(i, total_frames - 1) for i in clip_indices]
+        clip_frames = all_frames[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 = 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(),
@@ -69,11 +123,6 @@ def main():
     ])
     clip_input = torch.stack([clip_transform(f) for f in clip_frames]).unsqueeze(0).to(device)
 
-    # Extract Sync frames (25fps, 224x224)
-    sync_indices = [int(i * fps / args.sync_fps) for i in range(int(duration * args.sync_fps))]
-    sync_indices = [min(i, total_frames - 1) for i in sync_indices]
-    sync_frames = vr.get_batch(sync_indices).asnumpy()
-
     sync_transform = T.Compose([
         T.ToPILImage(),
         T.Resize(args.sync_size),
@@ -82,30 +131,39 @@ def main():
         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)
 
-    # Extract features
+    # ------------------------------------------------------------------
-    print("[PrismAudio] Encoding text with T5-Gemma...")
+    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)}")
 
-    print("[PrismAudio] Encoding video with VideoPrism...")
+    # ------------------------------------------------------------------
+    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)}")
 
-    print("[PrismAudio] Encoding video with Synchformer...")
+    # ------------------------------------------------------------------
+    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)}")
 
-    # Save as .npz
+    # ------------------------------------------------------------------
+    t0 = time.perf_counter()
+    print(f"[extract] Saving features to {args.output} ...", flush=True)
     np.savez(
         args.output,
-        video_features=video_features.cpu().numpy(),
+        video_features=video_features.cpu().float().numpy(),
-        global_video_features=global_video_features.cpu().numpy(),
+        global_video_features=global_video_features.cpu().float().numpy(),
-        text_features=text_features.cpu().numpy(),
+        text_features=text_features.cpu().float().numpy(),
-        global_text_features=global_text_features.cpu().numpy(),
+        global_text_features=global_text_features.cpu().float().numpy(),
-        sync_features=sync_features.cpu().numpy(),
+        sync_features=sync_features.cpu().float().numpy(),
         caption_cot=args.cot_text,
         duration=duration,
     )
-    print(f"[PrismAudio] Features saved to {args.output}")
+    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__":

scripts/install_extract_env.sh

@@ -27,7 +27,7 @@ echo "[PrismAudio] Installing PyTorch stack..."
 
 echo "[PrismAudio] Installing feature-extraction dependencies..."
 "${PIP}" install \
-    "tensorflow-cpu==2.15.0" \
+    "tensorflow-cpu>=2.16.0" \
     "jax[cpu]" \
     "jaxlib" \
     "transformers" \

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
+}