- Full Flux2 TP1 H100 pipeline parity is exact.

- L40S TP2/TP4 has bounded BF16 TP drift.
- H100 image artifacts were generated at 1024x1024 and look good.
- Full model image conditioning/caption upsampling remains out of scope.

Author: macthecadillac

examples/inference/basic/basic_flux2.py

@@ -0,0 +1,102 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Run full Flux2 text-to-image generation through FastVideo.
+
+User story:
+    "I have a local or HF Diffusers-format full Flux2 checkpoint and want a
+    minimal text-to-image generation command that uses embedded guidance."
+"""
+import argparse
+import os
+from pathlib import Path
+
+from fastvideo import VideoGenerator
+from fastvideo.api.sampling_param import SamplingParam
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run full Flux2 text-to-image generation.")
+    parser.add_argument(
+        "--model-path",
+        default="black-forest-labs/FLUX.2-dev",
+        help="HF id or local diffusers-format full Flux2 weights directory.",
+    )
+    parser.add_argument(
+        "--output",
+        default="outputs/flux2/flux2.png",
+        help="Output PNG path.",
+    )
+    parser.add_argument(
+        "--prompt",
+        default="a photo of a banana on a wooden table, studio lighting",
+        help="Text prompt.",
+    )
+    parser.add_argument("--height", type=int, default=1024)
+    parser.add_argument("--width", type=int, default=1024)
+    parser.add_argument("--steps", type=int, default=50)
+    parser.add_argument("--guidance-scale", type=float, default=4.0)
+    parser.add_argument("--max-sequence-length", type=int, default=None)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--num-gpus", type=int, default=1)
+    parser.add_argument("--tp-size", type=int, default=None)
+    parser.add_argument("--sp-size", type=int, default=None)
+    parser.add_argument(
+        "--backend",
+        default=None,
+        help="Set FASTVIDEO_ATTENTION_BACKEND, for example TORCH_SDPA.",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    if args.backend:
+        os.environ["FASTVIDEO_ATTENTION_BACKEND"] = args.backend
+
+    output = Path(args.output)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    tp_size = args.tp_size if args.tp_size is not None else (
+        args.num_gpus if args.num_gpus > 1 else 1
+    )
+    sp_size = args.sp_size if args.sp_size is not None else (
+        1 if args.num_gpus > 1 else args.num_gpus
+    )
+
+    generator = VideoGenerator.from_pretrained(
+        args.model_path,
+        num_gpus=args.num_gpus,
+        tp_size=tp_size,
+        sp_size=sp_size,
+        workload_type="t2i",
+        use_fsdp_inference=False,
+        dit_cpu_offload=False,
+        vae_cpu_offload=True,
+        text_encoder_cpu_offload=True,
+        pin_cpu_memory=False,
+        override_pipeline_cls_name="Flux2Pipeline",
+    )
+    try:
+        sampling = SamplingParam.from_pretrained(args.model_path)
+        sampling.prompt = args.prompt
+        sampling.height = args.height
+        sampling.width = args.width
+        sampling.num_frames = 1
+        sampling.fps = 1
+        sampling.num_inference_steps = args.steps
+        sampling.guidance_scale = args.guidance_scale
+        sampling.max_sequence_length = args.max_sequence_length
+        sampling.seed = args.seed
+        sampling.output_path = str(output)
+        sampling.save_video = True
+        sampling.return_frames = False
+
+        generator.generate_video(
+            args.prompt,
+            sampling_param=sampling,
+            output_path=str(output),
+        )
+    finally:
+        generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/basic_flux2_klein.py

@@ -52,6 +52,7 @@ def main() -> None:
     generator = VideoGenerator.from_pretrained(
         args.model_path,
         num_gpus=args.num_gpus,
+        workload_type="t2i",
         use_fsdp_inference=False,
         dit_cpu_offload=False,
         vae_cpu_offload=True,

fastvideo/api/sampling_param.py

@@ -29,6 +29,10 @@ class SamplingParam:
     # Video inputs
     video_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
+
     # Action control inputs (Matrix-Game)
     mouse_cond: Any | None = None  # Shape: (B, T, 2)
     keyboard_cond: Any | None = None  # Shape: (B, T, K)
@@ -64,6 +68,7 @@ class SamplingParam:
     # 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

fastvideo/configs/models/encoders/__init__.py

@@ -7,6 +7,7 @@
 from fastvideo.configs.models.encoders.siglip import SiglipVisionConfig
 from fastvideo.configs.models.encoders.reason1 import Reason1ArchConfig, Reason1Config
 from fastvideo.configs.models.encoders.gemma import LTX2GemmaConfig
+from fastvideo.configs.models.encoders.mistral3 import Mistral3TextConfig
 from fastvideo.configs.models.encoders.qwen3 import Qwen3TextConfig
 from fastvideo.configs.models.encoders.stable_audio_conditioner import (StableAudioConditionerArchConfig,
                                                                         StableAudioConditionerConfig)
@@ -16,5 +17,5 @@
     "EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig", "BaseEncoderOutput", "CLIPTextConfig",
     "CLIPVisionConfig", "WAN2_1ControlCLIPVisionConfig", "LlamaConfig", "T5Config", "T5LargeConfig", "Qwen2_5_VLConfig",
     "Reason1ArchConfig", "Reason1Config", "LTX2GemmaConfig", "SiglipVisionConfig", "StableAudioConditionerArchConfig",
-    "StableAudioConditionerConfig", "T5GemmaEncoderConfig", "Qwen3TextConfig"
+    "StableAudioConditionerConfig", "T5GemmaEncoderConfig", "Qwen3TextConfig", "Mistral3TextConfig"
 ]

fastvideo/configs/models/encoders/mistral3.py

@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Mistral3 text encoder configuration for full Flux2."""
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.encoders.base import (
+    TextEncoderArchConfig,
+    TextEncoderConfig,
+)
+
+
+@dataclass
+class Mistral3TextArchConfig(TextEncoderArchConfig):
+    """Architecture config for the Mistral3 text encoder used by full Flux2."""
+
+    architectures: list[str] = field(default_factory=lambda: ["Mistral3ForConditionalGeneration"])
+    hidden_size: int = 5120
+    num_hidden_layers: int = 40
+    text_len: int = 512
+    output_hidden_states: bool = True
+
+    def __post_init__(self) -> None:
+        self.tokenizer_kwargs = {
+            "padding": "max_length",
+            "truncation": True,
+            "max_length": self.text_len,
+            "return_tensors": "pt",
+        }
+
+
+@dataclass
+class Mistral3TextConfig(TextEncoderConfig):
+    """Top-level config for the Mistral3 full Flux2 text encoder."""
+
+    arch_config: TextEncoderArchConfig = field(default_factory=Mistral3TextArchConfig)
+    prefix: str = "mistral3"
+    is_chat_model: bool = True

fastvideo/configs/pipelines/flux_2.py

@@ -9,6 +9,7 @@
 from fastvideo.configs.models.dits.flux_2 import Flux2Config
 from fastvideo.configs.models.encoders import BaseEncoderOutput
 from fastvideo.configs.models.encoders.base import EncoderArchConfig
+from fastvideo.configs.models.encoders.mistral3 import Mistral3TextConfig
 from fastvideo.configs.models.encoders.qwen3 import Qwen3TextConfig
 from fastvideo.configs.models.vaes.flux2vae import Flux2VAEConfig
 from fastvideo.configs.pipelines.base import PipelineConfig, preprocess_text
@@ -17,37 +18,34 @@
 @dataclass
 class Flux2PipelineConfig(PipelineConfig):
     """Configuration for Flux2 image generation pipeline."""
-    
+
     # Flux2-specific parameters
-    embedded_cfg_scale: float = 4.0
-    
+    embedded_cfg_scale: float | None = 4.0
+    flux2_text_encoder_type: str = "mistral3"
+    text_encoder_out_layers: tuple[int, ...] = (10, 20, 30)
+
     # DiT configuration
     dit_config: DiTConfig = field(default_factory=Flux2Config)
     dit_precision: str = "bf16"
-    
+
     # VAE configuration
     vae_config: VAEConfig = field(default_factory=Flux2VAEConfig)
     vae_precision: str = "fp32"
     vae_tiling: bool = False  # Flux2 is image model, disable tiling by default
     vae_sp: bool = False
-    
-    # Text encoder configuration (Flux2 uses Mistral/Qwen)
-    text_encoder_configs: tuple[EncoderConfig, ...] = field(
-        default_factory=lambda: (EncoderConfig(),)
-    )
-    text_encoder_precisions: tuple[str, ...] = field(
-        default_factory=lambda: ("bf16",)
-    )
-    
+
+    # Text encoder configuration (full Flux2 uses Mistral3)
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (Mistral3TextConfig(), ))
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", ))
+
     # Default postprocess function (can be overridden)
     @staticmethod
     def default_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
         """Default text postprocessing for Flux2."""
         return outputs.last_hidden_state
-    
-    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor], ...] = field(
-        default_factory=lambda: (Flux2PipelineConfig.default_postprocess_text,)
-    )
+
+    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
+                                  ...] = field(default_factory=lambda: (Flux2PipelineConfig.default_postprocess_text, ))
 
 
 def flux2_klein_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
@@ -57,9 +55,7 @@ def flux2_klein_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
         raise ValueError("Flux2 Klein requires output_hidden_states=True from text encoder")
     out = torch.stack([outputs.hidden_states[k] for k in hidden_states_layers], dim=1)
     batch_size, num_channels, seq_len, hidden_dim = out.shape
-    prompt_embeds = out.permute(0, 2, 1, 3).reshape(
-        batch_size, seq_len, num_channels * hidden_dim
-    )
+    prompt_embeds = out.permute(0, 2, 1, 3).reshape(batch_size, seq_len, num_channels * hidden_dim)
     return prompt_embeds
 
 
@@ -79,15 +75,10 @@ class Flux2KleinTextEncoderConfig(EncoderConfig):
 class Flux2KleinPipelineConfig(Flux2PipelineConfig):
     """Configuration for Flux2 Klein (distilled, 4-step, no guidance)."""
     embedded_cfg_scale: float | None = None  # Klein distilled: no guidance embedding (matches Diffusers)
-    text_encoder_configs: tuple[EncoderConfig, ...] = field(
-        default_factory=lambda: (Qwen3TextConfig(),)
-    )
-    text_encoder_precisions: tuple[str, ...] = field(
-        default_factory=lambda: ("bf16",)
-    )
-    preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(
-        default_factory=lambda: (preprocess_text,)
-    )
-    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor], ...] = field(
-        default_factory=lambda: (flux2_klein_postprocess_text,)
-    )
+    flux2_text_encoder_type: str = "qwen3"
+    text_encoder_out_layers: tuple[int, ...] = (9, 18, 27)
+    text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=lambda: (Qwen3TextConfig(), ))
+    text_encoder_precisions: tuple[str, ...] = field(default_factory=lambda: ("bf16", ))
+    preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(default_factory=lambda: (preprocess_text, ))
+    postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor],
+                                  ...] = field(default_factory=lambda: (flux2_klein_postprocess_text, ))

fastvideo/entrypoints/video_generator.py

@@ -383,6 +383,7 @@ def _generate_video_impl(
         if grid_sizes is not None:
             kwargs['grid_sizes'] = grid_sizes
 
+        prompt_embeds = kwargs.pop("prompt_embeds", None)
         sampling_param.update(kwargs)
 
         if fastvideo_args.prompt_txt is not None or sampling_param.prompt_path is not None:
@@ -432,6 +433,8 @@ def _generate_video_impl(
             raise ValueError("Either prompt or prompt_txt must be provided")
         output_path = self._prepare_output_path(sampling_param.output_path, prompt)
         kwargs["output_path"] = output_path
+        if prompt_embeds is not None:
+            kwargs["prompt_embeds"] = prompt_embeds
         return self._generate_single_video(
             prompt=prompt,
             sampling_param=sampling_param,
@@ -534,6 +537,7 @@ def _generate_single_video(
         prompt = prompt.strip()
         sampling_param = deepcopy(sampling_param)
         output_path = kwargs["output_path"]
+        prompt_embeds = kwargs.get("prompt_embeds")
         sampling_param.prompt = prompt
         # Process negative prompt
         if sampling_param.negative_prompt is not None:
@@ -581,12 +585,8 @@ def _generate_single_video(
             VSA_sparsity=fastvideo_args.VSA_sparsity,
         )
         # Allow precomputed prompt_embeds (e.g. from diffusers) to skip text encoding
-        if kwargs.get("prompt_embeds") is not None:
-            batch.prompt_embeds = (
-                list(kwargs["prompt_embeds"])
-                if isinstance(kwargs["prompt_embeds"], (list, tuple))
-                else [kwargs["prompt_embeds"]]
-            )
+        if prompt_embeds is not None:
+            batch.prompt_embeds = (list(prompt_embeds) if isinstance(prompt_embeds, list | tuple) else [prompt_embeds])
 
         # Run inference
         start_time = time.perf_counter()

fastvideo/models/dits/flux_2.py

@@ -988,6 +988,8 @@ def forward(
         hidden_states: torch.Tensor,
         encoder_hidden_states: torch.Tensor = None,
         timestep: torch.LongTensor = None,
+        img_ids: torch.Tensor = None,
+        txt_ids: torch.Tensor = None,
         guidance: torch.Tensor = None,
         freqs_cis: torch.Tensor = None,
         joint_attention_kwargs: Optional[Dict[str, Any]] = None,
@@ -1047,21 +1049,22 @@ def forward(
 
         # 3. Compute RoPE positional embeddings when not provided externally
         if freqs_cis is None:
-            if input_was_5d:
-                img_h, img_w = h, w
-            else:
-                img_seq_len = hidden_states.shape[1]
-                img_h = img_w = int(img_seq_len ** 0.5)
-
-            txt_len = encoder_hidden_states.shape[1]
-            txt_ids = torch.cartesian_prod(
-                torch.arange(1), torch.arange(1),
-                torch.arange(1), torch.arange(txt_len),
-            ).to(device=hidden_states.device)
-            img_ids = torch.cartesian_prod(
-                torch.arange(1), torch.arange(img_h),
-                torch.arange(img_w), torch.arange(1),
-            ).to(device=hidden_states.device)
+            if txt_ids is None or img_ids is None:
+                if input_was_5d:
+                    img_h, img_w = h, w
+                else:
+                    img_seq_len = hidden_states.shape[1]
+                    img_h = img_w = int(img_seq_len ** 0.5)
+
+                txt_len = encoder_hidden_states.shape[1]
+                txt_ids = torch.cartesian_prod(
+                    torch.arange(1), torch.arange(1),
+                    torch.arange(1), torch.arange(txt_len),
+                ).to(device=hidden_states.device)
+                img_ids = torch.cartesian_prod(
+                    torch.arange(1), torch.arange(img_h),
+                    torch.arange(img_w), torch.arange(1),
+                ).to(device=hidden_states.device)
             freqs_cis = compute_flux2_freqs_cis_from_ids(
                 self.rotary_emb, txt_ids, img_ids, device=hidden_states.device,
             )