[model] Flux2 Klein Port (#1349)

Co-authored-by: Gnav3852 <63612880+Gnav3852@users.noreply.github.com>
Co-authored-by: Mac Lee <macthecadillac@gmail.com>

Author: 3 people

.gitignore

@@ -34,6 +34,8 @@ env
 *.log
 weights/
 logs/
+official_weights/
+converted_weights/
 
 # SSIM test outputs
 fastvideo/tests/ssim/generated_videos/

examples/inference/basic/basic_flux2.py

@@ -0,0 +1,124 @@
+# 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 import (
+    ComponentConfig,
+    EngineConfig,
+    GenerationRequest,
+    GeneratorConfig,
+    OffloadConfig,
+    OutputConfig,
+    ParallelismConfig,
+    PipelineSelection,
+    SamplingConfig,
+)
+
+
+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_config = GeneratorConfig(
+        model_path=args.model_path,
+        engine=EngineConfig(
+            num_gpus=args.num_gpus,
+            parallelism=ParallelismConfig(tp_size=tp_size, sp_size=sp_size),
+            use_fsdp_inference=False,
+            offload=OffloadConfig(
+                dit=False,
+                vae=True,
+                text_encoder=True,
+                pin_cpu_memory=False,
+            ),
+        ),
+        pipeline=PipelineSelection(
+            workload_type="t2i",
+            components=ComponentConfig(override_pipeline_cls_name="Flux2Pipeline"),
+        ),
+    )
+
+    generator = VideoGenerator.from_config(generator_config)
+    try:
+        sampling = SamplingConfig(
+            height=args.height,
+            width=args.width,
+            num_frames=1,
+            fps=1,
+            num_inference_steps=args.steps,
+            guidance_scale=args.guidance_scale,
+            seed=args.seed,
+        )
+        extensions = {}
+        if args.max_sequence_length is not None:
+            extensions["max_sequence_length"] = args.max_sequence_length
+
+        request = GenerationRequest(
+            prompt=args.prompt,
+            sampling=sampling,
+            output=OutputConfig(
+                output_path=str(output),
+                save_video=True,
+                return_frames=False,
+            ),
+            extensions=extensions,
+        )
+        generator.generate(request)
+    finally:
+        generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/basic_flux2_klein.py

@@ -0,0 +1,98 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Run Flux2 Klein text-to-image generation through FastVideo.
+
+User story:
+    "I need a short local smoke for the Flux2 Klein checkpoint before wiring it
+    into an image workflow. Use the model's distilled four-step defaults and
+    write a single PNG so I can compare the output against the reference."
+"""
+
+import argparse
+import os
+
+from fastvideo import VideoGenerator
+from fastvideo.api import (
+    EngineConfig,
+    GenerationRequest,
+    GeneratorConfig,
+    OffloadConfig,
+    OutputConfig,
+    PipelineSelection,
+    SamplingConfig,
+)
+
+
+DEFAULT_PROMPT = "a brushed steel espresso machine on a marble counter, morning window light"
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run Flux2 Klein text-to-image generation.")
+    parser.add_argument(
+        "--model-path",
+        default="black-forest-labs/FLUX.2-klein-4B",
+        help="HF id or local diffusers-format Flux2 Klein weights directory.",
+    )
+    parser.add_argument(
+        "--output-path",
+        default="outputs/flux2/flux2_klein.png",
+        help="PNG output path or output directory.",
+    )
+    parser.add_argument("--prompt", default=DEFAULT_PROMPT, help="Prompt text.")
+    parser.add_argument("--seed", type=int, default=0, help="Generation seed.")
+    parser.add_argument("--height", type=int, default=1024, help="Output image height.")
+    parser.add_argument("--width", type=int, default=1024, help="Output image width.")
+    parser.add_argument("--steps", type=int, default=4, help="Number of denoising steps.")
+    parser.add_argument("--num-gpus", type=int, default=1, help="Number of GPUs to use.")
+    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
+
+    generator_config = GeneratorConfig(
+        model_path=args.model_path,
+        engine=EngineConfig(
+            num_gpus=args.num_gpus,
+            use_fsdp_inference=False,
+            offload=OffloadConfig(
+                dit=False,
+                vae=True,
+                text_encoder=True,
+                pin_cpu_memory=False,
+            ),
+        ),
+        pipeline=PipelineSelection(workload_type="t2i"),
+    )
+
+    generator = VideoGenerator.from_config(generator_config)
+    try:
+        request = GenerationRequest(
+            prompt=args.prompt,
+            sampling=SamplingConfig(
+                height=args.height,
+                width=args.width,
+                num_frames=1,
+                fps=1,
+                num_inference_steps=args.steps,
+                guidance_scale=1.0,
+                seed=args.seed,
+            ),
+            output=OutputConfig(
+                output_path=args.output_path,
+                save_video=True,
+            ),
+        )
+        generator.generate(request)
+    finally:
+        generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

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/dits/__init__.py

@@ -1,5 +1,6 @@
 from fastvideo.configs.models.dits.cosmos import CosmosVideoConfig
 from fastvideo.configs.models.dits.cosmos2_5 import Cosmos25VideoConfig
+from fastvideo.configs.models.dits.flux_2 import Flux2Config
 from fastvideo.configs.models.dits.hunyuangamecraft import HunyuanGameCraftConfig
 from fastvideo.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
 from fastvideo.configs.models.dits.hunyuanvideo15 import HunyuanVideo15Config
@@ -14,5 +15,5 @@
 __all__ = [
     "HunyuanVideoConfig", "HunyuanVideo15Config", "HunyuanGameCraftConfig", "WanVideoConfig", "CosmosVideoConfig",
     "Cosmos25VideoConfig", "LongCatVideoConfig", "LTX2VideoConfig", "HYWorldConfig", "Kandinsky5VideoConfig",
-    "MagiHumanVideoConfig", "StableAudioConfig"
+    "MagiHumanVideoConfig", "StableAudioConfig", "Flux2Config"
 ]

fastvideo/configs/models/dits/base.py

@@ -14,6 +14,11 @@ class DiTArchConfig(ArchConfig):
     param_names_mapping: dict = field(default_factory=dict)
     reverse_param_names_mapping: dict = field(default_factory=dict)
     lora_param_names_mapping: dict = field(default_factory=dict)
+    # When True, the denoising stage casts text/prompt embeddings to the DiT's
+    # working dtype before the diffusion loop. Flux2 requires this (BFL casts ctx
+    # to bf16 before denoising); models with fp32 text encoders (Wan, Hunyuan15,
+    # SD3.5) leave it False to preserve full-precision embeddings.
+    cast_prompt_embeds_to_dit_dtype: bool = False
     _supported_attention_backends: tuple[AttentionBackendEnum,
                                          ...] = (AttentionBackendEnum.SAGE_ATTN, AttentionBackendEnum.FLASH_ATTN,
                                                  AttentionBackendEnum.TORCH_SDPA,

fastvideo/configs/models/dits/flux_2.py

@@ -0,0 +1,77 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copied and adapted from: https://github.com/sglang-ai/sglang
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class Flux2ArchConfig(DiTArchConfig):
+    """Architecture configuration for Flux2 transformer model."""
+
+    cast_prompt_embeds_to_dit_dtype: bool = True
+
+    # Flux2-specific architecture parameters
+    patch_size: int = 1
+    in_channels: int = 64
+    out_channels: int | None = None
+    num_layers: int = 19  # Number of double-stream transformer blocks
+    num_single_layers: int = 38  # Number of single-stream transformer blocks
+    attention_head_dim: int = 128
+    num_attention_heads: int = 24
+    joint_attention_dim: int = 4096  # Dimension for text encoder output
+    timestep_guidance_channels: int = 256  # Dimension for timestep embedding
+    mlp_ratio: float = 3.0
+    axes_dims_rope: tuple[int, ...] = (32, 32, 32, 32)  # RoPE dimensions per axis (match diffusers Flux2)
+    rope_theta: int = 2000  # Base frequency for RoPE (match diffusers Flux2)
+    eps: float = 1e-6
+    guidance_embeds: bool = True  # Whether to use guidance embeddings
+    # When True, compute SwiGLU in fp32 inside ``ff_context`` only (bf16 noise mitigation).
+    ff_context_swiglu_fp32: bool = False
+
+    # Parameter name mapping for loading HuggingFace checkpoints
+    param_names_mapping: dict = field(default_factory=lambda: {
+        r"transformer\.(\w*)\.(.*)$": r"\1.\2",
+    })
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        self.out_channels = self.out_channels or self.in_channels
+        self.hidden_size = self.num_attention_heads * self.attention_head_dim
+        self.num_channels_latents = self.out_channels
+
+    def update_from_weight_keys(self, all_keys: set[str]) -> None:
+        """Infer num_layers and num_single_layers from checkpoint weight keys so the model is built with the same number of blocks as the weights."""
+        if not all_keys:
+            return
+        num_layers = 0
+        num_single_layers = 0
+        for k in all_keys:
+            if "single_transformer_blocks." not in k and "transformer_blocks." in k:
+                parts = k.split("transformer_blocks.")[-1].split(".")
+                if parts[0].isdigit():
+                    num_layers = max(num_layers, int(parts[0]) + 1)
+            if "single_transformer_blocks." in k:
+                parts = k.split("single_transformer_blocks.")[-1].split(".")
+                if parts[0].isdigit():
+                    num_single_layers = max(num_single_layers, int(parts[0]) + 1)
+        if num_layers > 0:
+            self.num_layers = num_layers
+            logger.info("Inferred num_layers=%s from checkpoint keys", num_layers)
+        if num_single_layers > 0:
+            self.num_single_layers = num_single_layers
+            logger.info("Inferred num_single_layers=%s from checkpoint keys", num_single_layers)
+        if num_layers > 0 or num_single_layers > 0:
+            self.__post_init__()
+
+
+@dataclass
+class Flux2Config(DiTConfig):
+    """Configuration for Flux2 transformer model."""
+
+    arch_config: DiTArchConfig = field(default_factory=Flux2ArchConfig)
+
+    prefix: str = "Flux"

fastvideo/configs/models/encoders/__init__.py

@@ -7,6 +7,8 @@
 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)
 from fastvideo.configs.models.encoders.t5gemma import T5GemmaEncoderConfig
@@ -15,5 +17,5 @@
     "EncoderConfig", "TextEncoderConfig", "ImageEncoderConfig", "BaseEncoderOutput", "CLIPTextConfig",
     "CLIPVisionConfig", "WAN2_1ControlCLIPVisionConfig", "LlamaConfig", "T5Config", "T5LargeConfig", "Qwen2_5_VLConfig",
     "Reason1ArchConfig", "Reason1Config", "LTX2GemmaConfig", "SiglipVisionConfig", "StableAudioConditionerArchConfig",
-    "StableAudioConditionerConfig", "T5GemmaEncoderConfig"
+    "StableAudioConditionerConfig", "T5GemmaEncoderConfig", "Qwen3TextConfig", "Mistral3TextConfig"
 ]

fastvideo/configs/models/encoders/base.py

@@ -36,6 +36,11 @@ class TextEncoderArchConfig(EncoderArchConfig):
         default_factory=list)  # mapping from huggingface weight names to custom names
     tokenizer_kwargs: dict[str, Any] = field(default_factory=dict)
     _fsdp_shard_conditions: list = field(default_factory=lambda: [])
+    # When True, the tokenizer loader prefers AutoProcessor over AutoTokenizer
+    # for encoders whose tokenizer dir ships a processor_config.json (e.g. Flux2
+    # full's Mistral3 multimodal processor). Default False keeps every existing
+    # encoder on the historical AutoTokenizer path.
+    require_processor: bool = False
 
     def __post_init__(self) -> None:
         self.tokenizer_kwargs = {