--- title: flux_stages ---

fastvideo.pipelines.stages.flux_stages.get_local_torch_device

get_local_torch_device() -> torch.device

Return 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: bool

Whether 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) -> FlexibleArgumentParser

fastvideo.pipelines.stages.flux_stages.FastVideoArgs.from_cli_args

from_cli_args(args: argparse.Namespace) -> FastVideoArgs

fastvideo.pipelines.stages.flux_stages.FastVideoArgs.from_kwargs

from_kwargs(**kwargs: Any) -> FastVideoArgs

fastvideo.pipelines.stages.flux_stages.FastVideoArgs.check_fastvideo_args

check_fastvideo_args() -> None

Validate inference arguments for consistency

fastvideo.pipelines.stages.flux_stages.FastVideoArgs.finalize_device_offload_policy

finalize_device_offload_policy(device_id: int = 0) -> bool

Apply 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) -> bool

Disable 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) -> _FastvideoLogger

The 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) -> VerificationResult

Verify 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: ForwardBatch

The current batch information.

fastvideo_args: FastVideoArgs

The inference arguments.

Returns

VerificationResult

A VerificationResult containing the verification status.

fastvideo.pipelines.stages.flux_stages.PipelineStage.verify_output

verify_output(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult

Verify the output for the stage.

Parameters

batch: ForwardBatch

The current batch information.

fastvideo_args: FastVideoArgs

The inference arguments.

Returns

VerificationResult

A VerificationResult containing the verification status.

fastvideo.pipelines.stages.flux_stages.PipelineStage.device

device: torch.device

Get 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: bool

Whether to enable logging.

fastvideo.pipelines.stages.flux_stages.PipelineStage.forward

forward(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch

Forward pass of the stage's processing.

This method should be implemented by subclasses to provide the forward processing logic for the stage.

Parameters

batch: ForwardBatch

The current batch information.

fastvideo_args: FastVideoArgs

The inference arguments.

Returns

ForwardBatch

The 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) -> ForwardBatch

Validate and prepare inputs.

Parameters

batch: ForwardBatch

The current batch information.

fastvideo_args: FastVideoArgs

The inference arguments.

Returns

ForwardBatch

The validated batch information.

fastvideo.pipelines.stages.flux_stages.InputValidationStage.verify_input

verify_input(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult

Verify input validation stage inputs.

fastvideo.pipelines.stages.flux_stages.InputValidationStage.verify_output

verify_output(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult

Verify 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) -> ForwardBatch

Prepare timesteps for the diffusion process.

Parameters

batch: ForwardBatch

The current batch information.

fastvideo_args: FastVideoArgs

The inference arguments.

Returns

ForwardBatch

The batch with prepared timesteps.

fastvideo.pipelines.stages.flux_stages.TimestepPreparationStage.verify_input

verify_input(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult

Verify timestep preparation stage inputs.

fastvideo.pipelines.stages.flux_stages.TimestepPreparationStage.verify_output

verify_output(batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult

Verify 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)

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

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

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

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

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

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

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

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

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

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