akhaliq/Dramabox

1#!/usr/bin/env python32"""3Warm TTS server — loads models once, accepts requests via stdin or function call.4 5The key insight: inference.py spends 11s on Gemma + 8s on model load every call.6This server loads everything once and keeps it warm.7 8We import and call the same code paths as inference.py but cache the heavy objects.9"""10import json11import logging12import os13import re14import sys15import time16from pathlib import Path17 18import torch19import torchaudio20 21# Setup paths22APP_DIR = Path(__file__).parent.parent23sys.path.insert(0, str(APP_DIR / "ltx2"))24sys.path.insert(0, str(APP_DIR / "src"))25 26logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")27 28from audio_conditioning import AudioConditionByReferenceLatent29from ltx_core.components.noisers import GaussianNoiser30from ltx_core.components.patchifiers import AudioPatchifier31from ltx_core.components.guiders import MultiModalGuider, MultiModalGuiderParams32from ltx_core.components.schedulers import LTX2Scheduler33from ltx_core.components.diffusion_steps import EulerDiffusionStep34from ltx_core.loader import DummyRegistry35from ltx_core.loader.single_gpu_model_builder import SingleGPUModelBuilder as Builder36from ltx_core.loader.sd_ops import SDOps37from ltx_core.model.transformer.model import LTXModel, LTXModelType, X0Model38from ltx_core.model.transformer.rope import LTXRopeType39from ltx_core.model.transformer.text_projection import create_caption_projection40from ltx_core.model.transformer.attention import AttentionFunction41from ltx_core.model.model_protocol import ModelConfigurator42from ltx_core.tools import AudioLatentTools43from ltx_core.types import Audio, AudioLatentShape, VideoPixelShape44from ltx_core.model.audio_vae import encode_audio as vae_encode_audio45from ltx_pipelines.utils.blocks import AudioConditioner, AudioDecoder, PromptEncoder46from ltx_pipelines.utils.media_io import decode_audio_from_file47from ltx_pipelines.utils.denoisers import GuidedDenoiser48from ltx_pipelines.utils.samplers import euler_denoising_loop49from safetensors import safe_open50 51 52DEFAULT_NEG = "worst quality, inconsistent, robotic, distorted, noise, static, muffled, unclear, unnatural, monotone"53 54 55def estimate_duration(prompt, multiplier=1.1):56    """Defer to the richer CLI estimator (sentence-aware + non-verbal action57    budget) so warm-server outputs match the lengths of the per-call CLI runs."""58    from inference import estimate_speech_duration59    base = estimate_speech_duration(prompt)60    return max(3.0, round(base * multiplier, 1))61 62 63def auto_rescale_for_cfg(cfg: float) -> float:64    """CFG-aware std-rescale schedule that prevents output clipping at high cfg.65 66    The CFG formula `pred = cond + (cfg-1)*(cond - uncond)` makes pred.std()67    grow roughly linearly with cfg, which the audio VAE+vocoder render as68    progressively louder waveforms. By cfg≈3 the output starts hard-clipping69    at 0 dBFS — and clipped information is unrecoverable in post.70 71    Empirical sweep on the blues prompt with the back-porch-boogie ref72    (rescale_scale needed for ≥1 dB peak headroom):73        cfg=2.5 → 0.2 ;  cfg=3 → 0.6 ; cfg=4 → 0.8 ; cfg=5–8 → 0.8 ; cfg=10 → 1.074 75    Piecewise-linear fit through those points; returns 0 below cfg=2 (no CFG76    even applied at cfg=1), plateaus at 0.8 between cfg=4 and cfg=8 to77    preserve the "extra punch" of high-CFG generations, and ramps to 1.0 by78    cfg=10.79    """80    if cfg <= 2.0:81        return 0.082    if cfg <= 3.0:83        return 0.6 * (cfg - 2.0)               # 0 → 0.684    if cfg <= 4.0:85        return 0.6 + 0.2 * (cfg - 3.0)         # 0.6 → 0.886    if cfg <= 8.0:87        return 0.8                              # plateau88    return min(1.0, 0.8 + 0.1 * (cfg - 8.0))   # 0.8 → 1.0 at cfg=1089 90 91class TTSServer:92    def __init__(self, checkpoint=None, full_checkpoint=None, gemma_root=None,93                 device="cuda", dtype="bf16", compile_model=True, bnb_4bit=True):94        MODELS = APP_DIR / "models"95        self.checkpoint = checkpoint or str(MODELS / "ltx-2.3-22b-dev-audio-only-v13-merged.safetensors")96        self.full_checkpoint = full_checkpoint or os.environ.get(97            "LTX_FULL_CHECKPOINT", "/mnt/persistent0/manmay/models/ltx23/ltx-2.3-22b-dev.safetensors")98        if gemma_root is None and not os.environ.get("GEMMA_DIR"):99            from model_downloader import get_gemma_path100            gemma_root = get_gemma_path()101        self.gemma_root = gemma_root or os.environ["GEMMA_DIR"]102        self.device = torch.device(device)103        self.dtype = torch.float16 if dtype == "fp16" else torch.bfloat16104        self.compile_model = compile_model105        self.bnb_4bit = bnb_4bit106        self.patchifier = AudioPatchifier(patch_size=1)107 108        # Cached models109        self._prompt_encoder = None110        self._velocity_model = None111        self._audio_conditioner = None112        self._audio_decoder = None113 114        logging.info(f"TTSServer loading on {device}...")115        t0 = time.time()116        self._load_all()117        logging.info(f"All models loaded in {time.time()-t0:.1f}s — ready for requests")118 119    def _load_all(self):120        # 1. Prompt encoder (Gemma + embeddings processor kept warm)121        t0 = time.time()122        self._prompt_encoder = PromptEncoder(123            checkpoint_path=self.full_checkpoint,124            gemma_root=self.gemma_root,125            dtype=self.dtype, device=self.device,126            warm=True,127            use_bnb_4bit=self.bnb_4bit,128            audio_only=True,129        )130        logging.info(f"  PromptEncoder (warm): {time.time()-t0:.1f}s")131 132        # 2. Audio conditioner (VAE encoder kept warm)133        t0 = time.time()134        self._audio_conditioner = AudioConditioner(135            checkpoint_path=self.full_checkpoint,136            dtype=self.dtype, device=self.device,137            warm=True,138        )139        logging.info(f"  AudioConditioner (warm): {time.time()-t0:.1f}s")140 141        # 3. Transformer142        t0 = time.time()143        with safe_open(self.checkpoint, framework="pt") as f:144            config = json.loads(f.metadata()["config"])145 146        t = config.get("transformer", {})147 148        class AudioOnlyConfigurator(ModelConfigurator[LTXModel]):149            @classmethod150            def from_config(cls, cfg):151                t = cfg.get("transformer", {})152                cp = None153                if not t.get("caption_proj_before_connector", False):154                    with torch.device("meta"):155                        cp = create_caption_projection(t, audio=True)156                return LTXModel(157                    model_type=LTXModelType.AudioOnly,158                    audio_num_attention_heads=t.get("audio_num_attention_heads", 32),159                    audio_attention_head_dim=t.get("audio_attention_head_dim", 64),160                    audio_in_channels=t.get("audio_in_channels", 128),161                    audio_out_channels=t.get("audio_out_channels", 128),162                    num_layers=t.get("num_layers", 48),163                    audio_cross_attention_dim=t.get("audio_cross_attention_dim", 2048),164                    norm_eps=t.get("norm_eps", 1e-6),165                    attention_type=AttentionFunction(t.get("attention_type", "default")),166                    positional_embedding_theta=10000.0,167                    audio_positional_embedding_max_pos=[20.0],168                    timestep_scale_multiplier=t.get("timestep_scale_multiplier", 1000),169                    use_middle_indices_grid=t.get("use_middle_indices_grid", True),170                    rope_type=LTXRopeType(t.get("rope_type", "interleaved")),171                    double_precision_rope=t.get("frequencies_precision", False) == "float64",172                    apply_gated_attention=t.get("apply_gated_attention", False),173                    audio_caption_projection=cp,174                    cross_attention_adaln=t.get("cross_attention_adaln", False),175                )176 177        audio_sd_ops = SDOps("AO").with_matching(prefix="model.diffusion_model.").with_replacement(178            "model.diffusion_model.", "")179        builder = Builder(180            model_path=self.checkpoint,181            model_class_configurator=AudioOnlyConfigurator,182            model_sd_ops=audio_sd_ops,183            registry=DummyRegistry(),184        )185        self._velocity_model = builder.build(device=self.device, dtype=self.dtype).to(self.device).eval()186        n_params = sum(p.numel() for p in self._velocity_model.parameters()) / 1e9187        vram_gb = sum(p.numel() * p.element_size() for p in self._velocity_model.parameters()) / 1e9188        logging.info(f"  Transformer: {time.time()-t0:.1f}s ({n_params:.1f}B params, {vram_gb:.1f}GB VRAM, {self.dtype})")189 190        # torch.compile for faster denoising191        if self.compile_model:192            t0 = time.time()193            logging.info("  Compiling transformer with torch.compile (default mode)...")194            self._velocity_model = torch.compile(self._velocity_model, mode="default", dynamic=True)195            logging.info(f"  Compiled: {time.time()-t0:.1f}s (first call triggers actual compilation)")196 197        # 4. Audio decoder (VAE decoder + vocoder kept warm)198        t0 = time.time()199        self._audio_decoder = AudioDecoder(200            checkpoint_path=self.full_checkpoint,201            dtype=self.dtype, device=self.device,202            warm=True,203        )204        logging.info(f"  AudioDecoder (warm): {time.time()-t0:.1f}s")205 206    @torch.inference_mode()207    def generate(self, prompt, voice_ref=None, cfg_scale=2.5, stg_scale=1.5,208                 duration_multiplier=1.1, seed=42, ref_duration=10.0,209                 rescale_scale="auto", gen_duration: float = 0.0):210        """Generate audio. Returns (waveform_path, duration_seconds).211 212        rescale_scale: latent-side CFG std-rescale that prevents clipping at213            high cfg. Set to "auto" (default) for the cfg-aware schedule, a214            float in [0, 1] for a fixed override, or 0 to disable.215        gen_duration: explicit target duration in seconds. 0 (default) → auto216            from prompt + duration_multiplier; >0 overrides everything else.217        """218        t_total = time.time()219 220        # Duration + target shape — explicit gen_duration wins over the estimator.221        if gen_duration and gen_duration > 0:222            gen_dur = float(gen_duration)223        else:224            gen_dur = estimate_duration(prompt, duration_multiplier)225        fps = 25.0226        n_frames = int(round(gen_dur * fps)) + 1227        n_frames = ((n_frames - 1 + 4) // 8) * 8 + 1228        pixel_shape = VideoPixelShape(batch=1, frames=n_frames, height=64, width=64, fps=fps)229        target_shape = AudioLatentShape.from_video_pixel_shape(pixel_shape)230        audio_tools = AudioLatentTools(patchifier=self.patchifier, target_shape=target_shape)231 232        # Initial state233        state = audio_tools.create_initial_state(device=self.device, dtype=self.dtype)234 235        # Voice ref conditioning236        if voice_ref and os.path.exists(voice_ref):237            t0 = time.time()238            voice = decode_audio_from_file(voice_ref, self.device, 0.0, ref_duration)239            w = voice.waveform240            if w.dim() == 2:241                if w.shape[0] == 1:242                    w = w.repeat(2, 1)243                w = w.unsqueeze(0)244            elif w.dim() == 3 and w.shape[1] == 1:245                w = w.repeat(1, 2, 1)246            target_samples = int(ref_duration * voice.sampling_rate)247            if w.shape[-1] < target_samples:248                w = w.repeat(1, 1, (target_samples // w.shape[-1]) + 1)249            w = w[..., :target_samples]250            peak = w.abs().max()251            if peak > 0:252                w = w * (10 ** (-4.0 / 20) / peak)253            voice = Audio(waveform=w, sampling_rate=voice.sampling_rate)254            ref_latent = self._audio_conditioner(lambda enc: vae_encode_audio(voice, enc, None))255            cond = AudioConditionByReferenceLatent(latent=ref_latent.to(self.device, self.dtype), strength=1.0)256            state = cond.apply_to(state, audio_tools)257            logging.info(f"Voice ref: {time.time()-t0:.2f}s")258 259        # Noise260        gen = torch.Generator(device=self.device).manual_seed(seed)261        noiser = GaussianNoiser(generator=gen)262        state = noiser(state, noise_scale=1.0)263 264        # Prompt encode265        t0 = time.time()266        prompts = [prompt, DEFAULT_NEG] if cfg_scale > 1.0 else [prompt]267        ctx = self._prompt_encoder(prompts, streaming_prefetch_count=None)268        a_ctx = ctx[0].audio_encoding269        a_ctx_neg = ctx[1].audio_encoding if cfg_scale > 1.0 else None270        logging.info(f"Prompt: {time.time()-t0:.2f}s")271 272        # Denoiser273        resc = auto_rescale_for_cfg(cfg_scale) if rescale_scale == "auto" else float(rescale_scale)274        if rescale_scale == "auto":275            logging.info(f"Auto rescale_scale = {resc:.2f} for cfg={cfg_scale}")276        guider = MultiModalGuider(277            params=MultiModalGuiderParams(278                cfg_scale=cfg_scale, stg_scale=stg_scale,279                stg_blocks=[29], rescale_scale=resc, modality_scale=1.0,280            ),281            negative_context=a_ctx_neg,282        )283        denoiser = GuidedDenoiser(284            v_context=None, a_context=a_ctx,285            video_guider=None, audio_guider=guider,286        )287 288        # Sigmas289        sigmas = LTX2Scheduler().execute(steps=30, latent=state.latent).to(self.device)290 291        # Denoise292        t0 = time.time()293        x0 = X0Model(self._velocity_model)294        _, audio_state = euler_denoising_loop(295            sigmas=sigmas, video_state=None, audio_state=state,296            stepper=EulerDiffusionStep(), transformer=x0, denoiser=denoiser,297        )298        logging.info(f"Denoise (30 steps): {time.time()-t0:.2f}s")299 300        # Strip + unpatchify + decode301        audio_state = audio_tools.clear_conditioning(audio_state)302        audio_state = audio_tools.unpatchify(audio_state)303 304        # End-of-clip silence-prior fix.305        # The base LTX-2.3 22B DiT was trained on audio clips ≤ ~20 s and306        # learned a strong "clip-end silence" prior that lands on the next307        # patchifier-aligned latent frame after 20 s — index 513 = 8*64+1.308        # When inference produces longer audio, this prior leaks through as a309        # high-norm latent burst at frame 513 (and adjacent 512), which the310        # audio VAE + vocoder render as a ~30 ms hard silence dip near 20.4 s.311        # Linear interpolation across the two affected frames removes the dip312        # cleanly without any retraining. Only runs when the latent is long313        # enough to actually contain the boundary.314        latent = audio_state.latent315        if latent.shape[2] > 513:316            f0, f1 = 511, 514          # neighbours used for interpolation317            n = f1 - f0                # = 3318            patched = latent.clone()319            for f in (512, 513):320                t = (f - f0) / n321                patched[:, :, f, :] = (1.0 - t) * latent[:, :, f0, :] + t * latent[:, :, f1, :]322            latent = patched323 324        t0 = time.time()325        decoded = self._audio_decoder(latent)326        logging.info(f"Decode: {time.time()-t0:.2f}s")327 328        total = time.time() - t_total329        dur = decoded.waveform.shape[-1] / decoded.sampling_rate330        logging.info(f"Total: {total:.2f}s for {dur:.1f}s audio")331        return decoded.waveform, decoded.sampling_rate332 333    def generate_to_file(self, prompt, output, watermark: bool = True, **kwargs):334        waveform, sr = self.generate(prompt, **kwargs)335        wav_cpu = waveform.cpu().float()336        if watermark:337            try:338                import numpy as np, perth339                if not hasattr(self, "_perth"):340                    self._perth = perth.PerthImplicitWatermarker()341                mono = wav_cpu.mean(dim=0).numpy() if wav_cpu.shape[0] > 1 else wav_cpu[0].numpy()342                mono_wm = self._perth.apply_watermark(mono, sample_rate=sr)343                mono_wm_t = torch.from_numpy(np.asarray(mono_wm, dtype=np.float32)).unsqueeze(0)344                wav_cpu = mono_wm_t if wav_cpu.shape[0] == 1 else mono_wm_t.repeat(wav_cpu.shape[0], 1)345            except Exception as e:346                logging.warning(f"Perth watermark skipped ({e})")347        torchaudio.save(output, wav_cpu, sr)348        logging.info(f"Saved: {output}")349        return output350 351 352if __name__ == "__main__":353    import argparse354    p = argparse.ArgumentParser()355    p.add_argument("--device", default="cuda")356    p.add_argument("--dtype", default="fp16", choices=["fp16", "bf16"])357    p.add_argument("--no-compile", action="store_true")358    p.add_argument("--no-bnb-4bit", action="store_true",359                   help="Disable bitsandbytes 4-bit path (default: on, since the default "360                        "unsloth Gemma checkpoint is pre-quantized).")361    args = p.parse_args()362 363    server = TTSServer(device=args.device, dtype=args.dtype, compile_model=not args.no_compile,364                       bnb_4bit=not args.no_bnb_4bit)365 366    # First call - includes any warmup367    logging.info("=== First request ===")368    server.generate_to_file(369        prompt='A woman speaks clearly, "The weather today will be sunny."',370        output="/tmp/warm_test1.wav",371        voice_ref="/mnt/persistent0/manmay/expressive/female_radio_nikole/female_radio_nikole.wav",372    )373 374    # Second call - should be much faster (models already warm)375    logging.info("\n=== Second request (warm) ===")376    server.generate_to_file(377        prompt='A man speaks excitedly, "This is amazing, I cannot believe it!"',378        output="/tmp/warm_test2.wav",379        voice_ref="/mnt/persistent0/manmay/expressive/male_arnie/male_arnie.mp3",380    )381