300 lines
11 KiB
Python
300 lines
11 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Inspired by SGLang: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/model_executor/forward_batch_info.py
|
|
"""
|
|
Data structures for functional pipeline processing.
|
|
|
|
This module defines the dataclasses used to pass state between pipeline components
|
|
in a functional manner, reducing the need for explicit parameter passing.
|
|
"""
|
|
|
|
import pprint
|
|
from dataclasses import asdict, dataclass, field
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
import PIL.Image
|
|
import torch
|
|
|
|
if TYPE_CHECKING:
|
|
from torchcodec.decoders import VideoDecoder
|
|
|
|
import time
|
|
from collections import OrderedDict
|
|
|
|
from fastvideo.attention import AttentionMetadata
|
|
|
|
|
|
class PipelineLoggingInfo:
|
|
"""Simple approach using OrderedDict to track stage metrics."""
|
|
|
|
def __init__(self):
|
|
# OrderedDict preserves insertion order and allows easy access
|
|
self.stages: OrderedDict[str, dict[str, Any]] = OrderedDict()
|
|
|
|
def add_stage_execution_time(self, stage_name: str, execution_time: float):
|
|
"""Add execution time for a stage."""
|
|
if stage_name not in self.stages:
|
|
self.stages[stage_name] = {}
|
|
self.stages[stage_name]['execution_time'] = execution_time
|
|
self.stages[stage_name]['timestamp'] = time.time()
|
|
|
|
def add_stage_metric(self, stage_name: str, metric_name: str, value: Any):
|
|
"""Add any metric for a stage."""
|
|
if stage_name not in self.stages:
|
|
self.stages[stage_name] = {}
|
|
self.stages[stage_name][metric_name] = value
|
|
|
|
def get_stage_info(self, stage_name: str) -> dict[str, Any]:
|
|
"""Get all info for a specific stage."""
|
|
return self.stages.get(stage_name, {})
|
|
|
|
def get_execution_order(self) -> list[str]:
|
|
"""Get stages in execution order."""
|
|
return list(self.stages.keys())
|
|
|
|
def get_total_execution_time(self) -> float:
|
|
"""Get total pipeline execution time."""
|
|
return sum(stage.get('execution_time', 0) for stage in self.stages.values())
|
|
|
|
|
|
@dataclass
|
|
class ForwardBatch:
|
|
"""
|
|
Complete state passed through the pipeline execution.
|
|
|
|
This dataclass contains all information needed during the diffusion pipeline
|
|
execution, allowing methods to update specific components without needing
|
|
to manage numerous individual parameters.
|
|
"""
|
|
# TODO(will): double check that args are separate from fastvideo_args
|
|
# properly. Also maybe think about providing an abstraction for pipeline
|
|
# specific arguments.
|
|
data_type: str
|
|
|
|
generator: torch.Generator | list[torch.Generator] | None = None
|
|
|
|
# Image inputs
|
|
image_path: str | None = None
|
|
image_embeds: list[torch.Tensor] = field(default_factory=list)
|
|
pil_image: torch.Tensor | PIL.Image.Image | None = None
|
|
preprocessed_image: torch.Tensor | None = None
|
|
|
|
# Text inputs
|
|
prompt: str | list[str] | None = None
|
|
negative_prompt: str | list[str] | None = None
|
|
prompt_path: str | None = None
|
|
output_path: str = "outputs/"
|
|
output_video_name: str | None = None
|
|
|
|
# Video inputs
|
|
video_path: str | None = None
|
|
video_latent: torch.Tensor | None = None
|
|
|
|
# Refine inputs (LongCat)
|
|
refine_from: str | None = None
|
|
t_thresh: float = 0.5
|
|
spatial_refine_only: bool = False
|
|
num_cond_frames: int = 0
|
|
stage1_video: list[PIL.Image.Image] | None = None # Loaded frames from refine_from
|
|
|
|
# Primary encoder embeddings
|
|
prompt_embeds: list[torch.Tensor] = field(default_factory=list)
|
|
negative_prompt_embeds: list[torch.Tensor] | None = None
|
|
prompt_attention_mask: list[torch.Tensor] | None = None
|
|
negative_attention_mask: list[torch.Tensor] | None = None
|
|
clip_embedding_pos: list[torch.Tensor] | None = None
|
|
clip_embedding_neg: list[torch.Tensor] | None = None
|
|
|
|
# Additional text-related parameters
|
|
max_sequence_length: int | None = None
|
|
prompt_template: dict[str, Any] | None = None
|
|
do_classifier_free_guidance: bool = False
|
|
|
|
# Batch info
|
|
batch_size: int | None = None
|
|
num_videos_per_prompt: int = 1
|
|
seed: int | None = None
|
|
seeds: list[int] | None = None
|
|
|
|
# Tracking if embeddings are already processed
|
|
is_prompt_processed: bool = False
|
|
|
|
# Latent tensors
|
|
latents: torch.Tensor | None = None
|
|
lq_latents: torch.Tensor | None = None
|
|
raw_latent_shape: tuple[int, ...] | None = None
|
|
noise_pred: torch.Tensor | None = None
|
|
image_latent: torch.Tensor | None = None
|
|
|
|
# Action control inputs (Matrix-Game)
|
|
mouse_cond: torch.Tensor | None = None # Shape: (B, T, 2)
|
|
keyboard_cond: torch.Tensor | None = None # Shape: (B, T, K)
|
|
grid_sizes: torch.Tensor | None = None # Shape: (3,) [F,H,W]
|
|
|
|
# Camera control inputs (HYWorld)
|
|
pose: str | None = None # Camera trajectory: pose string (e.g., 'w-31') or JSON file path
|
|
|
|
# Camera/action control inputs (GameCraft)
|
|
camera_states: torch.Tensor | None = None # Plücker coordinates [B, T, 6, H, W]
|
|
gt_latents: torch.Tensor | None = None # Ground truth latents for conditioning [B, 16, T, H, W]
|
|
conditioning_mask: torch.Tensor | None = None # Mask for conditioning [B, 1, T, H, W]
|
|
camera_trajectory: str | None = None # Camera trajectory file/identifier
|
|
action_list: list[str] | None = None # List of actions (e.g., ['forward', 'left'])
|
|
action_speed_list: list[float] | None = None # Speed for each action
|
|
# Camera control inputs (LingBotWorld)
|
|
c2ws_plucker_emb: torch.Tensor | None = None # Plucker embedding: [B, C, F_lat, H_lat, W_lat]
|
|
|
|
# Latent dimensions
|
|
height_latents: list[int] | int | None = None
|
|
width_latents: list[int] | int | None = None
|
|
num_frames: list[int] | int = 1 # Default for image models
|
|
|
|
# Original dimensions (before VAE scaling)
|
|
height: list[int] | int | None = None
|
|
width: list[int] | int | None = None
|
|
height_sr: list[int] | int | None = None
|
|
width_sr: list[int] | int | None = None
|
|
fps: list[int] | int | None = None
|
|
|
|
# Timesteps
|
|
timesteps: torch.Tensor | None = None
|
|
timestep: torch.Tensor | float | int | None = None
|
|
step_index: int | None = None
|
|
boundary_ratio: float | None = None
|
|
|
|
# Scheduler parameters
|
|
num_inference_steps: int = 50
|
|
num_inference_steps_sr: int = 50
|
|
guidance_scale: float = 1.0
|
|
guidance_scale_2: float | None = None
|
|
guidance_rescale: float = 0.0
|
|
eta: float = 0.0
|
|
sigmas: list[float] | None = None
|
|
|
|
# LTX-2 multi-modal CFG parameters
|
|
ltx2_cfg_scale_video: float = 1.0
|
|
ltx2_cfg_scale_audio: float = 1.0
|
|
ltx2_modality_scale_video: float = 1.0
|
|
ltx2_modality_scale_audio: float = 1.0
|
|
ltx2_rescale_scale: float = 0.0
|
|
# STG (Spatio-Temporal Guidance) parameters
|
|
ltx2_stg_scale_video: float = 0.0
|
|
ltx2_stg_scale_audio: float = 0.0
|
|
ltx2_stg_blocks_video: list[int] = field(default_factory=list)
|
|
ltx2_stg_blocks_audio: list[int] = field(default_factory=list)
|
|
|
|
n_tokens: int | None = None
|
|
|
|
# Other parameters that may be needed by specific schedulers
|
|
extra_step_kwargs: dict[str, Any] = field(default_factory=dict)
|
|
|
|
# Component modules (populated by the pipeline)
|
|
modules: dict[str, Any] = field(default_factory=dict)
|
|
|
|
# Final output (after pipeline completion)
|
|
output: torch.Tensor | None = None
|
|
return_trajectory_latents: bool = False
|
|
return_trajectory_decoded: bool = False
|
|
trajectory_timesteps: list[torch.Tensor] | None = None
|
|
trajectory_latents: torch.Tensor | None = None
|
|
trajectory_decoded: list[torch.Tensor] | None = None
|
|
|
|
# Extra parameters that might be needed by specific pipeline implementations
|
|
extra: dict[str, Any] = field(default_factory=dict)
|
|
|
|
# Misc
|
|
save_video: bool = True
|
|
return_frames: bool = False
|
|
|
|
is_cfg_negative: bool = False
|
|
|
|
# VSA parameters
|
|
VSA_sparsity: float = 0.0
|
|
|
|
# Logging info
|
|
logging_info: PipelineLoggingInfo = field(default_factory=PipelineLoggingInfo)
|
|
|
|
def __post_init__(self):
|
|
"""Initialize dependent fields after dataclass initialization."""
|
|
|
|
# Enable CFG for standard guidance_scale and LTX-2 text CFG scales.
|
|
ltx2_text_cfg_enabled = (self.ltx2_cfg_scale_video != 1.0 or self.ltx2_cfg_scale_audio != 1.0)
|
|
if self.guidance_scale > 1.0 or ltx2_text_cfg_enabled:
|
|
self.do_classifier_free_guidance = True
|
|
if self.negative_prompt_embeds is None:
|
|
self.negative_prompt_embeds = []
|
|
if self.guidance_scale_2 is None:
|
|
self.guidance_scale_2 = self.guidance_scale
|
|
|
|
def __str__(self):
|
|
return pprint.pformat(asdict(self), indent=2, width=120)
|
|
|
|
|
|
@dataclass
|
|
class TrainingBatch:
|
|
current_timestep: int = 0
|
|
current_vsa_sparsity: float = 0.0
|
|
|
|
# Dataloader batch outputs
|
|
latents: torch.Tensor | None = None
|
|
raw_latent_shape: tuple[int, ...] | None = None
|
|
noise_latents: torch.Tensor | None = None
|
|
encoder_hidden_states: torch.Tensor | None = None
|
|
encoder_attention_mask: torch.Tensor | None = None
|
|
# LTX related audio inputs
|
|
audio_latents: torch.Tensor | None = None
|
|
audio_noisy_model_input: torch.Tensor | None = None
|
|
audio_timesteps: torch.Tensor | None = None
|
|
audio_noise: torch.Tensor | None = None
|
|
audio_encoder_hidden_states: torch.Tensor | None = None
|
|
audio_encoder_attention_mask: torch.Tensor | None = None
|
|
conditioning_mask: torch.Tensor | None = None
|
|
# i2v
|
|
preprocessed_image: torch.Tensor | None = None
|
|
image_embeds: torch.Tensor | None = None
|
|
image_latents: torch.Tensor | None = None
|
|
infos: list[dict[str, Any]] | None = None
|
|
mask_lat_size: torch.Tensor | None = None
|
|
|
|
# ODE trajectory supervision
|
|
trajectory_latents: torch.Tensor | None = None
|
|
trajectory_timesteps: torch.Tensor | None = None
|
|
|
|
# Transformer inputs
|
|
noisy_model_input: torch.Tensor | None = None
|
|
timesteps: torch.Tensor | None = None
|
|
sigmas: torch.Tensor | None = None
|
|
noise: torch.Tensor | None = None
|
|
|
|
attn_metadata_vsa: AttentionMetadata | None = None
|
|
attn_metadata: AttentionMetadata | None = None
|
|
|
|
# input kwargs
|
|
input_kwargs: dict[str, Any] | None = None
|
|
|
|
# Training loss
|
|
loss: torch.Tensor | None = None
|
|
|
|
# Training outputs
|
|
total_loss: float | None = None
|
|
grad_norm: float | None = None
|
|
|
|
# Distillation-specific attributes
|
|
encoder_hidden_states_neg: torch.Tensor | None = None
|
|
encoder_attention_mask_neg: torch.Tensor | None = None
|
|
conditional_dict: dict[str, Any] | None = None
|
|
unconditional_dict: dict[str, Any] | None = None
|
|
|
|
# Distillation losses
|
|
generator_loss: float = 0.0
|
|
fake_score_loss: float = 0.0
|
|
|
|
dmd_latent_vis_dict: dict[str, Any] = field(default_factory=dict)
|
|
latent_vis_dict: dict[str, Any] = field(default_factory=dict)
|
|
fake_score_latent_vis_dict: dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
|
@dataclass
|
|
class PreprocessBatch(ForwardBatch):
|
|
video_loader: list["VideoDecoder"] | list[str] = field(default_factory=list)
|
|
video_file_name: list[str] = field(default_factory=list)
|