akhaliq/Dramabox

1#!/usr/bin/env python32"""3LTX-2.3 TTS with IC-LoRA voice cloning.4 5Uses AudioConditionByReferenceLatent to append reference audio tokens to the6end of the target sequence.  Auto-detects distilled vs dev checkpoint and7selects the appropriate denoiser (SimpleDenoiser / GuidedDenoiser) and sigma8schedule.  Leverages the official euler_denoising_loop, AudioLatentTools,9GaussianNoiser, and X0Model wrapper throughout.10 11Usage (distilled):12    python tts_iclora.py \13        --voice-sample reference.wav \14        --prompt "A woman speaks clearly: The weather today will be sunny." \15        --output tts_output.wav16 17Usage (dev):18    python tts_iclora.py \19        --voice-sample reference.wav \20        --prompt "A woman speaks clearly: The weather today will be sunny." \21        --checkpoint ltx-2.3-22b-dev-audio-only.safetensors \22        --full-checkpoint ltx-2.3-22b-dev.safetensors \23        --output tts_output.wav24"""25 26import argparse27import json28import logging29import os30import re31import struct32import sys33import time34from pathlib import Path35 36import torch37import torchaudio38 39REPO_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))40sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "ltx2"))41# ltx-pipelines already on path via ltx2/42 43# Also add the local directory so audio_conditioning.py is importable44sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))45 46MODEL_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "models")47GEMMA_DIR = os.environ.get("GEMMA_DIR", "gemma-3-12b-it-qat-q4_0-unquantized")48 49 50# ---------------------------------------------------------------------------51# Helpers52# ---------------------------------------------------------------------------53 54 55def detect_model_type(checkpoint_path: str) -> str:56    """Detect if checkpoint is distilled or dev by checking filename and metadata."""57    path_lower = checkpoint_path.lower()58    if "distilled" in path_lower:59        return "distilled"60    if "dev" in path_lower:61        return "dev"62    # Fallback: try to read safetensors metadata63    try:64        with open(checkpoint_path, "rb") as f:65            header_size = struct.unpack("<Q", f.read(8))[0]66            header = json.loads(f.read(header_size).decode())67        metadata = header.get("__metadata__", {})68        version = metadata.get("model_version", "")69        if "distilled" in version.lower():70            return "distilled"71    except Exception:72        pass73    # Default to distilled (most common for audio-only)74    return "distilled"75 76 77_LAUGH_VERBS = {78    # base seconds per occurrence; gets scaled by the modifier found nearby.79    # Verb regex covers inflections: laugh/laughs/laughed/laughing.80    r"\blaugh(?:s|ed|ing)?\b": 1.5,81    r"\bcackl(?:e|es|ed|ing)\b": 1.5,82    r"\bchuckl(?:e|es|ed|ing)\b": 1.0,83    r"\bgiggl(?:e|es|ed|ing)\b": 1.0,84    r"\bsnicker(?:s|ed|ing)?\b": 0.8,85    r"\bcru?el laugh\b": 1.5,86}87 88 89def _contextual_laugh_duration(text: str) -> float:90    """Context-aware laugh budget.91 92    For each laugh verb in the prompt, look at the adjective/adverb that93    modifies it and scale the base duration:94      - short modifiers  (briefly, softly, once)     -> 0.4x base95      - long modifiers   (maniacally, heartily, ...) -> 1.2x base96      - default (no mod / neutral)                   -> 1.0x base97    Also reward phonetic repetition inside quotes -- 'Hahahahahaha' buys more98    time than 'Haha' -- at ~0.2s per extra repeated syllable.99    """100    # "softly" / "quietly" describe volume not length, so keep at default 1.0x.101    short_mod = re.compile(102        r"^\s*(?:[a-z]+ly )?(?:briefly|shortly|once|quickly)",103        re.IGNORECASE)104    long_mod = re.compile(105        r"^\s*(?:[a-z]+ly )?(?:maniacally|heartily|uproariously|uncontrollably|"106        r"hysterically|darkly|wickedly|evilly|loudly|long)"107        r"|^\s*between phrases", re.IGNORECASE)108 109    total = 0.0110    for pat, base_dur in _LAUGH_VERBS.items():111        for m in re.finditer(pat, text, re.IGNORECASE):112            ctx = text[m.end(): m.end() + 40]113            if short_mod.match(ctx):114                total += base_dur * 0.4115            elif long_mod.match(ctx):116                total += base_dur * 1.2117            else:118                total += base_dur119 120    # Phonetic laugh repetition inside quotes:121    #   'Haha' = 2 syllables (base, no bonus)122    #   'Hahahaha' = 4 syllables (+0.4s)123    #   'Hehehehahahahahahahaha' ~ 10 syllables (+1.6s)124    for q in re.findall(r'"([^"]+)"', text) + re.findall(r"'((?:[^']|'(?![\s.,!?)\]]))+)'", text):125        for run in re.findall(r"(?:h[ae]){3,}|(?:h[ae][ \-]?){3,}", q, re.IGNORECASE):126            syls = len(re.findall(r"h[ae]", run, re.IGNORECASE))127            total += 0.2 * max(syls - 2, 0)128    return total129 130 131def _estimate_nonverbal_duration(text: str) -> float:132    """Estimate extra duration for non-verbal sounds and actions in the prompt.133 134    Laugh-verb handling lives in ``_contextual_laugh_duration`` so cackle /135    chuckle / laugh budgets scale with the adjective ("maniacally" vs136    "briefly") and with the repetition length of 'Ha'/'He' tokens inside137    quotes.138    """139    PATTERNS = {140        # Breathing / sighs141        r'\bsighs?\b': 0.8, r'\bshaky breath\b': 1.0, r'\bbreathing deeply\b': 1.0,142        r'\bgasps?\b': 0.5, r'\bburps?\b': 0.5, r'\byawns?\b': 1.0,143        r'\bpants?\b': 0.8, r'\bwheezes?\b': 0.8, r'\bcoughs?\b': 0.8,144        r'\bsniffles?\b': 0.5, r'\bsnorts?\b': 0.3, r'\bgroans?\b': 0.8,145        # Pauses (trimmed; earlier values over-budgeted silence)146        r'\blong pause\b': 1.0, r'\bpauses? briefly\b': 0.3,147        r'\bpauses?\b': 0.5, r'\bsilence\b': 1.0,148        r'\blets? the .{1,20} hang\b': 1.0, r'\blets? .{1,20} sink in\b': 1.0,149        # Physical actions that produce sound150        r'\bslams?\b': 0.5, r'\bclaps?\b': 0.3,151        r'\bdraws? (?:his|her|a) sword\b': 0.5,152        r'\btakes? a (?:drag|swig|sip|drink)\b': 0.5,153        r'\bwhistles?\b': 1.0, r'\bhums?\b': 0.8,154        # Vocal actions (not in quotes but take time)155        r'\bmutters?\b': 1.5, r'\bmumbles?\b': 1.0, r'\bwhispers?\b': 0.0,156        r'\bclears? (?:his|her) throat\b': 0.5, r'\bgulps?\b': 0.5,157        r'\bswallows?\b': 0.5,158        # (laugh / chuckle / cackle / giggle / snicker handled by159        # _contextual_laugh_duration below -- modifier-aware, not flat.)160        # Emotional transitions161        r'\bvoice (?:breaks?|cracks?|trembles?|drops?|rises?)\b': 0.5,162        r'\bsteadies? (?:him|her)self\b': 1.0,163        r'\bcatches? (?:his|her) breath\b': 1.0,164        r'\bcomposes? (?:him|her)self\b': 0.8,165        # Scene transitions that imply time166        r'\bdemeanor shifts?\b': 0.5, r'\bsettles? in\b': 0.5,167        r'\bleans? in\b': 0.3, r'\bwipes? (?:his|her) eyes\b': 0.5,168    }169    extra = 0.0170    for pattern, dur in PATTERNS.items():171        extra += dur * len(re.findall(pattern, text, re.IGNORECASE))172    extra += _contextual_laugh_duration(text)173    return extra174 175 176def estimate_speech_duration(text: str, speed: float = 1.0) -> float:177    """Estimate speech duration from spoken content + non-verbal actions.178 179    Extracts spoken text by priority:180    1. Quoted text ('...' or "...") -- official prompt guide format181    2. Text after colon -- simple "Speaker: dialogue" format182    3. Full text -- fallback183 184    Also scans the full prompt for non-verbal cues (laughs, pauses, sighs,185    gasps, etc.) and adds estimated duration for each.186    """187    # Try double quotes first (clean, no contraction issues)188    quotes = re.findall(r'"([^"]+)"', text)189    if not quotes:190        # Single quotes: allow apostrophes in contractions (don't, can't, it's)191        # Match ' to ' but apostrophes NOT followed by space/punctuation are kept inside192        quotes = re.findall(r"'((?:[^']|'(?![\s.,!?)\]]))+)'", text)193        # Filter out short fragments (scene directions like "He pauses")194        quotes = [q for q in quotes if len(q.split()) > 3]195    if quotes:196        spoken = " ".join(quotes)197    elif ":" in text:198        spoken = text.split(":", 1)[1].strip()199    else:200        spoken = text201 202    CHARS_PER_SEC = 14.0203    text_len = len(spoken)204 205    if text_len < 40:206        chars_per_sec = CHARS_PER_SEC * 0.6207    elif text_len < 80:208        chars_per_sec = CHARS_PER_SEC * 0.8209    else:210        chars_per_sec = CHARS_PER_SEC211 212    chars_per_sec *= speed213    duration = text_len / chars_per_sec214 215    sentence_count = spoken.count(".") + spoken.count("!") + spoken.count("?")216    duration += sentence_count * 0.3217 218    # Add time for non-verbal sounds/actions in the full prompt219    duration += _estimate_nonverbal_duration(text)220 221    return max(3.0, round(duration + 2.0, 1))222 223 224def parse_args():225    p = argparse.ArgumentParser(description="LTX-2.3 TTS with IC-LoRA voice cloning")226 227    p.add_argument("--voice-sample", default=None, help="Voice reference WAV")228    p.add_argument("--no-ref", action="store_true", help="Skip voice reference conditioning (raw base model)")229    p.add_argument("--prompt", required=True, help="Text/scene description to synthesize")230    p.add_argument("--output", default="tts_output.wav")231 232    p.add_argument("--ref-duration", type=float, default=10.0, help="Seconds of voice reference to use")233    p.add_argument("--gen-duration", type=float, default=0.0,234                   help="Target output duration in seconds (0 = auto from prompt + multiplier). "235                        "Set explicitly for long-form prompts (e.g. --gen-duration 30 for music). "236                        "Outputs >20.5s automatically engage the end-of-clip silence-prior patch.")237    p.add_argument("--pad-start", type=float, default=0.0,238                   help="Prepend N seconds of silent padding, trimmed after decode (use 0 for clean starts)")239    p.add_argument("--speed", type=float, default=1.0)240    p.add_argument("--duration-multiplier", type=float, default=1.0,241                   help="Multiply auto-estimated duration by this factor (e.g. 1.1 for 10%% more breathing room)")242 243    p.add_argument("--checkpoint", default=os.path.join(MODEL_DIR, "ltx-2.3-audio-only.safetensors"))244    p.add_argument("--full-checkpoint", default=os.path.join(MODEL_DIR, "ltx-2.3-22b-distilled.safetensors"))245    p.add_argument("--gemma-root", default=GEMMA_DIR)246    p.add_argument("--bnb-4bit", dest="bnb_4bit", action="store_true", default=True,247                   help="Load Gemma text encoder via the bitsandbytes 4-bit path "248                        "(required for the default unsloth/gemma-3-12b-it-bnb-4bit "249                        "pre-quantized weights). Default: on.")250    p.add_argument("--no-bnb-4bit", dest="bnb_4bit", action="store_false",251                   help="Disable the bitsandbytes path (use only if --gemma-root "252                        "points at an unquantized Gemma checkpoint).")253    p.add_argument("--lora", default=None, help="Path to trained IC-LoRA .safetensors (audio-only)")254    p.add_argument("--lora-rank", type=int, default=128, help="LoRA rank (must match training)")255    p.add_argument("--id-guidance-scale", type=float, default=3.0, help="Identity guidance scale (0=disabled)")256    p.add_argument("--seed", type=int, default=42)257 258    # Auto-set based on model type but overridable259    p.add_argument("--no-watermark", action="store_true",260                   help="Skip Perth audio watermarking on the output (default: watermark on).")261    p.add_argument("--sampler", choices=["euler", "heun"], default="euler",262                   help="Denoising loop. 'heun' = jkass_quality 2nd-order predictor-corrector (~2x model calls, cleaner audio).")263    p.add_argument("--cfg-scale", type=float, default=None, help="CFG scale (auto: 1.0 distilled, 7.0 dev)")264    p.add_argument("--stg-scale", type=float, default=None, help="STG scale (auto: 0.0 distilled, 1.0 dev)")265    p.add_argument("--stg-block", type=int, default=29, help="Block index for STG perturbation")266    p.add_argument("--rescale-scale", type=float, default=None,267                   help="Latent CFG std-rescale (default auto: cfg-aware schedule that prevents "268                        "output clipping at high cfg; pass any float in [0,1] to override).")269    p.add_argument("--modality-scale", type=float, default=None, help="Modality (auto: 1.0 distilled, 3.0 dev)")270    p.add_argument("--cfg-clamp", type=float, default=0.0, help="Clamp guided pred std to N * cond std (0=disabled)")271    p.add_argument("--steps", type=int, default=None, help="Override steps (auto: distilled sigmas / 30 dev)")272    p.add_argument("--fps", type=float, default=None, help="FPS (auto: 24.0 distilled, 25.0 dev)")273    p.add_argument(274        "--negative-prompt",275        default=(276            "worst quality, inconsistent motion, blurry, jittery, distorted, "277            "robotic voice, echo, background noise, off-sync audio, repetitive speech"278        ),279        help="Negative prompt for CFG (dev model)",280    )281 282    return p.parse_args()283 284 285@torch.inference_mode()286def main():287    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")288    args = parse_args()289    t0 = time.time()290 291    # ---- Imports (deferred to avoid startup cost when checking --help) ----292    from audio_conditioning import AudioConditionByReferenceLatent293 294    from ltx_core.batch_split import BatchSplitAdapter295    from ltx_core.components.diffusion_steps import EulerDiffusionStep296    from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderParams297    from ltx_core.components.noisers import GaussianNoiser298    from ltx_core.components.patchifiers import AudioPatchifier299    from ltx_core.components.schedulers import LTX2Scheduler300    from ltx_core.loader.registry import DummyRegistry301    from ltx_core.loader.sd_ops import SDOps302    from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder303    from ltx_core.model.audio_vae import encode_audio as vae_encode_audio304    from ltx_core.model.model_protocol import ModelConfigurator305    from ltx_core.model.transformer.attention import AttentionFunction306    from ltx_core.model.transformer.model import LTXModel, LTXModelType, X0Model307    from ltx_core.model.transformer.rope import LTXRopeType308    from ltx_core.tools import AudioLatentTools309    from ltx_core.types import Audio, AudioLatentShape, LatentState, VideoPixelShape310    from ltx_pipelines.utils.blocks import AudioConditioner, AudioDecoder, PromptEncoder311    from ltx_pipelines.utils.constants import DISTILLED_SIGMA_VALUES312    from ltx_pipelines.utils.denoisers import GuidedDenoiser, SimpleDenoiser313    from ltx_pipelines.utils.gpu_model import gpu_model314    from ltx_pipelines.utils.media_io import decode_audio_from_file315    from ltx_pipelines.utils.samplers import euler_denoising_loop, heun_denoising_loop316 317    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")318    dtype = torch.bfloat16319    patchifier = AudioPatchifier(patch_size=1)320 321    # ---- Detect model type and set defaults ----322    model_type = detect_model_type(args.full_checkpoint)323    logging.info(f"Detected model type: {model_type}")324 325    is_distilled = model_type == "distilled"326 327    if args.cfg_scale is None:328        args.cfg_scale = 1.0 if is_distilled else 7.0329    if args.stg_scale is None:330        args.stg_scale = 0.0 if is_distilled else 1.0331    if args.rescale_scale is None:332        # Auto cfg-aware rescale: imported from inference_server to keep one source of truth.333        from inference_server import auto_rescale_for_cfg334        args.rescale_scale = 0.0 if is_distilled else auto_rescale_for_cfg(args.cfg_scale)335    if args.modality_scale is None:336        args.modality_scale = 1.0 if is_distilled else 3.0337    if args.fps is None:338        args.fps = 24.0 if is_distilled else 25.0339 340    logging.info(341        f"Params: cfg={args.cfg_scale}, stg={args.stg_scale}, rescale={args.rescale_scale}, "342        f"modality={args.modality_scale}, fps={args.fps}"343    )344 345    # ---- Auto duration ----346    if args.gen_duration <= 0:347        args.gen_duration = estimate_speech_duration(args.prompt, args.speed)348        if args.duration_multiplier != 1.0:349            args.gen_duration = round(args.gen_duration * args.duration_multiplier, 1)350        logging.info(f"Auto duration: {args.gen_duration}s for {len(args.prompt)} chars"351                     f"{f' (x{args.duration_multiplier})' if args.duration_multiplier != 1.0 else ''}")352 353    # ---- Compute target shape (include pad_start in duration) ----354    padded_duration = args.gen_duration + args.pad_start355    raw_frames = int(round(padded_duration * args.fps)) + 1356    num_frames = ((raw_frames - 1 + 4) // 8) * 8 + 1357    pixel_shape = VideoPixelShape(batch=1, frames=num_frames, height=64, width=64, fps=args.fps)358    tgt_shape = AudioLatentShape.from_video_pixel_shape(pixel_shape)359    logging.info(f"Target shape: {tgt_shape} ({args.gen_duration}s, {num_frames} frames)")360 361    # ---- AudioLatentTools for target ----362    audio_tools = AudioLatentTools(patchifier=patchifier, target_shape=tgt_shape)363 364    # ---- Create initial state ----365    state = audio_tools.create_initial_state(device, dtype)366    logging.info(367        f"Initial state: latent={state.latent.shape}, positions={state.positions.shape}, "368        f"denoise_mask={state.denoise_mask.shape}"369    )370 371    if not args.no_ref and args.voice_sample:372        # ---- Encode voice reference ----373        logging.info(f"Loading voice reference: {args.voice_sample}")374        voice = decode_audio_from_file(args.voice_sample, device, 0.0, args.ref_duration)375        if voice is None:376            raise ValueError(f"Could not load audio from {args.voice_sample}")377 378        w = voice.waveform379        if w.dim() == 2:380            if w.shape[0] == 1:381                w = w.repeat(2, 1)382            w = w.unsqueeze(0)383        elif w.dim() == 3 and w.shape[1] == 1:384            w = w.repeat(1, 2, 1)385 386        target_samples = int(args.ref_duration * voice.sampling_rate)387        if w.shape[-1] < target_samples:388            w = w.repeat(1, 1, (target_samples // w.shape[-1]) + 1)389        w = w[..., :target_samples]390 391        # Peak normalize reference392        peak = w.abs().max()393        if peak > 0:394            target_peak = 10 ** (-4.0 / 20)  # -4dB395            w = w * (target_peak / peak)396            logging.info(f"Normalized reference: peak {peak:.4f} -> {target_peak:.4f}")397 398        voice = Audio(waveform=w, sampling_rate=voice.sampling_rate)399 400        logging.info("Encoding voice through Audio VAE...")401        ac = AudioConditioner(checkpoint_path=args.full_checkpoint, dtype=dtype, device=device)402        ref_latent = ac(lambda enc: vae_encode_audio(voice, enc, None))403        del ac404        torch.cuda.empty_cache()405        logging.info(f"Reference latent: {ref_latent.shape}")406 407        # ---- Apply conditioning: append ref tokens to END ----408        conditioning = AudioConditionByReferenceLatent(latent=ref_latent.to(device, dtype), strength=1.0)409        state = conditioning.apply_to(latent_state=state, latent_tools=audio_tools)410        logging.info(411            f"After conditioning: latent={state.latent.shape}, positions={state.positions.shape}, "412            f"attention_mask={'None' if state.attention_mask is None else state.attention_mask.shape}"413        )414    else:415        logging.info("No voice reference — running raw base model")416 417    # ---- Apply noise ----418    generator = torch.Generator(device=device).manual_seed(args.seed)419    noiser = GaussianNoiser(generator=generator)420    noised_state = noiser(state, noise_scale=1.0)421    logging.info("Applied Gaussian noise to state")422 423    # ---- Encode prompt ----424    use_cfg = args.cfg_scale > 1.0425    logging.info("Encoding prompt...")426    pe = PromptEncoder(checkpoint_path=args.full_checkpoint, gemma_root=args.gemma_root, dtype=dtype, device=device,427                       use_bnb_4bit=args.bnb_4bit, warm=True)428    prompts_to_encode = [args.prompt]429    if use_cfg:430        prompts_to_encode.append(args.negative_prompt)431    ctx = pe(prompts_to_encode, streaming_prefetch_count=None)432    a_ctx = ctx[0].audio_encoding433    a_ctx_neg = ctx[1].audio_encoding if use_cfg else None434    del pe435    torch.cuda.empty_cache()436    logging.info(f"Prompt encoded: a_ctx={a_ctx.shape}" + (f", a_ctx_neg={a_ctx_neg.shape}" if a_ctx_neg is not None else ""))437 438    # ---- Build audio-only model ----439    logging.info("Building audio-only model...")440    audio_only_sd_ops = SDOps("AO").with_matching(prefix="model.diffusion_model.").with_replacement(441        "model.diffusion_model.", ""442    )443 444    class AudioOnlyConfigurator(ModelConfigurator[LTXModel]):445        @classmethod446        def from_config(cls, config):447            t = config.get("transformer", {})448            cp = None449            if not t.get("caption_proj_before_connector", False):450                from ltx_core.model.transformer.text_projection import create_caption_projection451 452                with torch.device("meta"):453                    cp = create_caption_projection(t, audio=True)454            return LTXModel(455                model_type=LTXModelType.AudioOnly,456                audio_num_attention_heads=t.get("audio_num_attention_heads", 32),457                audio_attention_head_dim=t.get("audio_attention_head_dim", 64),458                audio_in_channels=t.get("audio_in_channels", 128),459                audio_out_channels=t.get("audio_out_channels", 128),460                num_layers=t.get("num_layers", 48),461                audio_cross_attention_dim=t.get("audio_cross_attention_dim", 2048),462                norm_eps=t.get("norm_eps", 1e-6),463                attention_type=AttentionFunction(t.get("attention_type", "default")),464                positional_embedding_theta=10000.0,465                audio_positional_embedding_max_pos=[20.0],466                timestep_scale_multiplier=t.get("timestep_scale_multiplier", 1000),467                use_middle_indices_grid=t.get("use_middle_indices_grid", True),468                rope_type=LTXRopeType(t.get("rope_type", "interleaved")),469                double_precision_rope=t.get("frequencies_precision", False) == "float64",470                apply_gated_attention=t.get("apply_gated_attention", False),471                audio_caption_projection=cp,472                cross_attention_adaln=t.get("cross_attention_adaln", False),473            )474 475    builder = Builder(476        model_path=args.checkpoint,477        model_class_configurator=AudioOnlyConfigurator,478        model_sd_ops=audio_only_sd_ops,479        registry=DummyRegistry(),480    )481    velocity_model = builder.build(device=device, dtype=dtype).to(device).eval()482 483    # ---- Load LoRA weights (if provided) ----484    if args.lora and os.path.exists(args.lora):485        from peft import LoraConfig, get_peft_model486        from safetensors.torch import load_file as st_load487 488        logging.info(f"Loading LoRA: {args.lora}")489        lora_sd = st_load(args.lora)490 491        is_peft_format = any("base_model.model." in k for k in lora_sd.keys())492        is_original_idlora = any("diffusion_model." in k for k in lora_sd.keys())493 494        lora_config = LoraConfig(495            r=args.lora_rank,496            lora_alpha=args.lora_rank,497            lora_dropout=0.0,498            bias="none",499            target_modules=[500                "audio_attn1.to_k",501                "audio_attn1.to_q",502                "audio_attn1.to_v",503                "audio_attn1.to_out.0",504                "audio_attn2.to_k",505                "audio_attn2.to_q",506                "audio_attn2.to_v",507                "audio_attn2.to_out.0",508                "audio_ff.net.0.proj",509                "audio_ff.net.2",510            ],511        )512        velocity_model = get_peft_model(velocity_model, lora_config)513 514        if is_peft_format:515            mapped_sd = {}516            for k, v in lora_sd.items():517                new_key = k518                if ".lora_A.weight" in k and ".lora_A.default.weight" not in k:519                    new_key = k.replace(".lora_A.weight", ".lora_A.default.weight")520                if ".lora_B.weight" in k and ".lora_B.default.weight" not in k:521                    new_key = k.replace(".lora_B.weight", ".lora_B.default.weight")522                mapped_sd[new_key] = v523            missing, unexpected = velocity_model.load_state_dict(mapped_sd, strict=False)524            loaded = len(mapped_sd) - len(unexpected)525            logging.info(f"Loaded {loaded} LoRA weights (peft format)")526        elif is_original_idlora:527            audio_keys = {528                k: v529                for k, v in lora_sd.items()530                if "audio_attn1" in k or "audio_attn2" in k or "audio_ff" in k531            }532            mapped_sd = {}533            for k, v in audio_keys.items():534                new_key = k.replace("diffusion_model.", "base_model.model.")535                new_key = new_key.replace(".lora_A.weight", ".lora_A.default.weight")536                new_key = new_key.replace(".lora_B.weight", ".lora_B.default.weight")537                mapped_sd[new_key] = v538            missing, unexpected = velocity_model.load_state_dict(mapped_sd, strict=False)539            loaded = len(mapped_sd) - len(unexpected)540            logging.info(f"Loaded {loaded} LoRA weights (original ID-LoRA)")541 542        velocity_model = velocity_model.merge_and_unload()543        logging.info("Merged LoRA into model")544 545    logging.info(f"Model: {sum(p.numel() for p in velocity_model.parameters()) / 1e9:.1f}B params")546 547    # ---- Wrap velocity model in X0Model ----548    x0_model = X0Model(velocity_model)549 550    # ---- Build denoiser and sigmas ----551    stepper = EulerDiffusionStep()552 553    # ---- Sigma schedule ----554    if is_distilled:555        if args.steps is not None and args.steps > 0:556            sigmas = LTX2Scheduler().execute(steps=args.steps, latent=noised_state.latent).to(device)557            logging.info(f"Distilled with custom {args.steps}-step schedule")558        else:559            sigmas = torch.tensor(DISTILLED_SIGMA_VALUES, dtype=torch.float32, device=device)560            logging.info(f"Distilled {len(DISTILLED_SIGMA_VALUES) - 1}-step schedule")561    else:562        steps = args.steps if args.steps is not None and args.steps > 0 else 30563        sigmas = LTX2Scheduler().execute(steps=steps, latent=noised_state.latent).to(device)564        logging.info(f"Dev {steps}-step schedule")565 566    # ---- Denoiser: use GuidedDenoiser if any guidance is active, SimpleDenoiser otherwise ----567    needs_guidance = args.cfg_scale > 1.0 or args.stg_scale > 0.0 or args.modality_scale > 1.0568    if needs_guidance:569        audio_guider = MultiModalGuider(570            params=MultiModalGuiderParams(571                cfg_scale=args.cfg_scale,572                stg_scale=args.stg_scale,573                stg_blocks=[args.stg_block] if args.stg_scale > 0 else [],574                rescale_scale=args.rescale_scale,575                modality_scale=args.modality_scale,576                cfg_clamp_scale=args.cfg_clamp,577            ),578            negative_context=a_ctx_neg,579        )580        denoiser = GuidedDenoiser(581            v_context=None,582            a_context=a_ctx,583            video_guider=None,584            audio_guider=audio_guider,585        )586        logging.info(f"GuidedDenoiser: cfg={args.cfg_scale}, stg={args.stg_scale}, "587                     f"rescale={args.rescale_scale}, modality={args.modality_scale}")588    else:589        denoiser = SimpleDenoiser(v_context=None, a_context=a_ctx)590        logging.info("SimpleDenoiser (no guidance)")591 592    logging.info(f"Sigmas: {sigmas.tolist()}")593 594    # ---- Denoising loop ----595    logging.info(f"Running denoising loop ({len(sigmas) - 1} steps)...")596    with gpu_model(x0_model) as model:597        batched_model = BatchSplitAdapter(model, max_batch_size=1)598 599        denoise_fn = heun_denoising_loop if args.sampler == "heun" else euler_denoising_loop600        _, audio_state = denoise_fn(601            sigmas=sigmas,602            video_state=None,603            audio_state=noised_state,604            stepper=stepper,605            transformer=batched_model,606            denoiser=denoiser,607        )608 609    del velocity_model, x0_model610    torch.cuda.empty_cache()611 612    # ---- Strip ref tokens and unpatchify ----613    logging.info("Stripping conditioning and unpatchifying...")614    audio_state = audio_tools.clear_conditioning(audio_state)615    audio_state = audio_tools.unpatchify(audio_state)616    logging.info(f"Final latent shape: {audio_state.latent.shape}")617 618    # ---- End-of-clip silence-prior fix ----619    # Base LTX-2.3 22B was trained on audio clips ≤ ~20 s and learned a strong620    # "clip-end silence" prior at the next patchifier-aligned latent boundary621    # (frame 513 = 8 × 64 + 1). For longer outputs that prior leaks through as622    # a ~30 ms hard silence dip near 20.4 s. Linearly interpolating frames623    # 512–513 between their neighbours (511 and 514) removes the dip cleanly.624    latent_in = audio_state.latent625    if latent_in.shape[2] > 513:626        f0, f1 = 511, 514627        n = f1 - f0628        patched = latent_in.clone()629        for f in (512, 513):630            t = (f - f0) / n631            patched[:, :, f, :] = (1.0 - t) * latent_in[:, :, f0, :] + t * latent_in[:, :, f1, :]632        latent_in = patched633 634    # ---- Decode audio ----635    logging.info("Decoding audio...")636    ad = AudioDecoder(checkpoint_path=args.full_checkpoint, dtype=dtype, device=device)637    decoded = ad(latent_in)638    del ad639    torch.cuda.empty_cache()640 641    wav = decoded.waveform642    if wav.dim() == 1:643        wav = wav.unsqueeze(0)644    sr = decoded.sampling_rate645 646    # Trim leading pad if --pad-start was used647    if args.pad_start > 0:648        trim_samples = int(args.pad_start * sr)649        wav = wav[..., trim_samples:]650        logging.info(f"Trimmed {args.pad_start}s ({trim_samples} samples) of start padding")651 652    # Apply Perth (Perceptual Threshold) imperceptible neural watermark — see653    # https://github.com/resemble-ai/perth. Mono waveform required; if stereo,654    # we average to mono for the watermark and broadcast back. Skip on655    # --no-watermark for debugging.656    wav_cpu = wav.float().cpu()657    if not getattr(args, "no_watermark", False):658        try:659            import perth660            import numpy as np661            wm = perth.PerthImplicitWatermarker()662            mono = wav_cpu.mean(dim=0).numpy() if wav_cpu.shape[0] > 1 else wav_cpu[0].numpy()663            mono_wm = wm.apply_watermark(mono, sample_rate=sr)664            mono_wm_t = torch.from_numpy(np.asarray(mono_wm, dtype=np.float32)).unsqueeze(0)665            wav_cpu = mono_wm_t if wav_cpu.shape[0] == 1 else mono_wm_t.repeat(wav_cpu.shape[0], 1)666        except Exception as e:667            logging.warning(f"Perth watermark skipped ({e})")668 669    os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)670    torchaudio.save(args.output, wav_cpu, sr)671 672    elapsed = time.time() - t0673    logging.info(f"Output: {args.output} ({wav.shape[-1] / sr:.1f}s)")674    logging.info(f"Total time: {elapsed:.1f}s")675 676 677if __name__ == "__main__":678    main()679