flux_stages
fastvideo.pipelines.stages.flux_stages.get_local_torch_device
get_local_torch_device() -> torch.deviceReturn the torch device for the current rank.
fastvideo.pipelines.stages.flux_stages.FastVideoArgs
FastVideoArgs(model_path: str, mode: ExecutionMode = ExecutionMode.INFERENCE, workload_type: WorkloadType = WorkloadType.T2V, distributed_executor_backend: str = 'mp', enable_streaming_ipc_queues: bool = False, ray_placement_group: PlacementGroup | None = None, ray_runtime_env: RuntimeEnv | None = None, inference_mode: bool = True, trust_remote_code: bool = False, revision: str | None = None, num_gpus: int = 1, tp_size: int = -1, sp_size: int = -1, hsdp_replicate_dim: int = 1, hsdp_shard_dim: int = -1, dist_timeout: int | None = None, pipeline_config: PipelineConfig = PipelineConfig(), preprocess_config: PreprocessConfig | None = None, lora_path: str | None = None, lora_nickname: str = 'default', lora_strength: float = 1.0, lora_target_modules: list[str] | None = None, output_type: str = 'pil', attention_backend: str | None = None, dit_cpu_offload: bool = True, use_fsdp_inference: bool = False, dit_layerwise_offload: bool = True, text_encoder_cpu_offload: bool = True, image_encoder_cpu_offload: bool = True, vae_cpu_offload: bool = True, pin_cpu_memory: bool = True, h3_sequential_load: bool | None = None, h3_encoder_split: bool | None = None, h3_encoder_nodes: int = -1, h3_encoder_workers: int = 0, video_decode_backend: str = 'h3-vae', taeh3_checkpoint: str | None = None, taeh3_chunk_size: int = 5, h3_vae_nvfp4_decoder: str | None = None, lazy_module_load: bool | None = None, vae_parallel_decode: bool = False, vae_parallel_encode: bool = False, vae_parallel_decode_strategy: str | None = None, enable_torch_compile: bool = False, regional_compile: bool = False, enable_torch_compile_text_encoder: bool = False, enable_torch_compile_vae: bool = False, enable_torch_compile_audio_vae: bool = False, torch_compile_kwargs: dict[str, Any] = dict(), torch_compile_kwargs_dit: dict[str, Any] = dict(), torch_compile_kwargs_text_encoder: dict[str, Any] = dict(), torch_compile_kwargs_vae: dict[str, Any] = dict(), torch_compile_kwargs_audio_vae: dict[str, Any] = dict(), inference_torch_compile: bool = False, h3_vsa_fp4: dict[str, Any] = dict(), disable_autocast: bool = False, VSA_sparsity: float | None = None, VSA_tile_size: int | None = None, moba_config_path: str | None = None, moba_config: dict[str, Any] = dict(), master_port: int | None = None, enable_stage_verification: bool = True, prompt_txt: str | None = None, ltx2_vae_tiling: bool | None = None, ltx2_vae_spatial_tile_size_in_pixels: int | None = None, ltx2_vae_spatial_tile_overlap_in_pixels: int | None = None, ltx2_vae_temporal_tile_size_in_frames: int | None = None, ltx2_vae_temporal_tile_overlap_in_frames: int | None = None, ltx2_initial_latent_path: str | None = None, ltx2_audio_latent_path: str | None = None, refine_enabled: bool | None = None, refine_upsampler_path: str | None = None, refine_transformer_path: str | None = None, refine_lora_path: str | None = None, refine_num_inference_steps: int | None = None, refine_guidance_scale: float | None = None, refine_add_noise: bool | None = None, refine_noise_path: str | None = None, refine_audio_noise_path: str | None = None, ltx2_refine_enabled: bool = False, ltx2_refine_upsampler_path: str | None = None, ltx2_refine_transformer_path: str | None = None, ltx2_refine_lora_path: str | None = None, ltx2_refine_num_inference_steps: int = 3, ltx2_refine_guidance_scale: float = 1.0, ltx2_refine_add_noise: bool = True, ltx2_refine_noise_path: str | None = None, ltx2_refine_audio_noise_path: str | None = None, ltx2_legacy_native_noise_order: bool = False, ltx2_use_distilled_sigmas: bool = True, model_paths: dict[str, str] = dict(), model_loaded: dict[str, bool] = (lambda: {'transformer': True, 'vae': True, 'upsampler': True})(), override_text_encoder_safetensors: str | None = None, override_text_encoder_quant: QuantizationMethods = None, transformer_quant: Any | None = None, override_transformer_cls_name: str | None = None, init_weights_from_safetensors: str = '', init_weights_from_safetensors_2: str = '', override_pipeline_cls_name: str | None = None, boundary_ratio: float | None = 0.875)fastvideo.pipelines.stages.flux_stages.FastVideoArgs.is_output_rank
is_output_rank: boolWhether this process may materialize user-facing outputs.
This is runtime state assigned by an SPMD executor, not a public configuration field, so dataclass serialization and typed config parsing deliberately ignore it.
fastvideo.pipelines.stages.flux_stages.FastVideoArgs.add_cli_args
add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParserfastvideo.pipelines.stages.flux_stages.FastVideoArgs.from_cli_args
from_cli_args(args: argparse.Namespace) -> FastVideoArgsfastvideo.pipelines.stages.flux_stages.FastVideoArgs.from_kwargs
from_kwargs(**kwargs: Any) -> FastVideoArgsfastvideo.pipelines.stages.flux_stages.FastVideoArgs.check_fastvideo_args
check_fastvideo_args() -> NoneValidate inference arguments for consistency
fastvideo.pipelines.stages.flux_stages.FastVideoArgs.finalize_device_offload_policy
finalize_device_offload_policy(device_id: int = 0) -> boolApply device-local memory policy, then resolve incompatible modes.
fastvideo.pipelines.stages.flux_stages.FastVideoArgs.disable_offload_on_unified_memory
disable_offload_on_unified_memory(device_id: int = 0, *, offload_flag: str | None = None) -> boolDisable host offload after a worker has selected its device.
CUDA's unified-memory probe reads runtime device properties and may
initialize a CUDA context. Callers must therefore use this only inside
a device-owning process, after selecting and binding device_id.
Returning the classification lets direct component-loader callers
apply the same policy to explicit per-call overrides. When
offload_flag is given, the return value says whether this policy
covers that component role.
fastvideo.pipelines.stages.flux_stages.set_forward_context
set_forward_context(current_timestep, attn_metadata, forward_batch: Optional[ForwardBatch] = None)A context manager that stores the current forward context, can be attention metadata, etc. Here we can inject common logic for every model forward pass.
fastvideo.pipelines.stages.flux_stages.init_logger
init_logger(name: str) -> _FastvideoLoggerThe main purpose of this function is to ensure that loggers are retrieved in such a way that we can be sure the root fastvideo logger has already been configured.
fastvideo.pipelines.stages.flux_stages.ForwardBatch
ForwardBatch(data_type: str, generator: torch.Generator | list[torch.Generator] | None = None, image_path: str | None = None, image_embeds: list[torch.Tensor] = list(), pil_image: torch.Tensor | PIL.Image.Image | None = None, last_image: torch.Tensor | PIL.Image.Image | None = None, references: list[Any] | None = None, preprocessed_image: torch.Tensor | None = None, 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_path: str | None = None, video_latent: torch.Tensor | None = None, 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, pose_latents: torch.Tensor | None = None, face_pixel_values: torch.Tensor | None = None, audio_path: str | None = None, audio_embeds: torch.Tensor | None = None, t_thresh: float = 0.5, num_cond_frames: int = 0, prompt_embeds: list[torch.Tensor] = 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, max_sequence_length: int | None = None, prompt_template: dict[str, Any] | None = None, do_classifier_free_guidance: bool = False, use_embedded_guidance: bool = False, true_cfg_scale: float = 1.0, batch_size: int | None = None, num_videos_per_prompt: int = 1, seed: int | None = None, seeds: list[int] | None = None, is_prompt_processed: bool = False, latents: torch.Tensor | None = None, audio_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, first_frame_latent: torch.Tensor | None = None, helios_latent_chunks: list[torch.Tensor] | None = None, mouse_cond: torch.Tensor | None = None, keyboard_cond: torch.Tensor | None = None, grid_sizes: torch.Tensor | None = None, num_iterations: int | None = None, use_base_model: bool = False, pose: str | None = None, c2ws_plucker_emb: torch.Tensor | None = None, action_path: str | None = None, trajectory_type: str | None = None, movement_distance: float | None = None, camera_rotation: str | None = None, height_latents: list[int] | int | None = None, width_latents: list[int] | int | None = None, num_frames: list[int] | int = 1, 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: torch.Tensor | None = None, timestep: torch.Tensor | float | int | None = None, step_index: int | None = None, boundary_ratio: float | None = None, 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, guidance_rescale: float = 0.0, eta: float = 0.0, sigmas: list[float] | None = None, 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, use_zero_init: bool = True, zero_steps: int = 1, is_amplify_first_chunk: bool = False, enable_teacache: bool = False, 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] = list(), ltx2_stg_blocks_audio: list[int] = list(), ltx2_use_ancestral_sampler: bool = False, ltx2_images: list[tuple[str, int, float]] | None = None, ltx2_image_crf: float = 33.0, ltx2_conditioning_latent_stage1: torch.Tensor | None = None, ltx2_conditioning_latent_stage2: torch.Tensor | None = None, ltx2_video_conditions: list[tuple[list[str], int, float]] | None = None, audio_start_in_s: float | None = None, audio_end_in_s: float | None = None, init_audio: Any = None, init_audio_strength: float | None = None, init_noise_level: float | None = None, inpaint_audio: Any = None, inpaint_mask: Any = None, n_tokens: int | None = None, extra_step_kwargs: dict[str, Any] = dict(), modules: dict[str, Any] = dict(), 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, continuation_state: ContinuationState | None = None, return_continuation_state: bool = False, extra: dict[str, Any] = dict(), save_video: bool = True, return_frames: bool = False, is_cfg_negative: bool = False, VSA_sparsity: float = 0.0, logging_info: PipelineLoggingInfo = PipelineLoggingInfo(), *, return_samples: bool = False)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.
fastvideo.pipelines.stages.flux_stages.PipelineStage
PipelineStage()Abstract base class for all pipeline stages.
A pipeline stage represents a discrete step in the diffusion process that can be composed with other stages to create a complete pipeline. Each stage is responsible for a specific part of the process, such as prompt encoding, latent preparation, etc.
fastvideo.pipelines.stages.flux_stages.PipelineStage.verify_input
verify_input(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResultVerify the input for the stage.
Example
from fastvideo.pipelines.stages.validators import V, VerificationResult
def verify_input(self, batch, fastvideo_args): result = VerificationResult() result.add_check("height", batch.height, V.positive_int_divisible(8)) result.add_check("width", batch.width, V.positive_int_divisible(8)) result.add_check("image_latent", batch.image_latent, V.is_tensor) return result
Parameters
batch: ForwardBatchThe current batch information.
fastvideo_args: FastVideoArgsThe inference arguments.
Returns
VerificationResultA VerificationResult containing the verification status.
fastvideo.pipelines.stages.flux_stages.PipelineStage.verify_output
verify_output(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResultVerify the output for the stage.
Parameters
batch: ForwardBatchThe current batch information.
fastvideo_args: FastVideoArgsThe inference arguments.
Returns
VerificationResultA VerificationResult containing the verification status.
fastvideo.pipelines.stages.flux_stages.PipelineStage.device
device: torch.deviceGet the device for this stage.
fastvideo.pipelines.stages.flux_stages.PipelineStage.set_logging
set_logging(enable: bool)Enable or disable logging for this stage.
Parameters
enable: boolWhether to enable logging.
fastvideo.pipelines.stages.flux_stages.PipelineStage.forward
forward(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatchForward pass of the stage's processing.
This method should be implemented by subclasses to provide the forward processing logic for the stage.
Parameters
batch: ForwardBatchThe current batch information.
fastvideo_args: FastVideoArgsThe inference arguments.
Returns
ForwardBatchThe updated batch information after this stage's processing.
fastvideo.pipelines.stages.flux_stages.InputValidationStage
InputValidationStage()Stage for validating and preparing inputs for diffusion pipelines.
This stage validates that all required inputs are present and properly formatted before proceeding with the diffusion process.
fastvideo.pipelines.stages.flux_stages.InputValidationStage.forward
forward(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatchValidate and prepare inputs.
Parameters
batch: ForwardBatchThe current batch information.
fastvideo_args: FastVideoArgsThe inference arguments.
Returns
ForwardBatchThe validated batch information.
fastvideo.pipelines.stages.flux_stages.InputValidationStage.verify_input
verify_input(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResultVerify input validation stage inputs.
fastvideo.pipelines.stages.flux_stages.InputValidationStage.verify_output
verify_output(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResultVerify input validation stage outputs.
fastvideo.pipelines.stages.flux_stages.TimestepPreparationStage
TimestepPreparationStage(scheduler)Stage for preparing timesteps for the diffusion process.
This stage handles the preparation of the timestep sequence that will be used during the diffusion process.
fastvideo.pipelines.stages.flux_stages.TimestepPreparationStage.forward
forward(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatchPrepare timesteps for the diffusion process.
Parameters
batch: ForwardBatchThe current batch information.
fastvideo_args: FastVideoArgsThe inference arguments.
Returns
ForwardBatchThe batch with prepared timesteps.
fastvideo.pipelines.stages.flux_stages.TimestepPreparationStage.verify_input
verify_input(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResultVerify timestep preparation stage inputs.
fastvideo.pipelines.stages.flux_stages.TimestepPreparationStage.verify_output
verify_output(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResultVerify timestep preparation stage outputs.
fastvideo.pipelines.stages.flux_stages.FluxInputValidationStage
FluxInputValidationStage()Require height/width divisible by 16 (VAE scale × 2 for FLUX packing).
Source
class FluxInputValidationStage(InputValidationStage):
"""Require height/width divisible by 16 (VAE scale × 2 for FLUX packing)."""
def forward(
self,
batch: ForwardBatch,
fastvideo_args: FastVideoArgs,
) -> ForwardBatch:
if (batch.height is not None and batch.width is not None and (batch.height % 16 != 0 or batch.width % 16 != 0)):
raise ValueError("FLUX expects height and width divisible by 16 "
f"(VAE latent grid × 2× packing); got {batch.height}×{batch.width}.")
return super().forward(batch, fastvideo_args)</code></pre></details>
fastvideo.pipelines.stages.flux_stages.FluxInputValidationStage.forward
forward(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch
Source
def forward(
self,
batch: ForwardBatch,
fastvideo_args: FastVideoArgs,
) -> ForwardBatch:
if (batch.height is not None and batch.width is not None and (batch.height % 16 != 0 or batch.width % 16 != 0)):
raise ValueError("FLUX expects height and width divisible by 16 "
f"(VAE latent grid × 2× packing); got {batch.height}×{batch.width}.")
return super().forward(batch, fastvideo_args)
fastvideo.pipelines.stages.flux_stages.FluxConditioningStage
FluxConditioningStage()
Build CLIP pooled + T5 sequence + text_ids (and optional negative for true CFG).
Source
class FluxConditioningStage(PipelineStage):
"""Build CLIP pooled + T5 sequence + ``text_ids`` (and optional negative for true CFG)."""
@torch.no_grad()
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
if len(batch.prompt_embeds) < 2:
raise ValueError("FluxConditioningStage expects 2 prompt_embeds (CLIP pooled, T5 sequence), "
f"got {len(batch.prompt_embeds)}")
device = get_local_torch_device()
target_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]
pooled = batch.prompt_embeds[0].to(device=device, dtype=target_dtype)
enc = batch.prompt_embeds[1].to(device=device, dtype=target_dtype)
seq_len = enc.shape[1]
text_ids = torch.zeros(seq_len, 3, device=device, dtype=torch.long)
batch.extra["flux_pooled_projections"] = pooled
batch.extra["flux_encoder_hidden_states"] = enc
batch.extra["flux_text_ids"] = text_ids
if batch.do_classifier_free_guidance:
if not batch.negative_prompt_embeds or len(batch.negative_prompt_embeds) < 2:
raise ValueError("True CFG requires two negative_prompt_embeds (CLIP, T5).")
neg_pooled = batch.negative_prompt_embeds[0].to(device=device, dtype=target_dtype)
neg_enc = batch.negative_prompt_embeds[1].to(device=device, dtype=target_dtype)
batch.extra["flux_negative_pooled_projections"] = neg_pooled
batch.extra["flux_negative_encoder_hidden_states"] = neg_enc
return batch</code></pre></details>
fastvideo.pipelines.stages.flux_stages.FluxConditioningStage.forward
forward(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch
Source
@torch.no_grad()
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
if len(batch.prompt_embeds) < 2:
raise ValueError("FluxConditioningStage expects 2 prompt_embeds (CLIP pooled, T5 sequence), "
f"got {len(batch.prompt_embeds)}")
device = get_local_torch_device()
target_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]
pooled = batch.prompt_embeds[0].to(device=device, dtype=target_dtype)
enc = batch.prompt_embeds[1].to(device=device, dtype=target_dtype)
seq_len = enc.shape[1]
text_ids = torch.zeros(seq_len, 3, device=device, dtype=torch.long)
batch.extra["flux_pooled_projections"] = pooled
batch.extra["flux_encoder_hidden_states"] = enc
batch.extra["flux_text_ids"] = text_ids
if batch.do_classifier_free_guidance:
if not batch.negative_prompt_embeds or len(batch.negative_prompt_embeds) < 2:
raise ValueError("True CFG requires two negative_prompt_embeds (CLIP, T5).")
neg_pooled = batch.negative_prompt_embeds[0].to(device=device, dtype=target_dtype)
neg_enc = batch.negative_prompt_embeds[1].to(device=device, dtype=target_dtype)
batch.extra["flux_negative_pooled_projections"] = neg_pooled
batch.extra["flux_negative_encoder_hidden_states"] = neg_enc
return batch</code></pre></details>
fastvideo.pipelines.stages.flux_stages.FluxTimestepPreparationStage
FluxTimestepPreparationStage(scheduler)
Flow Match with resolution-dependent mu from packed image sequence length.
Source
class FluxTimestepPreparationStage(TimestepPreparationStage):
"""Flow Match with resolution-dependent ``mu`` from packed image sequence length."""
@staticmethod
def _calculate_mu(
image_seq_len: int,
base_seq_len: int = 256,
max_seq_len: int = 4096,
base_shift: float = 0.5,
max_shift: float = 1.15,
) -> float:
m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
b = base_shift - m * base_seq_len
return float(image_seq_len) * m + b
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
sig = inspect.signature(self.scheduler.set_timesteps)
if "mu" not in sig.parameters:
logger.warning(
"FLUX timestep prep: scheduler %s.set_timesteps does not accept 'mu'; falling back to the base "
"timestep schedule. FLUX expects a FlowMatchEulerDiscreteScheduler with resolution-dependent "
"dynamic shifting — output quality may degrade.",
type(self.scheduler).__name__)
return super().forward(batch, fastvideo_args)
cfg = getattr(self.scheduler, "config", None)
use_dynamic = bool(getattr(cfg, "use_dynamic_shifting", False)) if cfg is not None else False
if not use_dynamic:
logger.warning(
"FLUX timestep prep: scheduler has use_dynamic_shifting=False; falling back to the base timestep "
"schedule and skipping the resolution-dependent 'mu' shift. FLUX requires dynamic shifting for "
"correct timesteps — output quality may degrade.")
return super().forward(batch, fastvideo_args)
if batch.height is None or batch.width is None:
raise ValueError("height/width must be set before FluxTimestepPreparationStage")
vae_arch = fastvideo_args.pipeline_config.vae_config.arch_config
spatial_ratio = int(getattr(vae_arch, "spatial_compression_ratio", 8))
h_lat = batch.height // spatial_ratio
w_lat = batch.width // spatial_ratio
if h_lat % 2 != 0 or w_lat % 2 != 0:
raise ValueError(
f"Latent spatial dims must be even for FLUX packing; got {h_lat}×{w_lat} from {batch.height}×{batch.width}."
)
image_seq_len = (h_lat // 2) * (w_lat // 2)
base_seq_len = int(getattr(cfg, "base_image_seq_len", 256))
max_seq_len = int(getattr(cfg, "max_image_seq_len", 4096))
base_shift = float(getattr(cfg, "base_shift", 0.5))
max_shift = float(getattr(cfg, "max_shift", 1.15))
device = get_local_torch_device()
mu = self._calculate_mu(
image_seq_len=image_seq_len,
base_seq_len=base_seq_len,
max_seq_len=max_seq_len,
base_shift=base_shift,
max_shift=max_shift,
)
self.scheduler.set_timesteps(batch.num_inference_steps, device=device, mu=mu)
batch.timesteps = self.scheduler.timesteps
return batch</code></pre></details>
fastvideo.pipelines.stages.flux_stages.FluxTimestepPreparationStage.forward
forward(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch
Source
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
sig = inspect.signature(self.scheduler.set_timesteps)
if "mu" not in sig.parameters:
logger.warning(
"FLUX timestep prep: scheduler %s.set_timesteps does not accept 'mu'; falling back to the base "
"timestep schedule. FLUX expects a FlowMatchEulerDiscreteScheduler with resolution-dependent "
"dynamic shifting — output quality may degrade.",
type(self.scheduler).__name__)
return super().forward(batch, fastvideo_args)
cfg = getattr(self.scheduler, "config", None)
use_dynamic = bool(getattr(cfg, "use_dynamic_shifting", False)) if cfg is not None else False
if not use_dynamic:
logger.warning(
"FLUX timestep prep: scheduler has use_dynamic_shifting=False; falling back to the base timestep "
"schedule and skipping the resolution-dependent 'mu' shift. FLUX requires dynamic shifting for "
"correct timesteps — output quality may degrade.")
return super().forward(batch, fastvideo_args)
if batch.height is None or batch.width is None:
raise ValueError("height/width must be set before FluxTimestepPreparationStage")
vae_arch = fastvideo_args.pipeline_config.vae_config.arch_config
spatial_ratio = int(getattr(vae_arch, "spatial_compression_ratio", 8))
h_lat = batch.height // spatial_ratio
w_lat = batch.width // spatial_ratio
if h_lat % 2 != 0 or w_lat % 2 != 0:
raise ValueError(
f"Latent spatial dims must be even for FLUX packing; got {h_lat}×{w_lat} from {batch.height}×{batch.width}."
)
image_seq_len = (h_lat // 2) * (w_lat // 2)
base_seq_len = int(getattr(cfg, "base_image_seq_len", 256))
max_seq_len = int(getattr(cfg, "max_image_seq_len", 4096))
base_shift = float(getattr(cfg, "base_shift", 0.5))
max_shift = float(getattr(cfg, "max_shift", 1.15))
device = get_local_torch_device()
mu = self._calculate_mu(
image_seq_len=image_seq_len,
base_seq_len=base_seq_len,
max_seq_len=max_seq_len,
base_shift=base_shift,
max_shift=max_shift,
)
self.scheduler.set_timesteps(batch.num_inference_steps, device=device, mu=mu)
batch.timesteps = self.scheduler.timesteps
return batch</code></pre></details>
fastvideo.pipelines.stages.flux_stages.FluxLatentPreparationStage
FluxLatentPreparationStage(scheduler)
Source
class FluxLatentPreparationStage(PipelineStage):
def __init__(self, scheduler) -> None:
super().__init__()
self.scheduler = scheduler
@torch.no_grad()
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
if batch.height is None or batch.width is None:
raise ValueError("height/width required for FluxLatentPreparationStage")
if isinstance(batch.prompt, list):
batch_size = len(batch.prompt)
elif batch.prompt is not None:
batch_size = 1
else:
if not batch.prompt_embeds:
raise ValueError("prompt or prompt_embeds must be provided")
batch_size = batch.prompt_embeds[0].shape[0]
batch_size *= batch.num_videos_per_prompt
if isinstance(batch.generator, list) and len(batch.generator) != batch_size:
raise ValueError(f"generator list length {len(batch.generator)} does not match batch_size {batch_size}")
arch = fastvideo_args.pipeline_config.dit_config.arch_config
in_channels = int(getattr(arch, "in_channels", 64))
num_channels_latents = in_channels // 4
vae_arch = fastvideo_args.pipeline_config.vae_config.arch_config
spatial_ratio = int(getattr(vae_arch, "spatial_compression_ratio", 8))
h_lat = batch.height // spatial_ratio
w_lat = batch.width // spatial_ratio
dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]
device = get_local_torch_device()
shape = (batch_size, num_channels_latents, h_lat, w_lat)
latents = batch.latents
if latents is None:
latents = randn_tensor(shape, generator=batch.generator, device=device, dtype=dtype)
if hasattr(self.scheduler, "init_noise_sigma"):
latents = latents * self.scheduler.init_noise_sigma
else:
latents = latents.to(device=device, dtype=dtype)
if latents.shape != shape:
raise ValueError(f"Expected latents shape {shape}, got {tuple(latents.shape)}")
if hasattr(self.scheduler, "init_noise_sigma"):
latents = latents * self.scheduler.init_noise_sigma
packed = _pack_latents(latents, batch_size, num_channels_latents, h_lat, w_lat)
patch_h, patch_w = h_lat // 2, w_lat // 2
img_ids = _prepare_latent_image_ids(patch_h, patch_w, device, dtype=torch.long)
batch.latents = packed
batch.raw_latent_shape = shape
batch.extra["flux_h_lat"] = h_lat
batch.extra["flux_w_lat"] = w_lat
batch.extra["flux_num_channels_latents"] = num_channels_latents
batch.extra["flux_latent_image_ids"] = img_ids
return batch</code></pre></details>
fastvideo.pipelines.stages.flux_stages.FluxLatentPreparationStage.forward
forward(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch
Source
@torch.no_grad()
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
if batch.height is None or batch.width is None:
raise ValueError("height/width required for FluxLatentPreparationStage")
if isinstance(batch.prompt, list):
batch_size = len(batch.prompt)
elif batch.prompt is not None:
batch_size = 1
else:
if not batch.prompt_embeds:
raise ValueError("prompt or prompt_embeds must be provided")
batch_size = batch.prompt_embeds[0].shape[0]
batch_size *= batch.num_videos_per_prompt
if isinstance(batch.generator, list) and len(batch.generator) != batch_size:
raise ValueError(f"generator list length {len(batch.generator)} does not match batch_size {batch_size}")
arch = fastvideo_args.pipeline_config.dit_config.arch_config
in_channels = int(getattr(arch, "in_channels", 64))
num_channels_latents = in_channels // 4
vae_arch = fastvideo_args.pipeline_config.vae_config.arch_config
spatial_ratio = int(getattr(vae_arch, "spatial_compression_ratio", 8))
h_lat = batch.height // spatial_ratio
w_lat = batch.width // spatial_ratio
dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]
device = get_local_torch_device()
shape = (batch_size, num_channels_latents, h_lat, w_lat)
latents = batch.latents
if latents is None:
latents = randn_tensor(shape, generator=batch.generator, device=device, dtype=dtype)
if hasattr(self.scheduler, "init_noise_sigma"):
latents = latents * self.scheduler.init_noise_sigma
else:
latents = latents.to(device=device, dtype=dtype)
if latents.shape != shape:
raise ValueError(f"Expected latents shape {shape}, got {tuple(latents.shape)}")
if hasattr(self.scheduler, "init_noise_sigma"):
latents = latents * self.scheduler.init_noise_sigma
packed = _pack_latents(latents, batch_size, num_channels_latents, h_lat, w_lat)
patch_h, patch_w = h_lat // 2, w_lat // 2
img_ids = _prepare_latent_image_ids(patch_h, patch_w, device, dtype=torch.long)
batch.latents = packed
batch.raw_latent_shape = shape
batch.extra["flux_h_lat"] = h_lat
batch.extra["flux_w_lat"] = w_lat
batch.extra["flux_num_channels_latents"] = num_channels_latents
batch.extra["flux_latent_image_ids"] = img_ids
return batch</code></pre></details>
fastvideo.pipelines.stages.flux_stages.FluxDenoisingStage
FluxDenoisingStage(transformer, scheduler)
Source
class FluxDenoisingStage(PipelineStage):
def __init__(self, transformer, scheduler) -> None:
super().__init__()
self.transformer = transformer
self.scheduler = scheduler
@staticmethod
def _step_kwargs(scheduler_step, batch: ForwardBatch) -> dict[str, Any]:
kwargs: dict[str, Any] = {}
sig = inspect.signature(scheduler_step)
if "generator" in sig.parameters:
gen = batch.generator[0] if isinstance(batch.generator, list) else batch.generator
kwargs["generator"] = gen
return kwargs
@torch.no_grad()
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
if batch.timesteps is None:
raise ValueError("timesteps must be set before FluxDenoisingStage")
if batch.latents is None:
raise ValueError("latents must be set before FluxDenoisingStage")
packed = batch.latents
timesteps = batch.timesteps
pooled = batch.extra["flux_pooled_projections"]
enc = batch.extra["flux_encoder_hidden_states"]
txt_ids = batch.extra["flux_text_ids"]
img_ids = batch.extra["flux_latent_image_ids"]
neg_pooled = batch.extra.get("flux_negative_pooled_projections")
neg_enc = batch.extra.get("flux_negative_encoder_hidden_states")
true_cfg_scale = float(batch.true_cfg_scale)
use_true_cfg = batch.do_classifier_free_guidance and true_cfg_scale > 1.0
# Prefer the loaded transformer's arch (HF ``guidance_embeds``), not static pipeline defaults.
tr_arch = self.transformer.fastvideo_config.arch_config
guidance_embeds = bool(getattr(tr_arch, "guidance_embeds", False))
device = get_local_torch_device()
target_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]
autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast
bs = packed.shape[0]
if guidance_embeds:
guidance = torch.full((bs, ), float(batch.guidance_scale), device=device, dtype=torch.float32)
else:
guidance = None
step_extras = self._step_kwargs(self.scheduler.step, batch)
for t in timesteps:
# Stop if interrupted
if getattr(self, "interrupt", False):
break
t_scalar = t
if not isinstance(t_scalar, torch.Tensor):
t_scalar = torch.tensor([t_scalar], device=device, dtype=torch.float32)
t_scalar = t_scalar.to(device=device, dtype=torch.float32)
timestep_model = t_scalar.expand(bs).float() / 1000.0
timestep_model = timestep_model.to(dtype=target_dtype)
ts_ctx = int(t_scalar.reshape(-1)[0].item())
with (
torch.autocast(
device_type="cuda",
dtype=target_dtype,
enabled=autocast_enabled and device.type == "cuda",
),
set_forward_context(
current_timestep=ts_ctx,
attn_metadata=None,
forward_batch=batch,
),
):
if use_true_cfg:
assert neg_enc is not None and neg_pooled is not None
n_neg = self.transformer(
hidden_states=packed,
encoder_hidden_states=neg_enc,
pooled_projections=neg_pooled,
timestep=timestep_model,
guidance=guidance,
txt_ids=txt_ids,
img_ids=img_ids,
return_dict=False,
)[0]
n_pos = self.transformer(
hidden_states=packed,
encoder_hidden_states=enc,
pooled_projections=pooled,
timestep=timestep_model,
guidance=guidance,
txt_ids=txt_ids,
img_ids=img_ids,
return_dict=False,
)[0]
noise_pred = n_neg + true_cfg_scale * (n_pos - n_neg)
else:
noise_pred = self.transformer(
hidden_states=packed,
encoder_hidden_states=enc,
pooled_projections=pooled,
timestep=timestep_model,
guidance=guidance,
txt_ids=txt_ids,
img_ids=img_ids,
return_dict=False,
)[0]
packed = self.scheduler.step(
noise_pred,
t_scalar,
packed,
return_dict=False,
**step_extras,
)[0]
batch.latents = packed
return batch</code></pre></details>
fastvideo.pipelines.stages.flux_stages.FluxDenoisingStage.forward
forward(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch
Source
@torch.no_grad()
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
if batch.timesteps is None:
raise ValueError("timesteps must be set before FluxDenoisingStage")
if batch.latents is None:
raise ValueError("latents must be set before FluxDenoisingStage")
packed = batch.latents
timesteps = batch.timesteps
pooled = batch.extra["flux_pooled_projections"]
enc = batch.extra["flux_encoder_hidden_states"]
txt_ids = batch.extra["flux_text_ids"]
img_ids = batch.extra["flux_latent_image_ids"]
neg_pooled = batch.extra.get("flux_negative_pooled_projections")
neg_enc = batch.extra.get("flux_negative_encoder_hidden_states")
true_cfg_scale = float(batch.true_cfg_scale)
use_true_cfg = batch.do_classifier_free_guidance and true_cfg_scale > 1.0
# Prefer the loaded transformer's arch (HF ``guidance_embeds``), not static pipeline defaults.
tr_arch = self.transformer.fastvideo_config.arch_config
guidance_embeds = bool(getattr(tr_arch, "guidance_embeds", False))
device = get_local_torch_device()
target_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.dit_precision]
autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast
bs = packed.shape[0]
if guidance_embeds:
guidance = torch.full((bs, ), float(batch.guidance_scale), device=device, dtype=torch.float32)
else:
guidance = None
step_extras = self._step_kwargs(self.scheduler.step, batch)
for t in timesteps:
# Stop if interrupted
if getattr(self, "interrupt", False):
break
t_scalar = t
if not isinstance(t_scalar, torch.Tensor):
t_scalar = torch.tensor([t_scalar], device=device, dtype=torch.float32)
t_scalar = t_scalar.to(device=device, dtype=torch.float32)
timestep_model = t_scalar.expand(bs).float() / 1000.0
timestep_model = timestep_model.to(dtype=target_dtype)
ts_ctx = int(t_scalar.reshape(-1)[0].item())
with (
torch.autocast(
device_type="cuda",
dtype=target_dtype,
enabled=autocast_enabled and device.type == "cuda",
),
set_forward_context(
current_timestep=ts_ctx,
attn_metadata=None,
forward_batch=batch,
),
):
if use_true_cfg:
assert neg_enc is not None and neg_pooled is not None
n_neg = self.transformer(
hidden_states=packed,
encoder_hidden_states=neg_enc,
pooled_projections=neg_pooled,
timestep=timestep_model,
guidance=guidance,
txt_ids=txt_ids,
img_ids=img_ids,
return_dict=False,
)[0]
n_pos = self.transformer(
hidden_states=packed,
encoder_hidden_states=enc,
pooled_projections=pooled,
timestep=timestep_model,
guidance=guidance,
txt_ids=txt_ids,
img_ids=img_ids,
return_dict=False,
)[0]
noise_pred = n_neg + true_cfg_scale * (n_pos - n_neg)
else:
noise_pred = self.transformer(
hidden_states=packed,
encoder_hidden_states=enc,
pooled_projections=pooled,
timestep=timestep_model,
guidance=guidance,
txt_ids=txt_ids,
img_ids=img_ids,
return_dict=False,
)[0]
packed = self.scheduler.step(
noise_pred,
t_scalar,
packed,
return_dict=False,
**step_extras,
)[0]
batch.latents = packed
return batch</code></pre></details>
fastvideo.pipelines.stages.flux_stages.FluxDecodingStage
FluxDecodingStage(vae)
Unpack latents, apply VAE scaling/shift, decode to pixels (5D output B×3×1×H×W).
Source
class FluxDecodingStage(PipelineStage):
"""Unpack latents, apply VAE scaling/shift, decode to pixels (5D output ``B×3×1×H×W``)."""
def __init__(self, vae) -> None:
super().__init__()
self.vae = vae
@staticmethod
def _denormalize_latents(latents: torch.Tensor, vae: Any) -> torch.Tensor:
cfg = getattr(vae, "config", None)
sf = getattr(cfg, "scaling_factor", None) if cfg is not None else None
sh = getattr(cfg, "shift_factor", None) if cfg is not None else None
if sf is None and hasattr(vae, "scaling_factor"):
sf = vae.scaling_factor
if sh is None and hasattr(vae, "shift_factor"):
sh = vae.shift_factor
if sf is not None:
latents = latents / (sf.to(latents.device, latents.dtype) if isinstance(sf, torch.Tensor) else sf)
if sh is not None:
latents = latents + (sh.to(latents.device, latents.dtype) if isinstance(sh, torch.Tensor) else sh)
return latents
@torch.no_grad()
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
if not fastvideo_args.is_output_rank:
# Flux decoding has no collectives. Avoid duplicating the VAE
# forward and device-to-host pixel copy on SPMD worker ranks whose
# result is discarded by the executor.
batch.output = torch.empty((0, 3, 0, 0, 0), device="cpu", dtype=torch.float32)
return batch
packed = batch.latents
if packed is None:
raise ValueError("latents must be set before FluxDecodingStage")
h_lat = int(batch.extra["flux_h_lat"])
w_lat = int(batch.extra["flux_w_lat"])
num_ch = int(batch.extra["flux_num_channels_latents"])
raw_shape = batch.raw_latent_shape
if raw_shape is None:
raise ValueError("raw_latent_shape missing; FluxLatentPreparationStage must run first.")
batch_size = int(raw_shape[0])
infer_device = get_local_torch_device()
packed = packed.to(infer_device)
latents_4d = _unpack_latents(packed, batch_size, num_ch, h_lat, w_lat)
latents_4d = self._denormalize_latents(latents_4d, self.vae)
vae_device = next(self.vae.parameters()).device
latents_4d = latents_4d.to(device=vae_device)
vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
autocast_enabled = (vae_dtype != torch.float32) and not fastvideo_args.disable_autocast
use_cuda_autocast = autocast_enabled and vae_device.type == "cuda"
with torch.autocast(
device_type="cuda",
dtype=vae_dtype,
enabled=use_cuda_autocast,
):
if not autocast_enabled:
latents_4d = latents_4d.to(dtype=vae_dtype)
dec = self.vae.decode(latents_4d)
image = dec.sample if hasattr(dec, "sample") else dec[0]
image = (image / 2 + 0.5).clamp(0, 1)
batch.output = image.unsqueeze(2).detach().float().cpu()
return batch</code></pre></details>
fastvideo.pipelines.stages.flux_stages.FluxDecodingStage.forward
forward(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch
Source
@torch.no_grad()
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
if not fastvideo_args.is_output_rank:
# Flux decoding has no collectives. Avoid duplicating the VAE
# forward and device-to-host pixel copy on SPMD worker ranks whose
# result is discarded by the executor.
batch.output = torch.empty((0, 3, 0, 0, 0), device="cpu", dtype=torch.float32)
return batch
packed = batch.latents
if packed is None:
raise ValueError("latents must be set before FluxDecodingStage")
h_lat = int(batch.extra["flux_h_lat"])
w_lat = int(batch.extra["flux_w_lat"])
num_ch = int(batch.extra["flux_num_channels_latents"])
raw_shape = batch.raw_latent_shape
if raw_shape is None:
raise ValueError("raw_latent_shape missing; FluxLatentPreparationStage must run first.")
batch_size = int(raw_shape[0])
infer_device = get_local_torch_device()
packed = packed.to(infer_device)
latents_4d = _unpack_latents(packed, batch_size, num_ch, h_lat, w_lat)
latents_4d = self._denormalize_latents(latents_4d, self.vae)
vae_device = next(self.vae.parameters()).device
latents_4d = latents_4d.to(device=vae_device)
vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
autocast_enabled = (vae_dtype != torch.float32) and not fastvideo_args.disable_autocast
use_cuda_autocast = autocast_enabled and vae_device.type == "cuda"
with torch.autocast(
device_type="cuda",
dtype=vae_dtype,
enabled=use_cuda_autocast,
):
if not autocast_enabled:
latents_4d = latents_4d.to(dtype=vae_dtype)
dec = self.vae.decode(latents_4d)
image = dec.sample if hasattr(dec, "sample") else dec[0]
image = (image / 2 + 0.5).clamp(0, 1)
batch.output = image.unsqueeze(2).detach().float().cpu()
return batch</code></pre></details>