From d2ffec5cce359dc6d6215af7fdfb9c36311b0da7 Mon Sep 17 00:00:00 2001 From: Junda Su <90978028+Davids048@users.noreply.github.com> Date: Fri, 15 May 2026 18:10:47 -0700 Subject: [PATCH] [perf] Dreamverse 14/14: Add LTX2 profile speedups (#1337) Co-authored-by: SolitaryThinker Co-authored-by: Matthew Noto <99706358+RandNMR73@users.noreply.github.com> Co-authored-by: XOR-op <17672363+XOR-op@users.noreply.github.com> Co-authored-by: Zhang Peiyuan <42993249+jzhang38@users.noreply.github.com> --- .../basic_ltx2_distilled_fast_profile.py | 336 ++++++++++++++++++ fastvideo/entrypoints/video_generator.py | 284 ++++++++++++++- fastvideo/models/vaes/ltx2vae.py | 34 ++ fastvideo/pipelines/composed_pipeline_base.py | 4 + fastvideo/pipelines/stages/base.py | 7 +- 5 files changed, 657 insertions(+), 8 deletions(-) create mode 100644 examples/inference/basic/basic_ltx2_distilled_fast_profile.py diff --git a/examples/inference/basic/basic_ltx2_distilled_fast_profile.py b/examples/inference/basic/basic_ltx2_distilled_fast_profile.py new file mode 100644 index 000000000..e5f284746 --- /dev/null +++ b/examples/inference/basic/basic_ltx2_distilled_fast_profile.py @@ -0,0 +1,336 @@ +# SPDX-License-Identifier: Apache-2.0 +import json +import os +import time +from collections import OrderedDict +from pathlib import Path + +import torch +import torch._inductor.config +from fastvideo import VideoGenerator +from fastvideo.configs.pipelines.base import PipelineConfig +from fastvideo.layers.quantization.nvfp4_config import NVFP4Config +from fastvideo.utils import maybe_download_model + +VALIDATION_JSON = ( + Path(__file__).resolve().parents[2] / "training" / "finetune" / "ltx2" / "validation.json" +) + +# Override with a local snapshot or converted directory when needed, e.g. +# export LTX2_MODEL_PATH=/raid/$USER/hf/FastVideo/LTX2-Distilled-Diffusers +MODEL_ID = os.path.expandvars( + os.path.expanduser(os.getenv("LTX2_MODEL_PATH", "FastVideo/LTX2-Distilled-Diffusers")) +) +OUTPUT_DIR = Path("outputs_video/ltx2_distilled_fast_profile") + +os.environ["FASTVIDEO_ATTENTION_BACKEND"] = "FLASH_ATTN" +os.environ["FASTVIDEO_STAGE_LOGGING"] = "1" + +# Tune Inductor flags +config = torch._inductor.config +config.conv_1x1_as_mm = True # treat 1x1 convolutions as matrix muls +config.coordinate_descent_tuning = True +config.coordinate_descent_check_all_directions = True +config.epilogue_fusion = False # do not fuse pointwise ops into matmuls + + +def load_validation_entries(path: Path) -> list[dict]: + with path.open(encoding="utf-8") as f: + data = json.load(f) + + if isinstance(data, dict) and isinstance(data.get("data"), list): + return [entry for entry in data["data"] if isinstance(entry, dict)] + + raise ValueError(f"Unsupported validation format in {path}. Expected {{'data': [...]}}.") + + +def print_stage_breakdown( + result: dict, + run_idx: int, + num_runs: int, +) -> float | None: + logging_info = result.get("logging_info") + if logging_info is None: + print(f"[{run_idx}/{num_runs}] Stage breakdown unavailable: no logging_info") + return None + + stages = getattr(logging_info, "stages", None) + if not stages: + print(f"[{run_idx}/{num_runs}] Stage breakdown unavailable: no stage timings") + return None + + print(f"[{run_idx}/{num_runs}] Stage breakdown:") + total = 0.0 + for stage_name, stage_metrics in stages.items(): + exec_time = float(stage_metrics.get("execution_time", 0.0)) + total += exec_time + print(f" - {stage_name}: {exec_time:.3f}s") + print(f" - total(stage sum): {total:.3f}s") + return total + + +def extract_sr_forward_latency( + result: dict, +) -> tuple[float | None, list[tuple[str, float]], list[str]]: + logging_info = result.get("logging_info") + if logging_info is None: + return None, [], [] + + stages = getattr(logging_info, "stages", None) + if not stages: + return None, [], [] + + stage_names = list(stages.keys()) + sr_match_substr = os.getenv("FASTVIDEO_SR_LATENCY_STAGE_SUBSTR", "").strip().lower() + + sr_stage_entries: list[tuple[str, float]] = [] + for stage_name, stage_metrics in stages.items(): + stage_name_l = stage_name.lower() + if sr_match_substr: + is_sr_stage = sr_match_substr in stage_name_l + else: + is_sr_stage = ( + "srdenoisingstage" in stage_name_l + or "sr_denoising" in stage_name_l + or "upsample" in stage_name_l + or ("refine" in stage_name_l and "denois" in stage_name_l) + ) + if not is_sr_stage: + continue + exec_time = float(stage_metrics.get("execution_time", 0.0)) + sr_stage_entries.append((stage_name, exec_time)) + + if not sr_stage_entries: + return None, [], stage_names + return sum(x[1] for x in sr_stage_entries), sr_stage_entries, stage_names + + +def collect_stage_times( + result: dict, + stage_times: dict[str, list[float]], + stage_order: OrderedDict[str, None], +) -> None: + logging_info = result.get("logging_info") + if logging_info is None: + return + stages = getattr(logging_info, "stages", None) + if not stages: + return + for stage_name, stage_metrics in stages.items(): + stage_order.setdefault(stage_name, None) + exec_time = float(stage_metrics.get("execution_time", 0.0)) + stage_times.setdefault(stage_name, []).append(exec_time) + + +def print_stage_averages( + stage_times: dict[str, list[float]], + stage_order: OrderedDict[str, None], + measured_runs: int, +) -> None: + if measured_runs <= 0: + return + if not stage_times: + print("No stage timings collected for measured runs.") + return + + print(f"Average stage times over {measured_runs} measured runs:") + total_avg = 0.0 + for stage_name in stage_order.keys(): + times = stage_times.get(stage_name, []) + if not times: + continue + avg = sum(times) / len(times) + total_avg += avg + print(f" - {stage_name}: {avg:.3f}s") + print(f" - total(stage sum avg): {total_avg:.3f}s") + + +def resolve_refine_upsampler_path(model_root: str) -> Path: + root = Path(model_root) + candidates = [ + root / "spatial_upscaler", + root / "spatial_upsampler", + ] + + env_path = os.getenv("LTX2_REFINE_UPSAMPLER_PATH") + if env_path: + candidates.insert(0, Path(os.path.expandvars(os.path.expanduser(env_path)))) + + for candidate in candidates: + if (candidate / "config.json").is_file(): + return candidate + + checked = "\n".join(f" - {candidate}" for candidate in candidates) + raise FileNotFoundError( + "Could not find an LTX2 refine upsampler directory.\n" + "Checked:\n" + f"{checked}" + ) + + +def main() -> None: + if not VALIDATION_JSON.exists(): + raise FileNotFoundError(f"Validation file not found: {VALIDATION_JSON}") + + validation_entries = load_validation_entries(VALIDATION_JSON) + if not validation_entries: + raise ValueError(f"No validation entries found in {VALIDATION_JSON}") + + benchmark_entry = validation_entries[0] + prompt = benchmark_entry.get("caption") + if not isinstance(prompt, str) or not prompt.strip(): + raise ValueError("First validation entry is missing a usable caption") + + num_runs = 12 + warmup_runs = 2 + avg_window = num_runs - warmup_runs + measured_start_idx = max(warmup_runs, num_runs - avg_window) + + model_root = maybe_download_model(MODEL_ID) + refine_upsampler_path = resolve_refine_upsampler_path(model_root) + print(f"Using refine upsampler: {refine_upsampler_path}") + + pipeline_config = PipelineConfig.from_pretrained(model_root) + pipeline_config.dit_config.quant_config = NVFP4Config() + OUTPUT_DIR.mkdir(parents=True, exist_ok=True) + torch_compile_kwargs = { + "backend": "inductor", + "fullgraph": True, + # Uncomment for final best-performance profiling. It is disabled + # for faster development iteration because autotuning is slow. + # "mode": "max-autotune-no-cudagraphs", + "dynamic": False, + } + + generator = VideoGenerator.from_pretrained( + model_root, + num_gpus=1, + ltx2_refine_enabled=True, + ltx2_refine_upsampler_path=str(refine_upsampler_path), + refine_lora_path="", # keep refine LoRA disabled in this repo's typed adapter + ltx2_refine_lora_path="", # keep refine LoRA disabled for distilled model + ltx2_refine_num_inference_steps=2, + ltx2_refine_guidance_scale=1.0, + ltx2_refine_add_noise=True, + pipeline_config=pipeline_config, + enable_torch_compile=True, + enable_torch_compile_text_encoder=True, + enable_torch_compile_vae=True, + torch_compile_kwargs=torch_compile_kwargs, + torch_compile_kwargs_vae=torch_compile_kwargs, + dit_cpu_offload=False, + text_encoder_cpu_offload=False, + vae_cpu_offload=False, + ltx2_vae_tiling=False, + ) + + run_times: list[float] = [] + e2e_times: list[float] = [] + sr_forward_times: list[float] = [] + non_stage_overhead_times: list[float] = [] + stage_times: dict[str, list[float]] = {} + stage_order: OrderedDict[str, None] = OrderedDict() + + try: + for i in range(num_runs): + output_path = OUTPUT_DIR / f"output_ltx2_basic_t2v_run_{i + 1}.mp4" + if output_path.exists(): + output_path.unlink() + print(f"[{i + 1}/{num_runs}] Removed existing file: {output_path}") + + print(f"[{i + 1}/{num_runs}] Generating: {output_path}") + if os.environ.get("FASTVIDEO_STAGE_LOGGING") == "0" and torch.cuda.is_available(): + torch.cuda.synchronize() + + start = time.perf_counter() + result = generator.generate_video( + prompt=prompt, + output_path=str(output_path), + fps=24, + seed=10, + save_video=True, + guidance_scale=1.0, + height=benchmark_entry.get("height", 1088), + width=benchmark_entry.get("width", 1920), + num_frames=121, + num_inference_steps=5, + # image_path="examples/inference/basic/prompt1.png", + # ltx2_image_crf=0.0 + ) + if os.environ.get("FASTVIDEO_STAGE_LOGGING") == "0": + torch.cuda.synchronize() + + elapsed = result.get("generation_time") if isinstance(result, dict) else None + e2e_elapsed = result.get("e2e_latency") if isinstance(result, dict) else None + if elapsed is None: + elapsed = time.perf_counter() - start + if e2e_elapsed is None: + e2e_elapsed = time.perf_counter() - start + + run_times.append(elapsed) + e2e_times.append(e2e_elapsed) + print(f"[{i + 1}/{num_runs}] Generation time: {elapsed:.2f}s") + print(f"[{i + 1}/{num_runs}] End-to-end latency: {e2e_elapsed:.2f}s") + + if isinstance(result, dict): + stage_sum = print_stage_breakdown(result, i + 1, num_runs) + if stage_sum is not None: + non_stage_overhead = e2e_elapsed - stage_sum + print(f"[{i + 1}/{num_runs}] Non-stage overhead (e2e - stage sum): {non_stage_overhead:.3f}s") + if i >= measured_start_idx: + non_stage_overhead_times.append(non_stage_overhead) + + sr_forward_total, sr_stage_entries, stage_names = extract_sr_forward_latency(result) + if sr_forward_total is None: + print(f"[{i + 1}/{num_runs}] SR forward latency unavailable") + if stage_names: + print(f" Available stage keys: {', '.join(stage_names)}") + print(" Tip: set FASTVIDEO_SR_LATENCY_STAGE_SUBSTR= to match your SR stage key.") + else: + print(f"[{i + 1}/{num_runs}] SR forward latency: {sr_forward_total:.3f}s") + for sr_stage_name, sr_exec_time in sr_stage_entries: + print(f" - {sr_stage_name}: {sr_exec_time:.3f}s") + if i >= measured_start_idx: + sr_forward_times.append(sr_forward_total) + + if i >= measured_start_idx: + collect_stage_times(result, stage_times, stage_order) + + measured_times = run_times[measured_start_idx:] + avg_time = sum(measured_times) / len(measured_times) + print( + f"Average video generation time over {len(measured_times)} runs " + f"(runs {measured_start_idx + 1}-{len(run_times)}, skipping first {warmup_runs} warmup runs): " + f"{avg_time:.2f}s" + ) + + measured_e2e_times = e2e_times[measured_start_idx:] + avg_e2e_time = sum(measured_e2e_times) / len(measured_e2e_times) + print( + f"Average end-to-end latency over {len(measured_e2e_times)} runs " + f"(runs {measured_start_idx + 1}-{len(e2e_times)}, skipping first {warmup_runs} warmup runs): " + f"{avg_e2e_time:.2f}s" + ) + + if sr_forward_times: + avg_sr_forward = sum(sr_forward_times) / len(sr_forward_times) + print(f"Average SR forward latency over {len(sr_forward_times)} runs: {avg_sr_forward:.3f}s") + else: + print("Average SR forward latency unavailable (no SR stages matched).") + + print_stage_averages(stage_times, stage_order, len(measured_times)) + + if non_stage_overhead_times: + avg_non_stage_overhead = sum(non_stage_overhead_times) / len(non_stage_overhead_times) + print( + "Average non-stage overhead over " + f"{len(non_stage_overhead_times)} measured runs: {avg_non_stage_overhead:.3f}s" + ) + else: + print("Average non-stage overhead unavailable (no stage timings).") + finally: + generator.shutdown() + + +if __name__ == "__main__": + main() diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index 94f55ce86..adf5bb7e7 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -9,11 +9,14 @@ diffusion models. import os import re import shutil +import subprocess import threading import time import tempfile +import types import warnings from collections.abc import Mapping +from contextlib import suppress from copy import deepcopy from typing import Any @@ -53,7 +56,14 @@ from fastvideo.pipelines import ForwardBatch from fastvideo.utils import align_to, shallow_asdict from fastvideo.worker.executor import Executor +fcntl: types.ModuleType | None +try: + import fcntl +except ImportError: + fcntl = None + logger = init_logger(__name__) +_FFMPEG_ENCODER_OPTION_CACHE: dict[tuple[str, str, str], bool] = {} _BATCH_EXTRA_PASSTHROUGH_KEYS: tuple[str, ...] = ( "ltx2_audio_latents", @@ -777,6 +787,7 @@ class VideoGenerator: is_latent_output = fastvideo_args.output_type == "latent" audio_only = bool(output_batch.extra.get("audio_only")) + postprocess_start = time.perf_counter() frames: list[np.ndarray] | None if is_latent_output or audio_only: frames = None if is_latent_output else [] @@ -787,35 +798,90 @@ class VideoGenerator: x = torchvision.utils.make_grid(x, nrow=6) x = x.permute(1, 2, 0).squeeze(-1) x = (x * 255).to(torch.uint8) - frames.append(x.cpu().numpy()) + frames.append(x.contiguous().cpu().numpy()) + postprocess_time = time.perf_counter() - postprocess_start + logger.info("PostDecodeFrameProcessStage completed in %.3f s", postprocess_time) + if logging_info is not None: + logging_info.add_stage_execution_time("PostDecodeFrameProcessStage", postprocess_time) save_to_disk = batch.save_video and not is_latent_output + save_video_time = 0.0 + audio_mux_time = 0.0 if save_to_disk: if audio_only: # Audio-only workload: write a standalone .wav rather than # muxing the audio into a placeholder mp4 (which forces # ffmpeg to round 8x8 placeholder frames up to 16x16). output_path = self._rewrite_extension(output_path, ".wav") + save_start = time.perf_counter() self._write_pcm_wav( output_path, output_batch.extra["audio"], int(output_batch.extra["audio_sample_rate"]), ) + save_video_time = time.perf_counter() - save_start logger.info("Saved audio to %s", output_path) elif self._is_image_workload(): # Image workloads (t2i, i2i, …): save the first frame as PNG. assert frames is not None # implied by save_to_disk and not audio_only + save_start = time.perf_counter() imageio.imwrite(output_path, frames[0]) + save_video_time = time.perf_counter() - save_start logger.info("Saved image to %s", output_path) else: assert frames is not None # implied by save_to_disk and not audio_only - imageio.mimsave(output_path, frames, fps=batch.fps, format="mp4") - logger.info("Saved video to %s", output_path) audio = output_batch.extra.get("audio") audio_sample_rate = output_batch.extra.get("audio_sample_rate") - if (audio is not None and audio_sample_rate is not None - and not self._mux_audio(output_path, audio, audio_sample_rate)): - logger.warning("Audio mux failed; saved video without audio.") + if audio is not None and audio_sample_rate is not None: + # Single-pass save path: encode video+audio once, avoiding + # second-pass remux overhead. AudioMuxStage remains 0.0 on + # success because audio is already present in the saved MP4. + save_start = time.perf_counter() + save_ok = self._save_video_with_audio_ffmpeg_pipe( + output_path=output_path, + frames=frames, + fps=batch.fps, + audio=audio, + sample_rate=int(audio_sample_rate), + ) + if not save_ok: + logger.warning("ffmpeg pipe save failed; trying PyAV single-pass save.") + save_ok = self._save_video_with_audio_single_pass( + output_path=output_path, + frames=frames, + fps=batch.fps, + audio=audio, + sample_rate=int(audio_sample_rate), + ) + save_video_time = time.perf_counter() - save_start + if save_ok: + audio_mux_time = 0.0 + else: + logger.warning("Single-pass save failed; falling back to two-step save/mux.") + save_start = time.perf_counter() + imageio.mimsave(output_path, frames, fps=batch.fps, format="mp4") + save_video_time = time.perf_counter() - save_start + mux_start = time.perf_counter() + mux_ok = self._mux_audio(output_path, audio, int(audio_sample_rate)) + audio_mux_time = time.perf_counter() - mux_start + if not mux_ok: + logger.warning("Audio mux failed; saved video without audio.") + else: + save_start = time.perf_counter() + imageio.mimsave(output_path, frames, fps=batch.fps, format="mp4") + save_video_time = time.perf_counter() - save_start + audio_mux_time = 0.0 + logger.info("Saved video to %s", output_path) + + logger.info("VideoSaveStage completed in %.3f s", save_video_time) + if logging_info is not None: + logging_info.add_stage_execution_time("VideoSaveStage", save_video_time) + logger.info("AudioMuxStage completed in %.3f s", audio_mux_time) + if logging_info is not None: + logging_info.add_stage_execution_time("AudioMuxStage", audio_mux_time) + + e2e_time = time.perf_counter() - start_time + logger.info("End-to-end latency: %.2f seconds", e2e_time) result: dict[str, Any] = { "prompts": prompt, @@ -829,6 +895,7 @@ class VideoGenerator: "ltx2_audio_latents": output_batch.extra.get("ltx2_audio_latents"), "size": (target_height, target_width, batch.num_frames), "generation_time": gen_time, + "e2e_latency": e2e_time, "logging_info": logging_info, "trajectory": output_batch.trajectory_latents, "trajectory_timesteps": output_batch.trajectory_timesteps, @@ -900,6 +967,211 @@ class VideoGenerator: f.writeframes(audio_int16.tobytes()) return num_channels + @classmethod + def _save_video_with_audio_single_pass( + cls, + output_path: str, + frames: list[np.ndarray], + fps: int, + audio: torch.Tensor | np.ndarray, + sample_rate: int, + ) -> bool: + """Encode video+audio into MP4 in one pass using PyAV.""" + try: + import av + except ImportError: + logger.warning("PyAV not installed; cannot use single-pass save.") + return False + + if not frames: + return False + + output = None + try: + audio_int16, num_channels = cls._audio_to_int16(audio) + layout = "stereo" if num_channels == 2 else "mono" + output = av.open(output_path, mode="w") + video_stream = output.add_stream("libx264", rate=fps) + video_stream.width = int(frames[0].shape[1]) + video_stream.height = int(frames[0].shape[0]) + video_stream.pix_fmt = "yuv420p" + video_stream.options = { + "preset": "ultrafast", + "tune": "zerolatency", + } + + audio_stream = output.add_stream("aac", rate=sample_rate, layout=layout) + + for frame_np in frames: + vframe = av.VideoFrame.from_ndarray(np.ascontiguousarray(frame_np), format="rgb24") + for packet in video_stream.encode(vframe): + output.mux(packet) + for packet in video_stream.encode(): + output.mux(packet) + + # AAC commonly uses 1024 samples/frame. Pad the tail frame to avoid + # shape errors in PyAV/FFmpeg. + chunk_size = 1024 + for start in range(0, audio_int16.shape[0], chunk_size): + chunk = audio_int16[start:start + chunk_size] + if chunk.shape[0] < chunk_size: + pad = np.zeros((chunk_size - chunk.shape[0], chunk.shape[1]), dtype=chunk.dtype) + chunk = np.concatenate([chunk, pad], axis=0) + chunk_planar = np.ascontiguousarray(chunk.T) + aframe = av.AudioFrame.from_ndarray(chunk_planar, format="s16p", layout=layout) + aframe.sample_rate = sample_rate + for packet in audio_stream.encode(aframe): + output.mux(packet) + for packet in audio_stream.encode(): + output.mux(packet) + + output.close() + return True + except Exception as e: + logger.warning("Single-pass video+audio save failed: %s", e) + if output is not None: + with suppress(Exception): + output.close() + return False + + @staticmethod + def _ffmpeg_encoder_supports_option(ffmpeg_bin: str, codec: str, option_name: str) -> bool: + """Best-effort check whether ffmpeg encoder supports an option.""" + cache_key = (ffmpeg_bin, codec, option_name) + cached = _FFMPEG_ENCODER_OPTION_CACHE.get(cache_key) + if cached is not None: + return cached + + try: + result = subprocess.run( + [ + ffmpeg_bin, + "-hide_banner", + "-h", + f"encoder={codec}", + ], + check=False, + capture_output=True, + text=True, + ) + if result.returncode != 0: + _FFMPEG_ENCODER_OPTION_CACHE[cache_key] = False + return False + haystack = f"{result.stdout}\n{result.stderr}" + supported = f"-{option_name}" in haystack + _FFMPEG_ENCODER_OPTION_CACHE[cache_key] = supported + return supported + except Exception: + _FFMPEG_ENCODER_OPTION_CACHE[cache_key] = False + return False + + @classmethod + def _save_video_with_audio_ffmpeg_pipe( + cls, + output_path: str, + frames: list[np.ndarray], + fps: int, + audio: torch.Tensor | np.ndarray, + sample_rate: int, + ) -> bool: + """Encode video+audio using ffmpeg via rawvideo stdin + WAV input.""" + ffmpeg_bin = shutil.which(os.getenv("FASTVIDEO_FFMPEG_BIN", "ffmpeg")) + if ffmpeg_bin is None: + logger.warning("ffmpeg not found; cannot use ffmpeg pipe save.") + return False + + if not frames: + return False + + height = int(frames[0].shape[0]) + width = int(frames[0].shape[1]) + codec = os.getenv("FASTVIDEO_VIDEO_CODEC", "libx264") + + try: + with tempfile.TemporaryDirectory() as tmpdir: + wav_path = os.path.join(tmpdir, "audio.wav") + try: + cls._write_pcm_wav(wav_path, audio, sample_rate) + except ValueError as e: + logger.warning("Unexpected audio tensor for ffmpeg pipe save: %s", e) + return False + + cmd = [ + ffmpeg_bin, + "-hide_banner", + "-loglevel", + "error", + "-y", + "-f", + "rawvideo", + "-pix_fmt", + "rgb24", + "-s:v", + f"{width}x{height}", + "-r", + str(fps), + "-i", + "pipe:0", + "-i", + wav_path, + "-c:v", + codec, + ] + + if codec.endswith("_nvenc"): + nvenc_options = [ + ("preset", os.getenv("FASTVIDEO_NVENC_PRESET", "p1")), + ("tune", os.getenv("FASTVIDEO_NVENC_TUNE", "ull")), + ("rc", os.getenv("FASTVIDEO_NVENC_RC", "constqp")), + ("qp", os.getenv("FASTVIDEO_NVENC_QP", "28")), + ("bf", os.getenv("FASTVIDEO_NVENC_BF", "0")), + ] + for option_name, option_value in nvenc_options: + if cls._ffmpeg_encoder_supports_option(ffmpeg_bin, codec, option_name): + cmd += [f"-{option_name}", option_value] + else: + cmd += ["-preset", os.getenv("FASTVIDEO_X264_PRESET", "ultrafast")] + + cmd += [ + "-c:a", + "aac", + "-pix_fmt", + os.getenv("FASTVIDEO_OUTPUT_PIX_FMT", "yuv420p"), + "-shortest", + "-movflags", + "+faststart", + output_path, + ] + + proc = subprocess.Popen(cmd, stdin=subprocess.PIPE) + if proc.stdin is None: + proc.kill() + proc.wait() + logger.warning("ffmpeg stdin unavailable.") + return False + if fcntl is not None and hasattr(fcntl, "F_SETPIPE_SZ"): + with suppress(OSError): + fcntl.fcntl(proc.stdin.fileno(), fcntl.F_SETPIPE_SZ, 1048576) + + try: + ts_start = time.perf_counter() + for frame in frames: + proc.stdin.write(frame.tobytes()) + proc.stdin.close() + logger.info("Wrote frames to ffmpeg stdin in %.3f s", time.perf_counter() - ts_start) + rc = proc.wait() + if rc != 0: + logger.warning("ffmpeg pipe save failed with return code %d", rc) + return False + return True + except Exception: + proc.kill() + proc.wait() + raise + except Exception as e: + logger.warning("ffmpeg pipe save failed: %s", e) + return False + @classmethod def _mux_audio( cls, diff --git a/fastvideo/models/vaes/ltx2vae.py b/fastvideo/models/vaes/ltx2vae.py index 5f921dccf..2e510d310 100644 --- a/fastvideo/models/vaes/ltx2vae.py +++ b/fastvideo/models/vaes/ltx2vae.py @@ -7,6 +7,7 @@ from __future__ import annotations import itertools import math +import os from dataclasses import dataclass, replace from enum import Enum from typing import Any, Callable, Iterator, List, NamedTuple, Tuple @@ -19,6 +20,11 @@ from einops import rearrange from fastvideo.models.vaes.common import DiagonalGaussianDistribution +def _is_env_enabled(name: str, default: str = "") -> bool: + value = os.getenv(name, default) + return value.lower() in {"1", "true", "yes", "on"} + + # ============================================================================= # Enums # ============================================================================= @@ -1590,6 +1596,10 @@ class LTX2VideoDecoder(nn.Module): return self.model(sample, timestep=timestep, generator=generator) +def _is_ltx2_vae_codec(name: str, submodule: nn.Module) -> bool: + return name in {"encoder", "decoder"} and isinstance(submodule, (VideoEncoder, VideoDecoder)) + + class LTX2CausalVideoAutoencoder(nn.Module): """ LTX-2 VAE that exposes FastVideo's VAE encode/decode interface. @@ -1599,6 +1609,7 @@ class LTX2CausalVideoAutoencoder(nn.Module): # LTX-2 VAE scale factors TIME_SCALE: int = 8 SPATIAL_SCALE: int = 32 + _compile_conditions = [_is_ltx2_vae_codec] def __init__(self, config: dict[str, Any]): super().__init__() @@ -1606,8 +1617,30 @@ class LTX2CausalVideoAutoencoder(nn.Module): self.encoder = VideoEncoderConfigurator.from_config(config) self.decoder = VideoDecoderConfigurator.from_config(config) self._use_tiling: bool = False + self._use_channels_last_3d: bool = False + + if _is_env_enabled("FASTVIDEO_LTX2_VAE_CHANNELS_LAST_3D", default="1"): + self.enable_channels_last_3d() + + def _as_channels_last_3d(self, tensor: torch.Tensor) -> torch.Tensor: + if not self._use_channels_last_3d or tensor.ndim != 5 or tensor.is_contiguous(memory_format=torch.channels_last_3d): + return tensor + return tensor.contiguous(memory_format=torch.channels_last_3d) + + def enable_channels_last_3d(self) -> None: + """Enable channels-last layout for 3D VAE convolutions.""" + self._use_channels_last_3d = True + self.encoder.to(memory_format=torch.channels_last_3d) + self.decoder.to(memory_format=torch.channels_last_3d) + + def disable_channels_last_3d(self) -> None: + """Restore contiguous layout for 3D VAE convolutions.""" + self._use_channels_last_3d = False + self.encoder.to(memory_format=torch.contiguous_format) + self.decoder.to(memory_format=torch.contiguous_format) def encode(self, x: torch.Tensor) -> DiagonalGaussianDistribution: + x = self._as_channels_last_3d(x) means = self.encoder(x) zeros = torch.zeros_like(means) return DiagonalGaussianDistribution(torch.cat([means, zeros], dim=1), deterministic=True) @@ -1619,6 +1652,7 @@ class LTX2CausalVideoAutoencoder(nn.Module): generator: torch.Generator | None = None, ) -> torch.Tensor: """Decode latents to video, using tiling if enabled.""" + z = self._as_channels_last_3d(z) if self._use_tiling: # Collect all chunks from tiled decode and concatenate chunks = list(self.tiled_decode(z, TilingConfig.default(), timestep, generator)) diff --git a/fastvideo/pipelines/composed_pipeline_base.py b/fastvideo/pipelines/composed_pipeline_base.py index f1269e946..517d959d2 100644 --- a/fastvideo/pipelines/composed_pipeline_base.py +++ b/fastvideo/pipelines/composed_pipeline_base.py @@ -463,6 +463,10 @@ class ComposedPipelineBase(ABC): def add_stage(self, stage_name: str, stage: PipelineStage): assert self.modules is not None, "No modules are registered" + # Preserve the pipeline-unique stage key for structured metrics. + # Multiple stages can share the same class (for example LTX2 main + # denoise and refine denoise), so class-name keys would collide. + stage._pipeline_stage_name = stage_name self._stages.append(stage) self._stage_name_mapping[stage_name] = stage setattr(self, stage_name, stage) diff --git a/fastvideo/pipelines/stages/base.py b/fastvideo/pipelines/stages/base.py index 06b6bee27..4ea9c6f45 100644 --- a/fastvideo/pipelines/stages/base.py +++ b/fastvideo/pipelines/stages/base.py @@ -126,7 +126,9 @@ class PipelineStage(ABC): Returns: The updated batch information after this stage's processing. """ - stage_name = self.__class__.__name__ + stage_class_name = self.__class__.__name__ + stage_key = getattr(self, "_pipeline_stage_name", stage_class_name) + stage_name = f"{stage_key}|{stage_class_name}" # Check if verification is enabled (simple approach for prototype) enable_verification = getattr(fastvideo_args, 'enable_stage_verification', False) @@ -151,7 +153,8 @@ class PipelineStage(ABC): torch.cuda.synchronize() execution_time = time.perf_counter() - start_time logger.info("[%s] Execution completed in %s ms", stage_name, execution_time * 1000) - batch.logging_info.add_stage_execution_time(stage_name, execution_time) + batch.logging_info.add_stage_execution_time(stage_key, execution_time) + batch.logging_info.add_stage_metric(stage_key, "stage_class", stage_class_name) except Exception as e: torch.cuda.synchronize() execution_time = time.perf_counter() - start_time