nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

cleanup

Browse Source
This commit is contained in:
Will Lin
2026-01-21 14:09:40 -08:00
parent beae2943cf
commit deb901f6fc
11 changed files with 3 additions and 245 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/pipelines/stages/decoding.py
+2 -81
View File
@@ -3,7 +3,6 @@
Decoding stage for diffusion pipelines.
"""
import os
import weakref
import torch
@@ -22,45 +21,6 @@ from fastvideo.utils import PRECISION_TO_TYPE
logger = init_logger(__name__)
def _log_non_finite(name: str, tensor: torch.Tensor | None) -> bool:
if tensor is None:
return False
if torch.isfinite(tensor).all():
return False
non_finite = ~torch.isfinite(tensor)
count = int(non_finite.sum().item())
logger.error(
"[VAE] Non-finite detected in %s: count=%d shape=%s dtype=%s",
name,
count,
tuple(tensor.shape),
tensor.dtype,
)
finite_vals = tensor[~non_finite]
if finite_vals.numel() > 0:
logger.error(
"[VAE] %s finite range: min=%s max=%s",
name,
finite_vals.min().item(),
finite_vals.max().item(),
)
return True
def _log_tensor_stats(name: str, tensor: torch.Tensor | None) -> None:
if tensor is None:
return
logger.info(
"[VAE] %s stats: shape=%s dtype=%s min=%s max=%s mean=%s",
name,
tuple(tensor.shape),
tensor.dtype,
tensor.min().item(),
tensor.max().item(),
tensor.mean().item(),
)
class DecodingStage(PipelineStage):
"""
Stage for decoding latent representations into pixel space.
@@ -146,42 +106,16 @@ class DecodingStage(PipelineStage):
Decoded video tensor with shape (batch, channels, frames, height, width),
normalized to [0, 1] range and moved to CPU as float32
"""
debug_vae = os.getenv("LTX2_DEBUG_VAE_LOG", "0") == "1"
debug_vae_abort = os.getenv("LTX2_DEBUG_VAE_ABORT", "0") == "1"
debug_pipeline = os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1"
self.vae = self.vae.to(get_local_torch_device())
latents = latents.to(get_local_torch_device())
# Setup VAE precision
vae_dtype = PRECISION_TO_TYPE[
fastvideo_args.pipeline_config.vae_precision]
disable_vae_autocast = os.getenv("LTX2_DISABLE_VAE_AUTOCAST",
"1") == "1"
vae_autocast_enabled = (vae_dtype != torch.float32
) and not fastvideo_args.disable_autocast and (
not disable_vae_autocast)
if debug_vae:
_log_tensor_stats("latents_in", latents)
if _log_non_finite("latents_in", latents) and debug_vae_abort:
raise RuntimeError("Non-finite latents before VAE decode.")
if debug_pipeline:
_log_tensor_stats("latents_in", latents)
vae_autocast_enabled = (
vae_dtype != torch.float32) and not fastvideo_args.disable_autocast
latents = self._denormalize_latents(latents)
if debug_vae:
_log_tensor_stats("latents_denorm", latents)
if _log_non_finite("latents_denorm", latents) and debug_vae_abort:
raise RuntimeError("Non-finite latents after denormalize.")
if debug_pipeline:
_log_tensor_stats("latents_denorm", latents)
if os.getenv("LTX2_DISABLE_VAE_NOISE", "1") == "1":
decoder = getattr(self.vae, "decoder", None)
if decoder is not None and hasattr(decoder, "decode_noise_scale"):
decoder.decode_noise_scale = 0.0
elif hasattr(self.vae, "decode_noise_scale"):
self.vae.decode_noise_scale = 0.0
# Decode latents
with torch.autocast(device_type="cuda",
@@ -195,21 +129,8 @@ class DecodingStage(PipelineStage):
latents = latents.to(vae_dtype)
image = self.vae.decode(latents)
if debug_vae:
_log_tensor_stats("decoded_raw", image)
if _log_non_finite("decoded_raw", image) and debug_vae_abort:
raise RuntimeError("Non-finite decoded output.")
if debug_pipeline:
_log_tensor_stats("decoded_raw", image)
# Normalize image to [0, 1] range
image = (image / 2 + 0.5).clamp(0, 1)
if debug_vae:
_log_tensor_stats("decoded_norm", image)
if _log_non_finite("decoded_norm", image) and debug_vae_abort:
raise RuntimeError("Non-finite decoded output after clamp.")
if debug_pipeline:
_log_tensor_stats("decoded_norm", image)
return image
@torch.no_grad()
fastvideo/pipelines/stages/ltx2_denoising.py
-104
View File
@@ -37,58 +37,6 @@ DISTILLED_SIGMA_VALUES = [
logger = init_logger(__name__)
def _log_non_finite(
name: str,
tensor: torch.Tensor | None,
step_index: int | None = None,
) -> bool:
if tensor is None:
return False
if torch.isfinite(tensor).all():
return False
non_finite = ~torch.isfinite(tensor)
count = int(non_finite.sum().item())
step_info = f" step={step_index}" if step_index is not None else ""
logger.error(
"[LTX2] Non-finite detected in %s%s: count=%d shape=%s dtype=%s",
name,
step_info,
count,
tuple(tensor.shape),
tensor.dtype,
)
finite_vals = tensor[~non_finite]
if finite_vals.numel() > 0:
logger.error(
"[LTX2] %s%s finite range: min=%s max=%s",
name,
step_info,
finite_vals.min().item(),
finite_vals.max().item(),
)
return True
def _log_tensor_stats(
name: str,
tensor: torch.Tensor | None,
step_index: int | None = None,
) -> None:
if tensor is None:
return
step_info = f" step={step_index}" if step_index is not None else ""
logger.info(
"[LTX2] %s%s stats: shape=%s dtype=%s min=%s max=%s mean=%s",
name,
step_info,
tuple(tensor.shape),
tensor.dtype,
tensor.min().item(),
tensor.max().item(),
tensor.mean().item(),
)
def _ltx2_sigmas(
steps: int,
latent: torch.Tensor | None,
@@ -141,11 +89,6 @@ class LTX2DenoisingStage(PipelineStage):
batch: ForwardBatch,
fastvideo_args: FastVideoArgs,
) -> ForwardBatch:
debug_nans = os.getenv("LTX2_DEBUG_NANS", "0") == "1"
abort_on_nans = os.getenv("LTX2_DEBUG_NANS_ABORT", "0") == "1"
verbose_logs = os.getenv("LTX2_DEBUG_DENOISE_LOG", "0") == "1"
step_log_interval = int(
os.getenv("LTX2_DEBUG_DENOISE_LOG_INTERVAL", "1"))
if batch.latents is None:
raise ValueError("Latents must be provided before denoising.")
@@ -278,34 +221,9 @@ class LTX2DenoisingStage(PipelineStage):
tuple(latents.shape),
)
if debug_nans:
if _log_non_finite("latents_start", latents) and abort_on_nans:
raise RuntimeError("Non-finite latents before denoising.")
if _log_non_finite("prompt_embeds",
prompt_embeds) and abort_on_nans:
raise RuntimeError("Non-finite prompt embeddings.")
if (neg_prompt_embeds is not None and _log_non_finite(
"negative_prompt_embeds", neg_prompt_embeds)
and abort_on_nans):
raise RuntimeError("Non-finite negative prompt embeddings.")
if verbose_logs:
_log_tensor_stats("latents_start", latents)
_log_tensor_stats("prompt_embeds", prompt_embeds)
if neg_prompt_embeds is not None:
_log_tensor_stats("negative_prompt_embeds", neg_prompt_embeds)
for step_index in tqdm(range(len(sigmas) - 1)):
sigma = sigmas[step_index]
sigma_next = sigmas[step_index + 1]
if verbose_logs and step_log_interval > 0 and (
step_index % step_log_interval == 0):
logger.info(
"[LTX2] Step %d/%d sigma=%s sigma_next=%s",
step_index + 1,
len(sigmas) - 1,
sigma.item(),
sigma_next.item(),
)
timestep = timestep_template * sigma
audio_timestep = (audio_timestep_template * sigma
if audio_timestep_template is not None else None)
@@ -356,22 +274,6 @@ class LTX2DenoisingStage(PipelineStage):
pos_audio = pos_audio + (batch.guidance_scale -
1) * (pos_audio - neg_audio)
if debug_nans:
if (_log_non_finite("pos_denoised", pos_denoised, step_index)
and abort_on_nans):
raise RuntimeError(
f"Non-finite pos_denoised at step {step_index}.")
if (neg_prompt_embeds is not None and _log_non_finite(
"neg_denoised", neg_denoised, step_index)
and abort_on_nans):
raise RuntimeError(
f"Non-finite neg_denoised at step {step_index}.")
if verbose_logs and step_log_interval > 0 and (
step_index % step_log_interval == 0):
_log_tensor_stats("pos_denoised", pos_denoised, step_index)
if neg_prompt_embeds is not None:
_log_tensor_stats("neg_denoised", neg_denoised, step_index)
sigma_value = sigma.to(torch.float32) if isinstance(
sigma, torch.Tensor) else torch.tensor(
float(sigma),
@@ -389,12 +291,6 @@ class LTX2DenoisingStage(PipelineStage):
audio_latents = (audio_latents.float() +
audio_velocity.float() * dt).to(
audio_latents.dtype)
if (debug_nans and _log_non_finite("latents", latents, step_index)
and abort_on_nans):
raise RuntimeError(f"Non-finite latents at step {step_index}.")
if verbose_logs and step_log_interval > 0 and (
step_index % step_log_interval == 0):
_log_tensor_stats("latents", latents, step_index)
batch.latents = latents
batch.extra["ltx2_audio_latents"] = audio_latents
fastvideo/pipelines/stages/ltx2_latent_preparation.py
-22
View File
@@ -3,7 +3,6 @@
Latent preparation stage for LTX-2 pipelines.
"""
import os
from pathlib import Path
import torch
@@ -20,19 +19,6 @@ from fastvideo.pipelines.stages.validators import VerificationResult
logger = init_logger(__name__)
def _debug_log_line(message: str) -> None:
if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") != "1":
return
log_path = os.getenv("LTX2_PIPELINE_DEBUG_PATH", "")
if not log_path:
return
log_dir = os.path.dirname(log_path)
if log_dir:
os.makedirs(log_dir, exist_ok=True)
with open(log_path, "a", encoding="utf-8") as f:
f.write(message + "\n")
class LTX2LatentPreparationStage(PipelineStage):
"""Prepare initial LTX-2 latents without relying on a diffusers scheduler."""
@@ -102,14 +88,12 @@ class LTX2LatentPreparationStage(PipelineStage):
f"You have passed a list of generators of length {len(generator)}, "
f"but requested an effective batch size of {batch_size}.")
latents_source = "provided"
if latents is None:
if latent_path:
loaded_latents = self._load_initial_latent(
latent_path, device, dtype)
if loaded_latents is not None:
latents = loaded_latents
latents_source = "loaded"
else:
latents = randn_tensor(
shape,
@@ -118,7 +102,6 @@ class LTX2LatentPreparationStage(PipelineStage):
dtype=dtype,
)
self._save_initial_latent(latent_path, latents)
latents_source = "generated"
else:
latents = randn_tensor(
shape,
@@ -126,14 +109,9 @@ class LTX2LatentPreparationStage(PipelineStage):
device=device,
dtype=dtype,
)
latents_source = "generated"
else:
latents = latents.to(device)
if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1":
_debug_log_line(
f"fastvideo:initial_latent:source={latents_source} sum={latents.float().sum().item():.6f}"
)
batch.latents = latents
batch.raw_latent_shape = shape
return batch
fastvideo/pipelines/stages/text_encoding.py
-37
View File
@@ -5,34 +5,17 @@ Prompt encoding stages for diffusion pipelines.
This module contains implementations of prompt encoding stages for diffusion pipelines.
"""
import os
import torch
from typing import Any
from fastvideo.distributed import get_local_torch_device
from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.forward_context import set_forward_context
from fastvideo.logger import init_logger
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
from fastvideo.pipelines.stages.base import PipelineStage
from fastvideo.pipelines.stages.validators import StageValidators as V
from fastvideo.pipelines.stages.validators import VerificationResult
logger = init_logger(__name__)
def _debug_log_line(message: str) -> None:
if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") != "1":
return
log_path = os.getenv("LTX2_PIPELINE_DEBUG_PATH", "")
if not log_path:
return
log_dir = os.path.dirname(log_path)
if log_dir:
os.makedirs(log_dir, exist_ok=True)
with open(log_path, "a", encoding="utf-8") as f:
f.write(message + "\n")
class TextEncodingStage(PipelineStage):
"""
@@ -114,15 +97,6 @@ class TextEncodingStage(PipelineStage):
for nm in neg_masks_list:
batch.negative_attention_mask.append(nm)
if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1":
prompt_sum = prompt_embeds_list[0].float().sum().item()
neg_sum = 0.0
if batch.do_classifier_free_guidance and neg_embeds_list:
neg_sum = neg_embeds_list[0].float().sum().item()
_debug_log_line(
f"fastvideo:text_encode:prompt_sum={prompt_sum:.6f} "
f"neg_sum={neg_sum:.6f}")
return batch
def verify_input(self, batch: ForwardBatch,
@@ -278,17 +252,6 @@ class TextEncodingStage(PipelineStage):
input_ids = text_inputs["input_ids"]
attention_mask = text_inputs["attention_mask"]
if os.getenv("LTX2_PIPELINE_DEBUG_LOG", "0") == "1":
input_sum = input_ids.sum().item()
mask_sum = attention_mask.sum().item()
input_head = input_ids[0, :16].tolist()
mask_head = attention_mask[0, :16].tolist()
_debug_log_line(
"fastvideo:text_inputs"
f":encoder_index={i} input_sum={input_sum:.6f} "
f"mask_sum={mask_sum:.6f} input_shape={tuple(input_ids.shape)} "
f"mask_shape={tuple(attention_mask.shape)} "
f"input_head={input_head} mask_head={mask_head}")
with set_forward_context(current_timestep=0, attn_metadata=None):
outputs = text_encoder(
fastvideo/tests/encoders/test_ltx2_gemma_encoder.py → tests/local_tests/encoders/test_ltx2_gemma_encoder.py
View File
fastvideo/tests/encoders/test_ltx2_gemma_parity.py → tests/local_tests/encoders/test_ltx2_gemma_parity.py
View File
fastvideo/tests/transformers/test_ltx2.py → tests/local_tests/transformers/test_ltx2.py
View File
fastvideo/tests/transformers/test_ltx2_audio.py → tests/local_tests/transformers/test_ltx2_audio.py
+1 -1
View File
@@ -24,7 +24,7 @@ from fastvideo.fastvideo_args import FastVideoArgs
from fastvideo.forward_context import set_forward_context
from fastvideo.models.dits.ltx2 import Modality as FastVideoModality
from fastvideo.models.loader.component_loader import TransformerLoader
from fastvideo.tests.transformers.test_ltx2 import (
from .test_ltx2 import (
_attach_block_detail_logging,
_attach_block_sum_logging,
_infer_patch_params,
fastvideo/tests/vaes/test_ltx2_audio_vae.py → tests/local_tests/vaes/test_ltx2_audio_vae.py
View File
fastvideo/tests/vaes/test_ltx2_vae.py → tests/local_tests/vaes/test_ltx2_vae.py
View File
fastvideo/tests/vaes/test_ltx2_vae_official.py → tests/local_tests/vaes/test_ltx2_vae_official.py
View File