500 lines
19 KiB
Python
500 lines
19 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
from dataclasses import dataclass, field, fields
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from fastvideo.logger import init_logger
|
|
from fastvideo.utils import StoreBoolean
|
|
|
|
if TYPE_CHECKING:
|
|
from fastvideo.api.schema import ContinuationState
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class SamplingParam:
|
|
"""
|
|
Sampling parameters for video generation.
|
|
"""
|
|
# All fields below are copied from ForwardBatch
|
|
data_type: str = "video"
|
|
|
|
# Image inputs
|
|
image_path: str | None = None
|
|
pil_image: Any | None = None
|
|
last_image: Any | None = None
|
|
references: list[Any] | None = None
|
|
|
|
# Video inputs
|
|
video_path: str | None = None
|
|
|
|
# Wan-Animate driving inputs: paths to the *preprocessed* artifacts the
|
|
# official preprocessing pipeline produces (src_pose.mp4, src_face.mp4;
|
|
# replace mode adds src_bg.mp4 + src_mask.mp4). `animate_mode` selects
|
|
# animation (character on the reference image's background) vs replace
|
|
# (character composited into the background video).
|
|
pose_video_path: str | None = None
|
|
face_video_path: str | None = None
|
|
background_video_path: str | None = None
|
|
mask_video_path: str | None = None
|
|
animate_mode: str | None = None
|
|
|
|
# Audio inputs. `audio_path` drives audio-conditioned generation (Wan S2V):
|
|
# the waveform is encoded and cross-attended to per video frame. Distinct
|
|
# from the Stable Audio `init_audio` fields below, which seed audio output.
|
|
audio_path: str | None = None
|
|
|
|
# Optional pre-generated diffusion latents. Used by parity/debug harnesses
|
|
# and advanced callers that need deterministic latent reuse.
|
|
latents: Any | None = None
|
|
audio_latents: Any | None = None
|
|
|
|
# Action control inputs (Matrix-Game)
|
|
mouse_cond: Any | None = None # Shape: (B, T, 2)
|
|
keyboard_cond: Any | None = None # Shape: (B, T, K)
|
|
grid_sizes: Any | 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
|
|
prompt_attention_mask: list = field(default_factory=list)
|
|
negative_attention_mask: list = field(default_factory=list)
|
|
|
|
# Camera control inputs (LingBotWorld and LingBotWorld2)
|
|
c2ws_plucker_emb: Any | None = None # Plucker embedding: [B, C, F_lat, H_lat, W_lat]
|
|
action_path: str | None = None # Directory containing poses.npy and intrinsics.npy
|
|
|
|
# Refine / conditioning-frame inputs
|
|
t_thresh: float = 0.5 # Threshold for timestep scheduling in refinement
|
|
num_cond_frames: int = 0 # Number of conditioning frames
|
|
|
|
# Text inputs
|
|
prompt: str | list[str] | None = None
|
|
negative_prompt: str = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards"
|
|
max_sequence_length: int | None = None
|
|
prompt_path: str | None = None
|
|
output_path: str = "outputs/"
|
|
output_video_name: str | None = None
|
|
|
|
# Batch info
|
|
num_videos_per_prompt: int = 1
|
|
seed: int = 1024
|
|
|
|
# Original dimensions (before VAE scaling)
|
|
num_frames: int = 125
|
|
height: int = 720
|
|
width: int = 1280
|
|
height_sr: int = 1072
|
|
width_sr: int = 1920
|
|
fps: int = 24
|
|
|
|
# Denoising parameters
|
|
num_inference_steps: int = 50
|
|
num_inference_steps_sr: int = 50
|
|
guidance_scale: float = 1.0
|
|
batch_cfg: bool = False
|
|
guidance_scale_2: float | None = None
|
|
# Embedded guidance (FLUX): do not treat ``guidance_scale > 1`` as classic CFG.
|
|
use_embedded_guidance: bool = False
|
|
# Diffusers-style true CFG for FLUX when > 1 (requires negative prompt encoding).
|
|
true_cfg_scale: float = 1.0
|
|
guidance_rescale: float = 0.0
|
|
boundary_ratio: float | None = None
|
|
sigmas: list[float] | None = None
|
|
|
|
# Helios autoregressive spatial-pyramid sampling.
|
|
pyramid_num_inference_steps_list: list[int] | None = None
|
|
history_sizes: list[int] | None = None
|
|
num_latent_frames_per_chunk: int = 9
|
|
keep_first_frame: bool = True
|
|
is_skip_first_chunk: bool = False
|
|
# Diffusers-signature compatibility only: the pinned Helios-Distilled model
|
|
# index declares ``is_cfg_zero_star=false``, so neither value changes sampling.
|
|
use_zero_init: bool = True
|
|
zero_steps: int = 1
|
|
is_amplify_first_chunk: bool = False
|
|
|
|
# TeaCache parameters
|
|
enable_teacache: bool = False
|
|
|
|
# GEN3C camera control
|
|
trajectory_type: str | None = None
|
|
movement_distance: float | None = None
|
|
camera_rotation: str | None = None
|
|
|
|
# LTX-2 multi-modal CFG and STG.
|
|
# Class-level defaults match the *distilled* LTX-2 schedule
|
|
# (mirrors ``FastVideo-internal/.../LTX2DistilledSamplingParam``):
|
|
# the distilled model expects neutral guidance scales — modality 1,
|
|
# rescale 0, STG 0 — and explicit-CFG callers (full LTX-2) opt back
|
|
# in by selecting the ``LTX2_BASE`` preset, which overrides these
|
|
# to mod=3.0 / rescale=0.7 / stg=1.0 in its ``defaults`` dict.
|
|
# cfg_scale defaults stay at 1.0 (CFG off) so
|
|
# ``ForwardBatch.__post_init__`` doesn't force CFG on non-LTX-2
|
|
# models that never override these fields.
|
|
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
|
|
ltx2_stg_scale_video: float = 0.0
|
|
ltx2_stg_scale_audio: float = 0.0
|
|
ltx2_stg_blocks_video: list[int] = field(default_factory=lambda: [29])
|
|
ltx2_stg_blocks_audio: list[int] = field(default_factory=lambda: [29])
|
|
# LTX-2.5 distilled stage 1 uses the official variance-preserving
|
|
# ancestral Euler update. Older LTX generations remain deterministic.
|
|
ltx2_use_ancestral_sampler: bool = False
|
|
|
|
# LTX-2 image / video / continuation conditioning. These flow from
|
|
# generate_video(...) kwargs through ``sampling_param.update(kwargs)``
|
|
# onto the ForwardBatch fields of the same name. ``ltx2_image_crf``
|
|
# gates the conditioning-image H.264 re-encode; the streaming
|
|
# session controller passes ``ltx2_image_crf=0.0`` because it
|
|
# conditions on already-decoded VAE-quality frames.
|
|
ltx2_images: list[tuple[str, int, float]] | None = None
|
|
ltx2_image_crf: float = 33.0
|
|
ltx2_conditioning_latent_stage1: Any | None = None
|
|
ltx2_conditioning_latent_stage2: Any | None = None
|
|
ltx2_video_conditions: list[tuple[list[str], int, float]] | None = None
|
|
|
|
# Stable Audio (T2A): clip start/end in seconds. Honored by
|
|
# `StableAudioConditioningStage` + `StableAudioDecodingStage`. Other
|
|
# families ignore them.
|
|
audio_start_in_s: float | None = None
|
|
audio_end_in_s: float | None = None
|
|
|
|
# Stable Audio audio-to-audio (variation):
|
|
# `init_audio` -- a path or `[B, C, samples]` waveform at the model
|
|
# sample rate; the pipeline encodes it via the VAE
|
|
# and uses it as the starting latent.
|
|
# `init_audio_strength` -- 0..1, higher = closer to the reference
|
|
# (matches the convention of Stability's
|
|
# commercial Stable Audio 2.0 UI). 1.0 ~=
|
|
# VAE round-trip, 0.0 ~= plain T2A.
|
|
# `init_noise_level` -- legacy raw `sigma_max` override (0.3..500,
|
|
# higher = more freedom). Kept for callers
|
|
# that already use it; prefer `init_audio_strength`.
|
|
init_audio: Any = None
|
|
init_audio_strength: float | None = None
|
|
init_noise_level: float | None = None
|
|
|
|
# Stable Audio inpainting (RePaint-style): `inpaint_audio` is the
|
|
# reference clip, `inpaint_mask` is a [samples] tensor in {0, 1} where
|
|
# 1 means *keep the reference* and 0 means *regenerate*.
|
|
inpaint_audio: Any = None
|
|
inpaint_mask: Any = None
|
|
|
|
# Continuation state carried across streaming/multi-segment calls.
|
|
continuation_state: ContinuationState | None = None
|
|
# When True, the pipeline returns a ContinuationState on the result so
|
|
# the caller can resume from the generated segment.
|
|
return_continuation_state: bool = False
|
|
|
|
# Misc
|
|
save_video: bool = True
|
|
return_frames: bool = True
|
|
return_samples: bool = field(default=False, kw_only=True)
|
|
return_trajectory_latents: bool = False # returns all latents for each timestep
|
|
return_trajectory_decoded: bool = False # returns decoded latents for each timestep
|
|
|
|
def __post_init__(self) -> None:
|
|
self.data_type = "video" if self.num_frames > 1 else "image"
|
|
|
|
def check_sampling_param(self) -> None:
|
|
if self.prompt_path and not self.prompt_path.endswith(".txt"):
|
|
raise ValueError("prompt_path must be a txt file")
|
|
|
|
def update(self, source_dict: dict[str, Any]) -> None:
|
|
valid_fields = {f.name for f in fields(self)}
|
|
unknown = [key for key in source_dict if key not in valid_fields]
|
|
if unknown:
|
|
raise ValueError(f"{type(self).__name__}.update() received unknown field(s): "
|
|
f"{sorted(unknown)}. All kwargs must correspond to declared "
|
|
f"SamplingParam fields. If a kwarg is meant to flow into "
|
|
f"ForwardBatch.extra (e.g. LTX2 audio conditioning), route it "
|
|
f"via VideoGenerator._BATCH_EXTRA_PASSTHROUGH_KEYS instead.")
|
|
for key, value in source_dict.items():
|
|
setattr(self, key, value)
|
|
|
|
self.__post_init__()
|
|
|
|
@classmethod
|
|
def from_pretrained(cls, model_path: str) -> SamplingParam:
|
|
sampling_param = cls._from_preset(model_path)
|
|
if sampling_param is not None:
|
|
return sampling_param
|
|
|
|
logger.warning(
|
|
"Couldn't find a preset for %s."
|
|
" Using the default sampling param.",
|
|
model_path,
|
|
)
|
|
return cls()
|
|
|
|
@classmethod
|
|
def _from_preset(
|
|
cls,
|
|
model_path: str,
|
|
) -> SamplingParam | None:
|
|
"""Build a SamplingParam from preset defaults.
|
|
|
|
Returns ``None`` when no preset is configured for
|
|
*model_path*, letting the caller fall back to the legacy
|
|
subclass lookup.
|
|
"""
|
|
from fastvideo.registry import get_preset_selection
|
|
|
|
try:
|
|
preset_name, model_family = get_preset_selection(model_path)
|
|
except (ValueError, RuntimeError):
|
|
return None
|
|
if preset_name is None or model_family is None:
|
|
return None
|
|
|
|
from fastvideo.api.presets import get_preset
|
|
|
|
preset = get_preset(preset_name, model_family)
|
|
sp = cls()
|
|
valid_fields = {f.name for f in fields(cls)}
|
|
for key, value in preset.defaults.items():
|
|
if key in valid_fields:
|
|
setattr(sp, key, copy.deepcopy(value))
|
|
sp.__post_init__()
|
|
return sp
|
|
|
|
@staticmethod
|
|
def add_cli_args(parser: Any) -> Any:
|
|
"""Add CLI arguments for SamplingParam fields"""
|
|
parser.add_argument(
|
|
"--prompt",
|
|
type=str,
|
|
default=SamplingParam.prompt,
|
|
help="Text prompt for video generation",
|
|
)
|
|
parser.add_argument(
|
|
"--negative-prompt",
|
|
type=str,
|
|
default=SamplingParam.negative_prompt,
|
|
help="Negative text prompt for video generation",
|
|
)
|
|
parser.add_argument(
|
|
"--prompt-path",
|
|
type=str,
|
|
default=SamplingParam.prompt_path,
|
|
help="Path to a text file containing the prompt",
|
|
)
|
|
parser.add_argument(
|
|
"--output-path",
|
|
type=str,
|
|
default=SamplingParam.output_path,
|
|
help="Path to save the generated video",
|
|
)
|
|
parser.add_argument(
|
|
"--output-video-name",
|
|
type=str,
|
|
default=SamplingParam.output_video_name,
|
|
help="Name of the output video",
|
|
)
|
|
parser.add_argument(
|
|
"--num-videos-per-prompt",
|
|
type=int,
|
|
default=SamplingParam.num_videos_per_prompt,
|
|
help="Number of videos to generate per prompt",
|
|
)
|
|
parser.add_argument(
|
|
"--seed",
|
|
type=int,
|
|
default=SamplingParam.seed,
|
|
help="Random seed for generation",
|
|
)
|
|
parser.add_argument(
|
|
"--num-frames",
|
|
type=int,
|
|
default=SamplingParam.num_frames,
|
|
help="Number of frames to generate",
|
|
)
|
|
parser.add_argument(
|
|
"--height",
|
|
type=int,
|
|
default=SamplingParam.height,
|
|
help="Height of generated video",
|
|
)
|
|
parser.add_argument(
|
|
"--width",
|
|
type=int,
|
|
default=SamplingParam.width,
|
|
help="Width of generated video",
|
|
)
|
|
parser.add_argument(
|
|
"--fps",
|
|
type=int,
|
|
default=SamplingParam.fps,
|
|
help="Frames per second for saved video",
|
|
)
|
|
parser.add_argument(
|
|
"--num-inference-steps",
|
|
type=int,
|
|
default=SamplingParam.num_inference_steps,
|
|
help="Number of denoising steps",
|
|
)
|
|
parser.add_argument(
|
|
"--guidance-scale",
|
|
type=float,
|
|
default=SamplingParam.guidance_scale,
|
|
help="Classifier-free guidance scale",
|
|
)
|
|
parser.add_argument(
|
|
"--batch-cfg",
|
|
action=StoreBoolean,
|
|
default=SamplingParam.batch_cfg,
|
|
help="Evaluate conditional and unconditional CFG branches in one batch",
|
|
)
|
|
parser.add_argument(
|
|
"--guidance-rescale",
|
|
type=float,
|
|
default=SamplingParam.guidance_rescale,
|
|
help="Guidance rescale factor",
|
|
)
|
|
parser.add_argument(
|
|
"--use-embedded-guidance",
|
|
action="store_true",
|
|
default=SamplingParam.use_embedded_guidance,
|
|
help="Use embedded guidance scale (FLUX-style) instead of classic CFG",
|
|
)
|
|
parser.add_argument(
|
|
"--true-cfg-scale",
|
|
type=float,
|
|
default=SamplingParam.true_cfg_scale,
|
|
help="True CFG scale for FLUX when > 1 (requires negative prompt encoding)",
|
|
)
|
|
parser.add_argument(
|
|
"--boundary-ratio",
|
|
type=float,
|
|
default=SamplingParam.boundary_ratio,
|
|
help="Boundary timestep ratio",
|
|
)
|
|
parser.add_argument(
|
|
"--pyramid-num-inference-steps-list",
|
|
nargs=3,
|
|
type=int,
|
|
default=SamplingParam.pyramid_num_inference_steps_list,
|
|
help="Denoising steps for the three Helios pyramid stages",
|
|
)
|
|
parser.add_argument(
|
|
"--history-sizes",
|
|
nargs=3,
|
|
type=int,
|
|
default=SamplingParam.history_sizes,
|
|
help="Long, mid, and short Helios latent history sizes",
|
|
)
|
|
parser.add_argument(
|
|
"--num-latent-frames-per-chunk",
|
|
type=int,
|
|
default=SamplingParam.num_latent_frames_per_chunk,
|
|
help="Helios autoregressive latent frames per chunk",
|
|
)
|
|
parser.add_argument(
|
|
"--keep-first-frame",
|
|
action=StoreBoolean,
|
|
default=SamplingParam.keep_first_frame,
|
|
help="Keep the first Helios latent frame as prefix conditioning",
|
|
)
|
|
parser.add_argument(
|
|
"--is-skip-first-chunk",
|
|
action=StoreBoolean,
|
|
default=SamplingParam.is_skip_first_chunk,
|
|
help="Skip the first Helios autoregressive chunk",
|
|
)
|
|
parser.add_argument(
|
|
"--use-zero-init",
|
|
action=StoreBoolean,
|
|
default=SamplingParam.use_zero_init,
|
|
help="Enable Helios CFG zero initialization when supported",
|
|
)
|
|
parser.add_argument(
|
|
"--zero-steps",
|
|
type=int,
|
|
default=SamplingParam.zero_steps,
|
|
help="Number of Helios CFG zero-initialization steps",
|
|
)
|
|
parser.add_argument(
|
|
"--is-amplify-first-chunk",
|
|
action=StoreBoolean,
|
|
default=SamplingParam.is_amplify_first_chunk,
|
|
help="Use the amplified DMD schedule for the first Helios chunk",
|
|
)
|
|
parser.add_argument(
|
|
"--save-video",
|
|
action="store_true",
|
|
default=SamplingParam.save_video,
|
|
help="Whether to save the video to disk",
|
|
)
|
|
parser.add_argument(
|
|
"--no-save-video",
|
|
action="store_false",
|
|
dest="save_video",
|
|
help="Don't save the video to disk",
|
|
)
|
|
parser.add_argument(
|
|
"--return-frames",
|
|
action="store_true",
|
|
default=False,
|
|
help="Whether to return the raw frames",
|
|
)
|
|
parser.add_argument(
|
|
"--image-path",
|
|
type=str,
|
|
default=SamplingParam.image_path,
|
|
help="Path to input image for image-to-video generation",
|
|
)
|
|
parser.add_argument(
|
|
"--video-path",
|
|
type=str,
|
|
default=SamplingParam.video_path,
|
|
help="Path to input video for video-to-video generation",
|
|
)
|
|
parser.add_argument(
|
|
"--audio-path",
|
|
type=str,
|
|
default=SamplingParam.audio_path,
|
|
help="Path to input audio for speech-driven generation (Wan S2V)",
|
|
)
|
|
parser.add_argument(
|
|
"--t-thresh",
|
|
type=float,
|
|
default=SamplingParam.t_thresh,
|
|
help="Threshold for timestep scheduling in refinement (default: 0.5)",
|
|
)
|
|
parser.add_argument(
|
|
"--num-cond-frames",
|
|
type=int,
|
|
default=SamplingParam.num_cond_frames,
|
|
help="Number of conditioning frames for refinement",
|
|
)
|
|
parser.add_argument(
|
|
"--moba-config-path",
|
|
type=str,
|
|
default=None,
|
|
help="Path to a JSON file containing V-MoBA specific configurations.",
|
|
)
|
|
parser.add_argument(
|
|
"--return-trajectory-latents",
|
|
action="store_true",
|
|
default=SamplingParam.return_trajectory_latents,
|
|
help="Whether to return the trajectory",
|
|
)
|
|
parser.add_argument(
|
|
"--return-trajectory-decoded",
|
|
action="store_true",
|
|
default=SamplingParam.return_trajectory_decoded,
|
|
help="Whether to return the decoded trajectory",
|
|
)
|
|
return parser
|
|
|
|
|
|
@dataclass
|
|
class CacheParams:
|
|
cache_type: str = "none"
|