feat: auto-install pip venv for feature extraction on first use

PrismAudioFeatureExtractor now creates and populates a managed venv
(_extract_env/) automatically when python_env is left as the default
'python'. Also adds scripts/install_extract_env.sh for manual/Docker
setup without conda.

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

.gitignore

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

README.md

@@ -0,0 +1,106 @@
+# 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.
+
+## Installation
+
+Clone into your ComfyUI custom nodes directory:
+
+```bash
+cd ComfyUI/custom_nodes
+git clone -b prismaudio https://github.com/FunAudioLLM/ThinkSound 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.
+
+## 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`. |
+
+## Workflows
+
+### Quality Path (Video-to-Audio)
+
+```
+Video → PrismAudio Feature Extractor → PrismAudio Sampler → Save Audio
+```
+
+### Pre-computed Path
+
+```
+PrismAudio Feature Loader (.npz) → PrismAudio Sampler → Save Audio
+```
+
+### Text-Only
+
+```
+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).
+
+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.
+
+## Model Files
+
+Weights are auto-downloaded to `ComfyUI/models/prismaudio/`:
+
+| File | Size | Description |
+|------|------|-------------|
+| `prismaudio.ckpt` | ~2.7 GB | Diffusion model |
+| `vae.ckpt` | ~2.5 GB | Stable Audio 2.0 VAE |
+| `synchformer_state_dict.pth` | ~950 MB | Synchformer |
+
+T5-Gemma is cached in the standard HuggingFace cache directory (`~/.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
+```
+
+## 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.
+
+## Credits
+
+PrismAudio by [FunAudioLLM](https://github.com/FunAudioLLM) (ICLR 2026). [Paper & code](https://github.com/FunAudioLLM/ThinkSound).

nodes/feature_extractor.py

@@ -0,0 +1,138 @@
+import os
+import sys
+import hashlib
+import subprocess
+import tempfile
+import torch
+
+from .utils import PRISMAUDIO_CATEGORY
+from .feature_loader import PrismAudioFeatureLoader
+
+# Managed venv created automatically when python_env is left as default
+_PLUGIN_DIR = os.path.dirname(os.path.dirname(__file__))
+_MANAGED_VENV = os.path.join(_PLUGIN_DIR, "_extract_env")
+_MANAGED_PYTHON = os.path.join(_MANAGED_VENV, "bin", "python")
+
+_EXTRACT_PACKAGES = [
+    "torch", "torchaudio", "torchvision",
+    "tensorflow-cpu==2.15.0",
+    "jax[cpu]", "jaxlib",
+    "transformers", "decord", "einops", "numpy", "mediapy",
+    "git+https://github.com/google-deepmind/videoprism.git",
+]
+
+
+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)
+    subprocess.run([sys.executable, "-m", "venv", _MANAGED_VENV], check=True)
+
+    pip = os.path.join(_MANAGED_VENV, "bin", "pip")
+    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)
+
+    print("[PrismAudio] Feature-extraction env ready.")
+    return _MANAGED_PYTHON
+
+
+def _hash_inputs(video_tensor, cot_text):
+    """Create a hash of the inputs for caching."""
+    h = hashlib.sha256()
+    h.update(video_tensor.cpu().numpy().tobytes()[:1024 * 1024])  # First 1MB for speed
+    h.update(cot_text.encode())
+    return h.hexdigest()[:16]
+
+
+def _save_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)
+
+
+class PrismAudioFeatureExtractor:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "video": ("IMAGE",),
+                "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."}),
+                "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)"}),
+            },
+        }
+
+    RETURN_TYPES = ("PRISMAUDIO_FEATURES",)
+    RETURN_NAMES = ("features",)
+    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()
+
+        # Determine cache directory
+        if not cache_dir:
+            cache_dir = os.path.join(tempfile.gettempdir(), "prismaudio_features")
+        os.makedirs(cache_dir, exist_ok=True)
+
+        # Check cache
+        cache_hash = _hash_inputs(video, caption_cot)
+        cached_path = os.path.join(cache_dir, f"{cache_hash}.npz")
+        if os.path.exists(cached_path):
+            print(f"[PrismAudio] Using cached features: {cached_path}")
+            loader = PrismAudioFeatureLoader()
+            return loader.load_features(cached_path)
+
+        # Save video to temp file
+        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp:
+            tmp_video = tmp.name
+        _save_video_tensor_to_mp4(video, tmp_video)
+
+        # Build subprocess command
+        script_path = os.path.join(
+            os.path.dirname(os.path.dirname(__file__)),
+            "scripts", "extract_features.py"
+        )
+
+        cmd = [
+            python_env,
+            script_path,
+            "--video", tmp_video,
+            "--cot_text", caption_cot,
+            "--output", cached_path,
+        ]
+        if synchformer_ckpt:
+            cmd.extend(["--synchformer_ckpt", synchformer_ckpt])
+
+        print(f"[PrismAudio] Extracting features via subprocess...")
+        try:
+            result = subprocess.run(
+                cmd,
+                capture_output=True,
+                text=True,
+                timeout=600,  # 10 minute timeout
+            )
+            if result.returncode != 0:
+                raise RuntimeError(
+                    f"[PrismAudio] Feature extraction failed:\n{result.stderr}"
+                )
+            print(result.stdout)
+        finally:
+            if os.path.exists(tmp_video):
+                os.unlink(tmp_video)
+
+        # Load the extracted features
+        loader = PrismAudioFeatureLoader()
+        return loader.load_features(cached_path)

nodes/feature_loader.py

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

nodes/model_loader.py

@@ -0,0 +1,129 @@
+import os
+import json
+import torch
+import folder_paths
+import comfy.model_management as mm
+import comfy.utils
+
+from .utils import (
+    PRISMAUDIO_CATEGORY, get_prismaudio_model_dir, register_model_folder,
+    get_device, get_offload_device, determine_precision, determine_offload_strategy,
+    soft_empty_cache, resolve_hf_token,
+)
+
+# HuggingFace repo for auto-download
+HF_REPO_ID = "FunAudioLLM/PrismAudio"
+REQUIRED_FILES = {
+    "diffusion": "prismaudio.ckpt",
+    "vae": "vae.ckpt",
+    "synchformer": "synchformer_state_dict.pth",
+}
+
+
+def _download_if_missing(filename, model_dir, hf_token=None):
+    """Download a model file from HuggingFace if not present locally."""
+    filepath = os.path.join(model_dir, filename)
+    if os.path.exists(filepath):
+        return filepath
+
+    from huggingface_hub import hf_hub_download
+    print(f"[PrismAudio] Downloading {filename} from {HF_REPO_ID}...")
+    try:
+        downloaded = hf_hub_download(
+            repo_id=HF_REPO_ID,
+            filename=filename,
+            local_dir=model_dir,
+            token=hf_token or None,
+        )
+        return downloaded
+    except Exception as e:
+        if "401" in str(e) or "403" in str(e) or "gated" in str(e).lower():
+            raise RuntimeError(
+                f"[PrismAudio] Model '{filename}' requires license acceptance. "
+                f"Visit https://huggingface.co/{HF_REPO_ID} to accept the license, "
+                f"then set HF_TOKEN env var or run: huggingface-cli login"
+            ) from e
+        raise
+
+
+class PrismAudioModelLoader:
+    @classmethod
+    def INPUT_TYPES(cls):
+        register_model_folder()
+        return {
+            "required": {
+                "precision": (["auto", "fp32", "fp16", "bf16"],),
+                "offload_strategy": (["auto", "keep_in_vram", "offload_to_cpu"],),
+            },
+        }
+
+    RETURN_TYPES = ("PRISMAUDIO_MODEL",)
+    RETURN_NAMES = ("model",)
+    FUNCTION = "load_model"
+    CATEGORY = PRISMAUDIO_CATEGORY
+
+    def load_model(self, precision, offload_strategy):
+        device = get_device()
+        dtype = determine_precision(precision, device)
+        strategy = determine_offload_strategy(offload_strategy)
+        token = resolve_hf_token()
+        model_dir = get_prismaudio_model_dir()
+
+        # Auto-download missing files
+        for key, filename in REQUIRED_FILES.items():
+            _download_if_missing(filename, model_dir, hf_token=token)
+
+        # Load config
+        config_path = os.path.join(
+            os.path.dirname(os.path.dirname(__file__)),
+            "prismaudio_core", "configs", "prismaudio.json"
+        )
+        with open(config_path) as f:
+            model_config = json.load(f)
+
+        # Create model from config
+        from prismaudio_core.factory import create_model_from_config
+        model = create_model_from_config(model_config)
+
+        # Load diffusion weights
+        diffusion_path = os.path.join(model_dir, REQUIRED_FILES["diffusion"])
+        diffusion_state = comfy.utils.load_torch_file(diffusion_path)
+        # Handle wrapped state dicts: some ckpts wrap in {"state_dict": ...}
+        if "state_dict" in diffusion_state:
+            diffusion_state = diffusion_state["state_dict"]
+        model.load_state_dict(diffusion_state, strict=False)
+
+        # 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)
+        # Strip "autoencoder." prefix from keys
+        vae_state = {}
+        prefix = "autoencoder."
+        for k, v in vae_full_state.items():
+            if k.startswith(prefix):
+                vae_state[k[len(prefix):]] = v
+            else:
+                vae_state[k] = v
+        model.pretransform.load_state_dict(vae_state)
+
+        # Apply precision: DiT + conditioners in user-selected dtype,
+        # but keep VAE (pretransform) in fp32 to avoid NaN from snake activations in fp16
+        model.model.to(dtype)        # DiTWrapper
+        model.conditioner.to(dtype)  # MultiConditioner
+        # model.pretransform stays in fp32
+
+        if strategy == "keep_in_vram":
+            model = model.to(device)
+        else:
+            model = model.to(get_offload_device())
+
+        model.eval()
+
+        return ({
+            "model": model,
+            "dtype": dtype,
+            "strategy": strategy,
+            "config": model_config,
+            "model_dir": model_dir,
+        },)

nodes/sampler.py

@@ -0,0 +1,143 @@
+import torch
+import comfy.model_management as mm
+import comfy.utils
+
+from .utils import (
+    PRISMAUDIO_CATEGORY, SAMPLE_RATE, DOWNSAMPLING_RATIO, IO_CHANNELS,
+    get_device, get_offload_device, soft_empty_cache,
+)
+
+
+class PrismAudioSampler:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("PRISMAUDIO_MODEL",),
+                "features": ("PRISMAUDIO_FEATURES",),
+                "duration": ("FLOAT", {"default": 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"}),
+                "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFF}),
+            },
+        }
+
+    RETURN_TYPES = ("AUDIO",)
+    RETURN_NAMES = ("audio",)
+    FUNCTION = "generate"
+    CATEGORY = PRISMAUDIO_CATEGORY
+
+    def generate(self, model, features, duration, steps, cfg_scale, seed):
+        device = get_device()
+        dtype = model["dtype"]
+        strategy = model["strategy"]
+        diffusion = model["model"]
+
+        # Compute latent dimensions
+        latent_length = round(SAMPLE_RATE * duration / DOWNSAMPLING_RATIO)
+
+        # Note: no seq length config needed — the model adapts to input tensor shapes
+        # dynamically via its transformer architecture.
+
+        # Determine if video features are present (not all zeros)
+        has_video = features.get("video_features") is not None and features["video_features"].abs().sum() > 0
+
+        # 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),
+            "text_features": features["text_features"].to(device, dtype=dtype),
+            "sync_features": features["sync_features"].to(device, dtype=dtype),
+            "video_exist": torch.tensor(has_video),
+        }
+        metadata = (sample_meta,)
+
+        # Move model to device if offloaded
+        if strategy == "offload_to_cpu":
+            diffusion.model.to(device)
+            diffusion.conditioner.to(device)
+            soft_empty_cache()
+
+        with torch.no_grad(), torch.amp.autocast(device_type=device.type, dtype=dtype):
+            # Run conditioning
+            conditioning = diffusion.conditioner(metadata, device)
+
+            # Handle missing video: substitute learned empty embeddings
+            if not has_video:
+                _substitute_empty_features(diffusion, conditioning, device, dtype)
+
+            # Assemble conditioning inputs for the DiT
+            cond_inputs = diffusion.get_conditioning_inputs(conditioning)
+
+            # Generate noise from seed (MPS doesn't support torch.Generator)
+            gen_device = "cpu" if device.type == "mps" else device
+            generator = torch.Generator(device=gen_device).manual_seed(seed)
+            noise = torch.randn(
+                [1, IO_CHANNELS, latent_length],
+                generator=generator,
+                device=gen_device,
+            ).to(device=device, dtype=dtype)
+
+            # Sample with progress bar
+            pbar = comfy.utils.ProgressBar(steps)
+
+            from prismaudio_core.inference.sampling import sample_discrete_euler
+
+            def on_step(info):
+                pbar.update(1)
+
+            fakes = sample_discrete_euler(
+                diffusion.model,
+                noise,
+                steps,
+                callback=on_step,
+                **cond_inputs,
+                cfg_scale=cfg_scale,
+                batch_cfg=True,
+            )
+
+            # Offload diffusion model and conditioner before VAE decode
+            if strategy == "offload_to_cpu":
+                diffusion.model.to(get_offload_device())
+                diffusion.conditioner.to(get_offload_device())
+                soft_empty_cache()
+                diffusion.pretransform.to(device)
+
+            # VAE decode in fp32 (snake activations overflow in fp16)
+            with torch.amp.autocast(device_type=device.type, enabled=False):
+                audio = diffusion.pretransform.decode(fakes.float())
+
+            # Offload VAE
+            if strategy == "offload_to_cpu":
+                diffusion.pretransform.to(get_offload_device())
+                soft_empty_cache()
+
+        # Peak normalize then clamp (matching reference: div by max abs before clamp)
+        audio = audio.float()
+        peak = audio.abs().max().clamp(min=1e-8)
+        audio = (audio / peak).clamp(-1, 1)
+
+        # 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.
+
+    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.
+
+    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)
+    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)
+        conditioning['sync_features'][0] = empty_expanded
+        conditioning['sync_features'][1] = torch.ones(batch_size, 1, device=device)

nodes/text_only.py

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

prismaudio_core/factory.py

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

prismaudio_core/inference/__init__.py

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

prismaudio_core/inference/sampling.py

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

prismaudio_core/inference/utils.py

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

scripts/environment.yml

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

scripts/extract_features.py

@@ -0,0 +1,112 @@
+#!/usr/bin/env python3
+"""
+Standalone PrismAudio feature extraction script.
+Run in a separate conda env with JAX/TF installed.
+
+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 numpy as np
+import torch
+
+
+def main():
+    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("--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()
+
+    if not os.path.exists(args.video):
+        print(f"Error: Video not found: {args.video}")
+        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
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # Initialize feature extractor
+    feat_utils = FeaturesUtils(
+        vae_config_path=args.vae_config,
+        synchformer_ckpt=args.synchformer_ckpt,
+        device=device,
+    )
+
+    # Load and preprocess video
+    vr = VideoReader(args.video, ctx=cpu(0))
+    fps = vr.get_avg_fps()
+    total_frames = len(vr)
+    duration = total_frames / fps
+
+    # 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()
+
+    clip_transform = T.Compose([
+        T.ToPILImage(),
+        T.Resize(args.clip_size),
+        T.CenterCrop(args.clip_size),
+        T.ToTensor(),
+        T.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+    ])
+    clip_input = torch.stack([clip_transform(f) for f in clip_frames]).unsqueeze(0).to(device)
+
+    # 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),
+        T.CenterCrop(args.sync_size),
+        T.ToTensor(),
+        T.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+    ])
+    sync_input = torch.stack([sync_transform(f) for f in sync_frames]).unsqueeze(0).to(device)
+
+    # Extract features
+    print("[PrismAudio] Encoding text with T5-Gemma...")
+    text_features = feat_utils.encode_t5_text([args.cot_text])
+
+    print("[PrismAudio] 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])
+
+    print("[PrismAudio] Encoding video with Synchformer...")
+    sync_features = feat_utils.encode_video_with_sync(sync_input)
+
+    # Save as .npz
+    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(),
+        caption_cot=args.cot_text,
+        duration=duration,
+    )
+    print(f"[PrismAudio] Features saved to {args.output}")
+
+
+if __name__ == "__main__":
+    main()

scripts/install_extract_env.sh

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