fix: return fps from non-cache extraction path

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

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

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.
-
-For the **Feature Extractor** node (video feature extraction), a separate conda environment is required — see [Feature Extraction Environment](#feature-extraction-environment) below.
+**flash-attn** is optional — detected at runtime, falls back to PyTorch SDPA if unavailable.
 
 ## Nodes
 
-| Node | Description |
-|------|-------------|
-| **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). |
-| **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. |
-| **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`. |
+### PrismAudio Model Loader
+
+Loads the DiT diffusion model and VAE. Auto-downloads weights from HuggingFace on first use.
+
+| Input | Options | Description |
+|-------|---------|-------------|
+| `precision` | auto / fp32 / fp16 / bf16 | DiT and conditioner dtype. VAE is always fp32. |
+| `offload_strategy` | auto / keep_in_vram / offload_to_cpu | Memory management. |
+
+---
+
+### PrismAudio Feature Extractor
+
+Extracts video features (VideoPrism LvT, Synchformer) and text features (T5-Gemma) from a video in a subprocess. Results are cached on disk.
+
+| Input | Description |
+|-------|-------------|
+| `video` | IMAGE tensor from any ComfyUI video loader |
+| `caption_cot` | Chain-of-thought description of the audio scene |
+| `video_info` | *(optional)* `VHS_VIDEOINFO` from VHS LoadVideo — sets fps automatically |
+| `fps` | Source fps — ignored if `video_info` is connected |
+| `python_env` | `managed_env` (auto-created isolated venv, recommended) or `comfyui_env` (current Python, see warning below) |
+| `cache_dir` | Directory for cached `.npz` files. Empty = system temp dir. |
+| `hf_token` | HuggingFace token for gated models. Prefer `HF_TOKEN` env var instead. |
+
+**Outputs:** `features` (PRISMAUDIO_FEATURES), `fps` (FLOAT)
+
+**`managed_env`** auto-creates a venv at `_extract_env/` inside the plugin directory on first use and installs JAX, TF, VideoPrism, and Synchformer. This takes several minutes the first time.
+
+**`comfyui_env`** uses the current ComfyUI Python — JAX/TF/videoprism must already be installed. Installing them into the ComfyUI environment may conflict with existing packages.
+
+---
+
+### PrismAudio Feature Loader
+
+Loads a pre-computed `.npz` feature file. Use this to re-use extracted features without re-running the extractor.
+
+| Input | Description |
+|-------|-------------|
+| `npz_path` | Path to a `.npz` file produced by the Feature Extractor |
+
+---
+
+### PrismAudio Sampler
+
+Video-to-audio generation. Takes model + features, produces AUDIO.
+
+| Input | Description |
+|-------|-------------|
+| `model` | From Model Loader |
+| `features` | From Feature Extractor or Feature Loader |
+| `duration` | Audio duration in seconds. Set to `0` to use the video duration from features automatically. |
+| `steps` | Sampling steps (default: 100) |
+| `cfg_scale` | Classifier-free guidance scale (default: 7.0) |
+| `seed` | RNG seed |
+
+---
+
+### PrismAudio Text Only
+
+Text-to-audio generation without video. Uses the T5-Gemma encoder.
+
+| Input | Description |
+|-------|-------------|
+| `model` | From Model Loader |
+| `text_prompt` | Chain-of-thought audio scene description. Longer, more detailed prompts produce better results. |
+| `duration` | Audio duration in seconds |
+| `steps` | Sampling steps (default: 100) |
+| `cfg_scale` | Classifier-free guidance scale (default: 7.0) |
+| `seed` | RNG seed |
+
+---
 
 ## Workflows
 
-### 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 |
-|------|----------|
-| 24 GB+ | Keep all models in VRAM |
-| 12–24 GB | Sequential offload |
-| 8–12 GB | Aggressive offload + fp16 |
-| < 8 GB | May work with aggressive offload |
-
-## 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
-```
+| VRAM | Recommended settings |
+|------|----------------------|
+| 24 GB+ | `keep_in_vram`, any precision |
+| 12–24 GB | `offload_to_cpu`, bf16/fp16 |
+| 8–12 GB | `offload_to_cpu`, fp16 |
+| < 8 GB | May work with `offload_to_cpu` + fp16 |
 
 ## Troubleshooting
 
 - **Gated model errors** — Accept the license at <https://huggingface.co/FunAudioLLM/PrismAudio> and set `HF_TOKEN`.
-- **VRAM errors** — Switch `offload_strategy` to `offload_to_cpu`, or use `fp16` precision.
-- **flash-attn** — Purely optional. Auto-detected at runtime; falls back to PyTorch SDPA.
+- **VRAM errors** — Switch `offload_strategy` to `offload_to_cpu` and/or use `fp16` precision.
+- **Feature extraction fails** — Ensure `synchformer_state_dict.pth` is in `models/prismaudio/`. On first run with `managed_env`, installation takes several minutes.
+- **flash-attn** — Optional. Auto-detected at runtime; falls back to PyTorch SDPA.
 
 ## Credits
 
-PrismAudio by [FunAudioLLM](https://github.com/FunAudioLLM) (ICLR 2026). [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]

nodes/feature_extractor.py

@@ -15,28 +15,58 @@ _MANAGED_PYTHON = os.path.join(_MANAGED_VENV, "bin", "python")
 
 _EXTRACT_PACKAGES = [
     "torch", "torchaudio", "torchvision",
-    "tensorflow-cpu==2.15.0",
-    "jax[cpu]", "jaxlib",
+    # TF 2.15 only supports Python <=3.11; use >=2.16 for Python 3.12+
+    "tensorflow-cpu>=2.16.0",
+    # jax[cuda13] includes jaxlib; pip-managed CUDA libs (no local toolkit needed)
+    "jax[cuda13]", "flax",
     "transformers", "decord", "einops", "numpy", "mediapy",
     "git+https://github.com/google-deepmind/videoprism.git",
 ]
 
 
+def _pip_install(pip, *packages, label=None):
+    """Install one or more packages with visible output; raise on failure."""
+    tag = label or packages[0]
+    print(f"[PrismAudio]   installing {tag} ...", flush=True)
+    result = subprocess.run(
+        [pip, "install", "--progress-bar", "on"] + list(packages),
+        capture_output=False,
+    )
+    if result.returncode != 0:
+        raise RuntimeError(
+            f"[PrismAudio] Failed to install {tag} (exit {result.returncode}). "
+            "See pip output above for details."
+        )
+    print(f"[PrismAudio]   {tag} OK", flush=True)
+
+
 def _ensure_extract_env():
     """Create and populate the managed venv on first use."""
     if os.path.exists(_MANAGED_PYTHON):
         return _MANAGED_PYTHON
 
-    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)...")
-    subprocess.run([pip, "install"] + _EXTRACT_PACKAGES, check=True)
+    total = len(_EXTRACT_PACKAGES)
+    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
 
 
@@ -48,13 +78,14 @@ def _hash_inputs(video_tensor, cot_text):
     return h.hexdigest()[:16]
 
 
-def _save_video_tensor_to_mp4(video_tensor, output_path, fps=30):
-    """Save ComfyUI IMAGE tensor [T,H,W,C] to MP4."""
-    import torchvision.io as tvio
-    # ComfyUI IMAGE is [T,H,W,C] float32 [0,1]
-    frames = (video_tensor * 255).to(torch.uint8)
-    # torchvision write_video expects [T,H,W,C] uint8
-    tvio.write_video(output_path, frames, fps=fps)
+def _save_frames_to_npy(video_tensor, output_path):
+    """Save ComfyUI IMAGE tensor [T,H,W,C] float32 [0,1] to .npy as uint8.
+
+    Lossless — avoids H.264 encode/decode roundtrip.
+    """
+    import numpy as np
+    frames_np = (video_tensor.cpu().numpy() * 255).astype("uint8")
+    np.save(output_path, frames_np)
 
 
 class PrismAudioFeatureExtractor:
@@ -66,21 +97,31 @@ 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_NAMES = ("features",)
+    RETURN_TYPES = ("PRISMAUDIO_FEATURES", "FLOAT")
+    RETURN_NAMES = ("features", "fps")
     FUNCTION = "extract_features"
     CATEGORY = PRISMAUDIO_CATEGORY
 
-    def extract_features(self, video, caption_cot, python_env="python", cache_dir="", synchformer_ckpt=""):
-        # Resolve python binary — auto-install managed venv if using default
-        if python_env == "python":
-            python_env = _ensure_extract_env()
+    def extract_features(self, video, caption_cot, video_info=None, fps=30.0, python_env="managed_env", cache_dir="", hf_token=""):
+        # Resolve fps from VHS video_info if connected
+        if video_info is not None:
+            fps = video_info["loaded_fps"]
+
+        # Resolve python binary
+        if python_env == "comfyui_env":
+            print("[PrismAudio] WARNING: using ComfyUI Python env — JAX/TF/videoprism must already be installed. "
+                  "Installing them here may conflict with existing packages and destabilize ComfyUI.", flush=True)
+            python_bin = sys.executable
+        else:
+            python_bin = _ensure_extract_env()
 
         # Determine cache directory
         if not cache_dir:
@@ -93,12 +134,18 @@ class PrismAudioFeatureExtractor:
         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
-        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp:
+        # Save frames to temp file (lossless .npy, no codec roundtrip)
+        import time
+        t0 = time.perf_counter()
+        frames = video.shape[0]
+        print(f"[PrismAudio] Saving {frames} frames to .npy (fps={fps})...", flush=True)
+        with tempfile.NamedTemporaryFile(suffix=".npy", delete=False) as tmp:
             tmp_video = tmp.name
-        _save_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 +153,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,
-                text=True,
+                capture_output=False,
                 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/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,9 @@ 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"}),
-                "steps": ("INT", {"default": 24, "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"}),
+                "duration": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 30.0, "step": 0.1, "tooltip": "Audio duration in seconds. Set to 0 to use the video duration from features automatically."}),
+                "steps": ("INT", {"default": 100, "min": 1, "max": 100, "tooltip": "Number of sampling steps"}),
+                "cfg_scale": ("FLOAT", {"default": 7.0, "min": 1.0, "max": 20.0, "step": 0.1, "tooltip": "Classifier-free guidance scale"}),
                 "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFF}),
             },
         }
@@ -33,6 +33,13 @@ class PrismAudioSampler:
         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)
 
@@ -42,12 +49,15 @@ class PrismAudioSampler:
         # 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),
-            "sync_features": features["sync_features"].to(device, dtype=dtype),
+            "sync_features":  sync_feat,
             "video_exist": torch.tensor(has_video),
         }
         metadata = (sample_meta,)
@@ -96,6 +106,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 +118,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 +127,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
-    (predict.py/app.py) checks for 'metaclip_features' which doesn't exist in the
-    prismaudio.json config, so video substitution never runs. Cond_MLP with zero
-    input + bias-free linear layers naturally produces near-zero output.
+    empty_clip_feat and empty_sync_feat are learned null embeddings in the conditioner
+    output space (1024-dim). Passing zero features through bias-free Cond_MLP produces
+    near-zero activations, NOT the learned null signal the model was trained with.
 
     The conditioner returns {key: [tensor, mask]} where tensor is [B, seq, dim].
     """
     dit = diffusion.model.model if hasattr(diffusion.model, 'model') else diffusion.model
 
-    # 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)
-        cond_tensor = conditioning['sync_features'][0]
-        batch_size = cond_tensor.shape[0]
-        empty_expanded = empty.unsqueeze(0).expand(batch_size, -1, -1)
+        empty = dit.empty_sync_feat.to(device, dtype=dtype)           # [1, 1024]
+        batch_size = conditioning['sync_features'][0].shape[0]
+        empty_expanded = empty.unsqueeze(0).expand(batch_size, -1, -1)  # [B, 1, 1024]
         conditioning['sync_features'][0] = empty_expanded
         conditioning['sync_features'][1] = torch.ones(batch_size, 1, device=device)

nodes/text_only.py

@@ -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}),
-                "cfg_scale": ("FLOAT", {"default": 5.0, "min": 1.0, "max": 20.0, "step": 0.1}),
+                "steps": ("INT", {"default": 100, "min": 1, "max": 100}),
+                "cfg_scale": ("FLOAT", {"default": 7.0, "min": 1.0, "max": 20.0, "step": 0.1}),
                 "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFF}),
             },
         }
@@ -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,6 +919,12 @@ class ContinuousTransformer(nn.Module):
                 x = self.fusion_mlp(x)
 
             if sync_cond is not None:
+                # Resample sync_cond to match audio sequence length if needed
+                if sync_cond.shape[1] != x.shape[1]:
+                    sync_cond = torch.nn.functional.interpolate(
+                        sync_cond.transpose(1, 2), size=x.shape[1],
+                        mode='linear', align_corners=False,
+                    ).transpose(1, 2)
                 if self.sync_film_generator is not None:
                     scale, shift = self.sync_film_generator(sync_cond).chunk(2, dim=-1)
                     x = x * (1 + scale) + shift

requirements.txt

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

scripts/extract_features.py

@@ -1,64 +1,116 @@
 #!/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)
-    from data_utils.v2a_utils.feature_utils_288 import FeaturesUtils
-    import torchvision.transforms as T
-    from decord import VideoReader, cpu
-
+    print(f"[extract] Device     : {'cuda' if torch.cuda.is_available() else 'cpu'}", flush=True)
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-    # 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
+    # ------------------------------------------------------------------
+    t0 = _step(3, 6, "reading and preprocessing video")
+    if args.video.endswith(".npy"):
+        all_frames = np.load(args.video)   # [T, H, W, C] uint8
+        fps = args.source_fps
+        total_frames = all_frames.shape[0]
+        duration = total_frames / fps
+        print(f"[extract]   fps={fps:.3f}  frames={total_frames}  duration={duration:.2f}s", flush=True)
+
+        clip_indices = [int(i * fps / args.clip_fps) for i in range(int(duration * args.clip_fps))]
+        clip_indices = [min(i, total_frames - 1) for i in clip_indices]
+        clip_frames = all_frames[clip_indices]
+        print(f"[extract]   CLIP frames  : {len(clip_indices)} @ {args.clip_fps}fps → {args.clip_size}×{args.clip_size}", flush=True)
+
+        sync_indices = [int(i * fps / args.sync_fps) for i in range(int(duration * args.sync_fps))]
+        sync_indices = [min(i, total_frames - 1) for i in sync_indices]
+        sync_frames = all_frames[sync_indices]
+        print(f"[extract]   Sync frames  : {len(sync_indices)} @ {args.sync_fps}fps → {args.sync_size}×{args.sync_size}", flush=True)
+    else:
+        from decord import VideoReader, cpu
         vr = VideoReader(args.video, ctx=cpu(0))
         fps = vr.get_avg_fps()
         total_frames = len(vr)
         duration = total_frames / fps
+        print(f"[extract]   fps={fps:.3f}  frames={total_frames}  duration={duration:.2f}s", flush=True)
 
-    # Extract CLIP frames (4fps, 288x288)
         clip_indices = [int(i * fps / args.clip_fps) for i in range(int(duration * args.clip_fps))]
         clip_indices = [min(i, total_frames - 1) for i in clip_indices]
         clip_frames = vr.get_batch(clip_indices).asnumpy()
+        print(f"[extract]   CLIP frames  : {len(clip_indices)} @ {args.clip_fps}fps → {args.clip_size}×{args.clip_size}", flush=True)
+
+        sync_indices = [int(i * fps / args.sync_fps) for i in range(int(duration * args.sync_fps))]
+        sync_indices = [min(i, total_frames - 1) for i in sync_indices]
+        sync_frames = vr.get_batch(sync_indices).asnumpy()
+        print(f"[extract]   Sync frames  : {len(sync_indices)} @ {args.sync_fps}fps → {args.sync_size}×{args.sync_size}", flush=True)
 
     clip_transform = T.Compose([
         T.ToPILImage(),
@@ -69,11 +121,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 +129,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(),
-        global_video_features=global_video_features.cpu().numpy(),
-        text_features=text_features.cpu().numpy(),
-        global_text_features=global_text_features.cpu().numpy(),
-        sync_features=sync_features.cpu().numpy(),
+        video_features=video_features.cpu().float().numpy(),
+        global_video_features=global_video_features.cpu().float().numpy(),
+        text_features=text_features.cpu().float().numpy(),
+        global_text_features=global_text_features.cpu().float().numpy(),
+        sync_features=sync_features.cpu().float().numpy(),
         caption_cot=args.cot_text,
         duration=duration,
     )
-    print(f"[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
+}