BlueWaveSemi45/DramaboxCPU
0
1#!/usr/bin/env python32"""Warm validation runner — loads base dev + LoRA + all aux models ONCE,3then iterates every speaker in val_config generating each output.4 5Matches the same generation path as inference.py but keeps Gemma / audio VAE6/ velocity model / audio decoder resident across entries. Inference7settings default to the Gradio warm-server values (cfg=2.5, stg=1.5,8modality=1.0, rescale=0, 30 steps, fps=25) — use --inference-params to9override.10"""11import argparse12import logging13import os14import sys15import time16import traceback17 18import torch19import torchaudio20 21REPO_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))22MODEL_DIR = REPO_DIR23sys.path.insert(0, os.path.join(REPO_DIR, "ltx2"))24sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))25 26DEV_FULL_CKPT = os.environ.get(27 "LTX_FULL_CHECKPOINT",28 os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "ltx-2.3-22b-dev.safetensors"),29)30GEMMA_ROOT = os.environ.get(31 "GEMMA_ROOT",32 os.path.expanduser("~/.cache/dramabox/gemma-3-12b-it-bnb-4bit"),33)34 35 36def parse_args():37 p = argparse.ArgumentParser()38 p.add_argument("--val-config", required=True)39 p.add_argument("--output-dir", required=True)40 p.add_argument("--lora", default=None)41 p.add_argument("--lora-rank", type=int, default=128)42 p.add_argument("--full-checkpoint", default=DEV_FULL_CKPT)43 p.add_argument("--gemma-root", default=GEMMA_ROOT)44 p.add_argument("--cfg-scale", type=float, default=2.5)45 p.add_argument("--stg-scale", type=float, default=1.5)46 p.add_argument("--rescale-scale", type=float, default=0.0)47 p.add_argument("--modality-scale", type=float, default=1.0)48 p.add_argument("--steps", type=int, default=30)49 p.add_argument("--fps", type=float, default=25.0)50 p.add_argument("--stg-block", type=int, default=29)51 p.add_argument("--cfg-clamp", type=float, default=0.0)52 p.add_argument("--seed", type=int, default=42)53 p.add_argument("--duration-multiplier", type=float, default=1.1)54 # Match Gradio / inference_server.py DEFAULT_NEG exactly55 p.add_argument("--negative-prompt", default=(56 "worst quality, inconsistent, robotic, distorted, noise, static, "57 "muffled, unclear, unnatural, monotone"58 ))59 return p.parse_args()60 61 62def estimate_speech_duration(prompt: str, speed: float = 1.0) -> float:63 import re64 quoted = re.findall(r'"([^"]*)"', prompt) or re.findall(r"'([^']*)'", prompt)65 text = " ".join(quoted) if quoted else prompt66 duration = len(text) * 0.065 / max(speed, 0.1) + 1.567 return max(3.0, round(duration, 1))68 69 70class WarmValidator:71 def __init__(self, full_checkpoint, gemma_root, lora_path=None, lora_rank=128,72 device="cuda", dtype=torch.bfloat16):73 from audio_conditioning import AudioConditionByReferenceLatent # noqa: F401 (imported by inference.py)74 from ltx_core.components.patchifiers import AudioPatchifier75 from ltx_pipelines.utils.blocks import PromptEncoder, AudioConditioner, AudioDecoder76 77 self.device = torch.device(device)78 self.dtype = dtype79 self.full_checkpoint = full_checkpoint80 self.gemma_root = gemma_root81 self.patchifier = AudioPatchifier(patch_size=1)82 83 logging.info("Loading PromptEncoder (Gemma + embeddings_processor)...")84 t0 = time.time()85 self.prompt_encoder = PromptEncoder(86 checkpoint_path=full_checkpoint, gemma_root=gemma_root,87 dtype=dtype, device=self.device, warm=True, audio_only=True,88 )89 logging.info(f" PromptEncoder ready in {time.time()-t0:.1f}s")90 91 logging.info("Loading AudioConditioner (audio VAE encoder)...")92 t0 = time.time()93 self.audio_conditioner = AudioConditioner(94 checkpoint_path=full_checkpoint, dtype=dtype, device=self.device, warm=True,95 )96 logging.info(f" AudioConditioner ready in {time.time()-t0:.1f}s")97 98 logging.info("Loading AudioDecoder...")99 t0 = time.time()100 self.audio_decoder = AudioDecoder(101 checkpoint_path=full_checkpoint, dtype=dtype, device=self.device, warm=True,102 )103 logging.info(f" AudioDecoder ready in {time.time()-t0:.1f}s")104 105 logging.info("Building velocity model (audio-only from base dev)...")106 t0 = time.time()107 self.velocity_model = self._build_velocity_model(full_checkpoint, lora_path, lora_rank)108 logging.info(f" Velocity model ready in {time.time()-t0:.1f}s "109 f"({sum(p.numel() for p in self.velocity_model.parameters()) / 1e9:.1f}B params)")110 111 def _build_velocity_model(self, checkpoint_path, lora_path, lora_rank):112 from ltx_core.loader.registry import DummyRegistry113 from ltx_core.loader.sd_ops import SDOps114 from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder115 from ltx_core.model.model_protocol import ModelConfigurator116 from ltx_core.model.transformer.attention import AttentionFunction117 from ltx_core.model.transformer.model import LTXModel, LTXModelType118 from ltx_core.model.transformer.rope import LTXRopeType119 120 sd_ops = (121 SDOps("AO")122 .with_matching(prefix="model.diffusion_model.")123 .with_replacement("model.diffusion_model.", "")124 )125 126 class Cfg(ModelConfigurator[LTXModel]):127 @classmethod128 def from_config(cls, config):129 t = config.get("transformer", {})130 cp = None131 if not t.get("caption_proj_before_connector", False):132 from ltx_core.model.transformer.text_projection import create_caption_projection133 with torch.device("meta"):134 cp = create_caption_projection(t, audio=True)135 return LTXModel(136 model_type=LTXModelType.AudioOnly,137 audio_num_attention_heads=t.get("audio_num_attention_heads", 32),138 audio_attention_head_dim=t.get("audio_attention_head_dim", 64),139 audio_in_channels=t.get("audio_in_channels", 128),140 audio_out_channels=t.get("audio_out_channels", 128),141 num_layers=t.get("num_layers", 48),142 audio_cross_attention_dim=t.get("audio_cross_attention_dim", 2048),143 norm_eps=t.get("norm_eps", 1e-6),144 attention_type=AttentionFunction(t.get("attention_type", "default")),145 positional_embedding_theta=10000.0,146 audio_positional_embedding_max_pos=[20.0],147 timestep_scale_multiplier=t.get("timestep_scale_multiplier", 1000),148 use_middle_indices_grid=t.get("use_middle_indices_grid", True),149 rope_type=LTXRopeType(t.get("rope_type", "interleaved")),150 double_precision_rope=t.get("frequencies_precision", False) == "float64",151 apply_gated_attention=t.get("apply_gated_attention", False),152 audio_caption_projection=cp,153 cross_attention_adaln=t.get("cross_attention_adaln", False),154 )155 156 builder = Builder(157 model_path=checkpoint_path, model_class_configurator=Cfg,158 model_sd_ops=sd_ops, registry=DummyRegistry(),159 )160 velocity = builder.build(device=self.device, dtype=self.dtype).to(self.device).eval()161 162 if lora_path and os.path.exists(lora_path):163 from peft import LoraConfig, get_peft_model164 from safetensors.torch import load_file as st_load165 logging.info(f"Attaching LoRA: {lora_path}")166 lora_sd = st_load(lora_path)167 is_peft = any("base_model.model." in k for k in lora_sd.keys())168 is_iclora = any("diffusion_model." in k for k in lora_sd.keys())169 cfg = LoraConfig(170 r=lora_rank, lora_alpha=lora_rank, lora_dropout=0.0, bias="none",171 target_modules=[172 "audio_attn1.to_k", "audio_attn1.to_q",173 "audio_attn1.to_v", "audio_attn1.to_out.0",174 "audio_attn2.to_k", "audio_attn2.to_q",175 "audio_attn2.to_v", "audio_attn2.to_out.0",176 "audio_ff.net.0.proj", "audio_ff.net.2",177 ],178 )179 velocity = get_peft_model(velocity, cfg)180 181 if is_peft:182 mapped = {}183 for k, v in lora_sd.items():184 nk = k185 if ".lora_A.weight" in k and ".lora_A.default.weight" not in k:186 nk = k.replace(".lora_A.weight", ".lora_A.default.weight")187 if ".lora_B.weight" in k and ".lora_B.default.weight" not in k:188 nk = k.replace(".lora_B.weight", ".lora_B.default.weight")189 mapped[nk] = v190 _, unexpected = velocity.load_state_dict(mapped, strict=False)191 logging.info(f" Loaded {len(mapped) - len(unexpected)} LoRA weights (peft)")192 elif is_iclora:193 audio_keys = {k: v for k, v in lora_sd.items()194 if "audio_attn1" in k or "audio_attn2" in k or "audio_ff" in k}195 mapped = {}196 for k, v in audio_keys.items():197 nk = k.replace("diffusion_model.", "base_model.model.")198 nk = nk.replace(".lora_A.weight", ".lora_A.default.weight")199 nk = nk.replace(".lora_B.weight", ".lora_B.default.weight")200 mapped[nk] = v201 _, unexpected = velocity.load_state_dict(mapped, strict=False)202 logging.info(f" Loaded {len(mapped) - len(unexpected)} LoRA weights (iclora)")203 204 velocity = velocity.merge_and_unload()205 logging.info(" Merged LoRA into base weights")206 207 return velocity208 209 @torch.inference_mode()210 def generate(self, prompt, output_path, voice_ref=None, args=None):211 from audio_conditioning import AudioConditionByReferenceLatent212 from ltx_core.batch_split import BatchSplitAdapter213 from ltx_core.components.diffusion_steps import EulerDiffusionStep214 from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderParams215 from ltx_core.components.noisers import GaussianNoiser216 from ltx_core.components.schedulers import LTX2Scheduler217 from ltx_core.model.audio_vae import encode_audio as vae_encode_audio218 from ltx_core.model.transformer.model import X0Model219 from ltx_core.tools import AudioLatentTools220 from ltx_core.types import Audio, AudioLatentShape, VideoPixelShape221 from ltx_pipelines.utils.denoisers import GuidedDenoiser, SimpleDenoiser222 from ltx_pipelines.utils.gpu_model import gpu_model223 from ltx_pipelines.utils.media_io import decode_audio_from_file224 from ltx_pipelines.utils.samplers import euler_denoising_loop225 226 t_total = time.time()227 228 # ---- Duration + shape ----229 gen_dur = estimate_speech_duration(prompt) * args.duration_multiplier230 raw_frames = int(round(gen_dur * args.fps)) + 1231 num_frames = ((raw_frames - 1 + 4) // 8) * 8 + 1232 pixel_shape = VideoPixelShape(batch=1, frames=num_frames, height=64, width=64, fps=args.fps)233 tgt_shape = AudioLatentShape.from_video_pixel_shape(pixel_shape)234 audio_tools = AudioLatentTools(patchifier=self.patchifier, target_shape=tgt_shape)235 236 state = audio_tools.create_initial_state(self.device, self.dtype)237 238 # ---- Voice reference ----239 if voice_ref and os.path.exists(voice_ref):240 voice = decode_audio_from_file(voice_ref, self.device, 0.0, 10.0)241 if voice is not None:242 w = voice.waveform243 if w.dim() == 2:244 if w.shape[0] == 1:245 w = w.repeat(2, 1)246 w = w.unsqueeze(0)247 elif w.dim() == 3 and w.shape[1] == 1:248 w = w.repeat(1, 2, 1)249 target_samples = int(10.0 * voice.sampling_rate)250 if w.shape[-1] < target_samples:251 w = w.repeat(1, 1, (target_samples // w.shape[-1]) + 1)252 w = w[..., :target_samples]253 peak = w.abs().max()254 if peak > 0:255 w = w * (10 ** (-4.0 / 20) / peak)256 voice = Audio(waveform=w, sampling_rate=voice.sampling_rate)257 ref_latent = self.audio_conditioner(lambda enc: vae_encode_audio(voice, enc, None))258 cond = AudioConditionByReferenceLatent(259 latent=ref_latent.to(self.device, self.dtype), strength=1.0,260 )261 state = cond.apply_to(latent_state=state, latent_tools=audio_tools)262 263 # ---- Noise ----264 gen = torch.Generator(device=self.device).manual_seed(args.seed)265 noiser = GaussianNoiser(generator=gen)266 state = noiser(state, noise_scale=1.0)267 268 # ---- Prompt encode ----269 use_cfg = args.cfg_scale > 1.0270 prompts = [prompt, args.negative_prompt] if use_cfg else [prompt]271 ctx = self.prompt_encoder(prompts, streaming_prefetch_count=None)272 a_ctx = ctx[0].audio_encoding273 a_ctx_neg = ctx[1].audio_encoding if use_cfg else None274 275 # ---- Denoiser ----276 needs_guidance = args.cfg_scale > 1.0 or args.stg_scale > 0.0 or args.modality_scale > 1.0277 if needs_guidance:278 guider = MultiModalGuider(279 params=MultiModalGuiderParams(280 cfg_scale=args.cfg_scale, stg_scale=args.stg_scale,281 stg_blocks=[args.stg_block] if args.stg_scale > 0 else [],282 rescale_scale=args.rescale_scale,283 modality_scale=args.modality_scale,284 cfg_clamp_scale=args.cfg_clamp,285 ),286 negative_context=a_ctx_neg,287 )288 denoiser = GuidedDenoiser(289 v_context=None, a_context=a_ctx,290 video_guider=None, audio_guider=guider,291 )292 else:293 denoiser = SimpleDenoiser(v_context=None, a_context=a_ctx)294 295 sigmas = LTX2Scheduler().execute(steps=args.steps, latent=state.latent).to(self.device)296 297 # ---- Denoise ----298 # NOTE: don't wrap in gpu_model() — that context manager moves the299 # model back off GPU on exit, which breaks subsequent iterations of300 # our warm validator. We keep the velocity model resident.301 x0 = X0Model(self.velocity_model)302 batched = BatchSplitAdapter(x0, max_batch_size=1)303 _, audio_state = euler_denoising_loop(304 sigmas=sigmas, video_state=None, audio_state=state,305 stepper=EulerDiffusionStep(), transformer=batched, denoiser=denoiser,306 )307 308 audio_state = audio_tools.clear_conditioning(audio_state)309 audio_state = audio_tools.unpatchify(audio_state)310 decoded = self.audio_decoder(audio_state.latent)311 312 wav = decoded.waveform313 if wav.dim() == 1:314 wav = wav.unsqueeze(0)315 os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)316 torchaudio.save(output_path, wav.float().cpu(), decoded.sampling_rate)317 logging.info(f" -> {output_path} ({wav.shape[-1]/decoded.sampling_rate:.1f}s, "318 f"{time.time()-t_total:.1f}s)")319 320 321def main():322 logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")323 args = parse_args()324 import yaml325 with open(args.val_config) as f:326 val_cfg = yaml.safe_load(f)327 os.makedirs(args.output_dir, exist_ok=True)328 329 # Build validator once (models warm for all entries).330 validator = WarmValidator(331 full_checkpoint=args.full_checkpoint,332 gemma_root=args.gemma_root,333 lora_path=args.lora,334 lora_rank=args.lora_rank,335 device="cuda" if torch.cuda.is_available() else "cpu",336 dtype=torch.bfloat16,337 )338 339 n_ok = n_fail = 0340 t0 = time.time()341 for entry in val_cfg.get("speakers", []):342 name = entry["name"]343 out_path = os.path.join(args.output_dir, f"{name}.wav")344 try:345 validator.generate(346 prompt=entry["prompt"],347 output_path=out_path,348 voice_ref=entry.get("reference"),349 args=args,350 )351 n_ok += 1352 logging.info(f" [{name}] OK")353 except Exception as e:354 n_fail += 1355 logging.warning(f" [{name}] FAILED: {e}")356 traceback.print_exc()357 358 logging.info(f"Validation done: ok={n_ok} fail={n_fail} in {(time.time()-t0)/60:.1f}min "359 f"at {args.output_dir}")360 361 362if __name__ == "__main__":363 main()364