4 Commits

Author SHA1 Message Date
Ethanfel a315093743 feat: sync_strength control and temporal coverage diagnostic in sampler
Adds sync_strength (0.0–3.0, default 1.0) to PrismAudioSampler.
The scale is applied post-conditioner (after Sync_MLP) to the conditioning
tensor before it enters the DiT. Since CFG always uses zeros as the null
sync embedding, this cleanly scales the sync guidance signal:
  effective_sync_guidance = cfg_scale * (sync_strength * cond - 0)
Higher values tighten temporal audio-video alignment; 0.0 disables sync
guidance entirely (audio conditioned only by video + text features).
Not applied in T2A mode where sync is replaced by the learned empty_sync_feat.

Also logs sync temporal coverage vs audio target duration, with a warning
when they differ by more than 0.5s (stale or mismatched features).

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-03-28 16:23:41 +01:00
Ethanfel e49f760b77 fix: feature extractor CUDA detection, cache correctness, and short-video crash
- Detect CUDA version at venv creation time and install matching jax[cuda12/13]
  instead of hardcoded jax[cuda13] — was broken on CUDA 12.x (most systems)
- Include fps in cache hash: same video+caption at different fps previously
  returned stale cached features with wrong frame sampling
- Guard frame index lists with max(1,...)/max(8,...) to prevent torch.stack([])
  crash on very short input clips; sync minimum is 8 to match Synchformer's
  segment size requirement
- Remove mediapy from managed venv packages — not imported anywhere
- Warn when caption_cot is empty (produces degenerate text features)

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-03-28 16:00:05 +01:00
Ethanfel 4f40e15db3 fix: guard model cleanup in try/finally and fix DiTWrapper comments
- Wrap training loop in try/finally so _unapply_lora always runs.
  Without this, an exception mid-training would leave LoRALinear wrappers
  in the cached DiTWrapper; a subsequent training run would then apply LoRA
  on top of existing LoRA, silently doubling the effective rank.
- Fix misleading comment: diffusion.model is DiTWrapper (not DiffusionTransformer).
  DiffusionTransformer is at diffusion.model.model; _apply_lora reaches it
  recursively but the direct attribute is the wrapper.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-03-28 15:49:04 +01:00
Ethanfel 08d73773c5 feat: LoRA trainer and loader nodes for PrismAudio DiT fine-tuning
Adds PrismAudioLoRATrainer and PrismAudioLoRALoader nodes enabling
low-rank adaptation of the DiT on paired (video features + audio) datasets.

- LoRALinear wraps nn.Linear with trainable lora_A/lora_B matrices
- Rectified flow training loop with fp16 GradScaler, AdamW, cfg dropout
- Checkpoint saving every N steps + _config.json metadata alongside weights
- _unapply_lora restores base model state after training completes
- Weight-merge loader: delta_W added in-place, no deep copy overhead
- Three target presets: attn_only, attn_ffn (default), full

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-03-28 12:18:50 +01:00
7 changed files with 444 additions and 10 deletions
+2
View File
@@ -7,6 +7,8 @@ _NODES = {
"PrismAudioFeatureExtractor": (".feature_extractor", "PrismAudioFeatureExtractor", "PrismAudio Feature Extractor"), "PrismAudioFeatureExtractor": (".feature_extractor", "PrismAudioFeatureExtractor", "PrismAudio Feature Extractor"),
"PrismAudioSampler": (".sampler", "PrismAudioSampler", "PrismAudio Sampler"), "PrismAudioSampler": (".sampler", "PrismAudioSampler", "PrismAudio Sampler"),
"PrismAudioTextOnly": (".text_only", "PrismAudioTextOnly", "PrismAudio Text Only"), "PrismAudioTextOnly": (".text_only", "PrismAudioTextOnly", "PrismAudio Text Only"),
"PrismAudioLoRATrainer": (".lora_trainer", "PrismAudioLoRATrainer", "PrismAudio LoRA Trainer"),
"PrismAudioLoRALoader": (".lora_loader", "PrismAudioLoRALoader", "PrismAudio LoRA Loader"),
} }
for key, (module_path, class_name, display_name) in _NODES.items(): for key, (module_path, class_name, display_name) in _NODES.items():
+26 -5
View File
@@ -13,13 +13,29 @@ _PLUGIN_DIR = os.path.dirname(os.path.dirname(__file__))
_MANAGED_VENV = os.path.join(_PLUGIN_DIR, "_extract_env") _MANAGED_VENV = os.path.join(_PLUGIN_DIR, "_extract_env")
_MANAGED_PYTHON = os.path.join(_MANAGED_VENV, "bin", "python") _MANAGED_PYTHON = os.path.join(_MANAGED_VENV, "bin", "python")
def _jax_package():
"""Return the correct jax extra for the current CUDA version."""
try:
import torch
if torch.cuda.is_available():
cuda_ver = torch.version.cuda or ""
major = int(cuda_ver.split(".")[0]) if cuda_ver else 0
if major >= 13:
return "jax[cuda13]"
elif major >= 12:
return "jax[cuda12]"
except Exception:
pass
return "jax" # CPU fallback
_EXTRACT_PACKAGES = [ _EXTRACT_PACKAGES = [
"torch", "torchaudio", "torchvision", "torch", "torchaudio", "torchvision",
# TF 2.15 only supports Python <=3.11; use >=2.16 for Python 3.12+ # TF 2.15 only supports Python <=3.11; use >=2.16 for Python 3.12+
"tensorflow-cpu>=2.16.0", "tensorflow-cpu>=2.16.0",
# jax[cuda13] includes jaxlib; pip-managed CUDA libs (no local toolkit needed) # jax CUDA extra is resolved at install time based on detected CUDA version
"jax[cuda13]", "flax", _jax_package(), "flax",
"transformers", "decord", "einops", "numpy", "mediapy", "transformers", "decord", "einops", "numpy",
"git+https://github.com/google-deepmind/videoprism.git", "git+https://github.com/google-deepmind/videoprism.git",
] ]
@@ -70,11 +86,12 @@ def _ensure_extract_env():
return _MANAGED_PYTHON return _MANAGED_PYTHON
def _hash_inputs(video_tensor, cot_text): def _hash_inputs(video_tensor, cot_text, fps):
"""Create a hash of the inputs for caching.""" """Create a hash of the inputs for caching."""
h = hashlib.sha256() h = hashlib.sha256()
h.update(video_tensor.cpu().numpy().tobytes()[:1024 * 1024]) # First 1MB for speed h.update(video_tensor.cpu().numpy().tobytes()[:1024 * 1024]) # First 1MB for speed
h.update(cot_text.encode()) h.update(cot_text.encode())
h.update(str(fps).encode()) # fps affects frame sampling — must be part of the key
return h.hexdigest()[:16] return h.hexdigest()[:16]
@@ -115,6 +132,10 @@ class PrismAudioFeatureExtractor:
if video_info is not None: if video_info is not None:
fps = video_info["loaded_fps"] fps = video_info["loaded_fps"]
if not caption_cot.strip():
print("[PrismAudio] Warning: caption_cot is empty — text features will be degenerate. "
"Provide a descriptive chain-of-thought caption for best results.", flush=True)
# Resolve python binary # Resolve python binary
if python_env == "comfyui_env": if python_env == "comfyui_env":
print("[PrismAudio] WARNING: using ComfyUI Python env — JAX/TF/videoprism must already be installed. " print("[PrismAudio] WARNING: using ComfyUI Python env — JAX/TF/videoprism must already be installed. "
@@ -129,7 +150,7 @@ class PrismAudioFeatureExtractor:
os.makedirs(cache_dir, exist_ok=True) os.makedirs(cache_dir, exist_ok=True)
# Check cache # Check cache
cache_hash = _hash_inputs(video, caption_cot) cache_hash = _hash_inputs(video, caption_cot, fps)
cached_path = os.path.join(cache_dir, f"{cache_hash}.npz") cached_path = os.path.join(cache_dir, f"{cache_hash}.npz")
if os.path.exists(cached_path): if os.path.exists(cached_path):
print(f"[PrismAudio] Using cached features: {cached_path}") print(f"[PrismAudio] Using cached features: {cached_path}")
+106
View File
@@ -0,0 +1,106 @@
import os
import json
import torch
import torch.nn as nn
from .utils import PRISMAUDIO_CATEGORY
def _merge_lora_weights(dit: nn.Module, lora_state: dict, rank: int, alpha: float, strength: float):
"""Add LoRA delta weights directly into the base model's nn.Linear tensors.
delta_W = lora_B @ lora_A * scale * strength
applied as: linear.weight += delta_W
This is equivalent to LoRALinear at inference but requires no wrapper,
no extra memory, and no change to the model's forward call graph.
"""
scale = (alpha / rank) * strength
# Group saved keys by module path
a_map = {
k.replace(".lora_A.weight", ""): v
for k, v in lora_state.items() if k.endswith("lora_A.weight")
}
b_map = {
k.replace(".lora_B.weight", ""): v
for k, v in lora_state.items() if k.endswith("lora_B.weight")
}
merged = 0
for path, lora_A in a_map.items():
if path not in b_map:
print(f"[PrismAudio] LoRA merge: missing lora_B for {path}, skipping", flush=True)
continue
lora_B = b_map[path] # [out_features, rank]
# delta_W: [out_features, in_features]
delta_W = (lora_B.float() @ lora_A.float()) * scale
# Navigate to the parent module using PyTorch's get_submodule
*parent_parts, child_name = path.split(".")
try:
parent = dit.get_submodule(".".join(parent_parts)) if parent_parts else dit
except AttributeError as e:
print(f"[PrismAudio] LoRA merge: could not find module '{path}': {e}", flush=True)
continue
linear = getattr(parent, child_name, None)
if not isinstance(linear, nn.Linear):
print(f"[PrismAudio] LoRA merge: expected nn.Linear at '{path}', got {type(linear)}", flush=True)
continue
linear.weight.data.add_(delta_W.to(linear.weight.dtype))
merged += 1
print(f"[PrismAudio] LoRA merged {merged} layer(s) (strength={strength:.3f})", flush=True)
class PrismAudioLoRALoader:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": ("PRISMAUDIO_MODEL",),
"lora_path": ("STRING", {"default": "", "tooltip": "Path to .safetensors LoRA file produced by PrismAudio LoRA Trainer"}),
"strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05, "tooltip": "LoRA influence scale. 1.0 = full strength, 0.0 = base model only"}),
},
}
RETURN_TYPES = ("PRISMAUDIO_MODEL",)
RETURN_NAMES = ("model",)
FUNCTION = "load_lora"
CATEGORY = PRISMAUDIO_CATEGORY
def load_lora(self, model, lora_path, strength):
from safetensors.torch import load_file
if not os.path.exists(lora_path):
raise FileNotFoundError(f"[PrismAudio] LoRA file not found: {lora_path}")
config_path = lora_path.replace(".safetensors", "_config.json")
if not os.path.exists(config_path):
raise FileNotFoundError(
f"[PrismAudio] LoRA config not found: {config_path}\n"
"Expected a _config.json alongside the .safetensors file."
)
with open(config_path) as f:
config = json.load(f)
rank = config["rank"]
alpha = config["alpha"]
lora_state = load_file(lora_path)
# Merge LoRA weights in-place into the DiT's base linear layers.
# ComfyUI re-executes the upstream ModelLoader on the next queue run
# when inputs change, providing a fresh base model as needed.
dit = model["model"].model # DiTWrapper
if strength == 0.0:
print("[PrismAudio] LoRA strength=0.0 — skipping merge, base model unchanged.", flush=True)
return (model,)
_merge_lora_weights(dit, lora_state, rank, alpha, strength)
return (model,)
+284
View File
@@ -0,0 +1,284 @@
import os
import math
import json
import random
import torch
import torch.nn as nn
import torch.nn.functional as F
import comfy.utils
from .utils import (
PRISMAUDIO_CATEGORY, SAMPLE_RATE,
get_device, get_offload_device, soft_empty_cache,
)
# ---------------------------------------------------------------------------
# LoRA primitives
# ---------------------------------------------------------------------------
class LoRALinear(nn.Module):
"""Low-rank adapter wrapping a frozen nn.Linear."""
def __init__(self, linear: nn.Linear, rank: int, alpha: float):
super().__init__()
self.linear = linear
self.scale = alpha / rank
in_f, out_f = linear.in_features, linear.out_features
self.lora_A = nn.Linear(in_f, rank, bias=False)
self.lora_B = nn.Linear(rank, out_f, bias=False)
nn.init.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5))
nn.init.zeros_(self.lora_B.weight)
def forward(self, x):
return self.linear(x) + self.lora_B(self.lora_A(x)) * self.scale
_TARGET_MODULE_PRESETS = {
"attn_only": {"to_q", "to_kv", "to_qkv", "to_out"},
"attn_ffn": {"to_q", "to_kv", "to_qkv", "to_out", "proj"},
"full": {"to_q", "to_kv", "to_qkv", "to_out", "proj", "project_in", "project_out"},
}
def _apply_lora(module: nn.Module, target_attrs: set, rank: int, alpha: float):
"""Recursively replace matching nn.Linear layers with LoRALinear."""
for name, child in list(module.named_children()):
if isinstance(child, nn.Linear) and name in target_attrs:
setattr(module, name, LoRALinear(child, rank, alpha))
else:
_apply_lora(child, target_attrs, rank, alpha)
def _unapply_lora(module: nn.Module):
"""Replace LoRALinear back with the original frozen Linear (no weight merge)."""
for name, child in list(module.named_children()):
if isinstance(child, LoRALinear):
child.linear.weight.requires_grad_(False)
setattr(module, name, child.linear)
else:
_unapply_lora(child)
def _get_lora_state_dict(module: nn.Module) -> dict:
"""Return only LoRA parameter tensors from a module's state dict."""
return {k: v for k, v in module.state_dict().items()
if "lora_A" in k or "lora_B" in k}
# ---------------------------------------------------------------------------
# Dataset helpers
# ---------------------------------------------------------------------------
_AUDIO_EXTS = (".wav", ".flac", ".mp3")
def _scan_dataset(dataset_dir: str):
"""Return list of (npz_path, audio_path) pairs matched by stem."""
pairs = []
for fname in os.listdir(dataset_dir):
if not fname.endswith(".npz"):
continue
stem = os.path.join(dataset_dir, fname[:-4])
for ext in _AUDIO_EXTS:
audio_path = stem + ext
if os.path.exists(audio_path):
pairs.append((stem + ".npz", audio_path))
break
return sorted(pairs)
def _load_audio(audio_path: str, device: torch.device) -> torch.Tensor:
"""Load audio to [1, 2, samples] float32 tensor at SAMPLE_RATE."""
import torchaudio
waveform, sr = torchaudio.load(audio_path)
if sr != SAMPLE_RATE:
waveform = torchaudio.functional.resample(waveform, sr, SAMPLE_RATE)
if waveform.shape[0] == 1:
waveform = waveform.expand(2, -1)
elif waveform.shape[0] > 2:
waveform = waveform[:2]
return waveform.unsqueeze(0).to(device) # [1, 2, samples]
def _load_metadata(npz_path: str, device: torch.device, dtype: torch.dtype) -> dict:
"""Load .npz features into a conditioner metadata dict."""
import numpy as np
data = np.load(npz_path, allow_pickle=True)
video_feat = torch.from_numpy(data["video_features"]).float().to(device, dtype=dtype)
text_feat = torch.from_numpy(data["text_features"]).float().to(device, dtype=dtype)
sync_feat = torch.from_numpy(data["sync_features"]).float().to(device, dtype=dtype)
has_video = bool(video_feat.abs().sum() > 0)
return {
"video_features": video_feat,
"text_features": text_feat,
"sync_features": sync_feat,
"video_exist": torch.tensor(has_video),
}
# ---------------------------------------------------------------------------
# Trainer node
# ---------------------------------------------------------------------------
class PrismAudioLoRATrainer:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": ("PRISMAUDIO_MODEL",),
"dataset_dir": ("STRING", {"default": "", "tooltip": "Directory containing paired .npz feature files and .wav/.flac audio files (matched by filename stem)"}),
"output_path": ("STRING", {"default": "", "tooltip": "Save path for .safetensors weights. Empty = models/prismaudio/lora/"}),
"lora_rank": ("INT", {"default": 64, "min": 1, "max": 512}),
"lora_alpha": ("FLOAT", {"default": 64.0, "min": 1.0, "max": 1024.0}),
"target_modules": (["attn_ffn", "attn_only", "full"], {"tooltip": "attn_only: Q/K/V/out only. attn_ffn: + FFN input (recommended). full: + transformer I/O projections"}),
"learning_rate": ("FLOAT", {"default": 1e-4, "min": 1e-7, "max": 1e-2, "step": 1e-6}),
"train_steps": ("INT", {"default": 1000, "min": 1, "max": 100000}),
"cfg_dropout_prob": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 0.5, "step": 0.01, "tooltip": "Probability of dropping conditioning per step — preserves CFG ability at inference"}),
"save_every": ("INT", {"default": 500, "min": 1, "max": 100000, "tooltip": "Save a checkpoint every N steps (in addition to final save)"}),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFF}),
},
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("lora_path",)
FUNCTION = "train"
CATEGORY = PRISMAUDIO_CATEGORY
def train(self, model, dataset_dir, output_path, lora_rank, lora_alpha,
target_modules, learning_rate, train_steps, cfg_dropout_prob, save_every, seed):
from safetensors.torch import save_file
device = get_device()
dtype = model["dtype"]
diffusion = model["model"]
strategy = model["strategy"]
torch.manual_seed(seed)
random.seed(seed)
# Scan dataset
pairs = _scan_dataset(dataset_dir)
if not pairs:
raise RuntimeError(f"[PrismAudio] No (.npz + audio) pairs found in: {dataset_dir}")
print(f"[PrismAudio] LoRA training — {len(pairs)} sample(s), {train_steps} steps", flush=True)
# Resolve output path
if not output_path:
import folder_paths
out_dir = os.path.join(folder_paths.models_dir, "prismaudio", "lora")
os.makedirs(out_dir, exist_ok=True)
output_path = os.path.join(out_dir, f"prismaudio_lora_r{lora_rank}.safetensors")
# Move model to device
diffusion.model.to(device)
diffusion.conditioner.to(device)
diffusion.pretransform.to(device)
# Freeze all DiT params, then apply LoRA (adds trainable lora_A/lora_B)
dit = diffusion.model # DiTWrapper
for p in dit.parameters():
p.requires_grad_(False)
target_attrs = _TARGET_MODULE_PRESETS[target_modules]
_apply_lora(dit, target_attrs, lora_rank, lora_alpha)
# Cast LoRA params to model dtype and move to device
for m in dit.modules():
if isinstance(m, LoRALinear):
m.lora_A.to(device=device, dtype=dtype)
m.lora_B.to(device=device, dtype=dtype)
trainable = [p for p in dit.parameters() if p.requires_grad]
n_params = sum(p.numel() for p in trainable)
print(f"[PrismAudio] LoRA trainable params: {n_params:,} ({n_params/1e6:.2f}M)", flush=True)
diffusion.conditioner.eval()
diffusion.pretransform.eval()
dit.train()
optimizer = torch.optim.AdamW(trainable, lr=learning_rate)
# GradScaler for fp16 to prevent underflow
use_scaler = (dtype == torch.float16)
scaler = torch.cuda.amp.GradScaler() if use_scaler else None
pbar = comfy.utils.ProgressBar(train_steps)
try:
for step in range(1, train_steps + 1):
npz_path, audio_path = random.choice(pairs)
with torch.no_grad():
# Encode audio to latent space
audio = _load_audio(audio_path, device)
x0 = diffusion.pretransform.encode(audio.float()).to(dtype) # [1, 64, L]
# Build conditioning from features
metadata = (_load_metadata(npz_path, device, dtype),)
conditioning = diffusion.conditioner(metadata, device)
cond_inputs = diffusion.get_conditioning_inputs(conditioning)
# Rectified flow: interpolate between data and noise
t = torch.rand(x0.shape[0], device=device, dtype=dtype) # [1]
noise = torch.randn_like(x0)
# t expanded for broadcast: [1] -> [1, 1, 1]
t_bcast = t[:, None, None]
x_t = (1.0 - t_bcast) * x0 + t_bcast * noise
v_target = noise - x0
with torch.amp.autocast(device_type=device.type, dtype=dtype):
v_pred = dit(x_t, t,
cfg_scale=1.0,
cfg_dropout_prob=cfg_dropout_prob,
**cond_inputs)
loss = F.mse_loss(v_pred.float(), v_target.float())
if use_scaler:
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
else:
loss.backward()
optimizer.step()
optimizer.zero_grad()
if step % 50 == 0:
print(f"[PrismAudio] step {step}/{train_steps} loss={loss.item():.6f}", flush=True)
if step % save_every == 0:
ckpt_path = output_path.replace(".safetensors", f"_step{step}.safetensors")
save_file(_get_lora_state_dict(dit), ckpt_path)
print(f"[PrismAudio] Checkpoint: {ckpt_path}", flush=True)
pbar.update(1)
# Save final weights
save_file(_get_lora_state_dict(dit), output_path)
# Save config alongside weights so the loader knows the structure
config_path = output_path.replace(".safetensors", "_config.json")
with open(config_path, "w") as f:
json.dump({
"rank": lora_rank,
"alpha": lora_alpha,
"target_modules": sorted(target_attrs),
}, f, indent=2)
print(f"[PrismAudio] LoRA saved: {output_path}", flush=True)
finally:
# Always restore model to base state — even on exception.
# Without this, LoRA wrappers would persist in the cached model and
# subsequent training runs would apply LoRA on top of existing LoRA.
dit.eval()
_unapply_lora(dit)
if strategy == "offload_to_cpu":
diffusion.model.to(get_offload_device())
diffusion.conditioner.to(get_offload_device())
diffusion.pretransform.to(get_offload_device())
soft_empty_cache()
return (output_path,)
+19 -1
View File
@@ -18,6 +18,7 @@ class PrismAudioSampler:
"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."}), "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"}), "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"}), "cfg_scale": ("FLOAT", {"default": 7.0, "min": 1.0, "max": 20.0, "step": 0.1, "tooltip": "Classifier-free guidance scale"}),
"sync_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 3.0, "step": 0.05, "tooltip": "Scale factor for sync conditioning. Higher values tighten audio-visual sync at the cost of audio naturalness; 0.0 disables sync guidance entirely."}),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFF}), "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFF}),
}, },
} }
@@ -27,7 +28,7 @@ class PrismAudioSampler:
FUNCTION = "generate" FUNCTION = "generate"
CATEGORY = PRISMAUDIO_CATEGORY CATEGORY = PRISMAUDIO_CATEGORY
def generate(self, model, features, duration, steps, cfg_scale, seed): def generate(self, model, features, duration, steps, cfg_scale, sync_strength, seed):
device = get_device() device = get_device()
dtype = model["dtype"] dtype = model["dtype"]
strategy = model["strategy"] strategy = model["strategy"]
@@ -43,6 +44,16 @@ class PrismAudioSampler:
# Compute latent dimensions # Compute latent dimensions
latent_length = round(SAMPLE_RATE * duration / DOWNSAMPLING_RATIO) latent_length = round(SAMPLE_RATE * duration / DOWNSAMPLING_RATIO)
# Sync temporal coverage diagnostic
sync_frames = features["sync_features"].shape[0]
sync_duration_covered = sync_frames / 25.0 # Synchformer always extracts at 25fps
print(f"[PrismAudio] sync: {sync_frames} frames @ 25fps = {sync_duration_covered:.2f}s | "
f"audio target: {latent_length} latent frames = {duration:.2f}s", flush=True)
if abs(sync_duration_covered - duration) > 0.5:
print(f"[PrismAudio] Warning: sync coverage ({sync_duration_covered:.2f}s) differs from "
f"audio duration ({duration:.2f}s) by more than 0.5s — consider re-extracting features "
f"with the correct video duration.", flush=True)
# Note: no seq length config needed — the model adapts to input tensor shapes # Note: no seq length config needed — the model adapts to input tensor shapes
# dynamically via its transformer architecture. # dynamically via its transformer architecture.
@@ -76,6 +87,13 @@ class PrismAudioSampler:
if not has_video: if not has_video:
_substitute_empty_features(diffusion, conditioning, device, dtype) _substitute_empty_features(diffusion, conditioning, device, dtype)
# Scale sync conditioning after the conditioner MLP (clean linear scale,
# avoids SiLU nonlinearity in Sync_MLP). The CFG null path always uses zeros,
# so this directly scales the sync guidance magnitude: cfg_scale * (strength*cond - 0).
# Only applied when video is present — T2A uses learned empty_sync_feat, not raw sync.
if has_video and sync_strength != 1.0 and 'sync_features' in conditioning:
conditioning['sync_features'][0] = conditioning['sync_features'][0] * sync_strength
# Assemble conditioning inputs for the DiT # Assemble conditioning inputs for the DiT
cond_inputs = diffusion.get_conditioning_inputs(conditioning) cond_inputs = diffusion.get_conditioning_inputs(conditioning)
+1
View File
@@ -9,3 +9,4 @@ descript-audio-codec
vector-quantize-pytorch vector-quantize-pytorch
scipy scipy
tqdm tqdm
torchaudio
+6 -4
View File
@@ -85,12 +85,13 @@ def main():
duration = total_frames / fps duration = total_frames / fps
print(f"[extract] fps={fps:.3f} frames={total_frames} duration={duration:.2f}s", flush=True) 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 = [int(i * fps / args.clip_fps) for i in range(max(1, int(duration * args.clip_fps)))]
clip_indices = [min(i, total_frames - 1) for i in clip_indices] clip_indices = [min(i, total_frames - 1) for i in clip_indices]
clip_frames = all_frames[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) 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))] # Synchformer processes in segments of 8; ensure at least 8 frames
sync_indices = [int(i * fps / args.sync_fps) for i in range(max(8, int(duration * args.sync_fps)))]
sync_indices = [min(i, total_frames - 1) for i in sync_indices] sync_indices = [min(i, total_frames - 1) for i in sync_indices]
sync_frames = all_frames[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) print(f"[extract] Sync frames : {len(sync_indices)} @ {args.sync_fps}fps → {args.sync_size}×{args.sync_size}", flush=True)
@@ -102,12 +103,13 @@ def main():
duration = total_frames / fps duration = total_frames / fps
print(f"[extract] fps={fps:.3f} frames={total_frames} duration={duration:.2f}s", flush=True) 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 = [int(i * fps / args.clip_fps) for i in range(max(1, int(duration * args.clip_fps)))]
clip_indices = [min(i, total_frames - 1) for i in clip_indices] clip_indices = [min(i, total_frames - 1) for i in clip_indices]
clip_frames = vr.get_batch(clip_indices).asnumpy() 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) 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))] # Synchformer processes in segments of 8; ensure at least 8 frames
sync_indices = [int(i * fps / args.sync_fps) for i in range(max(8, int(duration * args.sync_fps)))]
sync_indices = [min(i, total_frames - 1) for i in sync_indices] sync_indices = [min(i, total_frames - 1) for i in sync_indices]
sync_frames = vr.get_batch(sync_indices).asnumpy() 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) print(f"[extract] Sync frames : {len(sync_indices)} @ {args.sync_fps}fps → {args.sync_size}×{args.sync_size}", flush=True)