ltx_driver.py

#!/usr/bin/env python3
"""ltx_driver.py — LTX-2.3-22b-distilled IC-LoRA colorization, LOAD-ONCE + temporeel naadloos.
STRUCTURELE FIX voor venster-naden (ml-engineer + ac64 consult, 28 jun 2026):
- Carry-over: laatste N frames van venster K -> conditioning op frame_idx 0..N-1 van venster K+1
  (via ic_lora ImageConditioningInput op willekeurige frame_idx) -> venster start vanuit vorig eind -> naadloos.
- Optioneel globaal anker-keyframe in elk venster (palet globaal pinnen).
- Optioneel 8-frame cosine-crossfade bij stitch voor rest-gladheid.
- Model laadt 1x (i.p.v. 67x in ltx_full_loop.sh) -> ~55min i.p.v. ~102min.

COMPUTE OPTIMALISATIES (A100-SXM4-80GB, 28 jun 2026):
- flash-attn 2: DIFFUSERS_ATTN_BACKEND=flash -> LTX dispatch_attention_fn gebruikt flash-attn 2.
  Op Ampere (sm80) is dit de snelste attentie-backend. SageAttention NIET gebruiken op A100
  (SM90-FP8 variant is Hopper-only; FP16 variant geeft geen winst boven flash-attn 2).
- TF32: torch.backends.cuda.matmul.allow_tf32=True (default op PyTorch>=1.7 maar hier expliciet).
- CFG-skip: guidance_scale=1.0 (distilled model heeft guidance ingebakken). Als de pipeline
  met guidance_scale>1 draait, doet het 2 DiT-forwards per stap = 2x de compute. Met 1.0: 1 forward.
  DIAGNOSE EERST: .venv/bin/python -m ltx_pipelines.ic_lora --help | grep -i guid
- torch.compile: mode='reduce-overhead', dynamic=False (vaste shapes 81f/960x544). Alleen DiT,
  niet VAE/Gemma. Eerste compile ~3-4 min, amortiseert over 67 vensters.
  max-autotune NIET gebruiken op 80GB (OOM-risico met 46GB model).

LET OP: de ltx_pipelines.ic_lora API-class-naam kan per repo-versie verschillen.
Valideer op de instance: `python -c "import ltx_pipelines.ic_lora as m; print(dir(m))"`.
Native import faalt -> subprocess-fallback (correct, minder efficient; compile dan niet actief).

Manifest (JSON): [{"gray":"wK.mp4","keyframe":"ref_X.png","kf_pos":40,"prompt":"...","out":"wK.mp4","cut_at_start":false}, ...]
Run: DIFFUSERS_ATTN_BACKEND=flash python ltx_driver.py manifest.json \
  --ck ... --lora ... --gemma ... --num-frames 81 --width 960 --height 544 --steps 8 \
  --guidance-scale 1.0 --compile --carry-over 2 --crossfade 8 --stitch-out full_seamless.mp4 --fps 30
"""
# EERST: flash-attn backend instellen VOOR elke diffusers/ltx import
import os
_attn_backend = os.environ.get('DIFFUSERS_ATTN_BACKEND', 'flash')
os.environ['DIFFUSERS_ATTN_BACKEND'] = _attn_backend

import sys, json, time, argparse, tempfile, subprocess, shutil
from pathlib import Path

import torch
# TF32: op A100 versnelt matmul ~1.5x bij bf16. Standaard aan in PyTorch>=1.7 maar hier expliciet.
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument('manifest'); p.add_argument('--ck', required=True); p.add_argument('--lora', required=True)
    p.add_argument('--gemma', required=True); p.add_argument('--num-frames', type=int, default=81)
    p.add_argument('--width', type=int, default=960); p.add_argument('--height', type=int, default=544)
    p.add_argument('--steps', type=int, default=8); p.add_argument('--seed', type=int, default=42)
    p.add_argument('--lora-scale', type=float, default=1.0); p.add_argument('--carry-over', type=int, default=2)
    p.add_argument('--crossfade', type=int, default=8); p.add_argument('--carry-weight-base', type=float, default=0.85)
    p.add_argument('--carry-weight-decay', type=float, default=0.20); p.add_argument('--global-anchor', default=None)
    p.add_argument('--skip-stage2', action='store_true', default=True); p.add_argument('--prompt', default=None)
    p.add_argument('--stitch-out', default=None); p.add_argument('--fps', type=int, default=30)
    # A100 compute opts:
    p.add_argument('--guidance-scale', type=float, default=1.0,
                   help='CFG guidance. 1.0=single DiT forward/stap (distilled). >1 = 2 forwards = 2x trager.')
    p.add_argument('--compile', action='store_true', default=False,
                   help='torch.compile de DiT (reduce-overhead, dynamic=False). Eerste venster ~3-4 min trager.')
    p.add_argument('--compile-mode', default='reduce-overhead',
                   help='torch.compile mode. reduce-overhead=veilig op 80GB. max-autotune=OOM-risico.')
    return p.parse_args()

def _apply_compile(pipe, mode):
    """Compileer de DiT-transformer. Niet de VAE of text encoder (wisselende shapes / 1x gebruik)."""
    # Zoek het DiT-object: common attributen in ltx_pipelines en diffusers
    dit = None
    for attr in ('transformer', 'dit', 'model', 'unet'):
        if hasattr(pipe, attr):
            candidate = getattr(pipe, attr)
            if callable(getattr(candidate, 'forward', None)):
                dit = candidate; break
    if dit is None:
        print("[DRIVER] compile: DiT niet gevonden in pipeline — compile overgeslagen", flush=True); return
    print(f"[DRIVER] compile: torch.compile({attr}, mode='{mode}', dynamic=False) ...", flush=True)
    t0 = time.time()
    setattr(pipe, attr, torch.compile(dit, mode=mode, dynamic=False))
    print(f"[DRIVER] compile klaar ({round(time.time()-t0)}s) — eerste venster triggert kernel-compilatie", flush=True)

def _check_cfg(pipe, a):
    """Diagnose: print de effectieve guidance_scale zodat CFG-skip verifieerbaar is."""
    gs = getattr(pipe, '_guidance_scale', None) or getattr(pipe, 'guidance_scale', None) or a.guidance_scale
    n_forwards = 2 if float(gs) > 1.0 else 1
    print(f"[DRIVER] CFG: guidance_scale={gs} -> {n_forwards} DiT forward(s)/stap "
          f"(bij 8 steps: {8*n_forwards} forwards/venster)", flush=True)
    if n_forwards == 2:
        print("[DRIVER] WAARSCHUWING: guidance_scale>1 gedetecteerd. Gebruik --guidance-scale 1.0 "
              "of voeg '--guidance-scale 1.0' toe aan CLI voor 2x speedup.", flush=True)

def load_pipeline(a):
    # Rapporteer welke attentie-backend actief is
    print(f"[DRIVER] DIFFUSERS_ATTN_BACKEND={os.environ.get('DIFFUSERS_ATTN_BACKEND','native')}", flush=True)
    try:
        from ltx_pipelines.ic_lora.inference import ICLoRAInference
        pipe = ICLoRAInference(distilled_checkpoint_path=a.ck, lora_path=a.lora, lora_scale=a.lora_scale,
                               gemma_root=a.gemma, skip_stage2=a.skip_stage2)
        print("[DRIVER] native ICLoRAInference geladen", flush=True)
        _check_cfg(pipe, a)
        if a.compile:
            _apply_compile(pipe, a.compile_mode)
        return pipe, 'native'
    except (ImportError, AttributeError) as e:
        print(f"[DRIVER] native import faalt ({e}) -> subprocess-fallback "
              f"(compile niet actief in subprocess-mode)", flush=True)
        return None, 'subprocess'

def run_native(pipe, w, conds, a):
    # guidance_scale=1.0 = single forward per stap (CFG-skip voor distilled model)
    kwargs = dict(prompt=w.get('prompt', a.prompt or ''), video_conditioning_path=w['gray'],
                  video_conditioning_weight=1.0, image_conditioning=conds, num_frames=a.num_frames,
                  width=a.width, height=a.height, seed=a.seed, num_inference_steps=a.steps,
                  output_path=w['out'])
    # Pas guidance_scale toe als de pipeline het accepteert
    try:
        pipe.run(guidance_scale=a.guidance_scale, **kwargs)
    except TypeError:
        # Oudere API zonder guidance_scale parameter
        pipe.run(**kwargs)

def run_subprocess(w, conds, a):
    cmd = [sys.executable, '-m', 'ltx_pipelines.ic_lora', '--distilled-checkpoint-path', a.ck,
           '--lora', a.lora, str(a.lora_scale), '--gemma-root', a.gemma,
           '--video-conditioning', w['gray'], '1.0', '--num-frames', str(a.num_frames),
           '--width', str(a.width), '--height', str(a.height), '--seed', str(a.seed),
           '--frame-rate', str(a.fps), '--output-path', w['out'], '--prompt', w.get('prompt', a.prompt or '')]
    if a.skip_stage2: cmd.append('--skip-stage-2')
    # CFG-skip: probe of de CLI --guidance-scale accepteert (eenmalig) en voeg toe
    if a.guidance_scale <= 1.0:
        cmd += ['--guidance-scale', str(a.guidance_scale)]
    for (path, pos, wt) in conds: cmd += ['--image', path, str(pos), str(wt)]
    # Geef flash-attn backend door aan subprocess via env
    env = os.environ.copy()
    env['DIFFUSERS_ATTN_BACKEND'] = os.environ.get('DIFFUSERS_ATTN_BACKEND', 'flash')
    result = subprocess.run(cmd, env=env)
    if result.returncode != 0:
        # Als --guidance-scale niet herkend: opnieuw zonder die flag
        cmd2 = [c for c in cmd if c not in ['--guidance-scale', str(a.guidance_scale)]]
        subprocess.run(cmd2, env=env, check=True)

def last_frames(video, n, td):
    pr = subprocess.run(['ffprobe','-v','quiet','-select_streams','v:0','-count_packets',
                         '-show_entries','stream=nb_read_packets','-of','csv=p=0',video], capture_output=True, text=True)
    try: total = int(pr.stdout.strip())
    except Exception: total = 81
    out = []
    for i in range(n):
        fi = total - n + i; png = os.path.join(td, f'carry_{i:03d}.png')
        subprocess.run(['ffmpeg','-y','-i',video,'-vf',f'select=eq(n\\,{fi})','-vsync','0','-frames:v','1',png],capture_output=True)
        if os.path.exists(png): out.append(png)
    return out

def stitch(paths, xf, fps, out):
    if xf == 0 or len(paths) == 1:
        lst = tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False)
        for p in paths: lst.write(f"file '{p}'\n")
        lst.close(); subprocess.run(['ffmpeg','-y','-f','concat','-safe','0','-i',lst.name,'-c','copy',out],check=True); os.unlink(lst.name); return
    dur = 81/fps; xfd = xf/fps; inputs = []
    for p in paths: inputs += ['-i', p]
    parts = []; cur = '[0:v]'
    for k in range(1, len(paths)):
        off = max(0, k*dur - k*xfd); nl = f'[v{k}]' if k < len(paths)-1 else ''
        parts.append(f"{cur}[{k}:v]xfade=transition=fade:duration={xfd:.4f}:offset={off:.4f}{nl}"); cur = nl
    subprocess.run(['ffmpeg','-y']+inputs+['-filter_complex','; '.join(parts),'-fps_mode','vfr',out],check=True)

def main():
    a = parse_args(); manifest = json.load(open(a.manifest)); td = tempfile.mkdtemp(prefix='ltx_co_')
    print(f"[DRIVER] {len(manifest)} vensters, carry-over={a.carry_over}, crossfade={a.crossfade}", flush=True)
    t0 = time.time(); pipe, mode = load_pipeline(a); print(f"[DRIVER] load {round(time.time()-t0)}s mode={mode}", flush=True)
    rendered = []; prev = None
    for idx, w in enumerate(manifest):
        tw = time.time(); conds = []
        if prev and not w.get('cut_at_start', False):
            cn = min(a.carry_over, len(prev))
            for ci, cf in enumerate(prev[-cn:]):
                conds.append((cf, cn-1-ci, max(0.3, a.carry_weight_base - ci*a.carry_weight_decay)))
        if a.global_anchor and (not prev or w.get('cut_at_start')): conds.append((a.global_anchor, 0, 0.4))
        conds.append((w['keyframe'], w.get('kf_pos', a.num_frames//2), 1.0))
        print(f"[DRIVER] venster {idx+1}/{len(manifest)} {Path(w['out']).name} ({len(conds)} conds)", flush=True)
        (run_native if mode=='native' else (lambda p,ww,c,aa: run_subprocess(ww,c,aa)))(pipe, w, conds, a) if mode=='native' else run_subprocess(w, conds, a)
        if a.carry_over > 0: prev = last_frames(w['out'], a.carry_over, td)
        rendered.append(w['out']); print(f"[DRIVER] venster {idx+1} klaar ({round(time.time()-tw)}s)", flush=True)
    if a.stitch_out: print(f"[DRIVER] stitch -> {a.stitch_out}", flush=True); stitch(rendered, a.crossfade, a.fps, a.stitch_out)
    shutil.rmtree(td, ignore_errors=True); print("[DRIVER] ALL_DONE", flush=True)

if __name__ == '__main__':
    main()