File: /home/kiwerkzeuge.de/ltx2-av-runtime/test_fp4_av_diffusers.py
#!/usr/bin/env python3
import argparse
import json
import os
import subprocess
import time
import traceback
from pathlib import Path
import numpy as np
import soundfile as sf
import torch
from PIL import Image
from diffusers import AutoencoderKLLTX2Audio, LTX2ImageToVideoPipeline
from diffusers.pipelines.ltx2 import LTX2TextConnectors
from diffusers.utils import export_to_video
from transformers import Gemma3ForConditionalGeneration
DEFAULT_CHECKPOINT = Path(
"/home/kiwerkzeuge.de/ltx2-av-runtime/cache/hf-official/"
"models--Lightricks--LTX-2/snapshots/39c3be3ccfcb1deeb33c52769d72cb9170dd4573/"
"ltx-2-19b-dev-fp4.safetensors"
)
DEFAULT_IMAGE = Path("/home/kiwerkzeuge.de/public_html/uploads/20260226125524-f66b6411-sd.png")
DEFAULT_OUT_DIR = Path("/home/kiwerkzeuge.de/ltx2-av-runtime/test-output")
DEFAULT_CACHE_DIR = Path("/home/kiwerkzeuge.de/ltx2-av-runtime/cache/hf-official")
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="FP4 LTX-2 AV feasibility test via diffusers.")
parser.add_argument("--checkpoint", type=Path, default=DEFAULT_CHECKPOINT)
parser.add_argument("--image", type=Path, default=DEFAULT_IMAGE)
parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUT_DIR)
parser.add_argument("--cache-dir", type=Path, default=DEFAULT_CACHE_DIR)
parser.add_argument("--width", type=int, default=384)
parser.add_argument("--height", type=int, default=224)
parser.add_argument("--num-frames", type=int, default=9)
parser.add_argument("--fps", type=float, default=4.0)
parser.add_argument("--steps", type=int, default=4)
parser.add_argument("--guidance", type=float, default=2.0)
parser.add_argument(
"--prompt",
type=str,
default="cinematic movement, subtle ambience, synchronized environmental audio",
)
parser.add_argument(
"--negative-prompt",
type=str,
default="glitch, distortion, artifacts, frozen frame, silent output",
)
parser.add_argument("--seed", type=int, default=0, help="0 = auto seed")
return parser.parse_args()
def divisible_or_raise(width: int, height: int, num_frames: int) -> None:
if width % 32 != 0 or height % 32 != 0:
raise ValueError(f"Width/Height must be divisible by 32, got {width}x{height}")
if (num_frames - 1) % 8 != 0:
raise ValueError(f"num_frames must satisfy 8n+1, got {num_frames}")
def to_audio_channels_first(audio_value) -> np.ndarray:
arr = audio_value.detach().float().cpu().numpy() if isinstance(audio_value, torch.Tensor) else np.asarray(audio_value)
arr = arr.astype(np.float32)
if arr.ndim == 1:
arr = arr[None, :]
elif arr.ndim == 2:
# Accept either [channels, samples] or [samples, channels]
if arr.shape[0] > arr.shape[1] and arr.shape[1] in (1, 2):
arr = arr.T
else:
arr = np.squeeze(arr)
if arr.ndim == 1:
arr = arr[None, :]
elif arr.ndim == 2 and arr.shape[0] > arr.shape[1] and arr.shape[1] in (1, 2):
arr = arr.T
elif arr.ndim != 2:
raise ValueError(f"Unexpected audio ndim={arr.ndim}")
if arr.shape[0] == 1:
arr = np.repeat(arr, 2, axis=0)
if arr.shape[0] != 2:
raise ValueError(f"Expected 2 channels after normalization, got shape {arr.shape}")
return np.clip(arr, -1.0, 1.0)
def ffprobe_streams(path: Path) -> str:
cmd = [
"ffprobe",
"-v",
"error",
"-show_streams",
"-of",
"json",
str(path),
]
res = subprocess.run(cmd, capture_output=True, text=True, check=False)
return res.stdout.strip() if res.returncode == 0 else res.stderr.strip()
def load_pipeline_with_component_fallbacks(args, result: dict) -> LTX2ImageToVideoPipeline:
load_started = time.time()
component_overrides = {}
fallback_components = []
for _ in range(4):
try:
pipe = LTX2ImageToVideoPipeline.from_single_file(
str(args.checkpoint),
config="Lightricks/LTX-2",
torch_dtype=torch.bfloat16,
cache_dir=str(args.cache_dir),
local_files_only=False,
**component_overrides,
)
result["timings_sec"]["load"] = round(time.time() - load_started, 2)
if fallback_components:
result["fallback_components"] = fallback_components
return pipe
except Exception as exc:
message = str(exc)
if "LTX2TextConnectors" in message and "connectors" not in component_overrides:
connectors = LTX2TextConnectors.from_pretrained(
"Lightricks/LTX-2",
subfolder="connectors",
torch_dtype=torch.float16,
cache_dir=str(args.cache_dir),
local_files_only=False,
)
component_overrides["connectors"] = connectors
fallback_components.append("connectors")
continue
if "AutoencoderKLLTX2Audio" in message and "audio_vae" not in component_overrides:
audio_vae = AutoencoderKLLTX2Audio.from_pretrained(
"Lightricks/LTX-2",
subfolder="audio_vae",
torch_dtype=torch.float16,
cache_dir=str(args.cache_dir),
local_files_only=False,
)
component_overrides["audio_vae"] = audio_vae
fallback_components.append("audio_vae")
continue
if "Gemma3ForConditionalGeneration" in message and "text_encoder" not in component_overrides:
text_encoder = Gemma3ForConditionalGeneration.from_pretrained(
"Lightricks/LTX-2",
subfolder="text_encoder",
torch_dtype=torch.float16,
cache_dir=str(args.cache_dir),
local_files_only=False,
device_map="cpu",
)
component_overrides["text_encoder"] = text_encoder
fallback_components.append("text_encoder")
continue
raise
raise RuntimeError("Pipeline loading exhausted fallback attempts.")
def main() -> int:
args = parse_args()
divisible_or_raise(args.width, args.height, args.num_frames)
if not args.checkpoint.exists():
raise FileNotFoundError(f"Checkpoint not found: {args.checkpoint}")
if not args.image.exists():
raise FileNotFoundError(f"Image not found: {args.image}")
args.output_dir.mkdir(parents=True, exist_ok=True)
os.environ.setdefault("HF_HOME", str(args.cache_dir))
os.environ.setdefault("HUGGINGFACE_HUB_CACHE", str(args.cache_dir))
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
seed = args.seed if args.seed > 0 else int(time.time()) % 2_147_483_647
stamp = time.strftime("%Y%m%d-%H%M%S")
base = f"ltx2-fp4-av-{stamp}-{seed}"
video_path = args.output_dir / f"{base}.mp4"
wav_path = args.output_dir / f"{base}.wav"
mux_path = args.output_dir / f"{base}-mux.mp4"
result = {
"ok": False,
"seed": seed,
"checkpoint": str(args.checkpoint),
"input_image": str(args.image),
"video_path": str(video_path),
"audio_wav_path": str(wav_path),
"mux_path": str(mux_path),
"timings_sec": {},
}
pipe = None
start_total = time.time()
try:
pipe = load_pipeline_with_component_fallbacks(args, result)
if torch.cuda.is_available():
pipe.enable_model_cpu_offload()
if hasattr(pipe, "vae") and pipe.vae is not None:
try:
pipe.vae.enable_tiling()
pipe.vae.enable_slicing()
except Exception:
pass
image = Image.open(args.image).convert("RGB").resize((args.width, args.height), Image.LANCZOS)
generator = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu").manual_seed(seed)
infer_started = time.time()
out = pipe(
image=image,
prompt=args.prompt,
negative_prompt=args.negative_prompt,
width=args.width,
height=args.height,
num_frames=args.num_frames,
frame_rate=args.fps,
num_inference_steps=args.steps,
guidance_scale=args.guidance,
generator=generator,
output_type="np",
return_dict=True,
)
result["timings_sec"]["infer"] = round(time.time() - infer_started, 2)
frames = out.frames[0]
audio_raw = out.audio[0]
export_started = time.time()
export_to_video(frames, str(video_path), fps=max(1, int(round(args.fps))))
audio = to_audio_channels_first(audio_raw)
audio_sr = int(getattr(pipe.vocoder.config, "output_sampling_rate", 24000))
sf.write(str(wav_path), audio.T, samplerate=audio_sr)
ffmpeg_cmd = [
"ffmpeg",
"-y",
"-i",
str(video_path),
"-i",
str(wav_path),
"-c:v",
"copy",
"-c:a",
"aac",
"-shortest",
str(mux_path),
]
subprocess.run(ffmpeg_cmd, check=True, capture_output=True, text=True)
result["timings_sec"]["export_mux"] = round(time.time() - export_started, 2)
result["audio_shape"] = list(audio.shape)
result["audio_sample_rate"] = audio_sr
result["ffprobe_mux_streams"] = ffprobe_streams(mux_path)
result["ok"] = True
except Exception as exc:
result["error"] = str(exc)
result["traceback"] = traceback.format_exc()
finally:
if torch.cuda.is_available():
try:
torch.cuda.empty_cache()
except Exception:
pass
if pipe is not None:
del pipe
result["timings_sec"]["total"] = round(time.time() - start_total, 2)
print(json.dumps(result, ensure_ascii=True, indent=2))
return 0 if result["ok"] else 1
if __name__ == "__main__":
raise SystemExit(main())