[perf] Dreamverse 14/14: Add LTX2 profile speedups (#1337)
Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com> 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>
This commit is contained in:
Junda Su
co-authored by
SolitaryThinker
Matthew Noto
XOR-op
Zhang Peiyuan
SolitaryThinker
Matthew Noto
XOR-op
Zhang Peiyuan
parent
cb12e88713
commit
d2ffec5cce
@@ -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=<substring> 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()
|
||||
@@ -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,
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user