# Copyright (c) 2026 Scenema AI # https://scenema.ai # SPDX-License-Identifier: MIT """Audio generation engine for Scenema Audio. Loads the LTX 2.3 audio-only checkpoint, Audio VAE encoder, and Gemma 3 12B text encoder. VRAM management is auto-detected: models are moved between GPU and CPU as needed per inference phase. """ import gc import json import logging import os import time from contextlib import contextmanager from dataclasses import dataclass, replace as dc_replace import numpy as np import psutil import torch import torchaudio from safetensors import safe_open from safetensors.torch import load_file from ltx_core.batch_split import BatchSplitAdapter, BatchedPerturbationConfig from ltx_core.components.diffusion_steps import EulerDiffusionStep from ltx_core.components.noisers import GaussianNoiser from ltx_core.components.patchifiers import AudioPatchifier, VideoLatentPatchifier from ltx_core.model.audio_vae.audio_vae import Audio, encode_audio from ltx_core.model.audio_vae.model_configurator import AudioEncoderConfigurator from ltx_core.model.transformer.model import X0Model from ltx_core.model.transformer.model_configurator import LTXModelConfigurator from ltx_core.model.transformer.transformer import BasicAVTransformerBlock, rms_norm from ltx_core.tools import AudioLatentTools, LatentState, VideoLatentTools from ltx_core.types import AudioLatentShape, VideoLatentShape, VideoPixelShape from ltx_pipelines.distilled import DISTILLED_SIGMAS, DistilledPipeline from ltx_pipelines.utils.blocks import ModalitySpec, _build_state from ltx_pipelines.utils.denoisers import SimpleDenoiser from ltx_pipelines.utils.samplers import euler_denoising_loop from ltx_pipelines.utils.types import OffloadMode from ltx_core.text_encoders.gemma.tokenizer import LTXVGemmaTokenizer import bitsandbytes # noqa: F401 from transformers import BitsAndBytesConfig, Gemma3ForConditionalGeneration from .audio_utils import extract_wav logger = logging.getLogger(__name__) FPS = 24 MAX_REF_SECONDS = 5 class _Int8Linear(torch.nn.Module): """Linear layer with INT8 weights, dequantized to input dtype during forward. Keeps weights as int8 buffers in VRAM (~50% of bf16). Dequantization happens per forward pass: weight = int8 * scale, then cast to input dtype. Ported from bench_full_quantized.py. """ def __init__(self, weight_int8, scale, bias=None): super().__init__() self.register_buffer("weight_int8", weight_int8) self.register_buffer("scale", scale) if bias is not None: self.register_parameter("bias", torch.nn.Parameter(bias)) else: self.bias = None def forward(self, x): w = self.weight_int8.float() * self.scale.unsqueeze(1) w = w.to(x.dtype) return torch.nn.functional.linear(x, w, self.bias) # VRAM threshold: cards with this much VRAM keep all models GPU-resident # (Gemma bf16 on GPU, no offloading, MelBandRoFormer + SeedVC preloaded). # Below this: Gemma streams from CPU, models load/unload per request. HIGH_VRAM_THRESHOLD_GB = 40 @dataclass class AudioResult: waveform_np: np.ndarray # (samples,) or (samples, channels) float32 sample_rate: int duration_s: float def _materialize_meta_tensors(module, device="cpu"): """Replace meta tensors with zeros on the specified device.""" for name, param in list(module.named_parameters()): if param.is_meta: parts = name.split(".") mod = module for p in parts[:-1]: mod = getattr(mod, p) mod._parameters[parts[-1]] = torch.nn.Parameter( torch.zeros(param.shape, dtype=torch.bfloat16, device=device) ) for name, buf in list(module.named_buffers()): if buf.is_meta: parts = name.split(".") mod = module for p in parts[:-1]: mod = getattr(mod, p) mod._buffers[parts[-1]] = torch.zeros( buf.shape, dtype=torch.bfloat16, device=device ) def _audio_only_forward(self, video, audio, perturbations=None): """Monkey-patched forward for audio-only transformer blocks. Skips all video computation (attn1, attn2, ff, audio_to_video_attn) and only runs audio self-attention, cross-attention, and feedforward. """ if video is None and audio is None: raise ValueError("Need at least one modality") batch_size = (video or audio).x.shape[0] if perturbations is None: perturbations = BatchedPerturbationConfig.empty(batch_size) vx = video.x if video is not None else None ax = audio.x if audio is not None else None run_ax = audio is not None and audio.enabled and ax.numel() > 0 if run_ax: ashift_msa, ascale_msa, agate_msa = self.get_ada_values( self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(0, 3) ) norm_ax = rms_norm(ax, eps=self.norm_eps) * (1 + ascale_msa) + ashift_msa del ashift_msa, ascale_msa ax = ( ax + self.audio_attn1( norm_ax, pe=audio.positional_embeddings, mask=audio.self_attention_mask ) * agate_msa ) del agate_msa, norm_ax ax = ax + self._apply_text_cross_attention( ax, audio.context, self.audio_attn2, self.audio_scale_shift_table, getattr(self, "audio_prompt_scale_shift_table", None), audio.timesteps, audio.prompt_timestep, audio.context_mask, cross_attention_adaln=self.cross_attention_adaln, ) ashift_ff, ascale_ff, agate_ff = self.get_ada_values( self.audio_scale_shift_table, ax.shape[0], audio.timesteps, slice(3, 6) ) norm_ax_ff = rms_norm(ax, eps=self.norm_eps) * (1 + ascale_ff) + ashift_ff del ashift_ff, ascale_ff ax = ax + self.audio_ff(norm_ax_ff) * agate_ff del agate_ff, norm_ax_ff if video is not None: object.__setattr__(video, "x", vx) if audio is not None: object.__setattr__(audio, "x", ax) return video, audio # ── VRAM Manager ──────────────────────────────────────────────────────── class VRAMManager: """Manages model placement between GPU and CPU based on available VRAM. Tracks which models are on GPU and moves them as needed per inference phase. Offloading is determined by comparing total registered model size against available VRAM. If all models fit, no offloading occurs. """ def __init__(self, vram_gb: float): self.vram_gb = vram_gb self._models: dict[str, torch.nn.Module] = {} self._model_sizes: dict[str, float] = {} # GB per model self._on_gpu: set[str] = set() self.needs_offload = False # Determined after all models registered def register(self, name: str, model: torch.nn.Module, on_gpu: bool = True) -> None: """Register a model for VRAM management. Args: name: Identifier for the model. model: The PyTorch module. on_gpu: Whether the model is currently on GPU. """ self._models[name] = model size_gb = sum(p.numel() * p.element_size() for p in model.parameters()) / 1e9 self._model_sizes[name] = size_gb if on_gpu: self._on_gpu.add(name) def finalize(self) -> None: """Determine offloading strategy based on total model size vs VRAM. Call after all models are registered. Sets needs_offload based on whether all registered models fit in VRAM simultaneously with headroom for activations and pipeline overhead (~5GB). """ total_model_gb = sum(self._model_sizes.values()) # Gemma overhead depends on quantization mode: # bf16 streaming: ~16GB peak (13GB Gemma + 2GB embeddings + 1GB safety) # NF4: ~11GB peak (8GB NF4 model on GPU + 2GB embeddings + 1GB safety) gemma_nf4 = os.environ.get("GEMMA_QUANTIZE", "").lower() == "nf4" gemma_overhead_gb = 11.0 if gemma_nf4 else 16.0 self.needs_offload = (total_model_gb + gemma_overhead_gb) > self.vram_gb logger.info( "VRAM strategy: %.1f GB models + %.1f GB Gemma overhead (%s) vs %.1f GB VRAM -> offload=%s", total_model_gb, gemma_overhead_gb, "nf4" if gemma_nf4 else "bf16", self.vram_gb, "yes" if self.needs_offload else "no", ) def to_gpu(self, *names: str) -> None: """Move specified models to GPU, offloading others if needed. If offloading is required (VRAM < 40GB), all models NOT in the requested set are moved to CPU first to free VRAM. Args: names: Model names that should be on GPU for the current phase. """ if not self.needs_offload: # High VRAM: just ensure requested models are on GPU for name in names: if name not in self._on_gpu and name in self._models: self._models[name].cuda() self._on_gpu.add(name) return # Offload models that shouldn't be on GPU needed = set(names) to_offload = self._on_gpu - needed for name in to_offload: if name in self._models: self._models[name].cpu() self._on_gpu.discard(name) logger.debug("Offloaded %s to CPU", name) torch.cuda.empty_cache() # Load requested models to GPU for name in names: if name not in self._on_gpu and name in self._models: self._models[name].cuda() self._on_gpu.add(name) logger.debug("Loaded %s to GPU", name) def free_all(self) -> None: """Move all models to CPU.""" for name in list(self._on_gpu): if name in self._models: self._models[name].cpu() self._on_gpu.clear() torch.cuda.empty_cache() @contextmanager def phase(self, *names: str): """Context manager for a VRAM phase. Ensures specified models are on GPU for the duration, then returns to previous state on exit. Args: names: Model names needed on GPU for this phase. """ prev_on_gpu = set(self._on_gpu) self.to_gpu(*names) try: yield finally: # Restore previous state only if offloading is needed if self.needs_offload: to_restore = prev_on_gpu - set(names) to_remove = set(names) - prev_on_gpu for name in to_remove: if name in self._models and name in self._on_gpu: self._models[name].cpu() self._on_gpu.discard(name) for name in to_restore: if name in self._models and name not in self._on_gpu: self._models[name].cuda() self._on_gpu.add(name) torch.cuda.empty_cache() # ── Audio Engine ──────────────────────────────────────────────────────── class AudioEngine: """LTX 2.3 audio-only generation engine. Loads the baked audio checkpoint, Audio VAE encoder, and Gemma 3 12B text encoder. VRAM is managed automatically per inference phase. """ def __init__( self, audio_ckpt_path: str, vae_encoder_path: str, gemma_root: str, pipeline_ckpt_path: str | None = None, ): """Initialize AudioEngine. Args: audio_ckpt_path: Path to the audio-only transformer checkpoint. vae_encoder_path: Path to the standalone Audio VAE encoder checkpoint. gemma_root: Path to the Gemma 3 12B model directory. pipeline_ckpt_path: Path to checkpoint for DistilledPipeline. """ self.audio_ckpt_path = audio_ckpt_path self.vae_encoder_path = vae_encoder_path self.gemma_root = gemma_root self.pipeline_ckpt_path = pipeline_ckpt_path or audio_ckpt_path self._config = None self._mdl_wrapper = None self._audio_encoder = None self._pipeline = None self._vram: VRAMManager | None = None self._vae_sr = None self._loaded = False @property def vae_sample_rate(self) -> int: return self._vae_sr or 16000 def load(self) -> None: """Load all models. Call once at startup.""" if self._loaded: return vram_gb = torch.cuda.get_device_properties(0).total_memory / 1e9 ram_gb = psutil.virtual_memory().total / 1e9 logger.info( "System: %.1f GB VRAM, %.1f GB RAM, GPU: %s", vram_gb, ram_gb, torch.cuda.get_device_name(0), ) if vram_gb < 11: raise RuntimeError( f"Insufficient VRAM: {vram_gb:.0f}GB. Minimum 11GB required." ) if ram_gb < 24: raise RuntimeError( f"Insufficient RAM: {ram_gb:.0f}GB. Minimum 24GB required." ) self._vram = VRAMManager(vram_gb) self._load_audio_model() self._load_vae_encoder() self._patch_transformer_blocks() self._build_pipeline() # Determine offloading strategy based on actual model sizes vs VRAM self._vram.finalize() self._loaded = True logger.info("AudioEngine loaded") def _load_audio_model(self) -> None: """Load the audio-only checkpoint to GPU. Supports both bf16 and INT8 quantized checkpoints. INT8 checkpoints store weights as .weight.int8 (int8) + .weight.scale (float32) pairs. For INT8, nn.Linear layers are replaced with Int8Linear modules that keep weights quantized in VRAM (~5GB vs 9.8GB) and dequantize during the forward pass. """ t0 = time.time() with safe_open(self.audio_ckpt_path, framework="pt") as f: self._config = json.loads(f.metadata()["config"]) with torch.device("meta"): mdl = LTXModelConfigurator.from_config(self._config) sd = load_file(self.audio_ckpt_path, device="cpu") # Detect INT8 checkpoint format int8_map = { k.replace(".weight.int8", ""): k for k in sd if k.endswith(".weight.int8") } scale_map = { k.replace(".weight.scale", ""): k for k in sd if k.endswith(".weight.scale") } is_int8 = len(int8_map) > 0 if is_int8: # Load only non-quantized keys first (biases, norms, embeddings) regular_sd = { k: v for k, v in sd.items() if not k.endswith(".int8") and not k.endswith(".scale") } mdl_wrapper = X0Model(mdl) mdl_wrapper.load_state_dict(regular_sd, strict=False, assign=True) # Replace nn.Linear with Int8Linear for quantized weights n_replaced = 0 for name in int8_map: w_int8 = sd[int8_map[name]] w_scale = sd[scale_map[name]] parts = name.split(".") parent = mdl_wrapper for p in parts[:-1]: parent = getattr(parent, p) old = getattr(parent, parts[-1]) bias_key = name + ".bias" bias = sd.get(bias_key) if bias is None and hasattr(old, "bias") and old.bias is not None: bias = old.bias.data setattr(parent, parts[-1], _Int8Linear(w_int8, w_scale, bias)) n_replaced += 1 logger.info("INT8: replaced %d Linear layers with Int8Linear", n_replaced) else: mdl_wrapper = X0Model(mdl) mdl_wrapper.load_state_dict(sd, strict=False, assign=True) # Runtime INT8 quantization via BnB (bf16 checkpoint → INT8 on GPU) if os.environ.get("TRANSFORMER_QUANTIZE", "").lower() == "int8": import bitsandbytes as bnb n_quantized = 0 for name, module in list(mdl_wrapper.named_modules()): for cn, child in list(module.named_children()): if ( isinstance(child, torch.nn.Linear) and child.weight.numel() > 1_000_000 ): int8_layer = bnb.nn.Linear8bitLt( child.in_features, child.out_features, bias=child.bias is not None, has_fp16_weights=False, ) int8_layer.weight = bnb.nn.Int8Params( child.weight.data, requires_grad=False, has_fp16_weights=False, ) if child.bias is not None: int8_layer.bias = child.bias setattr(module, cn, int8_layer) n_quantized += 1 logger.info( "Runtime INT8: quantized %d Linear layers via BnB", n_quantized ) del sd gc.collect() for block in mdl.transformer_blocks: block.attn1 = torch.nn.Identity() block.attn2 = torch.nn.Identity() block.ff = torch.nn.Identity() block.audio_to_video_attn = torch.nn.Identity() gc.collect() _materialize_meta_tensors(mdl_wrapper) cross_pe = max( mdl.positional_embedding_max_pos[0], mdl.audio_positional_embedding_max_pos[0], ) mdl._init_preprocessors(cross_pe) self._mdl_wrapper = mdl_wrapper.cuda().eval() self._vram.register("audio_model", self._mdl_wrapper, on_gpu=True) logger.info( "Audio model loaded: %.1f GB, %.1fs", torch.cuda.memory_allocated() / 1e9, time.time() - t0, ) def _load_vae_encoder(self) -> None: """Load Audio VAE encoder from standalone checkpoint.""" t0 = time.time() avae_cfg = self._config["audio_vae"] preproc = avae_cfg["preprocessing"] self._vae_sr = preproc["audio"]["sampling_rate"] with torch.device("meta"): encoder = AudioEncoderConfigurator().from_config(avae_cfg) sd = load_file(self.vae_encoder_path, device="cpu") encoder.load_state_dict(sd, strict=False, assign=True) pcs = encoder.per_channel_statistics if "per_channel_statistics.std-of-means" in sd: pcs._buffers["std-of-means"] = sd["per_channel_statistics.std-of-means"] pcs._buffers["mean-of-means"] = sd["per_channel_statistics.mean-of-means"] del sd dd = avae_cfg["model"]["params"]["ddconfig"] encoder.mel_bins = dd["mel_bins"] encoder.mid.attn_1 = torch.nn.Identity() _materialize_meta_tensors(encoder, device="cpu") self._audio_encoder = encoder.cuda().eval().to(torch.bfloat16) self._vram.register("vae_encoder", self._audio_encoder, on_gpu=True) logger.info( "Audio VAE encoder loaded: %.1fM params, %.1fs", sum(p.numel() for p in self._audio_encoder.parameters()) / 1e6, time.time() - t0, ) def _patch_transformer_blocks(self) -> None: """Monkey-patch transformer blocks for audio-only forward pass.""" BasicAVTransformerBlock.forward = _audio_only_forward logger.info("Transformer blocks patched for audio-only forward") def _build_pipeline(self) -> None: """Build DistilledPipeline and cache Gemma + embeddings processor in CPU RAM. Caching eliminates the ~35s rebuild cost on every encode call. Gemma stays in CPU RAM permanently, streams to GPU layer-by-layer. Embeddings processor shuttles between CPU and GPU per call. """ t0 = time.time() mdl_wrapper = self._mdl_wrapper # Use NONE offload when VRAM is sufficient so Gemma stays GPU-resident # for fast encoding (~0.5s vs ~7s streaming). Fall back to CPU streaming # on smaller cards. offload = ( OffloadMode.NONE if self._vram.vram_gb >= HIGH_VRAM_THRESHOLD_GB else OffloadMode.CPU ) self._pipeline = DistilledPipeline( distilled_checkpoint_path=self.pipeline_ckpt_path, gemma_root=self.gemma_root, spatial_upsampler_path=None, loras=[], offload_mode=offload, ) @contextmanager def _gpu_ctx(**kw): yield mdl_wrapper self._pipeline.stage._transformer_ctx = _gpu_ctx pe = self._pipeline.prompt_encoder # Gemma loading strategy: # NF4: BitsAndBytes int4 quantization (~8GB on GPU, ~0.1s encode) # bf16 GPU: full precision on GPU (~24GB, ~1-2s encode) — when VRAM >= 40GB # bf16 streaming: streams from CPU RAM layer-by-layer (~7s encode) — when VRAM < 40GB self._gemma_nf4 = os.environ.get("GEMMA_QUANTIZE", "").lower() == "nf4" self._gemma_on_gpu = False if self._gemma_nf4: self._build_nf4_gemma() # NF4 needs its own embeddings processor and tokenizer self._cached_emb_proc = pe._embeddings_processor_builder.build( device="cuda", dtype=torch.bfloat16, ).eval() self._cached_tokenizer = LTXVGemmaTokenizer(self.gemma_root) logger.info("Embeddings processor cached on CUDA (NF4 mode)") elif self._vram.vram_gb >= HIGH_VRAM_THRESHOLD_GB: # Build pipeline's text encoder ONCE on GPU and keep it resident. # This uses the same builder as pipeline.prompt_encoder but # avoids the build/destroy cycle that makes each call ~30s. t_gemma = time.time() self._resident_text_encoder = pe._text_encoder_builder.build( device=torch.device("cuda"), dtype=torch.bfloat16, ).eval() self._cached_emb_proc = pe._embeddings_processor_builder.build( device="cuda", dtype=torch.bfloat16, ).eval() self._gemma_on_gpu = True vram_gb = torch.cuda.memory_allocated() / (1024**3) logger.info( "Gemma bf16 (pipeline encoder) GPU-resident: %.1fGB VRAM, %.1fs", vram_gb, time.time() - t_gemma, ) else: # Low VRAM: pipeline.prompt_encoder streams from CPU (~7s/encode) logger.info("Gemma managed by pipeline prompt_encoder (CPU streaming)") logger.info("Pipeline built: %.1fs", time.time() - t0) def _build_nf4_gemma(self) -> None: """Load Gemma 3 12B with BitsAndBytes NF4 quantization (~8GB on GPU). NF4 Gemma stays on GPU permanently. Encode is near-instant (~0.1s) since there's no CPU->GPU streaming. Slight quality tradeoff vs bf16 but acceptable for production use. """ t0 = time.time() quant_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_quant_type="nf4", ) self._nf4_gemma_model = Gemma3ForConditionalGeneration.from_pretrained( self.gemma_root, quantization_config=quant_config, device_map="cuda", dtype=torch.bfloat16, ).eval() # No streaming text encoder needed — _cached_text_encoder stays None self._cached_text_encoder = None vram_gb = torch.cuda.memory_allocated() / (1024**3) logger.info( "Gemma NF4 loaded on GPU: %.1fGB VRAM, %.1fs", vram_gb, time.time() - t0 ) def _build_bf16_gemma_gpu(self) -> None: """Load Gemma 3 12B bf16 directly on GPU (~24GB). For cards with >= 40GB VRAM. Gemma stays on GPU permanently. Encode is ~1-2s (pure inference, no CPU->GPU streaming). """ t0 = time.time() self._nf4_gemma_model = Gemma3ForConditionalGeneration.from_pretrained( self.gemma_root, device_map="cuda", torch_dtype=torch.bfloat16, ).eval() self._cached_text_encoder = None self._gemma_on_gpu = True vram_gb = torch.cuda.memory_allocated() / (1024**3) logger.info( "Gemma bf16 loaded on GPU: %.1fGB VRAM, %.1fs", vram_gb, time.time() - t0 ) def unload(self) -> None: """Free all GPU and CPU memory.""" if self._vram: self._vram.free_all() if ( hasattr(self, "_cached_text_encoder") and self._cached_text_encoder is not None ): self._cached_text_encoder.teardown() self._cached_text_encoder = None if hasattr(self, "_nf4_gemma_model"): del self._nf4_gemma_model self._nf4_gemma_model = None if hasattr(self, "_cached_emb_proc"): self._cached_emb_proc = None if hasattr(self, "_cached_tokenizer"): self._cached_tokenizer = None self._mdl_wrapper = None self._audio_encoder = None self._pipeline = None self._vram = None self._loaded = False gc.collect() torch.cuda.empty_cache() logger.info("AudioEngine unloaded") def encode_text(self, prompt: str): """Encode text prompt via Gemma 3 12B. Uses the pipeline's PromptEncoder which builds Gemma through the LTX-native builder. This ensures identical encoding to the reference pipeline (critical for SFX generation quality). Falls back to NF4/bf16 GPU-resident Gemma when available for speed, but routes through the pipeline encoder for correctness. Args: prompt: Compiled video-style text prompt. Returns: Tuple of (video_context, audio_context) tensors for denoising. """ t0 = time.time() with torch.inference_mode(): if self._gemma_nf4: # NF4: use BitsAndBytes quantized Gemma (fast, ~0.1s) tp = self._cached_tokenizer.tokenize_with_weights(prompt)["gemma"] ids = torch.tensor([[t[0] for t in tp]], device="cuda") mask = torch.tensor([[w[1] for w in tp]], device="cuda") out = self._nf4_gemma_model.model( input_ids=ids, attention_mask=mask, output_hidden_states=True, ) hs = out.hidden_states am = mask del out, ids emb = self._cached_emb_proc.process_hidden_states(hs, am) vc = emb.video_encoding ac = emb.audio_encoding del hs, am, emb elif self._gemma_on_gpu: # bf16 GPU-resident: use pipeline's text encoder (fast, ~0.5s) hs, am = self._resident_text_encoder.encode(prompt) emb = self._cached_emb_proc.process_hidden_states(hs, am) vc = emb.video_encoding ac = emb.audio_encoding del hs, am, emb else: # CPU streaming: use pipeline's prompt encoder (~7s) (emb,) = self._pipeline.prompt_encoder([prompt]) vc = emb.video_encoding ac = emb.audio_encoding logger.info("Gemma encode: %.1fs", time.time() - t0) return vc, ac def encode_reference(self, waveform_np: np.ndarray, sample_rate: int): """Encode reference audio to latent via Audio VAE encoder. Args: waveform_np: Audio samples, shape (samples,) or (samples, channels). sample_rate: Sample rate of the input audio in Hz. Returns: Reference latent tensor [B, C, T, F] on GPU. """ # Ensure VAE encoder is on GPU self._vram.to_gpu("vae_encoder") if waveform_np.ndim == 1: waveform_np = np.stack([waveform_np, waveform_np], axis=-1) if waveform_np.ndim == 2 and waveform_np.shape[1] == 2: wav = torch.from_numpy(waveform_np.T).float() else: wav = torch.from_numpy(waveform_np).float() if sample_rate != self._vae_sr: wav = torchaudio.functional.resample(wav, sample_rate, self._vae_sr) max_samples = MAX_REF_SECONDS * self._vae_sr if wav.shape[1] > max_samples: wav = wav[:, :max_samples] audio_obj = Audio(waveform=wav.unsqueeze(0), sampling_rate=self._vae_sr) with torch.inference_mode(): latent = encode_audio(audio_obj, self._audio_encoder) logger.info("Reference encoded: %s", latent.shape) return latent def generate( self, vc, ac, duration: float, seed: int, ref_latent=None, ) -> AudioResult: """Generate audio with optional A2V reference conditioning. Args: vc: Video context from encode_text(). ac: Audio context from encode_text(). duration: Target duration in seconds. seed: Random seed for reproducibility. ref_latent: Optional reference latent from encode_reference() for A2V voice conditioning. Returns: AudioResult with waveform numpy array and metadata. """ return self._generate_impl(vc, ac, duration, seed, ref_latent) @torch.inference_mode() def _generate_impl(self, vc, ac, duration, seed, ref_latent=None): # Ensure audio model is on GPU for denoising self._vram.to_gpu("audio_model") num_frames = ((int(duration * FPS) + 7) // 8) * 8 + 1 device = torch.device("cuda") gen = torch.Generator(device=device).manual_seed(seed) noiser = GaussianNoiser(generator=gen) sigmas = DISTILLED_SIGMAS.to(dtype=torch.float32, device=device) pixel_shape = VideoPixelShape( batch=1, frames=num_frames, width=64, height=64, fps=FPS ) v_shape = VideoLatentShape.from_pixel_shape(pixel_shape) video_tools = VideoLatentTools( VideoLatentPatchifier(patch_size=1), v_shape, fps=FPS ) video_state = _build_state( ModalitySpec(context=vc, conditionings=[]), video_tools, noiser, torch.bfloat16, device, ) a_shape = AudioLatentShape.from_video_pixel_shape(pixel_shape) audio_tools = AudioLatentTools(AudioPatchifier(patch_size=1), a_shape) audio_state = _build_state( ModalitySpec(context=ac), audio_tools, noiser, torch.bfloat16, device, ) ref_frames = 0 if ref_latent is not None: ref = ref_latent.to(device=device, dtype=torch.bfloat16) ref_frames = ref.shape[2] total_t = ref_frames + audio_state.latent.shape[1] ref_patchified = ref.permute(0, 2, 1, 3).reshape(1, ref_frames, -1) combined_latent = torch.cat([ref_patchified, audio_state.latent], dim=1) ref_mask = torch.zeros( 1, ref_frames, 1, device=device, dtype=audio_state.denoise_mask.dtype ) combined_mask = torch.cat([ref_mask, audio_state.denoise_mask], dim=1) combined_clean = torch.cat( [ref_patchified, torch.zeros_like(audio_state.clean_latent)], dim=1 ) combined_a_shape = AudioLatentShape( batch=1, channels=8, frames=total_t, mel_bins=16 ) combined_audio_tools = AudioLatentTools( AudioPatchifier(patch_size=1), combined_a_shape ) gen2 = torch.Generator(device=device).manual_seed(seed) noiser2 = GaussianNoiser(generator=gen2) tmp_state = _build_state( ModalitySpec(context=ac), combined_audio_tools, noiser2, torch.bfloat16, device, ) combined_positions = tmp_state.positions del tmp_state audio_state_final = LatentState( latent=combined_latent, denoise_mask=combined_mask, positions=combined_positions, clean_latent=combined_clean, attention_mask=None, ) else: audio_state_final = audio_state stepper = EulerDiffusionStep() with self._pipeline.stage._transformer_ctx() as transformer: wrapped = BatchSplitAdapter(transformer, max_batch_size=1) t0 = time.time() _, audio_state_out = euler_denoising_loop( sigmas=sigmas, video_state=video_state, audio_state=audio_state_final, stepper=stepper, transformer=wrapped, denoiser=SimpleDenoiser(vc, ac), ) logger.debug("Denoise: %.2fs", time.time() - t0) if ref_latent is not None and audio_state_out is not None and ref_frames > 0: audio_state_out = dc_replace( audio_state_out, latent=audio_state_out.latent[:, ref_frames:], denoise_mask=audio_state_out.denoise_mask[:, ref_frames:], positions=audio_state_out.positions[:, :, ref_frames:], clean_latent=( audio_state_out.clean_latent[:, ref_frames:] if audio_state_out.clean_latent is not None else None ), ) audio_state_out = audio_tools.clear_conditioning(audio_state_out) audio_state_out = audio_tools.unpatchify(audio_state_out) if torch.isnan(audio_state_out.latent).any(): logger.warning("NaN detected in denoised latent") # Offload audio model before VAE decode (pipeline handles decoder GPU usage) self._vram.to_gpu() audio = self._pipeline.audio_decoder(audio_state_out.latent) # Restore audio model after decode self._vram.to_gpu("audio_model") w, sr = extract_wav(audio) return AudioResult(waveform_np=w, sample_rate=sr, duration_s=w.shape[0] / sr)