[bugfix] Classify new config/sampling fields in schema parity inventory (#1446)

Co-authored-by: Aryan Kumar <aryan5v@users.noreply.github.com>

.agents/memory/codebase-map/README.md

@@ -117,30 +117,6 @@ FastVideo-WorldModel/
 | `TOKENIZERS_PARALLELISM` | Set `false` to avoid fork warnings |
 | `HF_HOME` | HuggingFace cache directory |
 
-## RL Reward Runtime Vendoring
-
-Commit `78b6995e` (`WIP: vendor HPSv3 and VideoAlign reward runtimes`) added
-vendored runtime code under `fastvideo/train/methods/rl/reward/HPSv3` and
-`fastvideo/train/methods/rl/reward/VideoAlign`, replacing broken gitlinks with
-normal tracked files.
-
-The provenance and porting rules live in the vendor package markers:
-`fastvideo/train/methods/rl/reward/HPSv3/__init__.py`,
-`fastvideo/train/methods/rl/reward/HPSv3/hpsv3/__init__.py`, and
-`fastvideo/train/methods/rl/reward/VideoAlign/__init__.py`.
-
-The FastVideo wrapper scripts `hpsv3.py` and `videoalign.py` depend on these
-vendored packages through explicit package imports. The copied runtime files
-are ported faithfully from upstream, with only import-path changes needed for
-package importability; this means some third-party code is unrelated to
-FastVideo internals but remains faithful to the source.
-
-Known follow-ups from commit `78b6995e`: the vendored code was not expected to
-pass pre-commit at that commit. VideoAlign also assumed checkpoint artifacts
-already existed under its checkpoints path or the configured
-`VIDEOALIGN_CHECKPOINT_PATH`; later work resolved the default checkpoint path by
-downloading the `KlingTeam/VideoReward` Hugging Face snapshot.
-
 ## Build & Test Commands
 
 ```bash

.gitignore

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

docker/Dockerfile.python3.10

@@ -55,12 +55,14 @@ RUN source $HOME/.local/bin/env && \
     echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \
     echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
 
-# Install FastVideo Unified Kernel
+# Install FastVideo Unified Kernel.
+# This build machine has no GPU, so target Hopper (sm_90a) explicitly instead
+# of probing a live device for the arch (matches the released kernel wheel).
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
     cd fastvideo-kernel && \
     git submodule update --init --recursive && \
-    ./build.sh
+    TORCH_CUDA_ARCH_LIST=9.0a ./build.sh
 
 
 EXPOSE 22

docker/Dockerfile.python3.11

@@ -55,12 +55,14 @@ RUN source $HOME/.local/bin/env && \
     echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \
     echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
 
-# Install FastVideo Unified Kernel
+# Install FastVideo Unified Kernel.
+# This build machine has no GPU, so target Hopper (sm_90a) explicitly instead
+# of probing a live device for the arch (matches the released kernel wheel).
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
     cd fastvideo-kernel && \
     git submodule update --init --recursive && \
-    ./build.sh
+    TORCH_CUDA_ARCH_LIST=9.0a ./build.sh
 
 
 EXPOSE 22

docker/Dockerfile.python3.12

@@ -55,11 +55,13 @@ RUN source $HOME/.local/bin/env && \
     echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \
     echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
 
-# Install FastVideo Unified Kernel
+# Install FastVideo Unified Kernel.
+# This build machine has no GPU, so target Hopper (sm_90a) explicitly instead
+# of probing a live device for the arch (matches the released kernel wheel).
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
     cd fastvideo-kernel && \
     git submodule update --init --recursive && \
-    ./build.sh
+    TORCH_CUDA_ARCH_LIST=9.0a ./build.sh
 
 EXPOSE 22

docker/Dockerfile.python3.12.cuda12.9.1

@@ -55,12 +55,14 @@ RUN source $HOME/.local/bin/env && \
     echo 'source /opt/venv/bin/activate' >> /root/.bashrc && \
     echo 'if [ -n "$ZSH_VERSION" ] && [ -f ~/.zshrc ]; then . ~/.zshrc; elif [ -f ~/.bashrc ]; then . ~/.bashrc; fi' > /root/.profile
 
-# Install FastVideo Unified Kernel
+# Install FastVideo Unified Kernel.
+# This build machine has no GPU, so target Hopper (sm_90a) explicitly instead
+# of probing a live device for the arch (matches the released kernel wheel).
 RUN source $HOME/.local/bin/env && \
     source /opt/venv/bin/activate && \
     cd fastvideo-kernel && \
     git submodule update --init --recursive && \
-    ./build.sh
+    TORCH_CUDA_ARCH_LIST=9.0a ./build.sh
 
 
 EXPOSE 22

docs/design/inference_schema_parity_inventory.yaml

@@ -108,6 +108,7 @@ surfaces:
       vae_sp: generator.pipeline.preset_overrides.vae_sp
       dmd_denoising_steps: generator.pipeline.preset_overrides.dmd_denoising_steps
       ti2v_task: generator.pipeline.preset_overrides.ti2v_task
+      lucy_edit_task: generator.pipeline.preset_overrides.lucy_edit_task
       boundary_ratio: generator.pipeline.preset_overrides.boundary_ratio
     compatibility_only:
       model_path: "Redundant with generator.model_path."
@@ -125,9 +126,18 @@ surfaces:
       text_encoder_configs: "Legacy internal component config object."
       preprocess_text_funcs: "Internal text preprocessing hooks."
       postprocess_text_funcs: "Internal text postprocessing hooks."
+      scheduler_step_in_fp32: "Runtime scheduler precision toggle; not part of the public typed inference API."
 
   pipeline_config_extensions:
     preset_owned:
+      flux2_text_encoder_type:
+        sources:
+          - fastvideo.configs.pipelines.flux_2.Flux2PipelineConfig
+          - fastvideo.configs.pipelines.flux_2.Flux2KleinPipelineConfig
+      text_encoder_out_layers:
+        sources:
+          - fastvideo.configs.pipelines.flux_2.Flux2PipelineConfig
+          - fastvideo.configs.pipelines.flux_2.Flux2KleinPipelineConfig
       conditioning_strategy:
         sources:
           - fastvideo.configs.pipelines.cosmos.CosmosConfig
@@ -491,6 +501,8 @@ surfaces:
       inpaint_mask: request.extensions.stable_audio.inpaint_mask
     internal_only:
       data_type: "Derived from the request shape and not a public input."
+      latents: "Pre-generated diffusion latents supplied by parity/debug harnesses; not a public input."
+      max_sequence_length: "Model-specific text-encoder sequence cap; not part of the public typed inference API."
 
   sampling_param_extensions: {}
 

docs/inference/support_matrix.md

@@ -58,6 +58,7 @@ pipeline initialization and sampling.
 | FastWan2.1 T2V 1.3B | `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` | 480P | ⭕ | ⭕ | ⭕ | ✅ | ⭕ |
 | FastWan2.2 TI2V 5B Full Attn* | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | 720P | ⭕ | ⭕ | ⭕ | ✅ | ⭕ |
 | Wan2.2 TI2V 5B | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 720P | ⭕ | ⭕ | ✅ | ⭕ | ⭕ |
+| Lucy Edit Dev 5B*** | `decart-ai/Lucy-Edit-Dev` | 480P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
 | Wan2.2 T2V A14B | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 480P<br>720P | ❌ | ❌ | ✅ | ⭕ | ⭕ |
 | Wan2.2 I2V A14B | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480P<br>720P | ❌ | ❌ | ✅ | ⭕ | ⭕ |
 | HunyuanVideo | `hunyuanvideo-community/HunyuanVideo` | 720px1280p<br>544px960p | ❌ | ✅ | ✅ | ⭕ | ⭕ |
@@ -78,6 +79,9 @@ pipeline initialization and sampling.
 
 **Note**: Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue.
 
+***Lucy Edit Dev uses a non-commercial model license. FastVideo support is
+focused on inference integration for video editing workflows.
+
 `Sliding Tile Attn (Legacy Branch)` entries refer to the archived
 `sta_do_not_delete` branch workflow, not active `main` inference wiring.
 

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

examples/inference/basic/basic_ltx2_3_distilled_i2v.py

@@ -0,0 +1,289 @@
+# SPDX-License-Identifier: Apache-2.0
+"""LTX-2.3 distilled image-to-video with torch.compile + timing breakdown.
+
+This example runs the LTX-2.3 distilled student model on a single GPU with
+torch.compile fully enabled, then prints a per-stage timing breakdown so the
+user can see where wall-time goes. It is meant as the canonical entry point
+for trying out the LTX-2.3 i2v path on `hao-ai-lab/FastVideo:main`.
+
+Quick start
+-----------
+    export LTX23_I2V_IMAGE=/path/to/your/portrait_or_product.jpg
+    # optional overrides:
+    #   export LTX23_I2V_PROMPT="a fashion model walks toward camera..."
+    #   export LTX23_OUTPUT_DIR=outputs_video/ltx2_3_distilled_i2v
+    python examples/inference/basic/basic_ltx2_3_distilled_i2v.py
+
+What the script does
+--------------------
+1. Loads FastVideo/LTX-2.3-Distilled-Diffusers (8 denoise + 3 refine steps,
+   CFG=1, no refine LoRA — the distilled production recipe).
+2. Compiles the DiT, text encoder, and VAE (fullgraph, Inductor default
+   mode — autotune adds ~7 min cold-compile here with no measurable
+   e2e gain).
+3. Runs 2 warmup calls (untimed) + 2 measured calls. Two warmups are kept
+   as a safety net — the first call pays cold compile + first-shape guard
+   work, and a second warmup ensures any residual recompiles settle before
+   we measure.
+4. Prints a per-stage breakdown and an average over the measured runs.
+
+Hardware notes
+--------------
+- Single-GPU example; for multi-GPU sequence-parallel see the gradio demo
+  under `examples/inference/gradio/local/gradio_local_demo_ltx2_3/`.
+- First-time compile takes ~30-40 min on GB200 (~20 min on H100; cached
+  in `$TORCHINDUCTOR_CACHE_DIR` afterwards). Subsequent invocations only
+  pay the one-time process load + a few seconds of dynamo trace.
+- On GB200 / Blackwell, run with `env -u LD_LIBRARY_PATH ...` to avoid a
+  system-cuBLAS / torch-cuBLAS mismatch that fails every GEMM. The
+  `_inductor.shape_padding = False` line below also avoids a pad_mm
+  landmine on the same generation of cards.
+"""
+from __future__ import annotations
+
+import os
+import time
+from collections import OrderedDict
+from pathlib import Path
+
+import torch._inductor.config as _inductor
+
+from fastvideo import VideoGenerator
+from fastvideo.configs.pipelines.base import PipelineConfig
+from fastvideo.utils import maybe_download_model
+
+# Env knobs (set BEFORE importing fastvideo where possible — but
+# FASTVIDEO_ATTENTION_BACKEND is fine here because the worker reads it
+# on generator construction).
+os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN")
+os.environ.setdefault("FASTVIDEO_STAGE_LOGGING", "1")
+
+# Inductor knobs. The first one (shape_padding=False) is mandatory on
+# Blackwell to avoid a cuBLAS INVALID_VALUE crash inside pad_mm during
+# the refine path. The rest are autotune-friendliness flags.
+_inductor.shape_padding = False
+_inductor.conv_1x1_as_mm = True
+_inductor.coordinate_descent_tuning = True
+_inductor.coordinate_descent_check_all_directions = True
+_inductor.epilogue_fusion = False
+
+MODEL_ID = os.path.expandvars(
+    os.path.expanduser(
+        os.getenv("LTX23_MODEL_PATH", "FastVideo/LTX-2.3-Distilled-Diffusers")
+    )
+)
+OUTPUT_DIR = Path(
+    os.getenv("LTX23_OUTPUT_DIR", "outputs_video/ltx2_3_distilled_i2v")
+)
+I2V_IMAGE = os.getenv("LTX23_I2V_IMAGE", "")
+DEFAULT_PROMPT = (
+    "A fashion model takes a slow step forward and shifts her weight, "
+    "the soft fabric of her clothing swaying and rippling with the "
+    "motion, her hair shifting gently, soft even studio lighting on a "
+    "clean light background, elegant slow-motion runway feel."
+)
+PROMPT = os.getenv("LTX23_I2V_PROMPT", DEFAULT_PROMPT)
+
+# Per-stage timing helpers --------------------------------------------------
+
+def _print_stage_breakdown(result: dict, label: str) -> float | None:
+    """Print stage execution times and return the sum, or None if missing."""
+    logging_info = result.get("logging_info")
+    stages = getattr(logging_info, "stages", None) if logging_info else None
+    if not stages:
+        print(f"  [{label}] stage breakdown unavailable")
+        return None
+    print(f"  [{label}] stage breakdown:")
+    total = 0.0
+    for name, metrics in stages.items():
+        exec_s = float(metrics.get("execution_time", 0.0))
+        total += exec_s
+        print(f"    - {name}: {exec_s:.3f}s")
+    print(f"    - stage_sum: {total:.3f}s")
+    return total
+
+
+def _collect_stage_times(
+    result: dict,
+    stage_times: dict[str, list[float]],
+    stage_order: OrderedDict[str, None],
+) -> None:
+    logging_info = result.get("logging_info")
+    stages = getattr(logging_info, "stages", None) if logging_info else None
+    if not stages:
+        return
+    for name, metrics in stages.items():
+        stage_order.setdefault(name, None)
+        stage_times.setdefault(name, []).append(
+            float(metrics.get("execution_time", 0.0))
+        )
+
+
+def _resolve_refine_upsampler(model_root: str) -> Path:
+    """LTX-2.3 distilled snapshots ship a `spatial_upscaler/` subdir."""
+    for name in ("spatial_upscaler", "spatial_upsampler"):
+        cand = Path(model_root) / name
+        if (cand / "config.json").is_file():
+            return cand
+    raise FileNotFoundError(
+        f"No refine upsampler directory under {model_root}. "
+        f"Expected `{model_root}/spatial_upscaler/config.json`."
+    )
+
+
+# Main ---------------------------------------------------------------------
+
+def main() -> None:
+    if not I2V_IMAGE:
+        raise SystemExit(
+            "LTX23_I2V_IMAGE is required for i2v. Example:\n"
+            "  export LTX23_I2V_IMAGE=/path/to/portrait_or_product.jpg\n"
+            "  python examples/inference/basic/basic_ltx2_3_distilled_i2v.py"
+        )
+    if not Path(I2V_IMAGE).is_file():
+        raise SystemExit(f"LTX23_I2V_IMAGE not found: {I2V_IMAGE}")
+
+    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+    model_root = maybe_download_model(MODEL_ID)
+    refine_upsampler_path = _resolve_refine_upsampler(model_root)
+    print(f"Model:           {model_root}")
+    print(f"Refine upsampler: {refine_upsampler_path}")
+    print(f"i2v image:       {I2V_IMAGE}")
+    print(f"Output dir:      {OUTPUT_DIR.resolve()}")
+
+    # mode="default" — Inductor's default schedule matches max-autotune on
+    # this pipeline (denoise/refine/decode all within ~5 ms, n=2) while
+    # saving ~7 min of cold compile on a single GB200.
+    torch_compile_kwargs = {
+        "backend": "inductor",
+        "fullgraph": True,
+        "mode": "default",
+        "dynamic": False,
+    }
+
+    # Loading the pipeline config *with model_path* binds model-specific
+    # tuning (notably VAE precision/decoder defaults) into the config. Without
+    # this, the generic pipeline config gives a substantially slower VAE
+    # decode stage. `basic_ltx2_distilled_fast_profile.py` uses the same
+    # pattern.
+    pipeline_config = PipelineConfig.from_pretrained(model_root)
+    pipeline_config.dit_config.quant_config = None
+
+    generator = VideoGenerator.from_pretrained(
+        model_root,
+        num_gpus=1,
+        # LTX-2.3 distilled uses the two-stage refine pipeline; the refine
+        # LoRA is intentionally empty for the distilled student.
+        ltx2_refine_enabled=True,
+        ltx2_refine_upsampler_path=str(refine_upsampler_path),
+        ltx2_refine_lora_path="",
+        ltx2_refine_num_inference_steps=3,
+        ltx2_refine_guidance_scale=1.0,
+        ltx2_refine_add_noise=True,
+        pipeline_config=pipeline_config,
+        enable_torch_compile=True,
+        enable_torch_compile_text_encoder=True,
+        # Compile the VAE codec submodules (encoder / decoder) too. The
+        # `LTX2CausalVideoAutoencoder` declares `_compile_conditions` so
+        # `_compile_with_conditions` targets just those submodules and
+        # leaves the surrounding tiling control flow eager — needed for
+        # fullgraph + dynamic=False to succeed. VAE eager decode is
+        # ~1.0s; compiling it brings the stage to ~0.3s.
+        enable_torch_compile_vae=True,
+        torch_compile_kwargs=torch_compile_kwargs,
+        torch_compile_kwargs_vae=torch_compile_kwargs,
+        # Keep everything resident — no CPU offload for serving-style runs.
+        dit_cpu_offload=False,
+        text_encoder_cpu_offload=False,
+        vae_cpu_offload=False,
+        ltx2_vae_tiling=False,
+    )
+
+    common_kwargs = dict(
+        prompt=PROMPT,
+        negative_prompt="",         # distilled is CFG-free; no negative needed
+        guidance_scale=1.0,         # CFG=1 for distilled
+        height=1280, width=832,     # portrait runway aspect
+        num_frames=121, fps=24,     # ~5s clip
+        num_inference_steps=8,      # distilled denoise steps
+        # i2v: anchor the input image at frame 0 with full strength.
+        # `ltx2_image_crf=0.0` skips an extra JPEG re-encode of an already
+        # JPEG conditioning image.
+        ltx2_images=[(I2V_IMAGE, 0, 1.0)],
+        ltx2_image_crf=0.0,
+        save_video=True,
+    )
+
+    warmup_runs = 2
+    measured_runs = 2
+    warmup_secs: list[float] = []
+    measured_secs: list[float] = []
+    stage_times: dict[str, list[float]] = {}
+    stage_order: OrderedDict[str, None] = OrderedDict()
+
+    try:
+        # Warmup: untimed (but we still wall-clock them so the first compile
+        # cost is visible to the reader).
+        for w in range(warmup_runs):
+            t0 = time.perf_counter()
+            print(f"\n[warmup {w + 1}/{warmup_runs}] compiling + generating…")
+            generator.generate_video(
+                output_path=str(OUTPUT_DIR / f"_warmup_{w + 1}.mp4"),
+                seed=7,
+                **common_kwargs,
+            )
+            dt = time.perf_counter() - t0
+            warmup_secs.append(dt)
+            print(f"[warmup {w + 1}/{warmup_runs}] wall={dt:.1f}s")
+
+        # Cleanup warmup artifacts so the user only sees measured outputs.
+        for w in range(warmup_runs):
+            (OUTPUT_DIR / f"_warmup_{w + 1}.mp4").unlink(missing_ok=True)
+
+        # Measured.
+        for m in range(measured_runs):
+            out_path = OUTPUT_DIR / f"output_ltx2_3_distilled_i2v_run_{m + 1}.mp4"
+            print(f"\n[measured {m + 1}/{measured_runs}] generating: {out_path}")
+            t0 = time.perf_counter()
+            result = generator.generate_video(
+                output_path=str(out_path),
+                seed=2002 + m,
+                **common_kwargs,
+            )
+            wall = time.perf_counter() - t0
+            e2e = (
+                result.get("e2e_latency")
+                if isinstance(result, dict) else None
+            ) or wall
+            measured_secs.append(e2e)
+            print(f"[measured {m + 1}/{measured_runs}] e2e={e2e:.2f}s wall={wall:.2f}s")
+            if isinstance(result, dict):
+                _print_stage_breakdown(result, f"measured {m + 1}")
+                _collect_stage_times(result, stage_times, stage_order)
+
+        # Summary.
+        print("\n=== summary ===")
+        print(f"warmup wall-times:      {[round(x, 1) for x in warmup_secs]}")
+        if measured_secs:
+            avg = sum(measured_secs) / len(measured_secs)
+            print(
+                f"measured e2e (n={len(measured_secs)}): "
+                f"{[round(x, 2) for x in measured_secs]} -> avg {avg:.2f}s"
+            )
+        if stage_times:
+            print(f"average stage times over {measured_runs} measured runs:")
+            avg_total = 0.0
+            for name in stage_order:
+                vals = stage_times.get(name) or []
+                if not vals:
+                    continue
+                avg_v = sum(vals) / len(vals)
+                avg_total += avg_v
+                print(f"  - {name}: {avg_v:.3f}s")
+            print(f"  - stage_sum_avg: {avg_total:.3f}s")
+    finally:
+        generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/basic_ltx2_3_distilled_i2v_typed.py

@@ -0,0 +1,350 @@
+# SPDX-License-Identifier: Apache-2.0
+"""LTX-2.3 distilled image-to-video — typed API (``from_config`` / ``generate``).
+
+Identical generation behavior to ``basic_ltx2_3_distilled_i2v.py``, but
+expressed through the newer typed surface (``GeneratorConfig`` /
+``GenerationRequest``) instead of the ``from_pretrained(**legacy_kwargs)``
+bridge. The typed API is now the preferred entry point — the legacy
+example still works but emits a ``DeprecationWarning`` for the LTX-2.3
+specific knobs.
+
+Quick start
+-----------
+    export LTX23_I2V_IMAGE=/path/to/your/portrait_or_product.jpg
+    # optional overrides:
+    #   export LTX23_I2V_PROMPT="a fashion model walks toward camera..."
+    #   export LTX23_OUTPUT_DIR=outputs_video/ltx2_3_distilled_i2v_typed
+    python examples/inference/basic/basic_ltx2_3_distilled_i2v_typed.py
+
+What the script does
+--------------------
+1. Loads FastVideo/LTX-2.3-Distilled-Diffusers (8 denoise + 3 refine
+   steps, CFG=1, no refine LoRA — the distilled production recipe).
+2. Compiles the DiT, text encoder, and VAE (fullgraph, Inductor default
+   mode — autotune adds ~7 min cold-compile here with no measurable
+   e2e gain).
+3. Runs 2 warmup calls (untimed) + 2 measured calls. Two warmups are
+   kept as a safety net — the first call pays cold compile + first-shape
+   guard work, and a second warmup ensures any residual recompiles
+   settle before we measure.
+4. Prints a per-stage breakdown and an average over the measured runs.
+
+Hardware notes
+--------------
+- Single-GPU example; for multi-GPU sequence-parallel see the gradio
+  demo under ``examples/inference/gradio/local/gradio_local_demo_ltx2_3/``.
+- First-time compile takes ~30-40 min on GB200 (~20 min on H100;
+  cached in ``$TORCHINDUCTOR_CACHE_DIR`` afterwards). Subsequent
+  invocations only pay the one-time process load + a few seconds of
+  dynamo trace.
+- On GB200 / Blackwell, run with ``env -u LD_LIBRARY_PATH ...`` to
+  avoid a system-cuBLAS / torch-cuBLAS mismatch that fails every GEMM.
+  The ``_inductor.shape_padding = False`` line below also avoids a
+  ``pad_mm`` landmine on the same generation of cards.
+
+Typed-API mapping (legacy kwarg ↔ typed field)
+----------------------------------------------
+- ``num_gpus``                      ↔ ``engine.num_gpus``
+- ``enable_torch_compile``          ↔ ``engine.compile.enabled``
+- ``enable_torch_compile_text_encoder`` ↔ ``engine.compile.text_encoder_enabled``
+- ``enable_torch_compile_vae``      ↔ ``engine.compile.vae_enabled``
+- ``torch_compile_kwargs``          ↔ ``engine.compile.backend/fullgraph/mode/dynamic``
+- ``torch_compile_kwargs_vae``      ↔ empty ``compile.vae_kwargs`` (inherits master)
+- ``dit_cpu_offload``               ↔ ``engine.offload.dit``
+- ``text_encoder_cpu_offload``      ↔ ``engine.offload.text_encoder``
+- ``vae_cpu_offload``               ↔ ``engine.offload.vae``
+- ``ltx2_vae_tiling``               ↔ ``pipeline.vae_tiling``
+- ``ltx2_refine_enabled``           ↔ ``pipeline.preset_overrides["refine"]["enabled"]``
+- ``ltx2_refine_upsampler_path``    ↔ ``pipeline.components.upsampler_weights``
+- ``ltx2_refine_lora_path``         ↔ ``pipeline.components.lora_path``
+- ``ltx2_refine_num_inference_steps`` ↔ ``pipeline.preset_overrides["refine"]["num_inference_steps"]``
+- ``ltx2_refine_guidance_scale``    ↔ ``pipeline.preset_overrides["refine"]["guidance_scale"]``
+- ``ltx2_refine_add_noise``         ↔ ``pipeline.preset_overrides["refine"]["add_noise"]``
+- ``pipeline_config=PipelineConfig.from_pretrained(model_root)`` ↔ (no-op — ``PipelineConfig.from_kwargs`` already resolves the model-specific class from ``model_path``)
+- ``pipeline_config.dit_config.quant_config = None`` ↔ leave ``engine.quantization`` unset
+- ``ltx2_images`` / ``ltx2_image_crf`` ↔ ``request.extensions`` (LTX-2 specific, no
+  first-class typed field yet)
+"""
+from __future__ import annotations
+
+import os
+import time
+from collections import OrderedDict
+from pathlib import Path
+
+import torch._inductor.config as _inductor
+
+from fastvideo import VideoGenerator
+from fastvideo.api import (
+    CompileConfig,
+    ComponentConfig,
+    EngineConfig,
+    GenerationRequest,
+    GeneratorConfig,
+    OffloadConfig,
+    OutputConfig,
+    PipelineSelection,
+    SamplingConfig,
+)
+from fastvideo.utils import maybe_download_model
+
+os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN")
+os.environ.setdefault("FASTVIDEO_STAGE_LOGGING", "1")
+
+# Inductor knobs. ``shape_padding=False`` is mandatory on Blackwell to
+# avoid a cuBLAS INVALID_VALUE crash inside pad_mm during the refine
+# path. The rest are autotune-friendliness flags.
+_inductor.shape_padding = False
+_inductor.conv_1x1_as_mm = True
+_inductor.coordinate_descent_tuning = True
+_inductor.coordinate_descent_check_all_directions = True
+_inductor.epilogue_fusion = False
+
+MODEL_ID = os.path.expandvars(
+    os.path.expanduser(
+        os.getenv("LTX23_MODEL_PATH", "FastVideo/LTX-2.3-Distilled-Diffusers")
+    )
+)
+OUTPUT_DIR = Path(
+    os.getenv(
+        "LTX23_OUTPUT_DIR", "outputs_video/ltx2_3_distilled_i2v_typed"
+    )
+)
+I2V_IMAGE = os.getenv("LTX23_I2V_IMAGE", "")
+DEFAULT_PROMPT = (
+    "A fashion model takes a slow step forward and shifts her weight, "
+    "the soft fabric of her clothing swaying and rippling with the "
+    "motion, her hair shifting gently, soft even studio lighting on a "
+    "clean light background, elegant slow-motion runway feel."
+)
+PROMPT = os.getenv("LTX23_I2V_PROMPT", DEFAULT_PROMPT)
+
+
+def _print_stage_breakdown(result, label: str) -> float | None:
+    logging_info = getattr(result, "logging_info", None)
+    stages = getattr(logging_info, "stages", None) if logging_info else None
+    if not stages:
+        print(f"  [{label}] stage breakdown unavailable")
+        return None
+    print(f"  [{label}] stage breakdown:")
+    total = 0.0
+    for name, metrics in stages.items():
+        exec_s = float(metrics.get("execution_time", 0.0))
+        total += exec_s
+        print(f"    - {name}: {exec_s:.3f}s")
+    print(f"    - stage_sum: {total:.3f}s")
+    return total
+
+
+def _collect_stage_times(
+    result,
+    stage_times: dict[str, list[float]],
+    stage_order: OrderedDict[str, None],
+) -> None:
+    logging_info = getattr(result, "logging_info", None)
+    stages = getattr(logging_info, "stages", None) if logging_info else None
+    if not stages:
+        return
+    for name, metrics in stages.items():
+        stage_order.setdefault(name, None)
+        stage_times.setdefault(name, []).append(
+            float(metrics.get("execution_time", 0.0))
+        )
+
+
+def _resolve_refine_upsampler(model_root: str) -> Path:
+    for name in ("spatial_upscaler", "spatial_upsampler"):
+        cand = Path(model_root) / name
+        if (cand / "config.json").is_file():
+            return cand
+    raise FileNotFoundError(
+        f"No refine upsampler directory under {model_root}. "
+        f"Expected `{model_root}/spatial_upscaler/config.json`."
+    )
+
+
+def main() -> None:
+    if not I2V_IMAGE:
+        raise SystemExit(
+            "LTX23_I2V_IMAGE is required for i2v. Example:\n"
+            "  export LTX23_I2V_IMAGE=/path/to/portrait_or_product.jpg\n"
+            "  python examples/inference/basic/"
+            "basic_ltx2_3_distilled_i2v_typed.py"
+        )
+    if not Path(I2V_IMAGE).is_file():
+        raise SystemExit(f"LTX23_I2V_IMAGE not found: {I2V_IMAGE}")
+
+    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+    model_root = maybe_download_model(MODEL_ID)
+    refine_upsampler_path = _resolve_refine_upsampler(model_root)
+    print(f"Model:           {model_root}")
+    print(f"Refine upsampler: {refine_upsampler_path}")
+    print(f"i2v image:       {I2V_IMAGE}")
+    print(f"Output dir:      {OUTPUT_DIR.resolve()}")
+
+    # mode="default" — Inductor's default schedule matches max-autotune on
+    # this pipeline (denoise/refine/decode all within ~5 ms, n=2) while
+    # saving ~7 min of cold compile on a single GB200.
+    generator_config = GeneratorConfig(
+        model_path=model_root,
+        engine=EngineConfig(
+            num_gpus=1,
+            # Keep DiT / text encoder / VAE resident on GPU — no CPU offload
+            # for serving-style runs. ``image_encoder`` and
+            # ``pin_cpu_memory`` are left at their schema defaults
+            # (matches the legacy example, which only set these three).
+            offload=OffloadConfig(
+                dit=False,
+                text_encoder=False,
+                vae=False,
+            ),
+            compile=CompileConfig(
+                enabled=True,
+                text_encoder_enabled=True,
+                # ``vae_enabled`` triggers ``_compile_with_conditions`` on
+                # ``LTX2CausalVideoAutoencoder``, which compiles just the
+                # encoder/decoder submodules and leaves the surrounding
+                # tiling control flow eager (required for ``fullgraph``).
+                # Empty ``vae_kwargs`` → inherits the master kwargs below.
+                vae_enabled=True,
+                backend="inductor",
+                fullgraph=True,
+                mode="default",
+                dynamic=False,
+            ),
+        ),
+        pipeline=PipelineSelection(
+            # ``PipelineConfig.from_kwargs`` resolves the model-specific
+            # pipeline-config class from ``model_path`` automatically, so we
+            # don't need to set ``components.pipeline_config_path`` — the
+            # model-specific VAE precision / decoder defaults are picked up
+            # the same way the legacy example's
+            # ``PipelineConfig.from_pretrained(model_root)`` did them.
+            components=ComponentConfig(
+                upsampler_weights=str(refine_upsampler_path),
+                # Distilled has no refine LoRA — omit ``lora_path``.
+            ),
+            vae_tiling=False,
+            preset_overrides={
+                "refine": {
+                    "enabled": True,
+                    "num_inference_steps": 3,
+                    "guidance_scale": 1.0,
+                    "add_noise": True,
+                },
+            },
+        ),
+    )
+
+    generator = VideoGenerator.from_config(generator_config)
+
+    def build_request(out_path: Path, seed: int) -> GenerationRequest:
+        return GenerationRequest(
+            prompt=PROMPT,
+            # distilled is CFG-free; no negative prompt
+            negative_prompt="",
+            sampling=SamplingConfig(
+                num_videos_per_prompt=1,
+                seed=seed,
+                height=1280,
+                width=832,
+                num_frames=121,
+                fps=24,
+                num_inference_steps=8,
+                guidance_scale=1.0,
+            ),
+            output=OutputConfig(
+                output_path=str(out_path),
+                save_video=True,
+                return_frames=False,
+            ),
+            # LTX-2.3 i2v fields don't have first-class typed slots yet;
+            # extensions is the documented bridge. ``ltx2_image_crf=0.0``
+            # skips an extra JPEG re-encode of an already JPEG image.
+            extensions={
+                "ltx2_images": [(I2V_IMAGE, 0, 1.0)],
+                "ltx2_image_crf": 0.0,
+            },
+        )
+
+    warmup_runs = 2
+    measured_runs = 2
+    warmup_secs: list[float] = []
+    measured_secs: list[float] = []
+    stage_times: dict[str, list[float]] = {}
+    stage_order: OrderedDict[str, None] = OrderedDict()
+
+    try:
+        for w in range(warmup_runs):
+            print(f"\n[warmup {w + 1}/{warmup_runs}] compiling + generating…")
+            t0 = time.perf_counter()
+            generator.generate(
+                build_request(
+                    OUTPUT_DIR / f"_warmup_{w + 1}.mp4", seed=7
+                )
+            )
+            dt = time.perf_counter() - t0
+            warmup_secs.append(dt)
+            print(f"[warmup {w + 1}/{warmup_runs}] wall={dt:.1f}s")
+
+        for w in range(warmup_runs):
+            (OUTPUT_DIR / f"_warmup_{w + 1}.mp4").unlink(missing_ok=True)
+
+        for m in range(measured_runs):
+            out_path = (
+                OUTPUT_DIR
+                / f"output_ltx2_3_distilled_i2v_typed_run_{m + 1}.mp4"
+            )
+            print(
+                f"\n[measured {m + 1}/{measured_runs}] generating: {out_path}"
+            )
+            t0 = time.perf_counter()
+            result = generator.generate(
+                build_request(out_path, seed=2002 + m)
+            )
+            wall = time.perf_counter() - t0
+            # ``e2e_latency`` is currently surfaced via ``result.extra``;
+            # ``GenerationResult`` exposes ``generation_time`` as a
+            # first-class field but the LTX-2 pipeline only fills the
+            # legacy ``e2e_latency`` key. Prefer the explicit one, fall
+            # back to wall-clock.
+            e2e = (
+                result.extra.get("e2e_latency")
+                if hasattr(result, "extra") else None
+            ) or wall
+            measured_secs.append(e2e)
+            print(
+                f"[measured {m + 1}/{measured_runs}] "
+                f"e2e={e2e:.2f}s wall={wall:.2f}s"
+            )
+            _print_stage_breakdown(result, f"measured {m + 1}")
+            _collect_stage_times(result, stage_times, stage_order)
+
+        print("\n=== summary ===")
+        print(
+            f"warmup wall-times:      "
+            f"{[round(x, 1) for x in warmup_secs]}"
+        )
+        if measured_secs:
+            avg = sum(measured_secs) / len(measured_secs)
+            print(
+                f"measured e2e (n={len(measured_secs)}): "
+                f"{[round(x, 2) for x in measured_secs]} -> avg {avg:.2f}s"
+            )
+        if stage_times:
+            print(f"average stage times over {measured_runs} measured runs:")
+            avg_total = 0.0
+            for name in stage_order:
+                vals = stage_times.get(name) or []
+                if not vals:
+                    continue
+                avg_v = sum(vals) / len(vals)
+                avg_total += avg_v
+                print(f"  - {name}: {avg_v:.3f}s")
+            print(f"  - stage_sum_avg: {avg_total:.3f}s")
+    finally:
+        generator.shutdown()
+
+
+if __name__ == "__main__":
+    main()

examples/inference/basic/basic_lucy_edit.py

@@ -0,0 +1,38 @@
+from fastvideo import VideoGenerator
+
+OUTPUT_PATH = "video_samples_lucy_edit"
+
+
+def main():
+    generator = VideoGenerator.from_pretrained(
+        "decart-ai/Lucy-Edit-Dev",
+        num_gpus=1,
+        use_fsdp_inference=False,
+        dit_cpu_offload=True,
+        vae_cpu_offload=False,
+        text_encoder_cpu_offload=True,
+        pin_cpu_memory=True,
+    )
+
+    prompt = ("Change the apron and blouse to a classic clown costume: satin "
+              "polka-dot jumpsuit in bright primary colors, ruffled white collar, "
+              "oversized pom-pom buttons, white gloves, oversized red shoes, red "
+              "foam nose; soft window light from left, eye-level medium shot.")
+    video_path = "https://d2drjpuinn46lb.cloudfront.net/painter_original_edit.mp4"
+
+    generator.generate_video(
+        prompt,
+        negative_prompt="",
+        video_path=video_path,
+        output_path=OUTPUT_PATH,
+        save_video=True,
+        height=480,
+        width=832,
+        num_frames=81,
+        fps=24,
+        guidance_scale=5.0,
+    )
+
+
+if __name__ == "__main__":
+    main()

examples/inference/eval/eval_third_person_separation.py

@@ -0,0 +1,91 @@
+"""Run ``judge.third_person_separation`` (needs ``.[eval-judge]`` + a Gemini key)
+over each baseline and print the candidate's win-rate table — from a ``--manifest``
+of pairs, or by pairing ``--candidate-dir`` against each ``--reference`` dir by
+filename stem.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+from collections import defaultdict
+from pathlib import Path
+
+from fastvideo.eval import create_evaluator
+
+METRIC = "judge.third_person_separation"
+VIDEO_EXTS = {".mp4", ".avi", ".mov", ".mkv", ".webm"}
+IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp"}
+
+
+def _by_stem(directory: Path, exts: set[str]) -> dict[str, Path]:
+    """Map filename stem -> path for files with the given extensions."""
+    return {p.stem: p for p in sorted(directory.iterdir()) if p.suffix.lower() in exts}
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description=__doc__,
+                                formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--candidate-dir", type=Path, default=None,
+                   help="Directory of candidate clips (directory mode).")
+    p.add_argument("--reference", action="append", default=[], metavar="NAME=DIR",
+                   help="Baseline directory, repeatable: 'name=dir' or bare 'dir'.")
+    p.add_argument("--image-dir", type=Path, default=None,
+                   help="Optional first-frame images, matched to clips by stem.")
+    p.add_argument("--prompts-json", type=Path, default=None,
+                   help="Optional {stem: control-signal text} JSON.")
+    p.add_argument("--actions-json", type=Path, default=None,
+                   help="Optional {stem: action-label} JSON for the per-action breakdown.")
+    p.add_argument("--manifest", type=Path, default=None,
+                   help="JSON list of {baseline, video_path, reference_path, ...} rows.")
+    p.add_argument("--output", type=Path, default=None)
+    args = p.parse_args()
+
+    # Group path-only samples per baseline: {baseline: [sample dict, ...]}.
+    by_baseline: dict[str, list[dict]] = defaultdict(list)
+    if args.manifest is not None:
+        for row in json.loads(args.manifest.read_text()):
+            by_baseline[row.get("baseline", "baseline")].append(
+                {k: v for k, v in row.items() if k != "baseline"})
+    elif args.candidate_dir is not None and args.reference:
+        cands = _by_stem(args.candidate_dir, VIDEO_EXTS)
+        images = _by_stem(args.image_dir, IMAGE_EXTS) if args.image_dir else {}
+        prompts = json.loads(args.prompts_json.read_text()) if args.prompts_json else {}
+        actions = json.loads(args.actions_json.read_text()) if args.actions_json else {}
+        for spec in args.reference:
+            name, sep, ref_dir = spec.partition("=")
+            if not sep:
+                name, ref_dir = Path(spec).name, spec
+            refs = _by_stem(Path(ref_dir), VIDEO_EXTS)
+            for stem in sorted(cands.keys() & refs.keys()):
+                sample = {"video_path": str(cands[stem]), "reference_path": str(refs[stem])}
+                if stem in images:
+                    sample["image_path"] = str(images[stem])
+                if stem in prompts:
+                    sample["text_prompt"] = prompts[stem]
+                if stem in actions:
+                    sample["action"] = actions[stem]
+                by_baseline[name].append(sample)
+    else:
+        p.error("provide either --manifest, or --candidate-dir with at least one --reference")
+
+    ev = create_evaluator(metrics=[METRIC], device="cpu")
+    print("\n| Baseline | Candidate win-rate (excl. ties) | W / L / T | n |")
+    print("|---|---|---|---|")
+    rows = {}
+    for baseline, samples in by_baseline.items():
+        res = ev.evaluate(samples=samples).corpus[METRIC]
+        rows[baseline] = res
+        d = res.details
+        if res.score is None:
+            print(f"| {baseline} | — | — | 0 |")
+        else:
+            print(f"| {baseline} | {100 * res.score:.1f}% | {d['wins']}/{d['losses']}/{d['ties']} | {d['n']} |")
+
+    if args.output is not None:
+        payload = {b: {"score": r.score, "details": r.details} for b, r in rows.items()}
+        args.output.write_text(json.dumps(payload, indent=2))
+        print(f"\nWrote {args.output}")
+
+
+if __name__ == "__main__":
+    main()

examples/train/configs/genrl_wan2.1_t2v_1.3B_hpsv3_videoalign.yaml

@@ -1,164 +0,0 @@
-# GenRL / Video GRPO: Wan 2.1 T2V 1.3B with HPSv3 and VideoAlign rewards.
-#
-# Ported from GenRL/config/longcat.yaml.
-#
-# - Student: trainable full-parameter model by default.
-# - LoRA is still available via models.student.use_lora=true.
-# - Full fine-tuning with beta > 0 requires models.reference and much
-#   more memory; keep beta at 0.0 for the 4xH100 probe run.
-#
-# Usage:
-#   torchrun --nnodes=1 --nproc_per_node=4 \
-#     -m fastvideo.train.entrypoint.train \
-#     --config examples/train/configs/genrl_wan2.1_t2v_1.3B_hpsv3_videoalign.yaml
-
-models:
-  student:
-    _target_: fastvideo.train.models.wan.wan_genrl.GenRLWanModel
-    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-    trainable: true
-    # Set true for LoRA. LoRA can use method.beta with disable_adapter().
-    use_lora: false
-    lora_r: 128
-    lora_alpha: 64
-    lora_init_weights: gaussian
-    lora_path: null
-    lora_target_modules:
-      - to_k
-      - to_out
-      - to_q
-      - to_v
-      - ffn.fc_in
-      - ffn.fc_out
-    enable_gradient_checkpointing_type: full
-
-method:
-  _target_: fastvideo.train.methods.rl.genrl.GenRLMethod
-
-  # ---- Reward functions ----
-  reward_fn:
-    hpsv3_general: 1.0
-    hpsv3_percentile: 1.0
-    videoalign_mq: 1.0
-    videoalign_ta: 1.0
-  reward_module: null
-  reward_on_gpu: true
-
-  # ---- Data ----
-  prompt_dataset_path: GenRL/datasets/filtered_prompts
-  prompt_fn: filtered_prompts
-
-  # ---- Sampling ----
-  sample_batch_size: 4
-  eval_batch_size: 2
-  # Sample multiple rollout microbatches, average their PPO losses, then
-  # apply one optimizer update. This reduces reward/advantage variance.
-  num_batches_per_epoch: 4
-  accumulate_ppo_microbatches: true
-  eval_every_steps: 20
-  eval_num_batches: 1
-  eval_num_steps: 16
-  eval_guidance_scale: 4.5
-  num_inference_steps: 16
-  guidance_scale: 4.5
-  num_video_per_prompt: 4
-  noise_level: 1.0
-  sde_type: flow_sde
-  sde_window_size: 1
-  sde_window_range: [0, 6]
-  diffusion_clip: true
-  diffusion_clip_value: 0.45
-  kl_reward: 0
-  same_latent: true
-
-  # ---- Video dimensions ----
-  height: 480
-  width: 832
-  num_frames: 81
-
-  # ---- PPO training ----
-  train_batch_size: 4
-  num_inner_epochs: 1
-  clip_range: 1.0e-4
-  adv_clip_max: 5.0
-  # Official LoRA LongCat uses beta: 3.0e-4 with disable_adapter().
-  # For full fine-tuning on 4 H100s, avoid a second frozen Wan copy.
-  beta: 0.0
-  use_cfg: true
-  # Flash-GRPO-style temporal gradient rectification: avoid the large
-  # LongCat sigma/dt multiplier while debugging full-FT stability.
-  loss_reweighting: flash_tgr
-  loss_reweighting_clip: null
-  weight_advantages: true
-  # Match official GenRL PPO cadence. With sde_window_size: 1 this is
-  # equivalent to one optimizer step per sampled trajectory timestep.
-  optimizer_step_per_timestep: true
-  log_post_update_kl: true
-  max_grad_norm: 1.0
-  seed: 42
-
-  # ---- Advantage computation ----
-  per_prompt_stat_tracking: true
-  global_std: false
-  max_group_std: true
-
-training:
-  distributed:
-  
-    num_gpus: 4
-    sp_size: 1
-    tp_size: 1
-    # Full fine-tuning needs FSDP/HSDP sharding across all 4 GPUs.
-    hsdp_replicate_dim: 1
-    hsdp_shard_dim: 4
-
-  data:
-    # Not used by GenRL (prompt dataloaders are in method config)
-    # but required by the config parser.
-    data_path: ""
-    train_batch_size: 1
-    seed: 42
-    num_height: 480
-    num_width: 832
-    num_frames: 81
-
-  optimizer:
-    # Full-parameter visual GRPO is much more sensitive than LoRA.
-    # DanceGRPO reports 5e-6 to 2e-5 as the practical range.
-    learning_rate: 1.0e-5
-    betas: [0.9, 0.999]
-    weight_decay: 1.0e-4
-    lr_scheduler: constant
-    lr_warmup_steps: 0
-
-  loop:
-    # max_train_steps == num_epochs in GenRL terms.
-    max_train_steps: 100000
-    gradient_accumulation_steps: 1
-
-  checkpoint:
-    output_dir: outputs/genrl_longcat
-    training_state_checkpointing_steps: 100
-    checkpoints_total_limit: 3
-
-  tracker:
-    project_name: VideoRL
-    # Leave blank so W&B auto-generates a unique display name per run.
-    run_name: ""
-
-  model:
-    enable_gradient_checkpointing_type: full
-
-callbacks:
-  # Gradnorm call back Disabled; GenRLMethod clips internally.
-  ema:
-    decay: 0.9
-    start_iter: 0
-    update_interval: 8
-  log_rl_samples:
-    every_steps: 1
-    max_videos: 4
-    fps: 16
-
-pipeline:
-  flow_shift: 3.0

examples/train/configs/genrl_wan2.1_t2v_1.3B_ocr.yaml

@@ -1,113 +0,0 @@
-# GenRL / Video GRPO: Wan 2.1 T2V 1.3B — OCR reward, full finetune, 4 GPUs.
-#
-# Usage:
-#   torchrun --nnodes=1 --nproc_per_node=4 \
-#     -m fastvideo.train.entrypoint.train \
-#     --config examples/train/configs/genrl_wan2.1_t2v_1.3B_ocr.yaml
-
-models:
-  student:
-    _target_: fastvideo.train.models.wan.wan_genrl.GenRLWanModel
-    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
-    trainable: true
-    enable_gradient_checkpointing_type: full
-
-method:
-  _target_: fastvideo.train.methods.rl.genrl.GenRLMethod
-
-  # ---- Reward functions ----
-  reward_fn:
-    video_ocr: 1.0
-  reward_module: null
-
-  # ---- Data ----
-  prompt_dataset_path: GenRL/datasets/ocr
-  prompt_fn: general_ocr
-
-  # ---- Sampling ----
-  sample_batch_size: 4
-  eval_batch_size: 2
-  num_batches_per_epoch: 1
-  num_inference_steps: 16
-  guidance_scale: 4.5
-  num_video_per_prompt: 4
-  noise_level: 1.0
-  sde_type: flow_sde
-  sde_window_size: 1
-  sde_window_range: [0, 6]
-  diffusion_clip: true
-  diffusion_clip_value: 0.45
-  kl_reward: 0
-  same_latent: true
-
-  # ---- Video dimensions ----
-  height: 480
-  width: 832
-  num_frames: 81
-
-  # ---- PPO training ----
-  train_batch_size: 4
-  num_inner_epochs: 1
-  clip_range: 1.0e-3
-  adv_clip_max: 5.0
-  # No frozen reference model is configured in this launch.
-  beta: 0.0
-  use_cfg: true
-  loss_reweighting: longcat
-  weight_advantages: false
-  max_grad_norm: 1.0
-  seed: 42
-
-  # ---- Advantage computation ----
-  per_prompt_stat_tracking: true
-  global_std: false
-  max_group_std: true
-
-training:
-  distributed:
-    num_gpus: 4
-    sp_size: 1
-    tp_size: 1
-    hsdp_replicate_dim: 1
-    hsdp_shard_dim: 4
-
-  data:
-    data_path: ""
-    train_batch_size: 1
-    seed: 42
-    num_height: 480
-    num_width: 832
-    num_frames: 81
-
-  optimizer:
-    learning_rate: 1.0e-4
-    betas: [0.9, 0.999]
-    weight_decay: 1.0e-4
-    lr_scheduler: constant
-    lr_warmup_steps: 0
-
-  loop:
-    max_train_steps: 100000
-    gradient_accumulation_steps: 1
-
-  checkpoint:
-    output_dir: outputs/genrl_ocr
-    training_state_checkpointing_steps: 100
-    checkpoints_total_limit: 3
-
-  tracker:
-    project_name: VideoRL
-    run_name: wan_2_1_t2v_1_3b_ocr
-
-  model:
-    enable_gradient_checkpointing_type: full
-
-callbacks:
-  grad_clip:
-    max_grad_norm: 0.0  # Disabled; GenRLMethod clips internally.
-  ema:
-    decay: 0.9
-    start_iter: 0
-
-pipeline:
-  flow_shift: 3.0

examples/train/run_slurm.sh

@@ -107,7 +107,7 @@ srun torchrun \\
     --node_rank \$SLURM_PROCID \\
     --rdzv_backend=c10d \\
     --rdzv_endpoint="\$MASTER_ADDR:\$MASTER_PORT" \\
-    -m fastvideo.train.entrypoint.train \\
+    fastvideo/train/entrypoint/train.py \\
     --config ${CONFIG} \\
     --training.distributed.num_gpus ${TOTAL_GPUS} \\
     ${EXTRA_ARGS[*]:-}

fastvideo-kernel/build.sh

@@ -78,16 +78,26 @@ print(f'{mj}.{mn}')"
 }
 
 if [ "${GPU_BACKEND}" = "CUDA" ]; then
-    detected_cc="$(detect_with_torch)" || {
-        echo "ERROR: torch-based CUDA arch detection failed in uv environment." >&2
-        echo "       Ensure torch is installed and CUDA is available in the uv-selected Python." >&2
-        exit 1
-    }
-
-    cc_major="${detected_cc%%.*}"
-    cc_minor="${detected_cc##*.}"
+    # Compute capability drives the arch/TK defaults below. Prefer an explicit
+    # TORCH_CUDA_ARCH_LIST (works on GPU-less build machines such as CI/Docker);
+    # only probe a live GPU via torch when no arch was provided.
+    if [ -n "${TORCH_CUDA_ARCH_LIST:-}" ]; then
+        echo "Using TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST} (skipping torch GPU probe)"
+        first_arch="${TORCH_CUDA_ARCH_LIST%%[;, ]*}"   # first entry, e.g. 9.0a
+        first_arch="${first_arch%[af]}"                # strip trailing a/f suffix
+        cc_major="${first_arch%%.*}"
+        cc_minor="${first_arch##*.}"
+    else
+        detected_cc="$(detect_with_torch)" || {
+            echo "ERROR: torch-based CUDA arch detection failed and TORCH_CUDA_ARCH_LIST is unset." >&2
+            echo "       Set TORCH_CUDA_ARCH_LIST (e.g. 9.0a) for GPU-less builds, or build where CUDA is available." >&2
+            exit 1
+        }
+        cc_major="${detected_cc%%.*}"
+        cc_minor="${detected_cc##*.}"
+        echo "Detected compute capability via torch: ${detected_cc}"
+    fi
     cmake_arch="${cc_major}${cc_minor}"
-    echo "Detected compute capability via torch: ${detected_cc} (sm_${cmake_arch})"
 
     # Respect explicit overrides.
     if [ -z "${TORCH_CUDA_ARCH_LIST:-}" ]; then

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/attention/backends/video_sparse_attn.py

@@ -150,14 +150,11 @@ class VideoSparseAttentionMetadata(AttentionMetadata):
     # in postprocess_output().  Avoids materializing the intermediate
     # ``[B, len(non_pad_index), H, D]`` tensor on every layer.
     untile_combined_index: torch.LongTensor
-    # Per-step shared padded buffer used by tile().  Lazily populated on
-    # the first layer's call and reused by every subsequent VSA layer in
-    # the same denoising step.  Scoping to metadata (not class/instance)
-    # makes the reuse thread-safe across concurrent requests and keeps
-    # the "pad positions are zero" invariant trivially true (the buffer
-    # is freshly zeroed alongside ``non_pad_index`` so the index set
-    # cannot drift between calls).
+    # Per-step shared padded buffer used by tile().  Inference can reuse this
+    # across VSA layers, but training disables it so activation checkpointing
+    # can release the large tiled QKVG scratch tensor after each attention call.
     tile_buf: torch.Tensor | None = None
+    cache_tile_buf: bool = True
 
 
 class VideoSparseAttentionMetadataBuilder(AttentionMetadataBuilder):
@@ -175,6 +172,7 @@ class VideoSparseAttentionMetadataBuilder(AttentionMetadataBuilder):
         patch_size: tuple[int, int, int],
         VSA_sparsity: float,
         device: torch.device,
+        cache_tile_buf: bool = True,
         **kwargs: dict[str, Any],
     ) -> VideoSparseAttentionMetadata:
         patch_size = patch_size
@@ -201,7 +199,8 @@ class VideoSparseAttentionMetadataBuilder(AttentionMetadataBuilder):
             reverse_tile_partition_indices=reverse_tile_partition_indices,
             variable_block_sizes=variable_block_sizes,
             non_pad_index=non_pad_index,
-            untile_combined_index=untile_combined_index)
+            untile_combined_index=untile_combined_index,
+            cache_tile_buf=cache_tile_buf)
 
 
 class VideoSparseAttentionImpl(AttentionImpl):
@@ -237,6 +236,11 @@ class VideoSparseAttentionImpl(AttentionImpl):
         w_padded_size = num_tiles[2] * VSA_TILE_SIZE[2]
         target_shape = (x.shape[0], t_padded_size * h_padded_size * w_padded_size, x.shape[-2], x.shape[-1])
 
+        if not attn_metadata.cache_tile_buf:
+            buf = torch.zeros(target_shape, device=x.device, dtype=x.dtype)
+            buf[:, attn_metadata.non_pad_index] = x[:, attn_metadata.tile_partition_indices]
+            return buf
+
         # Reuse the per-step buffer stashed on metadata (lazily allocated
         # on the first VSA layer's call within a denoising step).  Pad
         # positions are zero from the initial torch.zeros and never

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 @@ from fastvideo.configs.models.dits.kandinsky5 import Kandinsky5VideoConfig
 __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.qwen2_5 import Qwen2_5_VLConfig
 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 @@ __all__ = [
     "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 = {

fastvideo/configs/models/encoders/mistral3.py

@@ -0,0 +1,38 @@
+# 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
+    # Mistral3 (full Flux2) ships a multimodal processor; load via AutoProcessor.
+    require_processor: 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/models/encoders/qwen3.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+# Ported from SGLang: python/sglang/multimodal_gen/configs/models/encoders/qwen3.py
+"""Qwen3 text encoder configuration for FastVideo diffusion models (e.g. Flux2 Klein)."""
+from dataclasses import dataclass, field
+from typing import Any
+
+from fastvideo.configs.models.encoders.base import (
+    TextEncoderArchConfig,
+    TextEncoderConfig,
+)
+
+
+def _is_transformer_layer(n: str, m: Any) -> bool:
+    return "layers" in n and str.isdigit(n.split(".")[-1])
+
+
+def _is_embeddings(n: str, m: Any) -> bool:
+    return n.endswith("embed_tokens")
+
+
+def _is_final_norm(n: str, m: Any) -> bool:
+    return n.endswith("norm")
+
+
+@dataclass
+class Qwen3TextArchConfig(TextEncoderArchConfig):
+    """Architecture config for Qwen3 text encoder.
+
+    Qwen3 is similar to LLaMA but with QK-Norm (RMSNorm on Q and K before attention).
+    Used by Flux2 Klein.
+    """
+
+    vocab_size: int = 151936
+    hidden_size: int = 2560
+    intermediate_size: int = 9728
+    num_hidden_layers: int = 36
+    num_attention_heads: int = 32
+    num_key_value_heads: int = 8
+    hidden_act: str = "silu"
+    max_position_embeddings: int = 40960
+    initializer_range: float = 0.02
+    rms_norm_eps: float = 1e-6
+    use_cache: bool = True
+    pad_token_id: int = 151643
+    bos_token_id: int = 151643
+    eos_token_id: int = 151645
+    tie_word_embeddings: bool = True
+    rope_theta: float = 1000000.0
+    rope_scaling: dict | None = None
+    attention_bias: bool = False
+    attention_dropout: float = 0.0
+    mlp_bias: bool = False
+    head_dim: int = 128
+    text_len: int = 512
+    output_hidden_states: bool = True  # Klein needs hidden states from layers 9, 18, 27
+
+    stacked_params_mapping: list[tuple[str, str, str | int]] = field(default_factory=lambda: [
+        (".qkv_proj", ".q_proj", "q"),
+        (".qkv_proj", ".k_proj", "k"),
+        (".qkv_proj", ".v_proj", "v"),
+        (".gate_up_proj", ".gate_proj", 0),
+        (".gate_up_proj", ".up_proj", 1),
+    ])
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [_is_transformer_layer, _is_embeddings, _is_final_norm])
+
+    def __post_init__(self) -> None:
+        self.tokenizer_kwargs = {
+            "padding": "max_length",
+            "truncation": True,
+            "max_length": self.text_len,
+            "return_tensors": "pt",
+        }
+
+
+@dataclass
+class Qwen3TextConfig(TextEncoderConfig):
+    """Top-level config for Qwen3 text encoder."""
+
+    arch_config: TextEncoderArchConfig = field(default_factory=Qwen3TextArchConfig)
+    prefix: str = "qwen3"
+    is_chat_model: bool = True

fastvideo/configs/models/vaes/__init__.py

@@ -6,6 +6,7 @@ from fastvideo.configs.models.vaes.hunyuanvae import HunyuanVAEConfig
 from fastvideo.configs.models.vaes.hunyuan15vae import Hunyuan15VAEConfig
 from fastvideo.configs.models.vaes.ltx2vae import LTX2VAEConfig
 from fastvideo.configs.models.vaes.oobleck import OobleckVAEArchConfig, OobleckVAEConfig
+from fastvideo.configs.models.vaes.flux2vae import Flux2VAEConfig
 from fastvideo.configs.models.vaes.wanvae import WanVAEConfig
 
 __all__ = [
@@ -19,4 +20,5 @@ __all__ = [
     "LTX2VAEConfig",
     "OobleckVAEArchConfig",
     "OobleckVAEConfig",
+    "Flux2VAEConfig",
 ]

fastvideo/configs/models/vaes/flux2vae.py

@@ -0,0 +1,58 @@
+# 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.vaes.base import VAEArchConfig, VAEConfig
+
+
+@dataclass
+class Flux2VAEArchConfig(VAEArchConfig):
+    """Architecture configuration for Flux2 VAE model."""
+
+    # Flux2 VAE-specific architecture parameters
+    in_channels: int = 3
+    out_channels: int = 3
+    down_block_types: tuple[str, ...] = (
+        "DownEncoderBlock2D",
+        "DownEncoderBlock2D",
+        "DownEncoderBlock2D",
+        "AttnDownEncoderBlock2D",
+    )
+    up_block_types: tuple[str, ...] = (
+        "AttnUpDecoderBlock2D",
+        "UpDecoderBlock2D",
+        "UpDecoderBlock2D",
+        "UpDecoderBlock2D",
+    )
+    block_out_channels: tuple[int, ...] = (128, 256, 512, 512)
+    layers_per_block: int = 2
+    act_fn: str = "silu"
+    latent_channels: int = 16
+    norm_num_groups: int = 32
+    sample_size: int = 512
+    force_upcast: bool = False
+    use_quant_conv: bool = True
+    use_post_quant_conv: bool = True
+    mid_block_add_attention: bool = True
+    batch_norm_eps: float = 1e-5
+    batch_norm_momentum: float = 0.1
+    patch_size: tuple[int, int] = (1, 1)
+
+    # Latent scaling for decode: avoid division-by-zero; match Flux/Flux2 convention (e.g. 0.13025)
+    scaling_factor: float = 0.13025
+
+    # Spatial compression (for images, this is typically 8)
+    spatial_compression_ratio: int = 8
+    temporal_compression_ratio: int = 1  # Images don't have temporal dimension
+
+
+@dataclass
+class Flux2VAEConfig(VAEConfig):
+    """Configuration for Flux2 VAE model."""
+
+    arch_config: Flux2VAEArchConfig = field(default_factory=Flux2VAEArchConfig)
+
+    # Flux2 is an image model, so disable temporal tiling
+    use_tiling: bool = False
+    use_temporal_tiling: bool = False
+    use_parallel_tiling: bool = False

fastvideo/configs/pipelines/__init__.py

@@ -9,12 +9,12 @@ from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig
 from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig
 from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig
 from fastvideo.registry import get_pipeline_config_cls_from_name
-from fastvideo.configs.pipelines.wan import (SelfForcingWanT2V480PConfig, WanI2V480PConfig, WanI2V720PConfig,
-                                             WanT2V480PConfig, WanT2V720PConfig)
+from fastvideo.configs.pipelines.wan import (LucyEditDevConfig, SelfForcingWanT2V480PConfig, WanI2V480PConfig,
+                                             WanI2V720PConfig, WanT2V480PConfig, WanT2V720PConfig)
 
 __all__ = [
     "HunyuanConfig", "FastHunyuanConfig", "HunyuanGameCraftPipelineConfig", "PipelineConfig", "Hunyuan15T2V480PConfig",
     "Hunyuan15T2V720PConfig", "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", "WanI2V720PConfig",
-    "SelfForcingWanT2V480PConfig", "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig", "HYWorldConfig",
-    "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name"
+    "SelfForcingWanT2V480PConfig", "LucyEditDevConfig", "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig",
+    "HYWorldConfig", "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name"
 ]

fastvideo/configs/pipelines/base.py

@@ -35,6 +35,11 @@ class PipelineConfig:
     flow_shift: float | None = None
     flow_shift_sr: float | None = None
     disable_autocast: bool = False
+    # When True, the scheduler's Euler update runs in fp32 outside the autocast
+    # block (Diffusers-style; avoids BF16 drift over multiple steps). Flux2 sets
+    # this True for reference parity; other models keep the legacy in-autocast
+    # behavior to preserve existing SSIM references.
+    scheduler_step_in_fp32: bool = False
     is_causal: bool = False
 
     # Model configuration
@@ -64,8 +69,9 @@ class PipelineConfig:
     # DMD parameters
     dmd_denoising_steps: list[int] | None = field(default=None)
 
-    # Wan2.2 TI2V parameters
+    # Wan2.2 task modifiers
     ti2v_task: bool = False
+    lucy_edit_task: bool = False
     boundary_ratio: float | None = None
 
     # Compilation

fastvideo/configs/pipelines/flux_2.py

@@ -0,0 +1,86 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copied and adapted from: https://github.com/sglang-ai/sglang
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+import torch
+
+from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
+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
+
+
+@dataclass
+class Flux2PipelineConfig(PipelineConfig):
+    """Configuration for Flux2 image generation pipeline."""
+
+    # Flux2-specific parameters
+    embedded_cfg_scale: float | None = 4.0
+    scheduler_step_in_fp32: bool = True
+    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 (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, ))
+
+
+def flux2_klein_postprocess_text(outputs: BaseEncoderOutput) -> torch.Tensor:
+    """Klein postprocess: hidden states from layers 9, 18, 27 (Qwen3)."""
+    hidden_states_layers: list[int] = [9, 18, 27]
+    if outputs.hidden_states is None:
+        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)
+    return prompt_embeds
+
+
+@dataclass
+class Flux2KleinEncoderArchConfig(EncoderArchConfig):
+    """Encoder arch config for Flux2 Klein (Qwen3); needs hidden states for layers 9, 18, 27."""
+    output_hidden_states: bool = True
+
+
+@dataclass
+class Flux2KleinTextEncoderConfig(EncoderConfig):
+    """Text encoder config for Flux2 Klein (Qwen3)."""
+    arch_config: EncoderArchConfig = field(default_factory=Flux2KleinEncoderArchConfig)
+
+
+@dataclass
+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)
+    scheduler_step_in_fp32: bool = True
+    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/configs/pipelines/wan.py

@@ -6,9 +6,11 @@ import torch
 
 from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
 from fastvideo.configs.models.dits import WanVideoConfig
+from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig
 from fastvideo.configs.models.encoders import (BaseEncoderOutput, CLIPVisionConfig, T5Config,
                                                WAN2_1ControlCLIPVisionConfig)
 from fastvideo.configs.models.vaes import WanVAEConfig
+from fastvideo.configs.models.vaes.wanvae import WanVAEArchConfig
 from fastvideo.configs.pipelines.base import PipelineConfig
 
 
@@ -120,6 +122,142 @@ class Wan2_2_TI2V_5B_Config(WanT2V480PConfig):
     expand_timesteps: bool = True
 
     def __post_init__(self) -> None:
+        assert not (self.ti2v_task and self.lucy_edit_task)
+        self.vae_config.load_encoder = True
+        self.vae_config.load_decoder = True
+        self.dit_config.expand_timesteps = self.expand_timesteps
+
+
+@dataclass
+class LucyEditDevConfig(Wan2_2_TI2V_5B_Config):
+    """Configuration for Decart Lucy Edit Dev video editing."""
+
+    dit_config: DiTConfig = field(default_factory=lambda: WanVideoConfig(arch_config=WanVideoArchConfig(
+        num_attention_heads=24,
+        in_channels=96,
+        out_channels=48,
+        ffn_dim=14336,
+        num_layers=30,
+    )))
+    vae_config: VAEConfig = field(default_factory=lambda: WanVAEConfig(arch_config=WanVAEArchConfig(
+        base_dim=160,
+        decoder_base_dim=256,
+        z_dim=48,
+        in_channels=12,
+        out_channels=12,
+        scale_factor_spatial=16,
+        patch_size=2,
+        is_residual=True,
+        clip_output=False,
+        latents_mean=(
+            -0.2289,
+            -0.0052,
+            -0.1323,
+            -0.2339,
+            -0.2799,
+            0.0174,
+            0.1838,
+            0.1557,
+            -0.1382,
+            0.0542,
+            0.2813,
+            0.0891,
+            0.1570,
+            -0.0098,
+            0.0375,
+            -0.1825,
+            -0.2246,
+            -0.1207,
+            -0.0698,
+            0.5109,
+            0.2665,
+            -0.2108,
+            -0.2158,
+            0.2502,
+            -0.2055,
+            -0.0322,
+            0.1109,
+            0.1567,
+            -0.0729,
+            0.0899,
+            -0.2799,
+            -0.1230,
+            -0.0313,
+            -0.1649,
+            0.0117,
+            0.0723,
+            -0.2839,
+            -0.2083,
+            -0.0520,
+            0.3748,
+            0.0152,
+            0.1957,
+            0.1433,
+            -0.2944,
+            0.3573,
+            -0.0548,
+            -0.1681,
+            -0.0667,
+        ),
+        latents_std=(
+            0.4765,
+            1.0364,
+            0.4514,
+            1.1677,
+            0.5313,
+            0.4990,
+            0.4818,
+            0.5013,
+            0.8158,
+            1.0344,
+            0.5894,
+            1.0901,
+            0.6885,
+            0.6165,
+            0.8454,
+            0.4978,
+            0.5759,
+            0.3523,
+            0.7135,
+            0.6804,
+            0.5833,
+            1.4146,
+            0.8986,
+            0.5659,
+            0.7069,
+            0.5338,
+            0.4889,
+            0.4917,
+            0.4069,
+            0.4999,
+            0.6866,
+            0.4093,
+            0.5709,
+            0.6065,
+            0.6415,
+            0.4944,
+            0.5726,
+            1.2042,
+            0.5458,
+            1.6887,
+            0.3971,
+            1.0600,
+            0.3943,
+            0.5537,
+            0.5444,
+            0.4089,
+            0.7468,
+            0.7744,
+        ),
+    )))
+    ti2v_task: bool = False
+    lucy_edit_task: bool = True
+
+    def __post_init__(self) -> None:
+        assert not (self.ti2v_task and self.lucy_edit_task)
+        # Lucy uses Wan2.2's enhanced 48-channel VAE latents. Denoising
+        # concatenates noise + video latents, matching the 96-channel
+        # transformer input declared above.
         self.vae_config.load_encoder = True
         self.vae_config.load_decoder = True
         self.dit_config.expand_timesteps = self.expand_timesteps

fastvideo/entrypoints/video_generator.py

@@ -517,6 +517,7 @@ class VideoGenerator:
             if _ek in kwargs:
                 extra_overrides[_ek] = kwargs.pop(_ek)
 
+        prompt_embeds = kwargs.pop("prompt_embeds", None)
         sampling_param.update(kwargs)
         kwargs["_extra_overrides"] = extra_overrides
 
@@ -567,6 +568,8 @@ class VideoGenerator:
             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,
@@ -669,6 +672,7 @@ class VideoGenerator:
         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:
@@ -715,6 +719,9 @@ class VideoGenerator:
             n_tokens=n_tokens,
             VSA_sparsity=fastvideo_args.VSA_sparsity,
         )
+        # Allow precomputed prompt_embeds (e.g. from diffusers) to skip text encoding
+        if prompt_embeds is not None:
+            batch.prompt_embeds = (list(prompt_embeds) if isinstance(prompt_embeds, list | tuple) else [prompt_embeds])
 
         extra_overrides = kwargs.pop("_extra_overrides", {})
         for _ek, _ev in extra_overrides.items():

fastvideo/eval/README.md

@@ -2,8 +2,9 @@
 
 In-process evaluation suite for video generations. Includes pixel
 metrics (SSIM, PSNR, LPIPS), Fréchet Video Distance (FVD), optical-flow
-comparisons, the full VBench suite, Physics-IQ, audio metrics, and a
-VLM scorer behind a single registry-driven API.
+comparisons, the full VBench suite, Physics-IQ, audio metrics, an
+absolute VLM scorer (`videoscore2`), and a pairwise VLM judge
+(`judge.third_person_separation`) — all behind a single registry-driven API.
 
 ## Install
 
@@ -159,6 +160,7 @@ fastvideo/
 │       ├── audio/                 # clap_score, audiobox_aesthetics, kl_divergence,
 │       │                          # frechet_distance, wer, desync, imagebind_score
 │       ├── videoscore2/           # VideoScore-2 (Qwen2.5-VL)
+│       ├── judge/                 # pairwise VLM judges (third_person_separation)
 │       ├── physics_iq/            # PhysicsIQ + sub-metrics
 │       └── vbench/                # adapter: sys.path bootstrap + shims
 │           ├── __init__.py
@@ -313,6 +315,40 @@ to control read/write behavior. The example script
 `examples/inference/eval/eval_fvd.py` demonstrates the full
 two-directory workflow.
 
+## `judge.third_person_separation` — pairwise VLM judge
+
+A **preference** metric (a judge, not an absolute score), and the suite's first
+remote-API one. For each pair the judge (Gemini) sees the shared first frame and
+two rollouts — a candidate and a reference model under the same control signal —
+and picks the one that better separates the third-person CHARACTER (foreground)
+from the BACKGROUND. The corpus score is the candidate's win-rate, excluding
+ties. Set-vs-set, motion-first; it reads native mp4s, so samples carry path
+strings, not decoded tensors.
+
+```bash
+uv pip install -e .[eval-judge]      # opt-in: needs network + an API key
+export GEMINI_API_KEY=...            # or GOOGLE_API_KEY, or ~/.gemini_token
+```
+
+```python
+from fastvideo.eval import create_evaluator
+
+ev = create_evaluator(metrics=["judge.third_person_separation"], device="cpu")
+result = ev.evaluate(samples=[
+    {"video_path": "cand/000.mp4", "reference_path": "base/000.mp4",
+     "image_path": "frames/000.png", "text_prompt": "W: moves forward", "action": "W"},
+    # ... more pairs ...
+]).corpus["judge.third_person_separation"]
+result.score    # candidate win-rate excl. ties; result.details has the breakdown
+```
+
+Only `video_path`/`reference_path` are required; `image_path`/`text_prompt`/
+`action` are optional. Verdicts are cached under `${FASTVIDEO_EVAL_CACHE}/eval/judge/`.
+The judge separates best when the control yields genuine parallax (e.g.
+translation); rigid whole-frame motion (e.g. pure camera rotation) is harder. To
+sweep several baselines into a table, see
+`examples/inference/eval/eval_third_person_separation.py`.
+
 ## Out of scope (follow-up PRs)
 
 - **MIND** metrics. Depend on a separate `vipe` upstream submodule.

fastvideo/eval/metrics/judge/third_person_separation/metric.py

@@ -0,0 +1,306 @@
+"""Pairwise VLM judge (Gemini): of two rollouts under the same control, which
+better separates the third-person character (foreground) from the background;
+the score is the candidate's win-rate over a reference.
+"""
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import time
+from pathlib import Path
+from typing import Any
+
+from fastvideo.eval.metrics.base import BaseMetric
+from fastvideo.eval.models import get_cache_dir
+from fastvideo.eval.registry import register
+from fastvideo.eval.types import MetricResult, Video
+
+# Part of the on-disk cache key; bump to invalidate cached verdicts.
+RUBRIC_ID = "v1"
+DEFAULT_MODEL = "gemini-2.5-pro"
+DEFAULT_K = 3
+
+SYSTEM_PROMPT = ("You are a strict comparative evaluator of third-person video-game "
+                 "rollouts. Two videos were generated by two different models from the "
+                 "SAME first frame and the SAME control signal. You judge which video "
+                 "better demonstrates that the model SEPARATES the third-person CHARACTER "
+                 "(foreground) from the BACKGROUND SCENE — i.e. the control animates the "
+                 "character as an INDEPENDENT AGENT with its own trajectory while the "
+                 "background moves with the camera. You do NOT reward whichever video "
+                 "merely looks cleaner, sharper, or higher-res.")
+
+RUBRIC = """\
+You are watching TWO generated video rollouts (Video 1, Video 2) played in full,
+from the same first frame under the same control signal. Pick the better
+third-person world-model rollout. Judge THREE things together, in this order:
+
+(C) MOTION / ACTION EXECUTION FIRST. The rollout must actually CARRY OUT the
+    control signal with substantial motion (the scene/character clearly moves as
+    commanded). A clip that is near-static, barely drifts, or only twitches has
+    NOT demonstrated controllable separation — it FAILS, no matter how clean it
+    looks. CRITICAL: do NOT reward a clip for looking "smoother" or "more
+    stable" when that smoothness is really just the ABSENCE OF MOTION. Less
+    motion is NOT better. If one clip executes the action with clear motion and
+    the other is comparatively static, the MOVING one wins (unless it fails B).
+
+(A) TEMPORAL COHERENCE — among clips that actually move, penalize GENUINE
+    corruption: flicker/strobing, texture boiling/crawling, geometry swimming,
+    the character or scene morphing/warping into mush, colors pulsing, or
+    progressive degradation into noise. Do NOT confuse LEGITIMATE large motion
+    (camera sweeping, character running, scene flowing past) with instability —
+    fast correct motion is GOOD, not a defect. Only true frame-to-frame
+    INCOHERENCE counts against a clip.
+
+(B) FOREGROUND/BACKGROUND SEPARATION — the character stays a distinct, coherent
+    entity with its own trajectory while the background responds to the control;
+    it does not dissolve/smear into the bg, and the whole frame does not slide
+    as one rigid sheet.
+
+Decision: among clips that genuinely execute the motion (C), pick the one that
+is both temporally coherent (A) and shows cleaner separation (B). A static or
+barely-moving clip loses to a moving one. Real motion is not instability. Do not
+reward resolution or placidity. "tie" only if truly equivalent on all three."""
+
+USER_TASK = ("Output a JSON object with four fields:\n"
+             " - video_1_analysis: FIRST, how much does Video 1 actually move — does it "
+             "execute the control with clear motion, or is it near-static / barely "
+             "drifting? THEN: among its motion, is there GENUINE corruption (flicker, "
+             "boiling, morphing into mush) as opposed to legitimate fast motion? THEN: "
+             "is the character a distinct coherent entity vs rigid-slide / dissolve?\n"
+             " - video_2_analysis: the same three checks for Video 2.\n"
+             " - comparison: apply (C) motion-first, then (A) genuine-coherence, then "
+             "(B) separation. A near-static clip loses to a moving one; legitimate large "
+             "motion is NOT a defect.\n"
+             " - winner: \"video_1\", \"video_2\", or \"tie\".")
+
+RESPONSE_SCHEMA = {
+    "type": "object",
+    "properties": {
+        "video_1_analysis": {
+            "type": "string"
+        },
+        "video_2_analysis": {
+            "type": "string"
+        },
+        "comparison": {
+            "type": "string"
+        },
+        "winner": {
+            "type": "string",
+            "enum": ["video_1", "video_2", "tie"]
+        },
+    },
+    "required": ["video_1_analysis", "video_2_analysis", "comparison", "winner"],
+    "propertyOrdering": ["video_1_analysis", "video_2_analysis", "comparison", "winner"],
+}
+
+
+def _path_of(sample: dict, key: str) -> str | None:
+    """Resolve a native-file path from a string key or a Video wrapper."""
+    p = sample.get(f"{key}_path")
+    if isinstance(p, str):
+        return p
+    v = sample.get(key)
+    if isinstance(v, Video) and isinstance(v.source, str):
+        return v.source
+    if isinstance(v, str):
+        return v
+    return None
+
+
+def _resolve_api_key() -> str:
+    for env in ("GEMINI_API_KEY", "GOOGLE_API_KEY"):
+        key = os.environ.get(env)
+        if key:
+            return key.strip()
+    token = Path("~/.gemini_token").expanduser()
+    if token.is_file():
+        return token.read_text().strip()
+    raise ValueError("judge.third_person_separation needs a Gemini API key. Set "
+                     "GEMINI_API_KEY (or GOOGLE_API_KEY), or write it to ~/.gemini_token.")
+
+
+@register("judge.third_person_separation")
+class ThirdPersonSeparationMetric(BaseMetric):
+    """Pairwise VLM judge of third-person fg/bg separation; corpus win-rate."""
+
+    name = "judge.third_person_separation"
+    requires_reference = True
+    higher_is_better = True
+    needs_gpu = False
+    is_set_metric = True
+    dependencies = ["google.genai"]
+
+    def __init__(self, model: str = DEFAULT_MODEL, k: int = DEFAULT_K) -> None:
+        super().__init__()
+        self.model = model
+        self.k = k
+        self._client: Any = None
+        self._files: dict[str, Any] = {}  # path -> uploaded Gemini file handle
+        self._records: list[dict] = []  # one per accumulated pair
+
+    # --- model / client -----------------------------------------------------
+    def setup(self) -> None:
+        if self._client is not None:
+            return
+        from google import genai
+        self._client = genai.Client(api_key=_resolve_api_key())
+
+    # --- set-vs-set protocol ------------------------------------------------
+    def reset(self) -> None:
+        self._records = []
+        self._files = {}
+
+    def accumulate(self, sample: dict) -> None:
+        cand = _path_of(sample, "video")
+        base = _path_of(sample, "reference")
+        if cand is None or base is None:
+            return  # nothing to compare
+
+        image = sample.get("image_path")
+        action_text = sample.get("text_prompt") or ""
+        action = sample.get("action")
+
+        rec = self._cached(cand, base, action_text)
+        if rec is None:
+            if self._client is None:
+                self.setup()
+            rec = self._judge_pair(cand, base, image, action_text)
+            self._write_cache(cand, base, action_text, rec)
+        rec = {**rec, "action": action}
+        self._records.append(rec)
+
+    def finalize(self) -> MetricResult:
+        recs = [r for r in self._records if r.get("verdict")]
+        if not recs:
+            return MetricResult(name=self.name, score=None, details={"skipped": "no pairs judged"})
+        wins = sum(r["verdict"] == "candidate" for r in recs)
+        losses = sum(r["verdict"] == "baseline" for r in recs)
+        ties = sum(r["verdict"] == "tie" for r in recs)
+        decided = wins + losses
+        score = wins / decided if decided else None
+
+        # Per-action win-rate, grouped by the raw label (no assumed control scheme).
+        per_action: dict[str, dict] = {}
+        labels: set[str] = {str(r["action"]) for r in recs if r.get("action")}
+        for action in sorted(labels):
+            gr = [r for r in recs if r.get("action") == action]
+            gw = sum(r["verdict"] == "candidate" for r in gr)
+            gl = sum(r["verdict"] == "baseline" for r in gr)
+            per_action[action] = {
+                "n": len(gr),
+                "wins": gw,
+                "losses": gl,
+                "ties": len(gr) - gw - gl,
+                "win_rate": (gw / (gw + gl)) if (gw + gl) else None,
+            }
+        return MetricResult(name=self.name,
+                            score=score,
+                            details={
+                                "wins": wins,
+                                "losses": losses,
+                                "ties": ties,
+                                "n": len(recs),
+                                "win_rate_excl_ties": score,
+                                "per_action": per_action,
+                            })
+
+    def merge_from(self, other: BaseMetric) -> None:
+        assert isinstance(other, ThirdPersonSeparationMetric)
+        self._records.extend(other._records)
+
+    # --- judging ------------------------------------------------------------
+    def _judge_pair(self, cand: str, base: str, image: str | None, action_text: str) -> dict:
+        """k counterbalanced comparisons → aggregated per-pair verdict."""
+        # Seed the A/B alternation from the pair itself so it is reproducible and
+        # independent of evaluation order or which subset is being run.
+        seed = int(hashlib.sha1(f"{cand}|{base}".encode()).hexdigest(), 16)
+        mapped: list[str] = []
+        for i in range(self.k):
+            cand_first = (seed + i) % 2 == 0
+            v1, v2 = (cand, base) if cand_first else (base, cand)
+            winner = self._one_call(image, v1, v2, action_text)
+            if winner == "tie":
+                mapped.append("tie")
+            elif (winner == "video_1") == cand_first:
+                mapped.append("candidate")
+            else:
+                mapped.append("baseline")
+        cand_w = mapped.count("candidate")
+        base_w = mapped.count("baseline")
+        verdict = ("candidate" if cand_w > base_w else "baseline" if base_w > cand_w else "tie")
+        return {
+            "verdict": verdict,
+            "candidate_wins": cand_w,
+            "baseline_wins": base_w,
+            "ties": mapped.count("tie"),
+            "k": self.k,
+            "rubric_id": RUBRIC_ID
+        }
+
+    def _one_call(self, image: str | None, vid1: str, vid2: str, action_text: str) -> str:
+        from google.genai import types
+        contents: list[Any] = [action_text or "Compare these two rollouts."]
+        if image is not None:
+            contents += ["\nFirst frame (input condition, shared by BOTH "
+                         "videos):", self._upload(image)]
+        contents += [
+            "\nVideo 1 (model A's full rollout — watch it in motion):",
+            self._upload(vid1),
+            "\nVideo 2 (model B's full rollout — watch it in motion):",
+            self._upload(vid2),
+            "\n" + RUBRIC + "\n\n" + USER_TASK,
+        ]
+        for attempt in range(6):
+            try:
+                resp = self._client.models.generate_content(model=self.model,
+                                                            contents=contents,
+                                                            config=types.GenerateContentConfig(
+                                                                system_instruction=SYSTEM_PROMPT,
+                                                                response_mime_type="application/json",
+                                                                response_schema=RESPONSE_SCHEMA,
+                                                                temperature=0.4))
+                return json.loads(resp.text).get("winner", "tie")
+            except Exception as exc:  # noqa: BLE001 - transient API errors
+                if attempt == 5:
+                    print(f"[judge] giving up after 6 attempts ({exc}); scoring this call a tie")
+                    break
+                is_429 = "429" in str(exc) or "RESOURCE_EXHAUSTED" in str(exc)
+                time.sleep(40 if is_429 else 2**attempt)
+        return "tie"
+
+    def _upload(self, path: str) -> Any:
+        f = self._files.get(path)
+        if f is not None:
+            return f
+        f = self._client.files.upload(file=path)
+        while f.state.name != "ACTIVE":
+            time.sleep(1)
+            f = self._client.files.get(name=f.name)
+            if f.state.name == "FAILED":
+                raise RuntimeError(f"Gemini upload failed for {path}")
+        self._files[path] = f
+        return f
+
+    # --- per-pair cache -----------------------------------------------------
+    def _cache_path(self, cand: str, base: str, action_text: str) -> Path:
+        # k is intentionally NOT in the key so a larger-k run can reuse an
+        # existing verdict with enough samples (see ``_cached``).
+        h = hashlib.sha1(f"{RUBRIC_ID}|{self.model}|{cand}|{base}|{action_text}".encode()).hexdigest()[:16]
+        return get_cache_dir() / "judge" / "third_person_separation" / f"{h}.json"
+
+    def _cached(self, cand: str, base: str, action_text: str) -> dict | None:
+        cp = self._cache_path(cand, base, action_text)
+        if not cp.is_file():
+            return None
+        try:
+            rec = json.loads(cp.read_text())
+        except Exception:
+            return None
+        return rec if rec.get("k", 0) >= self.k and "verdict" in rec else None
+
+    def _write_cache(self, cand: str, base: str, action_text: str, rec: dict) -> None:
+        cp = self._cache_path(cand, base, action_text)
+        cp.parent.mkdir(parents=True, exist_ok=True)
+        cp.write_text(json.dumps(rec, indent=2))

fastvideo/eval/registry.py

@@ -85,6 +85,8 @@ def _extra_for(metric_name: str) -> str:
         return "eval-physics-iq"
     if metric_name.startswith("audio."):
         return "eval-audio"
+    if metric_name.startswith("judge."):
+        return "eval-judge"
     return "eval"
 
 

fastvideo/layers/linear.py

@@ -219,6 +219,15 @@ class ReplicatedLinear(LinearBase):
                         (e.g. model.layers.0.qkv_proj)
     """
 
+    # Opt-in instrumentation: when ``enable_shape_tracking`` is set to True,
+    # ``forward`` records every unique ``(input_shape, output_shape)`` pair
+    # observed across all ``ReplicatedLinear`` instances, along with the
+    # subclass name that produced it. Used by upcoming QAT-aware backends
+    # to discover which GEMM shapes need quantized kernels. Defaults to
+    # False; default forward path is bit-identical to pre-slice behavior.
+    enable_shape_tracking = False
+    _shape_to_layer_types: dict[tuple[torch.Size, torch.Size], set[str]] = {}
+
     def __init__(
         self,
         input_size: int,
@@ -285,6 +294,8 @@ class ReplicatedLinear(LinearBase):
         bias = self.bias if not self.skip_bias_add else None
         assert self.quant_method is not None
         output = self.quant_method.apply(self, x, bias)
+        if self.enable_shape_tracking:
+            self._track_shape(x.shape, output.shape)
         output_bias = self.bias if self.skip_bias_add else None
         return output, output_bias
 
@@ -294,6 +305,41 @@ class ReplicatedLinear(LinearBase):
         s += f", bias={self.bias is not None}"
         return s
 
+    @classmethod
+    def get_shape_mapping(cls) -> dict:
+        """Get the mapping from (input_shape, output_shape) to layer types."""
+        return cls._shape_to_layer_types.copy()
+
+    @classmethod
+    def reset_shape_tracking(cls) -> None:
+        """Clear tracked shapes and layer type mappings."""
+        cls._shape_to_layer_types.clear()
+
+    def _track_shape(self, input_shape: torch.Size, output_shape: torch.Size) -> None:
+        shape_key = (input_shape, output_shape)
+        if shape_key not in self._shape_to_layer_types:
+            self._shape_to_layer_types[shape_key] = set()
+            logger.debug("Layer: %s | input shape: %s --> output shape: %s, Quant Method: %s", self.prefix, input_shape,
+                         output_shape, self.quant_method.__class__.__name__)
+        self._shape_to_layer_types[shape_key].add(self.__class__.__name__)
+
+    @classmethod
+    def print_shape_summary(cls) -> None:
+        """Log a summary of all unique shapes and their layer types."""
+        if not cls._shape_to_layer_types:
+            logger.info("No shapes have been processed yet.")
+            return
+
+        lines = [
+            "=== Matrix Multiplication Shape Summary ===",
+            f"Total unique shapes: {len(cls._shape_to_layer_types)}",
+        ]
+        for i, (shape_key, layer_types) in enumerate(cls._shape_to_layer_types.items(), 1):
+            input_shape, output_shape = shape_key
+            lines.append(f"{i}. Input: {input_shape} → Output: {output_shape}")
+            lines.append(f"   Layer types: {', '.join(sorted(layer_types))}")
+        logger.info("\n".join(lines))
+
 
 class ColumnParallelLinear(LinearBase):
     """Linear layer with column parallelism.

fastvideo/layers/lora/linear.py

@@ -317,25 +317,8 @@ class RowParallelLinearWithLoRA(BaseLayerWithLoRA):
             input_parallel = splitted_input[tp_rank].contiguous()
         output_parallel = self.base_layer.quant_method.apply(self.base_layer, input_parallel)
 
-        if not self.merged and not self.disable_lora:
-            lora_A = self.lora_A
-            lora_B = self.lora_B
-            if lora_A is None or lora_B is None:
-                raise RuntimeError("LoRA weights (lora_A, lora_B) must be initialized "
-                                   "before forward pass when LoRA is enabled.")
-            if isinstance(lora_A, DTensor):
-                lora_A = lora_A.to_local()
-            if isinstance(lora_B, DTensor):
-                lora_B = lora_B.to_local()
-
-            lora_A_sliced = self.slice_lora_a_weights(lora_A.to(input_parallel, non_blocking=True))
-            lora_B_sliced = self.slice_lora_b_weights(lora_B.to(output_parallel, non_blocking=True))
-            delta = input_parallel @ lora_A_sliced.T @ lora_B_sliced.T
-            if self.lora_alpha != self.lora_rank:
-                delta = delta * (
-                    self.lora_alpha / self.lora_rank  # type: ignore
-                )
-            output_parallel = output_parallel + delta
+        if self.set_lora:
+            output_parallel = self.apply_lora(output_parallel, input_parallel)
 
         if self.base_layer.reduce_results and self.base_layer.tp_size > 1:
             output_ = tensor_model_parallel_all_reduce(output_parallel)
@@ -351,17 +334,12 @@ class RowParallelLinearWithLoRA(BaseLayerWithLoRA):
         return output, output_bias
 
     def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor:
-        shard_size = self.base_layer.input_size_per_partition
-        # LoRA A gets its input size from base_layer.weight.shape[1].
-        # If that size is already input_size_per_partition, A is already
-        # sharded; otherwise it is a global tensor and needs TP slicing.
-        if A.shape[1] == shard_size:
-            return A.contiguous()
-
         tp_rank = get_tp_rank()
+        shard_size = self.base_layer.input_size_per_partition
         start_idx = tp_rank * shard_size
         end_idx = (tp_rank + 1) * shard_size
-        return A[:, start_idx:end_idx].contiguous()
+        A = A[:, start_idx:end_idx].contiguous()
+        return A
 
     def slice_lora_b_weights(self, B: torch.Tensor) -> torch.Tensor:
         return B

fastvideo/layers/mlp.py

@@ -5,6 +5,7 @@ import torch.nn as nn
 
 from fastvideo.layers.activation import get_act_fn
 from fastvideo.layers.linear import ReplicatedLinear
+from fastvideo.layers.quantization import QuantizationConfig
 
 
 class MLP(nn.Module):
@@ -21,18 +22,27 @@ class MLP(nn.Module):
         act_type: str = "gelu_pytorch_tanh",
         dtype: torch.dtype | None = None,
         prefix: str = "",
+        quant_config: QuantizationConfig | None = None,
     ):
         super().__init__()
         self.fc_in = ReplicatedLinear(
             input_dim,
             mlp_hidden_dim,  # For activation func like SiLU that need 2x width
             bias=bias,
-            params_dtype=dtype)
+            params_dtype=dtype,
+            quant_config=quant_config,
+            prefix=f"{prefix}.fc_in",
+        )
 
         self.act = get_act_fn(act_type)
         if output_dim is None:
             output_dim = input_dim
-        self.fc_out = ReplicatedLinear(mlp_hidden_dim, output_dim, bias=bias, params_dtype=dtype)
+        self.fc_out = ReplicatedLinear(mlp_hidden_dim,
+                                       output_dim,
+                                       bias=bias,
+                                       params_dtype=dtype,
+                                       quant_config=quant_config,
+                                       prefix=f"{prefix}.fc_out")
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         x, _ = self.fc_in(x)

fastvideo/layers/rotary_embedding.py

@@ -52,6 +52,7 @@ def apply_rotary_emb(
     freqs_cis: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
     use_real: bool = True,
     use_real_unbind_dim: int = -1,
+    sequence_dim: int = 2,
 ) -> torch.Tensor:
     """
     Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings
@@ -60,17 +61,23 @@ def apply_rotary_emb(
     tensors contain rotary embeddings and are returned as real tensors.
     Args:
         x (`torch.Tensor`):
-            Query or key tensor to apply rotary embeddings. [B, H, S, D] xk (torch.Tensor): Key tensor to apply
+            Query or key tensor to apply rotary embeddings. [B, H, S, D] if sequence_dim=2 else [B, S, H, D].
         freqs_cis (`Tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],)
+        sequence_dim: 1 = sequence at dim 1 (cos [1,S,1,D], x [B,S,H,D]); 2 = sequence at dim 2 (cos [1,1,S,D], x [B,H,S,D]).
     Returns:
         Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
     """
     if use_real:
         cos, sin = freqs_cis  # [S, D]
-        # Match Diffusers broadcasting (sequence_dim=2 case)
-        cos = cos[None, None, :, :]
-        sin = sin[None, None, :, :]
         cos, sin = cos.to(x.device), sin.to(x.device)
+        if sequence_dim == 2:
+            cos = cos[None, None, :, :]
+            sin = sin[None, None, :, :]
+        elif sequence_dim == 1:
+            cos = cos[None, :, None, :]
+            sin = sin[None, :, None, :]
+        else:
+            raise ValueError(f"sequence_dim must be 1 or 2, got {sequence_dim}")
 
         if use_real_unbind_dim == -1:
             # Used for flux, cogvideox, hunyuan-dit

fastvideo/models/dits/wanvideo.py

@@ -26,6 +26,7 @@ from fastvideo.layers.visual_embedding import (ModulateProjection, PatchEmbed,
 from fastvideo.logger import init_logger
 from fastvideo.models.dits.base import BaseDiT
 from fastvideo.platforms import AttentionBackendEnum, current_platform
+from fastvideo.layers.quantization import QuantizationConfig
 
 from fastvideo.distributed.parallel_state import get_sp_world_size
 
@@ -106,7 +107,9 @@ class WanSelfAttention(nn.Module):
                  window_size=(-1, -1),
                  qk_norm=True,
                  eps=1e-6,
-                 parallel_attention=False) -> None:
+                 parallel_attention=False,
+                 quant_config: QuantizationConfig | None = None,
+                 prefix: str = "") -> None:
         assert dim % num_heads == 0
         super().__init__()
         self.dim = dim
@@ -118,10 +121,10 @@ class WanSelfAttention(nn.Module):
         self.parallel_attention = parallel_attention
 
         # layers
-        self.to_q = ReplicatedLinear(dim, dim)
-        self.to_k = ReplicatedLinear(dim, dim)
-        self.to_v = ReplicatedLinear(dim, dim)
-        self.to_out = ReplicatedLinear(dim, dim)
+        self.to_q = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.to_q")
+        self.to_k = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.to_k")
+        self.to_v = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.to_v")
+        self.to_out = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.to_out")
         self.norm_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
         self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
 
@@ -194,13 +197,15 @@ class WanI2VCrossAttention(WanSelfAttention):
         qk_norm=True,
         eps=1e-6,
         supported_attention_backends: tuple[AttentionBackendEnum, ...]
-        | None = None
+        | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
     ) -> None:
         super().__init__(dim, num_heads, window_size, qk_norm, eps,
-                         supported_attention_backends)
+                         supported_attention_backends, quant_config=quant_config, prefix=prefix)
 
-        self.add_k_proj = ReplicatedLinear(dim, dim)
-        self.add_v_proj = ReplicatedLinear(dim, dim)
+        self.add_k_proj = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.add_k_proj")
+        self.add_v_proj = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.add_v_proj")
         self.norm_added_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
         self.norm_added_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
 
@@ -246,16 +251,17 @@ class WanTransformerBlock(nn.Module):
                  added_kv_proj_dim: int | None = None,
                  supported_attention_backends: tuple[AttentionBackendEnum, ...]
                  | None = None,
+                 quant_config: QuantizationConfig | None = None,
                  prefix: str = ""):
         super().__init__()
 
         # 1. Self-attention
         self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
-        self.to_q = ReplicatedLinear(dim, dim, bias=True)
-        self.to_k = ReplicatedLinear(dim, dim, bias=True)
-        self.to_v = ReplicatedLinear(dim, dim, bias=True)
+        self.to_q = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_q")
+        self.to_k = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_k")
+        self.to_v = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_v")
 
-        self.to_out = ReplicatedLinear(dim, dim, bias=True)
+        self.to_out = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_out")
         self.attn1 = DistributedAttention(
             num_heads=num_heads,
             head_size=dim // num_heads,
@@ -290,13 +296,17 @@ class WanTransformerBlock(nn.Module):
             self.attn2 = WanI2VCrossAttention(dim,
                                               num_heads,
                                               qk_norm=qk_norm,
-                                              eps=eps)
+                                              eps=eps,
+                                              quant_config=quant_config,
+                                              prefix=f"{prefix}.attn2")
         else:
             # T2V
             self.attn2 = WanT2VCrossAttention(dim,
                                               num_heads,
                                               qk_norm=qk_norm,
-                                              eps=eps)
+                                              eps=eps,
+                                              quant_config=quant_config,
+                                              prefix=f"{prefix}.attn2")
         self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(
             dim,
             norm_type="layer",
@@ -306,7 +316,7 @@ class WanTransformerBlock(nn.Module):
             compute_dtype=torch.float32)
 
         # 3. Feed-forward
-        self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh")
+        self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh", quant_config=quant_config, prefix=f"{prefix}.ffn")
         self.mlp_residual = ScaleResidual()
 
         self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
@@ -406,17 +416,17 @@ class WanTransformerBlock_VSA(nn.Module):
                  added_kv_proj_dim: int | None = None,
                  supported_attention_backends: tuple[AttentionBackendEnum, ...]
                  | None = None,
+                 quant_config: QuantizationConfig | None = None,
                  prefix: str = ""):
         super().__init__()
 
         # 1. Self-attention
         self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False)
-        self.to_q = ReplicatedLinear(dim, dim, bias=True)
-        self.to_k = ReplicatedLinear(dim, dim, bias=True)
-        self.to_v = ReplicatedLinear(dim, dim, bias=True)
-        self.to_gate_compress = ReplicatedLinear(dim, dim, bias=True)
-
-        self.to_out = ReplicatedLinear(dim, dim, bias=True)
+        self.to_q = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_q")
+        self.to_k = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_k")
+        self.to_v = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_v")
+        self.to_gate_compress = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_gate_compress")
+        self.to_out = ReplicatedLinear(dim, dim, bias=True, quant_config=quant_config, prefix=f"{prefix}.to_out")
         self.attn1 = DistributedAttention_VSA(
             num_heads=num_heads,
             head_size=dim // num_heads,
@@ -451,13 +461,17 @@ class WanTransformerBlock_VSA(nn.Module):
             self.attn2 = WanI2VCrossAttention(dim,
                                               num_heads,
                                               qk_norm=qk_norm,
-                                              eps=eps)
+                                              eps=eps,
+                                              quant_config=quant_config,
+                                              prefix=f"{prefix}.attn2")
         else:
             # T2V
             self.attn2 = WanT2VCrossAttention(dim,
                                               num_heads,
                                               qk_norm=qk_norm,
-                                              eps=eps)
+                                              eps=eps,
+                                              quant_config=quant_config,
+                                              prefix=f"{prefix}.attn2")
         self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift(
             dim,
             norm_type="layer",
@@ -467,7 +481,7 @@ class WanTransformerBlock_VSA(nn.Module):
             compute_dtype=torch.float32)
 
         # 3. Feed-forward
-        self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh")
+        self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh", quant_config=quant_config, prefix=f"{prefix}.ffn")
         self.mlp_residual = ScaleResidual()
 
         self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
@@ -556,6 +570,7 @@ class WanTransformer3DModel(BaseDiT):
     def __init__(self, config: WanVideoConfig, hf_config: dict[str,
                                                                Any]) -> None:
         super().__init__(config=config, hf_config=hf_config)
+        self.quant_config = config.quant_config
 
         inner_dim = config.num_attention_heads * config.attention_head_dim
         self.hidden_size = config.hidden_size
@@ -594,6 +609,7 @@ class WanTransformer3DModel(BaseDiT):
                               config.eps,
                               config.added_kv_proj_dim,
                               self._supported_attention_backends,
+                              quant_config=config.quant_config,
                               prefix=f"{config.prefix}.blocks.{i}")
             for i in range(config.num_layers)
         ])

fastvideo/models/encoders/mistral3.py

@@ -0,0 +1,48 @@
+# SPDX-License-Identifier: Apache-2.0
+"""HF-backed Mistral3 text encoder wrapper for full Flux2."""
+from typing import Any
+
+import torch
+from torch import nn
+
+from fastvideo.configs.models.encoders.mistral3 import Mistral3TextConfig
+from fastvideo.models.encoders.base import TextEncoder
+
+
+class Mistral3ForConditionalGeneration(TextEncoder):
+    """Loads the Transformers Mistral3 implementation for Flux2 text encoding."""
+
+    supports_hf_from_pretrained = True
+
+    def __init__(self, config: Mistral3TextConfig) -> None:
+        super().__init__(config)
+        self.config = config
+
+    @classmethod
+    def from_pretrained_local(
+        cls,
+        model_path: str,
+        model_config: Mistral3TextConfig,
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> nn.Module:
+        from transformers import AutoModelForImageTextToText
+
+        model = AutoModelForImageTextToText.from_pretrained(
+            model_path,
+            local_files_only=True,
+            torch_dtype=dtype,
+            low_cpu_mem_usage=True,
+        ).eval()
+        if device.type != "cpu":
+            model = model.to(device)
+        return model
+
+    def forward(self, *args: Any, **kwargs: Any) -> Any:
+        raise NotImplementedError(
+            "Mistral3ForConditionalGeneration is loaded through Transformers "
+            "via from_pretrained_local()."
+        )
+
+
+EntryClass = Mistral3ForConditionalGeneration

fastvideo/models/encoders/qwen3.py

@@ -0,0 +1,461 @@
+# SPDX-License-Identifier: Apache-2.0
+# Ported from SGLang: python/sglang/multimodal_gen/runtime/models/encoders/qwen3.py
+"""Qwen3 causal LM text encoder for FastVideo diffusion models (e.g. Flux2 Klein)."""
+from collections.abc import Iterable
+from typing import Any
+
+import torch
+from torch import nn
+
+from fastvideo.attention import LocalAttention
+from fastvideo.configs.models.encoders import BaseEncoderOutput
+from fastvideo.configs.models.encoders.qwen3 import Qwen3TextConfig
+from fastvideo.distributed import get_tp_world_size
+from fastvideo.layers.activation import SiluAndMul
+from fastvideo.layers.layernorm import RMSNorm
+from fastvideo.layers.linear import (
+    MergedColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from fastvideo.layers.quantization import QuantizationConfig
+from fastvideo.layers.rotary_embedding import get_rope
+from fastvideo.layers.vocab_parallel_embedding import VocabParallelEmbedding
+from fastvideo.models.encoders.base import TextEncoder
+from fastvideo.models.loader.weight_utils import (
+    default_weight_loader,
+    maybe_remap_kv_scale_name,
+)
+
+
+class Qwen3MLP(nn.Module):
+    """Qwen3 MLP with SwiGLU activation and tensor parallelism."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_act: str,
+        quant_config: QuantizationConfig | None = None,
+        bias: bool = False,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.gate_up_proj = MergedColumnParallelLinear(
+            input_size=hidden_size,
+            output_sizes=[intermediate_size] * 2,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.gate_up_proj",
+        )
+        self.down_proj = RowParallelLinear(
+            input_size=intermediate_size,
+            output_size=hidden_size,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.down_proj",
+        )
+        if hidden_act != "silu":
+            raise ValueError(
+                f"Unsupported activation: {hidden_act}. Only silu is supported."
+            )
+        self.act_fn = SiluAndMul()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x, _ = self.gate_up_proj(x)
+        x = self.act_fn(x)
+        x, _ = self.down_proj(x)
+        return x
+
+
+class Qwen3Attention(nn.Module):
+    """Qwen3 attention with QK-Norm and tensor parallelism.
+
+    Key difference from LLaMA: RMSNorm is applied to Q and K before attention.
+    """
+
+    def __init__(
+        self,
+        config: Qwen3TextConfig,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        rope_theta: float = 1000000.0,
+        rope_scaling: dict[str, Any] | None = None,
+        max_position_embeddings: int = 40960,
+        quant_config: QuantizationConfig | None = None,
+        bias: bool = False,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        tp_size = get_tp_world_size()
+        self.total_num_heads = num_heads
+        assert self.total_num_heads % tp_size == 0
+        self.num_heads = self.total_num_heads // tp_size
+        self.total_num_kv_heads = num_kv_heads
+        if self.total_num_kv_heads >= tp_size:
+            assert self.total_num_kv_heads % tp_size == 0
+        else:
+            assert tp_size % self.total_num_kv_heads == 0
+        self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
+
+        self.head_dim = getattr(
+            config, "head_dim", self.hidden_size // self.total_num_heads
+        )
+        self.rotary_dim = self.head_dim
+        self.q_size = self.num_heads * self.head_dim
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.scaling = self.head_dim**-0.5
+        self.rope_theta = rope_theta
+        self.max_position_embeddings = max_position_embeddings
+
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size=hidden_size,
+            head_size=self.head_dim,
+            total_num_heads=self.total_num_heads,
+            total_num_kv_heads=self.total_num_kv_heads,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+        )
+        self.o_proj = RowParallelLinear(
+            input_size=self.total_num_heads * self.head_dim,
+            output_size=hidden_size,
+            bias=bias,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        rms_norm_eps = getattr(config, "rms_norm_eps", 1e-6)
+        self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
+        self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
+
+        self.rotary_emb = get_rope(
+            self.head_dim,
+            rotary_dim=self.rotary_dim,
+            max_position=max_position_embeddings,
+            base=int(rope_theta),
+            rope_scaling=rope_scaling,
+            is_neox_style=True,
+        )
+
+        self.attn = LocalAttention(
+            self.num_heads,
+            self.head_dim,
+            self.num_kv_heads,
+            softmax_scale=self.scaling,
+            causal=True,
+            supported_attention_backends=config._supported_attention_backends,
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        qkv, _ = self.qkv_proj(hidden_states)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+
+        batch_size, seq_len = q.shape[0], q.shape[1]
+        q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim)
+        k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
+        v = v.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
+
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        q = q.reshape(batch_size, seq_len, -1)
+        k = k.reshape(batch_size, seq_len, -1)
+
+        q, k = self.rotary_emb(positions, q, k)
+
+        q = q.reshape(batch_size, seq_len, self.num_heads, self.head_dim)
+        k = k.reshape(batch_size, seq_len, self.num_kv_heads, self.head_dim)
+
+        if attention_mask is None:
+            attn_output = self.attn(q, k, v)
+        else:
+            q_sdpa = q.transpose(1, 2)
+            k_sdpa = k.transpose(1, 2)
+            v_sdpa = v.transpose(1, 2)
+            causal_mask = torch.ones(
+                seq_len,
+                seq_len,
+                device=q.device,
+                dtype=torch.bool,
+            ).tril()
+            key_mask = attention_mask.to(device=q.device, dtype=torch.bool)
+            attn_mask = causal_mask[None, None, :, :] & key_mask[:, None, None, :]
+            attn_output = torch.nn.functional.scaled_dot_product_attention(
+                q_sdpa,
+                k_sdpa,
+                v_sdpa,
+                attn_mask=attn_mask,
+                dropout_p=0.0,
+                is_causal=False,
+                scale=self.scaling,
+                enable_gqa=self.num_heads != self.num_kv_heads,
+            ).transpose(1, 2)
+
+        attn_output = attn_output.reshape(batch_size, seq_len, -1)
+
+        output, _ = self.o_proj(attn_output)
+        return output
+
+
+class Qwen3DecoderLayer(nn.Module):
+    """Qwen3 transformer decoder layer."""
+
+    def __init__(
+        self,
+        config: Qwen3TextConfig,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        rope_theta = getattr(config, "rope_theta", 1000000.0)
+        rope_scaling = getattr(config, "rope_scaling", None)
+        max_position_embeddings = getattr(config, "max_position_embeddings", 40960)
+        attention_bias = getattr(config, "attention_bias", False)
+
+        self.self_attn = Qwen3Attention(
+            config=config,
+            hidden_size=self.hidden_size,
+            num_heads=config.num_attention_heads,
+            num_kv_heads=getattr(
+                config, "num_key_value_heads", config.num_attention_heads
+            ),
+            rope_theta=rope_theta,
+            rope_scaling=rope_scaling,
+            max_position_embeddings=max_position_embeddings,
+            quant_config=quant_config,
+            bias=attention_bias,
+            prefix=f"{prefix}.self_attn",
+        )
+        self.mlp = Qwen3MLP(
+            hidden_size=self.hidden_size,
+            intermediate_size=config.intermediate_size,
+            hidden_act=config.hidden_act,
+            quant_config=quant_config,
+            bias=getattr(config, "mlp_bias", False),
+            prefix=f"{prefix}.mlp",
+        )
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor | None,
+        attention_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if residual is None:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+        else:
+            hidden_states, residual = self.input_layernorm(hidden_states, residual)
+
+        hidden_states = self.self_attn(
+            positions=positions,
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+        )
+
+        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+        hidden_states = self.mlp(hidden_states)
+        return hidden_states, residual
+
+
+class Qwen3ForCausalLM(TextEncoder):
+    """Qwen3 causal language model for text encoding in diffusion models (e.g. Flux2 Klein).
+
+    Features:
+    - Tensor parallelism support
+    - FlashAttention/SDPA support via LocalAttention
+    - QK-Norm for better training stability
+    - output_hidden_states for Klein (layers 9, 18, 27)
+    """
+
+    supports_hf_from_pretrained = True
+
+    def __init__(self, config: Qwen3TextConfig) -> None:
+        super().__init__(config)
+
+        self.config = config
+        self.quant_config = getattr(config, "quant_config", None)
+
+        if getattr(config, "lora_config", None) is not None:
+            max_loras = getattr(config.lora_config, "max_loras", 1)
+            lora_vocab_size = getattr(config.lora_config, "lora_extra_vocab_size", 1)
+            lora_vocab = lora_vocab_size * max_loras
+        else:
+            lora_vocab = 0
+        self.vocab_size = config.vocab_size + lora_vocab
+        self.org_vocab_size = config.vocab_size
+
+        self.embed_tokens = VocabParallelEmbedding(
+            self.vocab_size,
+            config.hidden_size,
+            org_num_embeddings=config.vocab_size,
+            quant_config=self.quant_config,
+        )
+
+        self.layers = nn.ModuleList(
+            [
+                Qwen3DecoderLayer(
+                    config=config,
+                    quant_config=self.quant_config,
+                    prefix=f"{config.prefix}.layers.{i}",
+                )
+                for i in range(config.num_hidden_layers)
+            ]
+        )
+
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    @classmethod
+    def from_pretrained_local(
+        cls,
+        model_path: str,
+        model_config: Qwen3TextConfig,
+        dtype: torch.dtype,
+        device: torch.device,
+    ) -> nn.Module:
+        from transformers import AutoModelForCausalLM
+
+        if device.type == "cpu" and torch.cuda.is_available():
+            from fastvideo.distributed import get_local_torch_device
+
+            device = get_local_torch_device()
+
+        return AutoModelForCausalLM.from_pretrained(
+            model_path,
+            local_files_only=True,
+            torch_dtype=dtype,
+            low_cpu_mem_usage=True,
+        ).eval().to(device)
+
+    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        **kwargs: Any,
+    ) -> BaseEncoderOutput:
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+
+        if inputs_embeds is not None:
+            hidden_states = inputs_embeds
+        else:
+            assert input_ids is not None
+            hidden_states = self.get_input_embeddings(input_ids)
+
+        residual = None
+
+        if position_ids is None:
+            position_ids = torch.arange(
+                0, hidden_states.shape[1], device=hidden_states.device
+            ).unsqueeze(0)
+
+        all_hidden_states: tuple[Any, ...] | None = (
+            () if output_hidden_states else None
+        )
+
+        for layer in self.layers:
+            if all_hidden_states is not None:
+                all_hidden_states += (
+                    (hidden_states,)
+                    if residual is None
+                    else (hidden_states + residual,)
+                )
+            hidden_states, residual = layer(
+                position_ids,
+                hidden_states,
+                residual,
+                attention_mask=attention_mask,
+            )
+
+        hidden_states, _ = self.norm(hidden_states, residual)
+
+        if all_hidden_states is not None:
+            all_hidden_states += (hidden_states,)
+
+        return BaseEncoderOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=all_hidden_states,
+        )
+
+    def load_weights(
+        self, weights: Iterable[tuple[str, torch.Tensor]]
+    ) -> set[str]:
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+
+        for name, loaded_weight in weights:
+            if name.startswith("model."):
+                name = name[6:]
+
+            if "rotary_emb.inv_freq" in name:
+                continue
+            if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name:
+                continue
+
+            if "scale" in name:
+                kv_scale_name: str | None = maybe_remap_kv_scale_name(
+                    name, params_dict
+                )
+                if kv_scale_name is None:
+                    continue
+                name = kv_scale_name
+
+            for (
+                param_name,
+                weight_name,
+                shard_id,
+            ) in self.config.arch_config.stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+
+                if name not in params_dict:
+                    continue
+
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+
+                if name not in params_dict:
+                    continue
+
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+
+            loaded_params.add(name)
+
+        return loaded_params
+
+
+EntryClass = Qwen3ForCausalLM

fastvideo/models/loader/component_loader.py

@@ -13,7 +13,7 @@ from typing import cast
 import torch
 import torch.distributed as dist
 import torch.nn as nn
-from safetensors.torch import load_file as safetensors_load_file
+from safetensors.torch import load_file as safetensors_load_file, safe_open
 from torch.distributed import init_device_mesh
 from transformers import AutoImageProcessor, AutoProcessor, AutoTokenizer
 from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
@@ -573,13 +573,53 @@ class TokenizerLoader(ComponentLoader):
                 # If parsing fails, fall through to AutoTokenizer below.
                 pass
 
-        tokenizer = AutoTokenizer.from_pretrained(
-            resolved_model_path,  # "<path to model>/tokenizer"
-            # in v0, this was same string as encoder_name "ClipTextModel"
-            # TODO(will): pass these tokenizer kwargs from inference args? Maybe
-            # other method of config?
-            local_files_only=os.path.isdir(resolved_model_path),
-        )
+        # Only Flux2 full's Mistral3 (require_processor=True) must load via
+        # AutoProcessor. Gate the processor_config.json shortcut on that flag so
+        # existing encoders (e.g. HunyuanVideo 1.5 / Qwen2.5-VL) stay on the
+        # historical AutoTokenizer path below even if their tokenizer dir happens
+        # to ship a processor_config.json.
+        require_processor = False
+        if hasattr(fastvideo_args.pipeline_config, "text_encoder_configs"):
+            try:
+                require_processor = any(
+                    getattr(getattr(cfg, "arch_config", None), "require_processor", False)
+                    for cfg in fastvideo_args.pipeline_config.text_encoder_configs)
+            except Exception:
+                require_processor = False
+
+        if require_processor and os.path.exists(os.path.join(resolved_model_path, "processor_config.json")):
+            processor = AutoProcessor.from_pretrained(
+                resolved_model_path,
+                local_files_only=os.path.isdir(resolved_model_path),
+                trust_remote_code=fastvideo_args.trust_remote_code,
+            )
+            logger.info(
+                "Loaded tokenizer/processor from %s: %s",
+                resolved_model_path,
+                processor.__class__.__name__,
+            )
+            return processor
+
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(
+                resolved_model_path,  # "<path to model>/tokenizer"
+                # in v0, this was same string as encoder_name "ClipTextModel"
+                # TODO(will): pass these tokenizer kwargs from inference args? Maybe
+                # other method of config?
+                local_files_only=os.path.isdir(resolved_model_path),
+            )
+        except (OSError, ValueError):
+            tokenizer = AutoProcessor.from_pretrained(
+                resolved_model_path,
+                local_files_only=os.path.isdir(resolved_model_path),
+                trust_remote_code=fastvideo_args.trust_remote_code,
+            )
+            logger.info(
+                "Loaded tokenizer/processor from %s: %s",
+                resolved_model_path,
+                tokenizer.__class__.__name__,
+            )
+            return tokenizer
         padding_side = None
         if hasattr(fastvideo_args.pipeline_config, "text_encoder_configs"):
             try:
@@ -864,6 +904,18 @@ class VocoderLoader(ComponentLoader):
         return vocoder.eval()
 
 
+def _collect_safetensors_keys(safetensors_list: list) -> set:
+    """Collect all weight keys from safetensors files."""
+    all_keys: set[str] = set()
+    for path in safetensors_list:
+        try:
+            with safe_open(path, framework="pt") as f:
+                all_keys.update(f.keys())
+        except Exception as e:
+            logger.warning("Could not read keys from %s: %s", path, e)
+    return all_keys
+
+
 class TransformerLoader(ComponentLoader):
     """Loader for transformer."""
 
@@ -899,6 +951,12 @@ class TransformerLoader(ComponentLoader):
         if not safetensors_list:
             raise ValueError(f"No safetensors files found in {model_path}")
 
+        # arch_config can infer architecture from weight keys (e.g. Flux2 layer counts)
+        update_fn = getattr(dit_config.arch_config, "update_from_weight_keys", None)
+        if callable(update_fn):
+            weight_keys = _collect_safetensors_keys(safetensors_list)
+            update_fn(weight_keys)
+
         # Check if we should use custom initialization weights
         custom_weights_path = getattr(
             fastvideo_args, "init_weights_from_safetensors", None

fastvideo/models/loader/fsdp_load.py

@@ -332,6 +332,7 @@ def load_model_from_full_model_state_dict(
         NotImplementedError: If got FSDP with more than 1D.
     """
     meta_sd = model.state_dict()
+    named_parameters = dict(model.named_parameters())
     sharded_sd = {}
     custom_param_sd, reverse_param_names_mapping = hf_to_custom_state_dict(
         full_sd_iterator, param_names_mapping)  # type: ignore
@@ -363,8 +364,25 @@ def load_model_from_full_model_state_dict(
             )
         if not hasattr(meta_sharded_param, "device_mesh"):
             full_tensor = full_tensor.to(device=device, dtype=param_dtype)
-            # In cases where parts of the model aren't sharded, some parameters will be plain tensors
-            sharded_tensor = full_tensor
+            target_param = named_parameters.get(target_param_name)
+            weight_loader = getattr(target_param, "weight_loader", None)
+            # Gated on a shape mismatch: only fused/stacked params with a custom
+            # weight_loader (e.g. Qwen3's merged QKV/gate-up) take this path.
+            # Existing models whose unsharded params match the checkpoint shape
+            # fall through to the original `sharded_tensor = full_tensor` below.
+            if target_param is not None and callable(weight_loader) and tuple(target_param.shape) != tuple(
+                    full_tensor.shape):
+                loaded_param = nn.Parameter(torch.empty(tuple(target_param.shape),
+                                                        device=device,
+                                                        dtype=param_dtype),
+                                            requires_grad=False)
+                for attr_name, attr_value in vars(target_param).items():
+                    setattr(loaded_param, attr_name, attr_value)
+                weight_loader(loaded_param, full_tensor)
+                sharded_tensor = loaded_param.data
+            else:
+                # In cases where parts of the model aren't sharded, some parameters will be plain tensors.
+                sharded_tensor = full_tensor
         else:
             full_tensor = full_tensor.to(device=device, dtype=param_dtype)
             sharded_tensor = distribute_tensor(

fastvideo/models/registry.py

@@ -42,6 +42,7 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
     "LingBotWorldTransformer3DModel": ("dits", "lingbotworld", "LingBotWorldTransformer3DModel"),
     "Gen3CTransformer3DModel": ("dits", "gen3c", "Gen3CTransformer3DModel"),
     "Kandinsky5Transformer3DModel": ("dits", "kandinsky5", "Kandinsky5Transformer3DModel"),
+    "Flux2Transformer2DModel": ("dits", "flux_2", "Flux2Transformer2DModel"),
 }
 
 _IMAGE_TO_VIDEO_DIT_MODELS = {
@@ -69,6 +70,9 @@ _TEXT_ENCODER_MODELS = {
     "Qwen2_5_VLForConditionalGeneration":
     ("encoders", "reason1", "Reason1TextEncoder"),
     "LTX2GemmaTextEncoderModel": ("encoders", "gemma", "LTX2GemmaTextEncoderModel"),
+    "Qwen3ForCausalLM": ("encoders", "qwen3", "Qwen3ForCausalLM"),
+    "Mistral3ForConditionalGeneration":
+    ("encoders", "mistral3", "Mistral3ForConditionalGeneration"),
 }
 
 _IMAGE_ENCODER_MODELS: dict[str, tuple] = {
@@ -90,6 +94,7 @@ _VAE_MODELS = {
     ("vaes", "gen3c_tokenizer_vae", "AutoencoderKLGen3CTokenizer"),
     "AutoencoderKLStepvideo": ("vaes", "stepvideovae", "AutoencoderKLStepvideo"),
     "CausalVideoAutoencoder": ("vaes", "ltx2vae", "LTX2CausalVideoAutoencoder"),
+    "AutoencoderKLFlux2": ("vaes", "flux2vae", "AutoencoderKLFlux2"),
     # `stable-audio-open-1.0/vae/config.json` ships `_class_name="AutoencoderOobleck"`
     # (Diffusers' name); FastVideo's class is `OobleckVAE`.
     "AutoencoderOobleck": ("vaes", "oobleck", "OobleckVAE"),

fastvideo/models/vaes/flux2_components.py

@@ -0,0 +1,532 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+# Adapted from: huggingface/diffusers `Encoder`/`Decoder` VAE components
+# at the installed 0.36.0 source surface used by Flux2 Klein.
+
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from fastvideo.models.vaes.common import DiagonalGaussianDistribution
+
+
+@dataclass
+class AutoencoderKLOutput:
+    latent_dist: DiagonalGaussianDistribution
+
+    def __getitem__(self, idx: int):
+        return (self.latent_dist,)[idx]
+
+    def __getattr__(self, name: str):
+        # Existing local Flux2 parity tests used `vae.encode(x).mean` while
+        # diffusers-style callers use `vae.encode(x).latent_dist.mean`.
+        return getattr(self.latent_dist, name)
+
+
+@dataclass
+class DecoderOutput:
+    sample: torch.Tensor
+    commit_loss: Optional[torch.Tensor] = None
+
+    def __getitem__(self, idx: int):
+        return (self.sample, self.commit_loss)[idx]
+
+
+def get_activation(act_fn: str) -> nn.Module:
+    if act_fn in ("swish", "silu"):
+        return nn.SiLU()
+    if act_fn == "mish":
+        return nn.Mish()
+    if act_fn == "gelu":
+        return nn.GELU()
+    if act_fn == "relu":
+        return nn.ReLU()
+    raise ValueError(f"Unsupported activation function: {act_fn}")
+
+
+class AttnProcessor:
+    def __call__(self, attn: "Attention", hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
+        return attn._forward(hidden_states, temb=temb)
+
+
+class AttnAddedKVProcessor(AttnProcessor):
+    pass
+
+
+ADDED_KV_ATTENTION_PROCESSORS = frozenset({AttnAddedKVProcessor})
+CROSS_ATTENTION_PROCESSORS = frozenset({AttnProcessor})
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        query_dim: int,
+        heads: int = 8,
+        dim_head: int = 64,
+        dropout: float = 0.0,
+        bias: bool = False,
+        upcast_softmax: bool = False,
+        norm_num_groups: Optional[int] = None,
+        spatial_norm_dim: Optional[int] = None,
+        out_bias: bool = True,
+        eps: float = 1e-5,
+        rescale_output_factor: float = 1.0,
+        residual_connection: bool = False,
+        _from_deprecated_attn_block: bool = False,
+        **_: object,
+    ):
+        super().__init__()
+        self.inner_dim = dim_head * heads
+        self.query_dim = query_dim
+        self.heads = heads
+        self.dim_head = dim_head
+        self.scale = dim_head**-0.5
+        self.upcast_softmax = upcast_softmax
+        self.rescale_output_factor = rescale_output_factor
+        self.residual_connection = residual_connection
+        self._from_deprecated_attn_block = _from_deprecated_attn_block
+        self.spatial_norm = None
+        if spatial_norm_dim is not None:
+            raise ValueError("Flux2 VAE does not use spatial attention norm in this port")
+        self.group_norm = (
+            nn.GroupNorm(num_channels=query_dim, num_groups=norm_num_groups, eps=eps, affine=True)
+            if norm_num_groups is not None
+            else None
+        )
+        self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias)
+        self.to_k = nn.Linear(query_dim, self.inner_dim, bias=bias)
+        self.to_v = nn.Linear(query_dim, self.inner_dim, bias=bias)
+        self.to_out = nn.ModuleList([nn.Linear(self.inner_dim, query_dim, bias=out_bias), nn.Dropout(dropout)])
+        self.processor = AttnProcessor()
+
+    def set_processor(self, processor: AttnProcessor) -> None:
+        self.processor = processor
+
+    def get_processor(self) -> AttnProcessor:
+        return self.processor
+
+    def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
+        return self.processor(self, hidden_states, temb=temb)
+
+    def _forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
+        residual = hidden_states
+        batch_size, channel, height, width = hidden_states.shape
+        if self.group_norm is not None:
+            hidden_states = self.group_norm(hidden_states)
+        hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+        query = self.to_q(hidden_states)
+        key = self.to_k(hidden_states)
+        value = self.to_v(hidden_states)
+
+        query = query.view(batch_size, -1, self.heads, self.dim_head).transpose(1, 2)
+        key = key.view(batch_size, -1, self.heads, self.dim_head).transpose(1, 2)
+        value = value.view(batch_size, -1, self.heads, self.dim_head).transpose(1, 2)
+
+        if self.upcast_softmax:
+            query = query.float()
+            key = key.float()
+            value = value.float()
+        hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, scale=self.scale)
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, height * width, self.inner_dim)
+        hidden_states = hidden_states.to(self.to_out[0].weight.dtype)
+        hidden_states = self.to_out[0](hidden_states)
+        hidden_states = self.to_out[1](hidden_states)
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, channel, height, width)
+        if self.residual_connection:
+            hidden_states = hidden_states + residual
+        return hidden_states / self.rescale_output_factor
+
+
+class Downsample2D(nn.Module):
+    def __init__(self, channels: int, use_conv: bool = False, out_channels: Optional[int] = None, padding: int = 1, name: str = "conv"):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.padding = padding
+        self.name = name
+        if use_conv:
+            conv = nn.Conv2d(self.channels, self.out_channels, kernel_size=3, stride=2, padding=padding)
+        else:
+            assert self.channels == self.out_channels
+            conv = nn.AvgPool2d(kernel_size=2, stride=2)
+        if name == "conv":
+            self.Conv2d_0 = conv
+            self.conv = conv
+        elif name == "Conv2d_0":
+            self.conv = conv
+        else:
+            self.conv = conv
+
+    def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+        assert hidden_states.shape[1] == self.channels
+        if self.use_conv and self.padding == 0:
+            hidden_states = F.pad(hidden_states, (0, 1, 0, 1), mode="constant", value=0)
+        return self.conv(hidden_states)
+
+
+class Upsample2D(nn.Module):
+    def __init__(self, channels: int, use_conv: bool = False, out_channels: Optional[int] = None, name: str = "conv"):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_conv = use_conv
+        self.use_conv_transpose = False
+        self.name = name
+        self.interpolate = True
+        conv = nn.Conv2d(self.channels, self.out_channels, kernel_size=3, padding=1) if use_conv else None
+        if name == "conv":
+            self.conv = conv
+        else:
+            self.Conv2d_0 = conv
+
+    def forward(self, hidden_states: torch.Tensor, output_size: Optional[int] = None, *args, **kwargs) -> torch.Tensor:
+        assert hidden_states.shape[1] == self.channels
+        dtype = hidden_states.dtype
+        if dtype == torch.bfloat16:
+            hidden_states = hidden_states.float()
+        if output_size is None:
+            hidden_states = F.interpolate(hidden_states, scale_factor=2.0, mode="nearest")
+        else:
+            hidden_states = F.interpolate(hidden_states, size=output_size, mode="nearest")
+        if dtype == torch.bfloat16:
+            hidden_states = hidden_states.to(dtype)
+        if self.use_conv:
+            hidden_states = self.conv(hidden_states) if self.name == "conv" else self.Conv2d_0(hidden_states)
+        return hidden_states
+
+
+class ResnetBlock2D(nn.Module):
+    def __init__(
+        self,
+        *,
+        in_channels: int,
+        out_channels: Optional[int] = None,
+        dropout: float = 0.0,
+        temb_channels: Optional[int] = 512,
+        groups: int = 32,
+        groups_out: Optional[int] = None,
+        eps: float = 1e-6,
+        non_linearity: str = "swish",
+        time_embedding_norm: str = "default",
+        output_scale_factor: float = 1.0,
+        use_in_shortcut: Optional[bool] = None,
+        conv_shortcut_bias: bool = True,
+        conv_2d_out_channels: Optional[int] = None,
+        **_: object,
+    ):
+        super().__init__()
+        if time_embedding_norm not in ("default", "scale_shift"):
+            raise ValueError(f"unknown time_embedding_norm: {time_embedding_norm}")
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.output_scale_factor = output_scale_factor
+        self.time_embedding_norm = time_embedding_norm
+        if groups_out is None:
+            groups_out = groups
+        self.norm1 = nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        if temb_channels is not None:
+            self.time_emb_proj = nn.Linear(temb_channels, out_channels if time_embedding_norm == "default" else 2 * out_channels)
+        else:
+            self.time_emb_proj = None
+        self.norm2 = nn.GroupNorm(num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True)
+        self.dropout = nn.Dropout(dropout)
+        conv_2d_out_channels = conv_2d_out_channels or out_channels
+        self.conv2 = nn.Conv2d(out_channels, conv_2d_out_channels, kernel_size=3, stride=1, padding=1)
+        self.nonlinearity = get_activation(non_linearity)
+        self.use_in_shortcut = self.in_channels != conv_2d_out_channels if use_in_shortcut is None else use_in_shortcut
+        self.conv_shortcut = None
+        if self.use_in_shortcut:
+            self.conv_shortcut = nn.Conv2d(in_channels, conv_2d_out_channels, kernel_size=1, stride=1, padding=0, bias=conv_shortcut_bias)
+
+    def forward(self, input_tensor: torch.Tensor, temb: Optional[torch.Tensor] = None, *args, **kwargs) -> torch.Tensor:
+        hidden_states = self.norm1(input_tensor)
+        hidden_states = self.nonlinearity(hidden_states)
+        hidden_states = self.conv1(hidden_states)
+        if self.time_emb_proj is not None and temb is not None:
+            temb = self.nonlinearity(temb)
+            temb = self.time_emb_proj(temb)[:, :, None, None]
+        if self.time_embedding_norm == "default":
+            if temb is not None:
+                hidden_states = hidden_states + temb
+            hidden_states = self.norm2(hidden_states)
+        else:
+            if temb is None:
+                raise ValueError("temb cannot be None for scale_shift")
+            time_scale, time_shift = torch.chunk(temb, 2, dim=1)
+            hidden_states = self.norm2(hidden_states)
+            hidden_states = hidden_states * (1 + time_scale) + time_shift
+        hidden_states = self.nonlinearity(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.conv2(hidden_states)
+        if self.conv_shortcut is not None:
+            input_tensor = self.conv_shortcut(input_tensor.contiguous() if self.training else input_tensor)
+        return (input_tensor + hidden_states) / self.output_scale_factor
+
+
+class UNetMidBlock2D(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        temb_channels: Optional[int],
+        dropout: float = 0.0,
+        num_layers: int = 1,
+        resnet_eps: float = 1e-6,
+        resnet_time_scale_shift: str = "default",
+        resnet_act_fn: str = "swish",
+        resnet_groups: int = 32,
+        attn_groups: Optional[int] = None,
+        resnet_pre_norm: bool = True,
+        add_attention: bool = True,
+        attention_head_dim: int = 1,
+        output_scale_factor: float = 1.0,
+    ):
+        super().__init__()
+        if resnet_time_scale_shift == "spatial":
+            raise ValueError("Flux2 VAE does not use spatial resnet conditioning in this port")
+        resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
+        self.add_attention = add_attention
+        if attn_groups is None:
+            attn_groups = resnet_groups
+        resnets = [
+            ResnetBlock2D(
+                in_channels=in_channels,
+                out_channels=in_channels,
+                temb_channels=temb_channels,
+                eps=resnet_eps,
+                groups=resnet_groups,
+                dropout=dropout,
+                time_embedding_norm=resnet_time_scale_shift,
+                non_linearity=resnet_act_fn,
+                output_scale_factor=output_scale_factor,
+                pre_norm=resnet_pre_norm,
+            )
+        ]
+        attentions = []
+        if attention_head_dim is None:
+            attention_head_dim = in_channels
+        for _ in range(num_layers):
+            attentions.append(
+                Attention(
+                    in_channels,
+                    heads=in_channels // attention_head_dim,
+                    dim_head=attention_head_dim,
+                    rescale_output_factor=output_scale_factor,
+                    eps=resnet_eps,
+                    norm_num_groups=attn_groups,
+                    residual_connection=True,
+                    bias=True,
+                    upcast_softmax=True,
+                    _from_deprecated_attn_block=True,
+                )
+                if self.add_attention
+                else None
+            )
+            resnets.append(
+                ResnetBlock2D(
+                    in_channels=in_channels,
+                    out_channels=in_channels,
+                    temb_channels=temb_channels,
+                    eps=resnet_eps,
+                    groups=resnet_groups,
+                    dropout=dropout,
+                    time_embedding_norm=resnet_time_scale_shift,
+                    non_linearity=resnet_act_fn,
+                    output_scale_factor=output_scale_factor,
+                    pre_norm=resnet_pre_norm,
+                )
+            )
+        self.attentions = nn.ModuleList(attentions)
+        self.resnets = nn.ModuleList(resnets)
+        self.gradient_checkpointing = False
+
+    def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
+        hidden_states = self.resnets[0](hidden_states, temb)
+        for attn, resnet in zip(self.attentions, self.resnets[1:]):
+            if attn is not None:
+                hidden_states = attn(hidden_states, temb=temb)
+            hidden_states = resnet(hidden_states, temb)
+        return hidden_states
+
+
+class DownEncoderBlock2D(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-6, resnet_act_fn: str = "swish", resnet_groups: int = 32, add_downsample: bool = True, downsample_padding: int = 1, **_: object):
+        super().__init__()
+        self.resnets = nn.ModuleList([
+            ResnetBlock2D(
+                in_channels=in_channels if i == 0 else out_channels,
+                out_channels=out_channels,
+                temb_channels=None,
+                eps=resnet_eps,
+                groups=resnet_groups,
+                dropout=dropout,
+                non_linearity=resnet_act_fn,
+            )
+            for i in range(num_layers)
+        ])
+        self.downsamplers = nn.ModuleList([Downsample2D(out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op")]) if add_downsample else None
+
+    def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+        for resnet in self.resnets:
+            hidden_states = resnet(hidden_states, temb=None)
+        if self.downsamplers is not None:
+            for downsampler in self.downsamplers:
+                hidden_states = downsampler(hidden_states)
+        return hidden_states
+
+
+class AttnDownEncoderBlock2D(DownEncoderBlock2D):
+    def __init__(self, in_channels: int, out_channels: int, attention_head_dim: int = 1, **kwargs: object):
+        super().__init__(in_channels=in_channels, out_channels=out_channels, **kwargs)
+        resnet_groups = int(kwargs.get("resnet_groups", 32))
+        resnet_eps = float(kwargs.get("resnet_eps", 1e-6))
+        output_scale_factor = float(kwargs.get("output_scale_factor", 1.0))
+        if attention_head_dim is None:
+            attention_head_dim = out_channels
+        self.attentions = nn.ModuleList([
+            Attention(out_channels, heads=out_channels // attention_head_dim, dim_head=attention_head_dim, rescale_output_factor=output_scale_factor, eps=resnet_eps, norm_num_groups=resnet_groups, residual_connection=True, bias=True, upcast_softmax=True, _from_deprecated_attn_block=True)
+            for _ in self.resnets
+        ])
+
+    def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+        for resnet, attn in zip(self.resnets, self.attentions):
+            hidden_states = resnet(hidden_states, temb=None)
+            hidden_states = attn(hidden_states)
+        if self.downsamplers is not None:
+            for downsampler in self.downsamplers:
+                hidden_states = downsampler(hidden_states)
+        return hidden_states
+
+
+class UpDecoderBlock2D(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-6, resnet_act_fn: str = "swish", resnet_groups: int = 32, add_upsample: bool = True, temb_channels: Optional[int] = None, **_: object):
+        super().__init__()
+        self.resnets = nn.ModuleList([
+            ResnetBlock2D(
+                in_channels=in_channels if i == 0 else out_channels,
+                out_channels=out_channels,
+                temb_channels=temb_channels,
+                eps=resnet_eps,
+                groups=resnet_groups,
+                dropout=dropout,
+                non_linearity=resnet_act_fn,
+            )
+            for i in range(num_layers)
+        ])
+        self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) if add_upsample else None
+        self.resolution_idx = None
+
+    def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
+        for resnet in self.resnets:
+            hidden_states = resnet(hidden_states, temb=temb)
+        if self.upsamplers is not None:
+            for upsampler in self.upsamplers:
+                hidden_states = upsampler(hidden_states)
+        return hidden_states
+
+
+class AttnUpDecoderBlock2D(UpDecoderBlock2D):
+    def __init__(self, in_channels: int, out_channels: int, attention_head_dim: int = 1, **kwargs: object):
+        super().__init__(in_channels=in_channels, out_channels=out_channels, **kwargs)
+        resnet_groups = int(kwargs.get("resnet_groups", 32))
+        resnet_eps = float(kwargs.get("resnet_eps", 1e-6))
+        output_scale_factor = float(kwargs.get("output_scale_factor", 1.0))
+        if attention_head_dim is None:
+            attention_head_dim = out_channels
+        self.attentions = nn.ModuleList([
+            Attention(out_channels, heads=out_channels // attention_head_dim, dim_head=attention_head_dim, rescale_output_factor=output_scale_factor, eps=resnet_eps, norm_num_groups=resnet_groups, residual_connection=True, bias=True, upcast_softmax=True, _from_deprecated_attn_block=True)
+            for _ in self.resnets
+        ])
+
+    def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor] = None) -> torch.Tensor:
+        for resnet, attn in zip(self.resnets, self.attentions):
+            hidden_states = resnet(hidden_states, temb=temb)
+            hidden_states = attn(hidden_states, temb=temb)
+        if self.upsamplers is not None:
+            for upsampler in self.upsamplers:
+                hidden_states = upsampler(hidden_states)
+        return hidden_states
+
+
+def get_down_block(down_block_type: str, **kwargs: object) -> nn.Module:
+    if down_block_type == "DownEncoderBlock2D":
+        return DownEncoderBlock2D(**kwargs)
+    if down_block_type == "AttnDownEncoderBlock2D":
+        return AttnDownEncoderBlock2D(**kwargs)
+    raise ValueError(f"Unsupported Flux2 VAE down block type: {down_block_type}")
+
+
+def get_up_block(up_block_type: str, **kwargs: object) -> nn.Module:
+    kwargs.pop("prev_output_channel", None)
+    kwargs.pop("resolution_idx", None)
+    if up_block_type == "UpDecoderBlock2D":
+        return UpDecoderBlock2D(**kwargs)
+    if up_block_type == "AttnUpDecoderBlock2D":
+        return AttnUpDecoderBlock2D(**kwargs)
+    raise ValueError(f"Unsupported Flux2 VAE up block type: {up_block_type}")
+
+
+class Encoder(nn.Module):
+    def __init__(self, in_channels: int = 3, out_channels: int = 3, down_block_types: Tuple[str, ...] = ("DownEncoderBlock2D",), block_out_channels: Tuple[int, ...] = (64,), layers_per_block: int = 2, norm_num_groups: int = 32, act_fn: str = "silu", double_z: bool = True, mid_block_add_attention: bool = True):
+        super().__init__()
+        self.layers_per_block = layers_per_block
+        self.conv_in = nn.Conv2d(in_channels, block_out_channels[0], kernel_size=3, stride=1, padding=1)
+        self.down_blocks = nn.ModuleList([])
+        output_channel = block_out_channels[0]
+        for i, down_block_type in enumerate(down_block_types):
+            input_channel = output_channel
+            output_channel = block_out_channels[i]
+            is_final_block = i == len(block_out_channels) - 1
+            self.down_blocks.append(get_down_block(down_block_type, num_layers=self.layers_per_block, in_channels=input_channel, out_channels=output_channel, add_downsample=not is_final_block, resnet_eps=1e-6, downsample_padding=0, resnet_act_fn=act_fn, resnet_groups=norm_num_groups, attention_head_dim=output_channel, temb_channels=None))
+        self.mid_block = UNetMidBlock2D(in_channels=block_out_channels[-1], resnet_eps=1e-6, resnet_act_fn=act_fn, output_scale_factor=1, resnet_time_scale_shift="default", attention_head_dim=block_out_channels[-1], resnet_groups=norm_num_groups, temb_channels=None, add_attention=mid_block_add_attention)
+        self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6)
+        self.conv_act = nn.SiLU()
+        conv_out_channels = 2 * out_channels if double_z else out_channels
+        self.conv_out = nn.Conv2d(block_out_channels[-1], conv_out_channels, 3, padding=1)
+        self.gradient_checkpointing = False
+
+    def forward(self, sample: torch.Tensor) -> torch.Tensor:
+        sample = self.conv_in(sample)
+        for down_block in self.down_blocks:
+            sample = down_block(sample)
+        sample = self.mid_block(sample)
+        sample = self.conv_norm_out(sample)
+        sample = self.conv_act(sample)
+        return self.conv_out(sample)
+
+
+class Decoder(nn.Module):
+    def __init__(self, in_channels: int = 3, out_channels: int = 3, up_block_types: Tuple[str, ...] = ("UpDecoderBlock2D",), block_out_channels: Tuple[int, ...] = (64,), layers_per_block: int = 2, norm_num_groups: int = 32, act_fn: str = "silu", norm_type: str = "group", mid_block_add_attention: bool = True):
+        super().__init__()
+        if norm_type != "group":
+            raise ValueError("Flux2 VAE Decoder only supports group norm in this port")
+        self.layers_per_block = layers_per_block
+        self.conv_in = nn.Conv2d(in_channels, block_out_channels[-1], kernel_size=3, stride=1, padding=1)
+        self.up_blocks = nn.ModuleList([])
+        self.mid_block = UNetMidBlock2D(in_channels=block_out_channels[-1], resnet_eps=1e-6, resnet_act_fn=act_fn, output_scale_factor=1, resnet_time_scale_shift="default", attention_head_dim=block_out_channels[-1], resnet_groups=norm_num_groups, temb_channels=None, add_attention=mid_block_add_attention)
+        reversed_block_out_channels = list(reversed(block_out_channels))
+        output_channel = reversed_block_out_channels[0]
+        for i, up_block_type in enumerate(up_block_types):
+            prev_output_channel = output_channel
+            output_channel = reversed_block_out_channels[i]
+            is_final_block = i == len(block_out_channels) - 1
+            self.up_blocks.append(get_up_block(up_block_type, num_layers=self.layers_per_block + 1, in_channels=prev_output_channel, out_channels=output_channel, prev_output_channel=prev_output_channel, add_upsample=not is_final_block, resnet_eps=1e-6, resnet_act_fn=act_fn, resnet_groups=norm_num_groups, attention_head_dim=output_channel, temb_channels=None, resnet_time_scale_shift=norm_type))
+        self.conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6)
+        self.conv_act = nn.SiLU()
+        self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, 3, padding=1)
+        self.gradient_checkpointing = False
+
+    def forward(self, sample: torch.Tensor, latent_embeds: Optional[torch.Tensor] = None) -> torch.Tensor:
+        sample = self.conv_in(sample)
+        sample = self.mid_block(sample, latent_embeds)
+        for up_block in self.up_blocks:
+            sample = up_block(sample, latent_embeds)
+        sample = self.conv_norm_out(sample)
+        sample = self.conv_act(sample)
+        return self.conv_out(sample)

fastvideo/models/vaes/flux2vae.py

@@ -0,0 +1,533 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copied and adapted from: https://github.com/sglang-ai/sglang
+import math
+from typing import Dict, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+from fastvideo.configs.models.vaes.flux2vae import Flux2VAEConfig
+from fastvideo.models.vaes.common import (
+    DiagonalGaussianDistribution,
+    ParallelTiledVAE,
+)
+from fastvideo.models.vaes.flux2_components import (
+    ADDED_KV_ATTENTION_PROCESSORS,
+    CROSS_ATTENTION_PROCESSORS,
+    Attention,
+    AttnAddedKVProcessor,
+    AttnProcessor,
+    AutoencoderKLOutput,
+    Decoder,
+    DecoderOutput,
+    Encoder,
+)
+
+AttentionProcessor = AttnProcessor
+
+
+class AutoencoderKLFlux2(nn.Module, ParallelTiledVAE):
+    r"""
+    A VAE model with KL loss for encoding images into latents and decoding latent representations into images.
+    
+    This model inherits from [`ParallelTiledVAE`] for tiling support and uses standard diffusers
+    Encoder/Decoder components for Flux2 image generation.
+    """
+
+    _supports_gradient_checkpointing = True
+    _no_split_modules = ["Attention", "ResnetBlock2D"]
+
+    def __init__(
+        self,
+        config: Flux2VAEConfig,
+    ):
+        nn.Module.__init__(self)
+        ParallelTiledVAE.__init__(self, config=config)
+
+        self.config = config
+        arch_config = config.arch_config
+
+        in_channels: int = arch_config.in_channels
+        out_channels: int = arch_config.out_channels
+        down_block_types: Tuple[str, ...] = arch_config.down_block_types
+        up_block_types: Tuple[str, ...] = arch_config.up_block_types
+        block_out_channels: Tuple[int, ...] = arch_config.block_out_channels
+        layers_per_block: int = arch_config.layers_per_block
+        act_fn: str = arch_config.act_fn
+        latent_channels: int = arch_config.latent_channels
+        norm_num_groups: int = arch_config.norm_num_groups
+        sample_size: int = arch_config.sample_size
+        force_upcast: bool = arch_config.force_upcast
+        use_quant_conv: bool = arch_config.use_quant_conv
+        use_post_quant_conv: bool = arch_config.use_post_quant_conv
+        mid_block_add_attention: bool = arch_config.mid_block_add_attention
+        batch_norm_eps: float = arch_config.batch_norm_eps
+        batch_norm_momentum: float = arch_config.batch_norm_momentum
+        patch_size: Tuple[int, int] = arch_config.patch_size
+        
+        # pass init params to Encoder
+        self.encoder = Encoder(
+            in_channels=in_channels,
+            out_channels=latent_channels,
+            down_block_types=down_block_types,
+            block_out_channels=block_out_channels,
+            layers_per_block=layers_per_block,
+            act_fn=act_fn,
+            norm_num_groups=norm_num_groups,
+            double_z=True,
+            mid_block_add_attention=mid_block_add_attention,
+        )
+
+        # pass init params to Decoder
+        self.decoder = Decoder(
+            in_channels=latent_channels,
+            out_channels=out_channels,
+            up_block_types=up_block_types,
+            block_out_channels=block_out_channels,
+            layers_per_block=layers_per_block,
+            norm_num_groups=norm_num_groups,
+            act_fn=act_fn,
+            mid_block_add_attention=mid_block_add_attention,
+        )
+
+        self.quant_conv = (
+            nn.Conv2d(2 * latent_channels, 2 * latent_channels, 1)
+            if use_quant_conv
+            else None
+        )
+        self.post_quant_conv = (
+            nn.Conv2d(latent_channels, latent_channels, 1)
+            if use_post_quant_conv
+            else None
+        )
+
+        self.bn = nn.BatchNorm2d(
+            math.prod(patch_size) * latent_channels,
+            eps=batch_norm_eps,
+            momentum=batch_norm_momentum,
+            affine=False,
+            track_running_stats=True,
+        )
+
+        self.use_slicing = False
+        self.use_tiling = False
+
+        # only relevant if vae tiling is enabled
+        self.tile_sample_min_size = sample_size
+        sample_size_val = (
+            sample_size[0]
+            if isinstance(sample_size, (list, tuple))
+            else sample_size
+        )
+        self.tile_latent_min_size = int(
+            sample_size_val / (2 ** (len(block_out_channels) - 1))
+        )
+        self.tile_overlap_factor = 0.25
+
+    @property
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.attn_processors
+    def attn_processors(self) -> Dict[str, AttentionProcessor]:
+        r"""
+        Returns:
+            `dict` of attention processors: A dictionary containing all attention processors used in the model with
+            indexed by its weight name.
+        """
+        # set recursively
+        processors = {}
+
+        def fn_recursive_add_processors(
+            name: str,
+            module: torch.nn.Module,
+            processors: Dict[str, AttentionProcessor],
+        ):
+            if hasattr(module, "get_processor"):
+                processors[f"{name}.processor"] = module.get_processor()
+
+            for sub_name, child in module.named_children():
+                fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
+
+            return processors
+
+        for name, module in self.named_children():
+            fn_recursive_add_processors(name, module, processors)
+
+        return processors
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_attn_processor
+    def set_attn_processor(
+        self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]
+    ):
+        r"""
+        Sets the attention processor to use to compute attention.
+
+        Parameters:
+            processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
+                The instantiated processor class or a dictionary of processor classes that will be set as the processor
+                for **all** `Attention` layers.
+
+                If `processor` is a dict, the key needs to define the path to the corresponding cross attention
+                processor. This is strongly recommended when setting trainable attention processors.
+
+        """
+        count = len(self.attn_processors.keys())
+
+        if isinstance(processor, dict) and len(processor) != count:
+            raise ValueError(
+                f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
+                f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
+            )
+
+        def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
+            if hasattr(module, "set_processor"):
+                if not isinstance(processor, dict):
+                    module.set_processor(processor)
+                else:
+                    module.set_processor(processor.pop(f"{name}.processor"))
+
+            for sub_name, child in module.named_children():
+                fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
+
+        for name, module in self.named_children():
+            fn_recursive_attn_processor(name, module, processor)
+
+    # Copied from diffusers.models.unets.unet_2d_condition.UNet2DConditionModel.set_default_attn_processor
+    def set_default_attn_processor(self):
+        """
+        Disables custom attention processors and sets the default attention implementation.
+        """
+        if all(
+            proc.__class__ in ADDED_KV_ATTENTION_PROCESSORS
+            for proc in self.attn_processors.values()
+        ):
+            processor = AttnAddedKVProcessor()
+        elif all(
+            proc.__class__ in CROSS_ATTENTION_PROCESSORS
+            for proc in self.attn_processors.values()
+        ):
+            processor = AttnProcessor()
+        else:
+            raise ValueError(
+                f"Cannot call `set_default_attn_processor` when attention processors are of type {next(iter(self.attn_processors.values()))}"
+            )
+
+        self.set_attn_processor(processor)
+
+    def _encode(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size, num_channels, height, width = x.shape
+
+        if self.use_tiling and (
+            width > self.tile_sample_min_size or height > self.tile_sample_min_size
+        ):
+            return self._tiled_encode(x)
+
+        enc = self.encoder(x)
+        if self.quant_conv is not None:
+            enc = self.quant_conv(enc)
+
+        return enc
+
+    def encode(
+        self, x: torch.Tensor, return_dict: bool = True
+    ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
+        """
+        Encode a batch of images into latents.
+
+        Args:
+            x (`torch.Tensor`): Input batch of images.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
+
+        Returns:
+                The latent representations of the encoded images. If `return_dict` is True, a
+                [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned.
+        """
+
+        if x.ndim == 5:
+            assert x.shape[2] == 1
+            x = x.squeeze(2)
+
+        if self.use_slicing and x.shape[0] > 1:
+            encoded_slices = [self._encode(x_slice) for x_slice in x.split(1)]
+            h = torch.cat(encoded_slices)
+        else:
+            h = self._encode(x)
+
+        posterior = DiagonalGaussianDistribution(h)
+        if not return_dict:
+            return (posterior,)
+
+        return AutoencoderKLOutput(latent_dist=posterior)
+
+    def _decode(
+        self, z: torch.Tensor, return_dict: bool = True
+    ) -> Union[DecoderOutput, torch.Tensor]:
+        if self.use_tiling and (
+            z.shape[-1] > self.tile_latent_min_size
+            or z.shape[-2] > self.tile_latent_min_size
+        ):
+            return self.tiled_decode(z, return_dict=return_dict)
+
+        if self.post_quant_conv is not None:
+            z = self.post_quant_conv(z)
+
+        dec = self.decoder(z)
+
+        if not return_dict:
+            return (dec,)
+
+        return DecoderOutput(sample=dec)
+
+    def decode(
+        self, z: torch.FloatTensor, return_dict: bool = True, generator=None
+    ) -> Union[DecoderOutput, torch.FloatTensor]:
+        """
+        Decode a batch of images.
+
+        Args:
+            z (`torch.Tensor`): Input batch of latent vectors.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~models.vae.DecoderOutput`] or `tuple`:
+                If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
+                returned.
+
+        """
+        if self.use_slicing and z.shape[0] > 1:
+            decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)]
+            decoded = torch.cat(decoded_slices)
+        else:
+            decoded = self._decode(z).sample
+
+        if not return_dict:
+            return (decoded,)
+
+        return DecoderOutput(sample=decoded)
+
+    def blend_v(
+        self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
+    ) -> torch.Tensor:
+        blend_extent = min(a.shape[2], b.shape[2], blend_extent)
+        for y in range(blend_extent):
+            b[:, :, y, :] = a[:, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[
+                :, :, y, :
+            ] * (y / blend_extent)
+        return b
+
+    def blend_h(
+        self, a: torch.Tensor, b: torch.Tensor, blend_extent: int
+    ) -> torch.Tensor:
+        blend_extent = min(a.shape[3], b.shape[3], blend_extent)
+        for x in range(blend_extent):
+            b[:, :, :, x] = a[:, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[
+                :, :, :, x
+            ] * (x / blend_extent)
+        return b
+
+    def _tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
+        r"""Encode a batch of images using a tiled encoder.
+
+        When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several
+        steps. This is useful to keep memory use constant regardless of image size. The end result of tiled encoding is
+        different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the
+        tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the
+        output, but they should be much less noticeable.
+
+        Args:
+            x (`torch.Tensor`): Input batch of images.
+
+        Returns:
+            `torch.Tensor`:
+                The latent representation of the encoded videos.
+        """
+
+        overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
+        row_limit = self.tile_latent_min_size - blend_extent
+
+        # Split the image into 512x512 tiles and encode them separately.
+        rows = []
+        for i in range(0, x.shape[2], overlap_size):
+            row = []
+            for j in range(0, x.shape[3], overlap_size):
+                tile = x[
+                    :,
+                    :,
+                    i : i + self.tile_sample_min_size,
+                    j : j + self.tile_sample_min_size,
+                ]
+                tile = self.encoder(tile)
+                if self.quant_conv is not None:
+                    tile = self.quant_conv(tile)
+                row.append(tile)
+            rows.append(row)
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :row_limit, :row_limit])
+            result_rows.append(torch.cat(result_row, dim=3))
+
+        enc = torch.cat(result_rows, dim=2)
+        return enc
+
+    def tiled_encode(
+        self, x: torch.Tensor, return_dict: bool = True
+    ) -> AutoencoderKLOutput:
+        r"""Encode a batch of images using a tiled encoder.
+
+        When this option is enabled, the VAE will split the input tensor into tiles to compute encoding in several
+        steps. This is useful to keep memory use constant regardless of image size. The end result of tiled encoding is
+        different from non-tiled encoding because each tile uses a different encoder. To avoid tiling artifacts, the
+        tiles overlap and are blended together to form a smooth output. You may still see tile-sized changes in the
+        output, but they should be much less noticeable.
+
+        Args:
+            x (`torch.Tensor`): Input batch of images.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.autoencoder_kl.AutoencoderKLOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~models.autoencoder_kl.AutoencoderKLOutput`] or `tuple`:
+                If return_dict is True, a [`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain
+                `tuple` is returned.
+        """
+        deprecation_message = (
+            "The tiled_encode implementation supporting the `return_dict` parameter is deprecated. In the future, the "
+            "implementation of this method will be replaced with that of `_tiled_encode` and you will no longer be able "
+            "to pass `return_dict`. You will also have to create a `DiagonalGaussianDistribution()` from the returned value."
+        )
+
+        overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
+        row_limit = self.tile_latent_min_size - blend_extent
+
+        # Split the image into 512x512 tiles and encode them separately.
+        rows = []
+        for i in range(0, x.shape[2], overlap_size):
+            row = []
+            for j in range(0, x.shape[3], overlap_size):
+                tile = x[
+                    :,
+                    :,
+                    i : i + self.tile_sample_min_size,
+                    j : j + self.tile_sample_min_size,
+                ]
+                tile = self.encoder(tile)
+                if self.quant_conv is not None:
+                    tile = self.quant_conv(tile)
+                row.append(tile)
+            rows.append(row)
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :row_limit, :row_limit])
+            result_rows.append(torch.cat(result_row, dim=3))
+
+        moments = torch.cat(result_rows, dim=2)
+        posterior = DiagonalGaussianDistribution(moments)
+
+        if not return_dict:
+            return (posterior,)
+
+        return AutoencoderKLOutput(latent_dist=posterior)
+
+    def tiled_decode(
+        self, z: torch.Tensor, return_dict: bool = True
+    ) -> Union[DecoderOutput, torch.Tensor]:
+        r"""
+        Decode a batch of images using a tiled decoder.
+
+        Args:
+            z (`torch.Tensor`): Input batch of latent vectors.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~models.vae.DecoderOutput`] or `tuple`:
+                If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
+                returned.
+        """
+        overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))
+        blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)
+        row_limit = self.tile_sample_min_size - blend_extent
+
+        # Split z into overlapping 64x64 tiles and decode them separately.
+        # The tiles have an overlap to avoid seams between tiles.
+        rows = []
+        for i in range(0, z.shape[2], overlap_size):
+            row = []
+            for j in range(0, z.shape[3], overlap_size):
+                tile = z[
+                    :,
+                    :,
+                    i : i + self.tile_latent_min_size,
+                    j : j + self.tile_latent_min_size,
+                ]
+                if self.post_quant_conv is not None:
+                    tile = self.post_quant_conv(tile)
+                decoded = self.decoder(tile)
+                row.append(decoded)
+            rows.append(row)
+        result_rows = []
+        for i, row in enumerate(rows):
+            result_row = []
+            for j, tile in enumerate(row):
+                # blend the above tile and the left tile
+                # to the current tile and add the current tile to the result row
+                if i > 0:
+                    tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+                if j > 0:
+                    tile = self.blend_h(row[j - 1], tile, blend_extent)
+                result_row.append(tile[:, :, :row_limit, :row_limit])
+            result_rows.append(torch.cat(result_row, dim=3))
+
+        dec = torch.cat(result_rows, dim=2)
+        if not return_dict:
+            return (dec,)
+
+        return DecoderOutput(sample=dec)
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+        sample_posterior: bool = False,
+        return_dict: bool = True,
+        generator: Optional[torch.Generator] = None,
+    ) -> Union[DecoderOutput, torch.Tensor]:
+        r"""
+        Args:
+            sample (`torch.Tensor`): Input sample.
+            sample_posterior (`bool`, *optional*, defaults to `False`):
+                Whether to sample from the posterior.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`DecoderOutput`] instead of a plain tuple.
+        """
+        x = sample
+        posterior = self.encode(x).latent_dist
+        if sample_posterior:
+            z = posterior.sample(generator=generator)
+        else:
+            z = posterior.mode()
+        dec = self.decode(z).sample
+
+        if not return_dict:
+            return (dec,)
+
+        return DecoderOutput(sample=dec)
+
+
+EntryClass = AutoencoderKLFlux2

fastvideo/pipelines/basic/flux_2/__init__.py

@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Flux2 pipeline module."""
+
+from fastvideo.pipelines.basic.flux_2.flux_2_pipeline import Flux2Pipeline
+from fastvideo.pipelines.basic.flux_2.flux_2_klein_pipeline import Flux2KleinPipeline
+
+__all__ = ["Flux2Pipeline", "Flux2KleinPipeline"]

fastvideo/pipelines/basic/flux_2/flux_2_klein_pipeline.py

@@ -0,0 +1,17 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copied and adapted from: https://github.com/sglang-ai/sglang
+"""
+Flux2 Klein image generation pipeline (distilled, 4-step, no guidance).
+"""
+
+from fastvideo.configs.pipelines.flux_2 import Flux2KleinPipelineConfig
+from fastvideo.pipelines.basic.flux_2.flux_2_pipeline import Flux2Pipeline
+
+
+class Flux2KleinPipeline(Flux2Pipeline):
+    """Flux2 Klein image diffusion pipeline (distilled, 4-step, no guidance)."""
+
+    pipeline_config_cls: type[Flux2KleinPipelineConfig] = Flux2KleinPipelineConfig
+
+
+EntryClass = Flux2KleinPipeline

fastvideo/pipelines/basic/flux_2/flux_2_latent_preparation.py

@@ -0,0 +1,138 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Flux2 latent preparation stage using packed 2x2 layout.
+
+Flux2 uses packed latents: transformer sees 128 channels (32*4) with half
+spatial resolution; after denoising we unpatchify to 32 channels and full
+spatial for VAE decode. This stage prepares (B, 128, T, H//2, W//2).
+"""
+
+import torch
+from diffusers.utils.torch_utils import randn_tensor
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.latent_preparation import LatentPreparationStage
+
+
+class Flux2LatentPreparationStage(LatentPreparationStage):
+    """
+    Latent preparation for Flux2: packed layout with half spatial dimensions.
+
+    Matches diffusers Flux2Pipeline.prepare_latents: shape is
+    (B, num_channels_latents, T, H_latent//2, W_latent//2) so the transformer
+    sees 128 channels and half spatial; after denoising we unpatchify to
+    (B, 32, H_latent, W_latent) before VAE.
+    """
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        """Prepare latents with Flux2 packed half-spatial shape."""
+        from fastvideo.distributed import get_local_torch_device
+
+        latent_num_frames = None
+        if hasattr(self, "adjust_video_length"):
+            latent_num_frames = self.adjust_video_length(batch, fastvideo_args)
+
+        if not batch.prompt_embeds:
+            if batch.keyboard_cond is not None:
+                batch_size = batch.keyboard_cond.shape[0]
+            elif batch.mouse_cond is not None:
+                batch_size = batch.mouse_cond.shape[0]
+            elif batch.image_embeds:
+                batch_size = batch.image_embeds[0].shape[0]
+            else:
+                batch_size = 1
+        elif isinstance(batch.prompt, list):
+            batch_size = len(batch.prompt)
+        elif batch.prompt is not None:
+            batch_size = 1
+        else:
+            batch_size = batch.prompt_embeds[0].shape[0]
+
+        batch_size *= batch.num_videos_per_prompt
+
+        if not batch.prompt_embeds:
+            transformer_dtype = next(self.transformer.parameters()).dtype
+            device = get_local_torch_device()
+            dummy_prompt = torch.zeros(
+                batch_size,
+                0,
+                self.transformer.hidden_size,
+                device=device,
+                dtype=transformer_dtype,
+            )
+            batch.prompt_embeds = [dummy_prompt]
+            batch.negative_prompt_embeds = []
+            batch.do_classifier_free_guidance = False
+
+        dtype = batch.prompt_embeds[0].dtype
+        device = get_local_torch_device()
+        generator = batch.generator
+        latents = batch.latents
+        num_frames = (latent_num_frames if latent_num_frames is not None else batch.num_frames)
+        height = batch.height
+        width = batch.width
+
+        if height is None or width is None:
+            raise ValueError("Height and width must be provided")
+
+        vae_arch = fastvideo_args.pipeline_config.vae_config.arch_config
+        scale = vae_arch.spatial_compression_ratio
+        # Flux2 packed: half spatial (2x2 patch packing)
+        latent_h = (height // scale) // 2
+        latent_w = (width // scale) // 2
+
+        if self.use_btchw_layout:
+            shape = (
+                batch_size,
+                num_frames,
+                self.transformer.num_channels_latents,
+                latent_h,
+                latent_w,
+            )
+            bcthw_shape = tuple(shape[i] for i in [0, 2, 1, 3, 4])
+        else:
+            shape = (
+                batch_size,
+                self.transformer.num_channels_latents,
+                num_frames,
+                latent_h,
+                latent_w,
+            )
+            bcthw_shape = shape
+
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(f"You have passed a list of generators of length {len(generator)}, "
+                             f"but requested an effective batch size of {batch_size}.")
+
+        if latents is None:
+            latents = randn_tensor(
+                shape,
+                generator=generator,
+                device=device,
+                dtype=dtype,
+            )
+            if hasattr(self.scheduler, "init_noise_sigma"):
+                latents = latents * self.scheduler.init_noise_sigma
+        else:
+            latents = latents.to(device)
+            is_longcat_refine = (batch.refine_from is not None or batch.stage1_video is not None)
+            if (not is_longcat_refine) and hasattr(self.scheduler, "init_noise_sigma"):
+                latents = latents * self.scheduler.init_noise_sigma
+
+        batch.latents = latents
+        batch.raw_latent_shape = bcthw_shape
+        latent_ids = torch.cartesian_prod(
+            torch.arange(num_frames, device=device),
+            torch.arange(latent_h, device=device),
+            torch.arange(latent_w, device=device),
+            torch.arange(1, device=device),
+        )
+        batch.extra["flux2_img_ids"] = latent_ids.unsqueeze(0).expand(batch_size, -1, -1)
+        # Flux2 mu depends on image_seq_len; use packed spatial size
+        batch.n_tokens = latent_h * latent_w
+        return batch

fastvideo/pipelines/basic/flux_2/flux_2_pipeline.py

@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copied and adapted from: https://github.com/sglang-ai/sglang
+"""
+Flux2 image generation pipeline implementation.
+
+This module contains an implementation of the Flux2 image diffusion pipeline
+using the modular pipeline architecture.
+"""
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline
+from fastvideo.pipelines.basic.flux_2.flux_2_latent_preparation import (
+    Flux2LatentPreparationStage, )
+from fastvideo.pipelines.basic.flux_2.flux_2_timestep_preparation import (
+    Flux2TimestepPreparationStage, )
+from fastvideo.pipelines.basic.flux_2.flux_2_text_encoding import (
+    Flux2TextEncodingStage, )
+from fastvideo.pipelines.stages import (
+    ConditioningStage,
+    DecodingStage,
+    DenoisingStage,
+    InputValidationStage,
+)
+
+logger = init_logger(__name__)
+
+
+class Flux2Pipeline(LoRAPipeline, ComposedPipelineBase):
+    """
+    Flux2 image diffusion pipeline with LoRA support.
+    """
+
+    _required_config_modules = [
+        "text_encoder",
+        "tokenizer",
+        "vae",
+        "transformer",
+        "scheduler",
+    ]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None:
+        """Set up pipeline stages with proper dependency injection."""
+
+        self.add_stage(
+            stage_name="input_validation_stage",
+            stage=InputValidationStage(),
+        )
+
+        self.add_stage(
+            stage_name="prompt_encoding_stage",
+            stage=Flux2TextEncodingStage(
+                text_encoders=[self.get_module("text_encoder")],
+                tokenizers=[self.get_module("tokenizer")],
+            ),
+        )
+
+        self.add_stage(
+            stage_name="conditioning_stage",
+            stage=ConditioningStage(),
+        )
+
+        self.add_stage(
+            stage_name="latent_preparation_stage",
+            stage=Flux2LatentPreparationStage(
+                scheduler=self.get_module("scheduler"),
+                transformer=self.get_module("transformer", None),
+            ),
+        )
+
+        self.add_stage(
+            stage_name="timestep_preparation_stage",
+            stage=Flux2TimestepPreparationStage(scheduler=self.get_module("scheduler"), ),
+        )
+
+        self.add_stage(
+            stage_name="denoising_stage",
+            stage=DenoisingStage(
+                transformer=self.get_module("transformer"),
+                scheduler=self.get_module("scheduler"),
+                vae=self.get_module("vae"),
+                pipeline=self,
+            ),
+        )
+
+        self.add_stage(
+            stage_name="decoding_stage",
+            stage=DecodingStage(
+                vae=self.get_module("vae"),
+                pipeline=self,
+            ),
+        )
+
+
+EntryClass = Flux2Pipeline

fastvideo/pipelines/basic/flux_2/flux_2_text_encoding.py

@@ -0,0 +1,161 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Flux2 text encoding stages."""
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.forward_context import set_forward_context
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.text_encoding import TextEncodingStage
+
+FLUX2_SYSTEM_MESSAGE = ("You are an AI that reasons about image descriptions. You give structured "
+                        "responses focusing on object relationships, object\nattribution and actions "
+                        "without speculation.")
+
+
+def _format_flux2_full_input(prompts: list[str], system_message: str) -> list[list[dict[str, Any]]]:
+    return [[
+        {
+            "role": "system",
+            "content": [{
+                "type": "text",
+                "text": system_message
+            }],
+        },
+        {
+            "role": "user",
+            "content": [{
+                "type": "text",
+                "text": prompt.replace("[IMG]", "")
+            }],
+        },
+    ] for prompt in prompts]
+
+
+def _prepare_flux2_text_ids(prompt_embeds: torch.Tensor) -> torch.Tensor:
+    batch_size, seq_len, _ = prompt_embeds.shape
+    text_ids = torch.cartesian_prod(
+        torch.arange(1, device=prompt_embeds.device),
+        torch.arange(1, device=prompt_embeds.device),
+        torch.arange(1, device=prompt_embeds.device),
+        torch.arange(seq_len, device=prompt_embeds.device),
+    )
+    return text_ids.unsqueeze(0).expand(batch_size, -1, -1)
+
+
+class Flux2TextEncodingStage(TextEncodingStage):
+    """Text encoding for Flux2 full and Klein variants."""
+
+    def _uses_embedded_guidance(self, fastvideo_args: FastVideoArgs) -> bool:
+        return getattr(fastvideo_args.pipeline_config, "embedded_cfg_scale", None) is not None
+
+    @torch.no_grad()
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        if self._uses_embedded_guidance(fastvideo_args):
+            batch.do_classifier_free_guidance = False
+            batch.negative_prompt_embeds = []
+
+        if batch.prompt_embeds is not None and len(batch.prompt_embeds) > 0:
+            if "flux2_txt_ids" not in batch.extra:
+                batch.extra["flux2_txt_ids"] = _prepare_flux2_text_ids(batch.prompt_embeds[0])
+            return batch
+
+        if getattr(fastvideo_args.pipeline_config, "flux2_text_encoder_type", "") != "mistral3":
+            return super().forward(batch, fastvideo_args)
+
+        assert batch.prompt is not None
+        prompt_embeds, attention_mask = self.encode_flux2_full_text(
+            batch.prompt,
+            fastvideo_args,
+            max_length=batch.max_sequence_length,
+        )
+        batch.prompt_embeds.append(prompt_embeds)
+        batch.extra["flux2_txt_ids"] = _prepare_flux2_text_ids(prompt_embeds)
+        if batch.prompt_attention_mask is not None:
+            batch.prompt_attention_mask.append(attention_mask)
+        return batch
+
+    @torch.no_grad()
+    def encode_flux2_full_text(
+        self,
+        text: str | list[str],
+        fastvideo_args: FastVideoArgs,
+        max_length: int | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        tokenizer = self.tokenizers[0]
+        text_encoder = self.text_encoders[0]
+        encoder_config = fastvideo_args.pipeline_config.text_encoder_configs[0]
+        arch_config = encoder_config.arch_config
+
+        prompts = [text] if isinstance(text, str) else text
+        max_sequence_length = max_length or getattr(arch_config, "text_len", 512) or 512
+        hidden_state_layers = getattr(
+            fastvideo_args.pipeline_config,
+            "text_encoder_out_layers",
+            (10, 20, 30),
+        )
+        system_message = getattr(
+            fastvideo_args.pipeline_config,
+            "flux2_system_message",
+            FLUX2_SYSTEM_MESSAGE,
+        )
+
+        messages = _format_flux2_full_input(prompts, system_message)
+        inputs = tokenizer.apply_chat_template(
+            messages,
+            add_generation_prompt=False,
+            tokenize=True,
+            return_dict=True,
+            return_tensors="pt",
+            padding="max_length",
+            truncation=True,
+            max_length=max_sequence_length,
+        )
+
+        try:
+            encoder_device = next(text_encoder.parameters()).device
+        except StopIteration:
+            encoder_device = get_local_torch_device()
+        encoder_dtype = getattr(text_encoder, "dtype", None)
+
+        input_ids = inputs["input_ids"].to(encoder_device)
+        attention_mask = inputs["attention_mask"].to(encoder_device)
+
+        forward_kwargs: dict[str, Any] = {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "output_hidden_states": True,
+            "use_cache": False,
+        }
+        if "pixel_values" in inputs:
+            forward_kwargs["pixel_values"] = inputs["pixel_values"].to(
+                device=encoder_device,
+                dtype=encoder_dtype or torch.bfloat16,
+            )
+        if "image_sizes" in inputs:
+            forward_kwargs["image_sizes"] = inputs["image_sizes"].to(encoder_device)
+
+        with set_forward_context(current_timestep=0, attn_metadata=None):
+            outputs = text_encoder(**forward_kwargs)
+
+        if outputs.hidden_states is None:
+            raise ValueError("Full Flux2 requires output_hidden_states=True from text encoder")
+
+        stacked = torch.stack([outputs.hidden_states[k] for k in hidden_state_layers], dim=1)
+        if encoder_dtype is not None:
+            stacked = stacked.to(dtype=encoder_dtype)
+        batch_size, num_layers, seq_len, hidden_dim = stacked.shape
+        prompt_embeds = stacked.permute(0, 2, 1, 3).reshape(
+            batch_size,
+            seq_len,
+            num_layers * hidden_dim,
+        )
+        return prompt_embeds, attention_mask

fastvideo/pipelines/basic/flux_2/flux_2_timestep_preparation.py

@@ -0,0 +1,111 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Flux2-specific timestep preparation."""
+
+import inspect
+
+import numpy as np
+import torch
+
+from fastvideo.distributed import get_local_torch_device
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
+from fastvideo.pipelines.stages.timestep_preparation import TimestepPreparationStage
+
+
+def compute_empirical_mu(image_seq_len: int, num_steps: int) -> float:
+    """
+    Resolution-dependent mu for Flux2 flow-match scheduler.
+    From Black Forest Labs flux2 official repo: sampling.compute_empirical_mu.
+    """
+    a1, b1 = 8.73809524e-05, 1.89833333
+    a2, b2 = 0.00016927, 0.45666666
+
+    if image_seq_len > 4300:
+        return float(a2 * image_seq_len + b2)
+
+    m_200 = a2 * image_seq_len + b2
+    m_10 = a1 * image_seq_len + b1
+    a = (m_200 - m_10) / 190.0
+    b = m_200 - 200.0 * a
+    return float(a * num_steps + b)
+
+
+class Flux2TimestepPreparationStage(TimestepPreparationStage):
+    """Flux2 timestep preparation matching the Diffusers Flux2 schedule."""
+
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+    ) -> ForwardBatch:
+        scheduler = self.scheduler
+        device = get_local_torch_device()
+        num_inference_steps = batch.num_inference_steps
+        timesteps = batch.timesteps
+        sigmas = batch.sigmas
+        n_tokens = batch.n_tokens
+
+        extra_set_timesteps_kwargs = {}
+        if n_tokens is not None and "n_tokens" in inspect.signature(scheduler.set_timesteps).parameters:
+            extra_set_timesteps_kwargs["n_tokens"] = n_tokens
+
+        # Flux2/BFL: Diffusers' Flux2 pipeline passes a custom sigma grid and
+        # always supplies the resolution-dependent mu when the scheduler accepts
+        # it.
+        scheduler_config = getattr(scheduler, "config", None)
+        use_flow_sigmas = (getattr(scheduler_config, "use_flow_sigmas", False) if scheduler_config else False)
+        if timesteps is None and sigmas is None and not use_flow_sigmas:
+            sigmas = np.linspace(1.0, 1.0 / num_inference_steps, num_inference_steps)
+
+        if "mu" in inspect.signature(scheduler.set_timesteps).parameters:
+            if batch.n_tokens is not None:
+                image_seq_len = batch.n_tokens
+            else:
+                h = (batch.height if isinstance(batch.height, int) else (batch.height[0] if batch.height else None))
+                w = (batch.width if isinstance(batch.width, int) else (batch.width[0] if batch.width else None))
+                vae_config = getattr(fastvideo_args.pipeline_config, "vae_config", None)
+                if vae_config is not None:
+                    arch = getattr(vae_config, "arch_config", None)
+                    scale = (getattr(arch, "spatial_compression_ratio", 8) if arch else 8)
+                else:
+                    scale = 8
+                image_seq_len = ((h // scale) * (w // scale) if h is not None and w is not None else 256)
+            extra_set_timesteps_kwargs["mu"] = compute_empirical_mu(image_seq_len, num_inference_steps)
+
+        if timesteps is not None and sigmas is not None:
+            raise ValueError("Only one of `timesteps` or `sigmas` can be passed. "
+                             "Please choose one to set custom values")
+
+        if timesteps is not None:
+            accepts_timesteps = "timesteps" in inspect.signature(scheduler.set_timesteps).parameters
+            if not accepts_timesteps:
+                raise ValueError(f"The current scheduler class {scheduler.__class__}'s "
+                                 f"`set_timesteps` does not support custom timestep schedules.")
+            timesteps_for_scheduler = (timesteps.cpu() if isinstance(timesteps, torch.Tensor) else timesteps)
+            scheduler.set_timesteps(
+                timesteps=timesteps_for_scheduler,
+                device=device,
+                **extra_set_timesteps_kwargs,
+            )
+            timesteps = scheduler.timesteps
+        elif sigmas is not None:
+            accept_sigmas = "sigmas" in inspect.signature(scheduler.set_timesteps).parameters
+            if not accept_sigmas:
+                raise ValueError(f"The current scheduler class {scheduler.__class__}'s "
+                                 f"`set_timesteps` does not support custom sigmas schedules.")
+            scheduler.set_timesteps(
+                sigmas=sigmas,
+                device=device,
+                **extra_set_timesteps_kwargs,
+            )
+            timesteps = scheduler.timesteps
+        else:
+            scheduler.set_timesteps(
+                num_inference_steps,
+                device=device,
+                **extra_set_timesteps_kwargs,
+            )
+            timesteps = scheduler.timesteps
+
+        batch.timesteps = timesteps
+        return batch

fastvideo/pipelines/basic/flux_2/presets.py

@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Flux2 model family pipeline presets.
+
+Each preset is a named inference preset that declares the user-facing
+stage topology, default sampling values, and which per-stage overrides
+are allowed.  Presets are registered explicitly from
+:func:`fastvideo.registry._register_presets`.
+"""
+from fastvideo.api.presets import InferencePreset, PresetStageSpec
+
+_DENOISE_STAGE = PresetStageSpec(
+    name="denoise",
+    kind="denoising",
+    description="Main denoising pass",
+    allowed_overrides=frozenset({
+        "num_inference_steps",
+        "guidance_scale",
+    }),
+)
+
+FLUX2_DEV = InferencePreset(
+    name="flux2_dev",
+    version=1,
+    model_family="flux2",
+    description="Flux2 full T2I with embedded guidance",
+    workload_type="t2i",
+    stage_schemas=(_DENOISE_STAGE, ),
+    defaults={
+        "height": 1024,
+        "width": 1024,
+        "num_frames": 1,
+        "fps": 1,
+        "seed": 0,
+        "guidance_scale": 4.0,
+        "num_inference_steps": 50,
+    },
+)
+
+FLUX2_KLEIN_4B = InferencePreset(
+    name="flux2_klein_4b",
+    version=1,
+    model_family="flux2",
+    description="Flux2 Klein 4B (distilled, 4-step, no guidance)",
+    workload_type="t2i",
+    stage_schemas=(_DENOISE_STAGE, ),
+    defaults={
+        "height": 1024,
+        "width": 1024,
+        "num_frames": 1,
+        "fps": 1,
+        "seed": 0,
+        "guidance_scale": 1.0,
+        "num_inference_steps": 4,
+    },
+)
+
+FLUX2_KLEIN_9B = InferencePreset(
+    name="flux2_klein_9b",
+    version=1,
+    model_family="flux2",
+    description="Flux2 Klein 9B (distilled, 4-step, no guidance)",
+    workload_type="t2i",
+    stage_schemas=(_DENOISE_STAGE, ),
+    defaults={
+        "height": 1024,
+        "width": 1024,
+        "num_frames": 1,
+        "fps": 1,
+        "seed": 0,
+        "guidance_scale": 1.0,
+        "num_inference_steps": 4,
+    },
+)
+
+ALL_PRESETS = (FLUX2_DEV, FLUX2_KLEIN_4B, FLUX2_KLEIN_9B)

fastvideo/pipelines/basic/wan/lucy_edit_pipeline.py

@@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Lucy Edit video editing pipeline.
+
+Lucy Edit uses a Wan2.2 5B transformer with an input video latent appended to
+the noisy latent channels. The stage topology is therefore closest to Wan V2V,
+but the model repo does not include CLIP image-encoder components.
+"""
+
+from fastvideo.fastvideo_args import FastVideoArgs
+from fastvideo.logger import init_logger
+from fastvideo.pipelines.basic.wan.wan_v2v_pipeline import WanVideoToVideoPipeline
+from fastvideo.pipelines.stages import (
+    ConditioningStage,
+    DecodingStage,
+    DenoisingStage,
+    InputValidationStage,
+    LatentPreparationStage,
+    TextEncodingStage,
+    TimestepPreparationStage,
+    VideoVAEEncodingStage,
+)
+
+logger = init_logger(__name__)
+
+
+class LucyEditPipeline(WanVideoToVideoPipeline):
+    """FastVideo pipeline for decart-ai/Lucy-Edit-Dev."""
+
+    _required_config_modules = [
+        "text_encoder",
+        "tokenizer",
+        "vae",
+        "transformer",
+        "scheduler",
+    ]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage())
+
+        self.add_stage(
+            stage_name="prompt_encoding_stage",
+            stage=TextEncodingStage(
+                text_encoders=[self.get_module("text_encoder")],
+                tokenizers=[self.get_module("tokenizer")],
+            ),
+        )
+
+        self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
+
+        self.add_stage(
+            stage_name="timestep_preparation_stage",
+            stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")),
+        )
+
+        self.add_stage(
+            stage_name="latent_preparation_stage",
+            stage=LatentPreparationStage(
+                scheduler=self.get_module("scheduler"),
+                transformer=self.get_module("transformer"),
+            ),
+        )
+
+        self.add_stage(
+            stage_name="video_latent_preparation_stage",
+            stage=VideoVAEEncodingStage(vae=self.get_module("vae")),
+        )
+
+        self.add_stage(
+            stage_name="denoising_stage",
+            stage=DenoisingStage(
+                transformer=self.get_module("transformer"),
+                transformer_2=self.get_module("transformer_2"),
+                scheduler=self.get_module("scheduler"),
+            ),
+        )
+
+        self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae")))
+
+
+EntryClass = LucyEditPipeline

fastvideo/pipelines/basic/wan/presets.py

@@ -268,6 +268,24 @@ FAST_WAN_2_2_TI2V_5B = InferencePreset(
     },
 )
 
+LUCY_EDIT_DEV = InferencePreset(
+    name="lucy_edit_dev",
+    version=1,
+    model_family="wan",
+    description="Lucy Edit Dev 5B video editing",
+    workload_type="t2v",
+    stage_schemas=(_DENOISE_STAGE, ),
+    defaults={
+        "height": 480,
+        "width": 832,
+        "num_frames": 81,
+        "fps": 24,
+        "guidance_scale": 5.0,
+        "num_inference_steps": 50,
+        "negative_prompt": "",
+    },
+)
+
 # -------------------------------------------------------------------
 # Self-Forcing (causal) presets
 # -------------------------------------------------------------------
@@ -341,6 +359,7 @@ ALL_PRESETS = (
     FAST_WAN_T2V_480P,
     WAN_2_2_TI2V_5B,
     FAST_WAN_2_2_TI2V_5B,
+    LUCY_EDIT_DEV,
     SF_WAN_T2V_1_3B,
     SF_WAN_2_2_T2V_A14B,
     SF_WAN_2_2_I2V_A14B,

fastvideo/pipelines/composed_pipeline_base.py

@@ -380,6 +380,8 @@ class ComposedPipelineBase(ABC):
         model_index.pop("boundary_ratio", None)
         # used by Wan2.2 ti2v
         model_index.pop("expand_timesteps", None)
+        # HF metadata (e.g. Flux2 Klein is_distilled); not a loadable module
+        model_index.pop("is_distilled", None)
 
         # some sanity checks
         assert len(model_index) > 1, "model_index.json must contain at least one pipeline module"

fastvideo/pipelines/stages/decoding.py

@@ -47,6 +47,23 @@ class DecodingStage(PipelineStage):
         result.add_check("output", batch.output, [V.is_tensor, V.with_dims(5)])
         return result
 
+    def _is_flux2_packed(self, latents: torch.Tensor) -> bool:
+        """Detect Flux2 packed latents by checking channel count against VAE geometry.
+
+        Flux2 packs latent_channels into 2x2 spatial patches, so the DiT
+        operates on ``latent_channels * 4`` channels at half spatial resolution.
+        The VAE's ``post_quant_conv`` input dimension equals ``latent_channels``.
+        """
+        if not hasattr(self.vae, "bn"):
+            return False
+        pqc = getattr(self.vae, "post_quant_conv", None)
+        if pqc is None:
+            return False
+        vae_latent_ch = pqc.weight.shape[1]
+        packed_ch = vae_latent_ch * 4  # 2x2 patch packing
+        ch_dim = 1 if latents.ndim >= 4 else -1
+        return latents.shape[ch_dim] == packed_ch
+
     def _denormalize_latents(self, latents: torch.Tensor) -> torch.Tensor:
         """Convert normalized latents into the VAE's expected latent space."""
         # Some VAEs handle latent (de)normalization internally.
@@ -77,6 +94,29 @@ class DecodingStage(PipelineStage):
 
         return latents
 
+    @staticmethod
+    def _unpatchify_latents(latents: torch.Tensor) -> torch.Tensor:
+        """Inverse of 2x2 patch packing: ``(B, C*4, H', W') -> (B, C, 2*H', 2*W')``."""
+        batch_size, num_channels, height, width = latents.shape
+        latents = latents.reshape(batch_size, num_channels // (2 * 2), 2, 2, height, width)
+        latents = latents.permute(0, 1, 4, 2, 5, 3)
+        latents = latents.reshape(batch_size, num_channels // (2 * 2), height * 2, width * 2)
+        return latents
+
+    def _flux2_bn_denorm_and_unpatchify(self, latents: torch.Tensor) -> torch.Tensor:
+        """BN denormalize then unpatchify packed latents for VAE decode.
+
+        Handles any channel count (e.g. 64->16, 128->32) via 2x2 spatial unpack.
+        """
+        running_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype)
+        running_var = self.vae.bn.running_var.view(1, -1, 1, 1).to(latents.device, latents.dtype)
+        cfg = getattr(self.vae, "config", None)
+        arch = getattr(cfg, "arch_config", None) if cfg else None
+        eps = getattr(arch, "batch_norm_eps", None) or getattr(cfg, "batch_norm_eps", 1e-5)
+        bn_std = torch.sqrt(torch.clamp(running_var + eps, min=1e-6))
+        latents = latents * bn_std + running_mean
+        return self._unpatchify_latents(latents)
+
     @torch.no_grad()
     def decode(self, latents: torch.Tensor, fastvideo_args: FastVideoArgs) -> torch.Tensor:
         """
@@ -100,7 +140,9 @@ class DecodingStage(PipelineStage):
         vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision]
         vae_autocast_enabled = (vae_dtype != torch.float32) and not fastvideo_args.disable_autocast
 
-        latents = self._denormalize_latents(latents)
+        # Flux2: skip denormalize on packed latents; BN denorm runs below instead
+        if not (latents.ndim == 5 and self._is_flux2_packed(latents)):
+            latents = self._denormalize_latents(latents)
 
         # Decode latents
         with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled):
@@ -110,7 +152,27 @@ class DecodingStage(PipelineStage):
             #     self.vae.enable_parallel()
             if not vae_autocast_enabled:
                 latents = latents.to(vae_dtype)
+            # Flux2's image VAE expects 4D (B, C, H, W); squeeze the singleton T
+            # only for Flux2 packed latents. Gated on `_is_flux2_packed` so video
+            # VAEs that legitimately decode 5D latents with T=1 are untouched.
+            squeezed_for_vae = False
+            if latents.ndim == 5 and latents.shape[2] == 1 and self._is_flux2_packed(latents):
+                latents = latents.squeeze(2)
+                squeezed_for_vae = True
+            # Flux2 packed: BN denorm + unpatchify for VAE decode.
+            # BN denorm is the complete inverse normalisation for Flux2 (no
+            # scaling_factor/shift_factor step), matching Diffusers.
+            if latents.ndim == 4 and self._is_flux2_packed(latents):
+                latents = self._flux2_bn_denorm_and_unpatchify(latents)
             image = self.vae.decode(latents)
+            # Unwrap diffusers-style DecoderOutput / tuple (Flux2 VAE returns a
+            # DecoderOutput). No-op for existing VAEs that return a plain tensor.
+            if hasattr(image, "sample"):
+                image = image.sample
+            elif isinstance(image, tuple | list):
+                image = image[0]
+            if squeezed_for_vae:
+                image = image.unsqueeze(2)
 
         # Normalize image to [0, 1] range
         image = (image / 2 + 0.5).clamp(0, 1)
@@ -201,7 +263,13 @@ class DecodingStage(PipelineStage):
                 pipeline.add_module("vae", self.vae)
             fastvideo_args.model_loaded["vae"] = True
 
-        frames = batch.latents if fastvideo_args.output_type == "latent" else self.decode(batch.latents, fastvideo_args)
+        if fastvideo_args.output_type == "latent":
+            frames = batch.latents
+            if frames.ndim == 5 and frames.shape[2] == 1 and self._is_flux2_packed(frames):
+                frames = self._flux2_bn_denorm_and_unpatchify(frames.squeeze(2))
+                frames = frames.unsqueeze(2)
+        else:
+            frames = self.decode(batch.latents, fastvideo_args)
 
         # decode trajectory latents if needed
         if batch.return_trajectory_decoded:

fastvideo/pipelines/stages/denoising.py

@@ -4,6 +4,7 @@ Denoising stage for diffusion pipelines.
 """
 
 import inspect
+import os
 import weakref
 from collections.abc import Iterable
 from typing import Any
@@ -103,7 +104,37 @@ class DenoisingStage(PipelineStage):
         # TODO(will): make the precision configurable for inference
         # target_dtype = PRECISION_TO_TYPE[fastvideo_args.precision]
         target_dtype = torch.bfloat16
-        autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast
+        # Flux2-only denoising compensations.
+        #
+        # `_is_flux` gates four behaviors that exist because Flux2's transformer
+        # forward() does things internally that the generic pipeline must undo or
+        # match. These are architectural facts about the Flux2 transformer, not
+        # tunable precision policies (the precision policies #5/#6 — prompt-embed
+        # casting and scheduler-step placement — were already moved to config:
+        # DiTArchConfig.cast_prompt_embeds_to_dit_dtype and
+        # PipelineConfig.scheduler_step_in_fp32).
+        #
+        # The four behaviors gated below:
+        #   1. env-var bf16-reduced-precision matmul disable (4-step Klein drift)
+        #   2. autocast disabled (Flux2 long-sequence attention breaks parity under autocast)
+        #   3. guidance: skip the external x1000 (Flux2 multiplies guidance by 1000 internally)
+        #   4. timestep: divide by 1000 with cast-before-divide (Flux2 multiplies timestep by 1000 internally)
+        #
+        # Contract: `prefix == "Flux"` is set ONLY by Flux2 (fastvideo/configs/
+        # models/dits/flux_2.py). No other model uses that prefix, so this exact
+        # match cannot false-positive. A future Flux variant that needs the same
+        # compensations must either set prefix == "Flux" too, OR (preferred) these
+        # gates should graduate to arch-config declarations like the precision
+        # policies above.
+        _is_flux = (getattr(fastvideo_args.pipeline_config.dit_config, "prefix", "") == "Flux")
+        if _is_flux and os.getenv("FASTVIDEO_FLUX2_DISABLE_BF16_REDUCED_PRECISION_REDUCTION",
+                                  "").lower() in {"1", "true", "yes"}:
+            # Gate 1: tighten bf16 matmul accumulation for the 4-step Klein model (opt-in via env var).
+            torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
+        # Gate 2: Flux2 runs its bf16 transformer WITHOUT autocast — autocast perturbs long-sequence attention enough to break 4-step latent parity.
+        autocast_enabled = ((target_dtype != torch.float32) and not fastvideo_args.disable_autocast and not _is_flux)
+        scheduler_fp32 = getattr(fastvideo_args.pipeline_config, "scheduler_step_in_fp32", False)
+        local_device = get_local_torch_device()
 
         # Get timesteps and calculate warmup steps
         timesteps = batch.timesteps
@@ -159,13 +190,40 @@ class DenoisingStage(PipelineStage):
             },
         )
 
+        for key in ("flux2_txt_ids", "flux2_img_ids"):
+            value = batch.extra.get(key)
+            if torch.is_tensor(value):
+                batch.extra[key] = value.to(device=local_device)
+
+        flux2_id_kwargs = self.prepare_extra_func_kwargs(
+            self.transformer.forward,
+            {
+                "txt_ids": batch.extra.get("flux2_txt_ids"),
+                "img_ids": batch.extra.get("flux2_img_ids"),
+            },
+        )
+
         # Get latents and embeddings
         latents = batch.latents
-        prompt_embeds = batch.prompt_embeds
+        cast_embeds = getattr(fastvideo_args.pipeline_config.dit_config, "cast_prompt_embeds_to_dit_dtype", False)
+        if cast_embeds:
+            prompt_embeds = [
+                embed.to(device=local_device, dtype=target_dtype) if torch.is_tensor(embed) else embed
+                for embed in batch.prompt_embeds
+            ]
+        else:
+            prompt_embeds = batch.prompt_embeds
         assert not torch.isnan(prompt_embeds[0]).any(), "prompt_embeds contains nan"
         if batch.do_classifier_free_guidance:
             neg_prompt_embeds = batch.negative_prompt_embeds
             assert neg_prompt_embeds is not None
+            if cast_embeds:
+                neg_prompt_embeds = [
+                    embed.to(device=local_device, dtype=target_dtype) if torch.is_tensor(embed) else embed
+                    for embed in neg_prompt_embeds
+                ]
+            else:
+                neg_prompt_embeds = batch.negative_prompt_embeds
             assert not torch.isnan(neg_prompt_embeds[0]).any(), "neg_prompt_embeds contains nan"
 
         # (Wan2.2) Calculate timestep to switch from high noise expert to low noise expert
@@ -212,23 +270,34 @@ class DenoisingStage(PipelineStage):
         # Initialize lists for ODE trajectory
         trajectory_timesteps: list[torch.Tensor] = []
         trajectory_latents: list[torch.Tensor] = []
+        is_lucy_edit = fastvideo_args.pipeline_config.lucy_edit_task
 
         # Hoisted out of the per-step loop: depends only on inputs that
         # are constant across denoising steps.
         use_meanflow = getattr(self.transformer.config, "use_meanflow", False)
+        # Gate 3: Flux2's transformer multiplies guidance by 1000 internally, so we
+        # skip the external *1000 pre-scaling for Flux models.
         embedded_cfg_scale = fastvideo_args.pipeline_config.embedded_cfg_scale
+        if _is_flux and embedded_cfg_scale is not None:
+            embedded_cfg_scale = batch.guidance_scale
         if embedded_cfg_scale is not None:
             guidance_expand = (torch.tensor(
                 [embedded_cfg_scale] * latents.shape[0],
                 dtype=torch.float32,
                 device=get_local_torch_device(),
-            ).to(target_dtype) * 1000.0)
+            ).to(target_dtype) * (1.0 if _is_flux else 1000.0))
         else:
             guidance_expand = None
         # V2V padding: zero-filled tensor concatenated with each step's
         # latent_model_input.  Shape is fixed by latents and is never
         # written to, so we allocate once.
         v2v_zero_pad = torch.zeros_like(latents) if batch.video_latent is not None else None
+        lucy_timestep_seq_len = None
+        if is_lucy_edit:
+            patch_size = fastvideo_args.pipeline_config.dit_config.arch_config.patch_size
+            assert patch_size[0] == 1, "Lucy Edit timestep expansion assumes temporal patch size 1"
+            lucy_timestep_seq_len = (latents.shape[2] * (latents.shape[3] // patch_size[1]) *
+                                     (latents.shape[4] // patch_size[2]))
 
         # CFG gating / stale-uncond reuse setup (Adaptive Guidance LinearAG
         # variant, Castillo et al. 2023).  When envs.FASTVIDEO_CFG_GATE_STEP
@@ -309,14 +378,25 @@ class DenoisingStage(PipelineStage):
                 # Expand latents for V2V/I2V
                 latent_model_input = latents.to(target_dtype)
                 if batch.video_latent is not None:
-                    latent_model_input = torch.cat([latent_model_input, batch.video_latent, v2v_zero_pad],
-                                                   dim=1).to(target_dtype)
+                    if is_lucy_edit:
+                        latent_model_input = torch.cat(
+                            [latent_model_input, batch.video_latent],
+                            dim=1,
+                        ).to(target_dtype)
+                    else:
+                        latent_model_input = torch.cat(
+                            [latent_model_input, batch.video_latent, v2v_zero_pad],
+                            dim=1,
+                        ).to(target_dtype)
                 elif batch.image_latent is not None:
                     assert not fastvideo_args.pipeline_config.ti2v_task, "image latents should not be provided for TI2V task"
                     latent_model_input = torch.cat([latent_model_input, batch.image_latent], dim=1).to(target_dtype)
 
                 assert not torch.isnan(latent_model_input).any(), "latent_model_input contains nan"
-                if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
+                if is_lucy_edit:
+                    assert lucy_timestep_seq_len is not None
+                    t_expand = t.repeat(latent_model_input.shape[0], lucy_timestep_seq_len)
+                elif fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
                     timestep = torch.stack([t]).to(get_local_torch_device())
                     temp_ts = (mask2[0][0][:, ::2, ::2] * timestep).flatten()
                     temp_ts = torch.cat([temp_ts, temp_ts.new_ones(seq_len - temp_ts.size(0)) * timestep])
@@ -324,7 +404,19 @@ class DenoisingStage(PipelineStage):
                     t_expand = timestep.repeat(latent_model_input.shape[0], 1)
                 else:
                     t_expand = t.repeat(latent_model_input.shape[0])
-                t_expand = t_expand.to(get_local_torch_device())
+                # Gate 4: Flux2 transformer multiplies timestep by 1000 internally, so
+                # the pipeline must pass timestep/1000 (matching Diffusers).
+                # Diffusers casts to the latent dtype before the division; doing
+                # the division in fp32 first changes BF16 rounding for the final
+                # Klein timestep and breaks latent parity.
+                if _is_flux:
+                    t_expand = t_expand.to(
+                        device=get_local_torch_device(),
+                        dtype=latent_model_input.dtype,
+                    )
+                    t_expand = t_expand / 1000.0
+                else:
+                    t_expand = t_expand.to(get_local_torch_device())
 
                 if use_meanflow:
                     if i == len(timesteps) - 1:
@@ -403,6 +495,7 @@ class DenoisingStage(PipelineStage):
                             **action_kwargs,
                             **camera_kwargs,
                             **timesteps_r_kwarg,
+                            **flux2_id_kwargs,
                         )
 
                     if batch.do_classifier_free_guidance:
@@ -444,6 +537,7 @@ class DenoisingStage(PipelineStage):
                                     **action_kwargs,
                                     **camera_kwargs,
                                     **timesteps_r_kwarg,
+                                    **flux2_id_kwargs,
                                 )
                             _cfg_gate_fresh_uncond += 1
 
@@ -467,12 +561,16 @@ class DenoisingStage(PipelineStage):
                                 noise_pred_text,
                                 guidance_rescale=batch.guidance_rescale,
                             )
-                    # Compute the previous noisy sample
+                if scheduler_fp32:
+                    # Diffusers-style: fp32 Euler update outside autocast avoids BF16 drift.
                     latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
-                    if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
-                        latents = latents.squeeze(0)
-                        latents = (1. - mask2[0]) * z + mask2[0] * latents
-                        # latents = latents.unsqueeze(0)
+                else:
+                    with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled):
+                        latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+                if fastvideo_args.pipeline_config.ti2v_task and batch.pil_image is not None:
+                    latents = latents.squeeze(0)
+                    latents = (1. - mask2[0]) * z + mask2[0] * latents
+                    # latents = latents.unsqueeze(0)
 
                 # save trajectory latents if needed
                 if batch.return_trajectory_latents:

fastvideo/pipelines/stages/image_encoding.py

@@ -629,9 +629,10 @@ class VideoVAEEncodingStage(ImageVAEEncodingStage):
             encoder_output = self.vae.encode(video_condition)
 
         generator = batch.generator
-        if generator is None:
-            raise ValueError("Generator must be provided")
-        latent_condition = self.retrieve_latents(encoder_output, generator)
+        sample_mode = "argmax" if fastvideo_args.pipeline_config.lucy_edit_task else "sample"
+        if sample_mode == "sample" and generator is None:
+            raise ValueError("Generator must be provided for sampled video VAE encoding")
+        latent_condition = self.retrieve_latents(encoder_output, generator, sample_mode=sample_mode)
 
         if (hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None):
             if isinstance(self.vae.shift_factor, torch.Tensor):

fastvideo/pipelines/stages/text_encoding.py

@@ -57,6 +57,10 @@ class TextEncodingStage(PipelineStage):
         assert len(self.tokenizers) == len(self.text_encoders)
         assert len(self.text_encoders) == len(fastvideo_args.pipeline_config.text_encoder_configs)
 
+        # Skip encoding if precomputed prompt_embeds were provided
+        if batch.prompt_embeds is not None and len(batch.prompt_embeds) > 0:
+            return batch
+
         # Encode positive prompt with all available encoders
         assert batch.prompt is not None
         prompt_text: str | list[str] = batch.prompt
@@ -218,7 +222,8 @@ class TextEncodingStage(PipelineStage):
                     # Qwen2-style tokenizers. Scoped via treat_empty_as_dot so
                     # models that legitimately use "" (e.g. negative_prompt="")
                     # are not affected.
-                    if not processed_text.strip() and getattr(encoder_config, "treat_empty_as_dot", False):
+                    if isinstance(processed_text, str) and not processed_text.strip() and getattr(
+                            encoder_config, "treat_empty_as_dot", False):
                         processed_text = "."
                     processed_texts.append(processed_text)
                 else:
@@ -235,7 +240,27 @@ class TextEncodingStage(PipelineStage):
             tok = getattr(tokenizer, "tokenizer", tokenizer)
 
             if encoder_config.is_chat_model:
-                text_inputs = tok.apply_chat_template(processed_texts, **tok_kwargs).to(target_device)
+                already_chat_formatted = bool(processed_texts) and isinstance(processed_texts[0], list)
+                if already_chat_formatted:
+                    # Existing chat models (e.g. HunyuanVideo 1.5 / Qwen2.5-VL)
+                    # pre-format prompts into message lists upstream and rely on
+                    # the inner tokenizer + full tokenizer_kwargs (which include
+                    # add_generation_prompt). Preserve that original path exactly.
+                    text_inputs = tok.apply_chat_template(processed_texts, **tok_kwargs).to(target_device)
+                else:
+                    # Two-step approach matching Diffusers: format with chat
+                    # template first, then tokenize the resulting strings.
+                    formatted_texts = []
+                    for pt in processed_texts:
+                        messages = [{"role": "user", "content": pt}]
+                        formatted = tokenizer.apply_chat_template(
+                            messages,
+                            tokenize=False,
+                            add_generation_prompt=True,
+                            enable_thinking=False,
+                        )
+                        formatted_texts.append(formatted)
+                    text_inputs = tokenizer(formatted_texts, **tok_kwargs).to(target_device)
             else:
                 text_inputs = tok(processed_texts, **tok_kwargs).to(target_device)
 

fastvideo/registry.py

@@ -30,6 +30,10 @@ from fastvideo.configs.pipelines.hyworld import HYWorldConfig
 from fastvideo.configs.pipelines.lingbotworld import LingBotWorldI2V480PConfig
 from fastvideo.configs.pipelines.longcat import LongCatT2V480PConfig
 from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig
+from fastvideo.configs.pipelines.flux_2 import (
+    Flux2KleinPipelineConfig,
+    Flux2PipelineConfig,
+)
 from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig
 from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig
 from fastvideo.configs.pipelines.turbodiffusion import (
@@ -40,6 +44,7 @@ from fastvideo.configs.pipelines.turbodiffusion import (
 from fastvideo.configs.pipelines.wan import (
     FastWan2_1_T2V_480P_Config,
     FastWan2_2_TI2V_5B_Config,
+    LucyEditDevConfig,
     SelfForcingWan2_2_T2V480PConfig,
     SelfForcingWanT2V480PConfig,
     WANV2VConfig,
@@ -324,6 +329,45 @@ def _register_configs() -> None:
         default_preset="stable_audio_open_small",
     )
 
+    def _is_flux2_klein(path: str) -> bool:
+        path_lower = path.lower()
+        return "flux.2-klein" in path_lower or "flux2-klein" in path_lower or "flux2klein" in path_lower
+
+    def _is_flux2_full(path: str) -> bool:
+        path_lower = path.lower()
+        is_flux2 = "flux.2" in path_lower or "flux2" in path_lower or "flux_2" in path_lower or "flux-2" in path_lower
+        return is_flux2 and "klein" not in path_lower
+
+    # Flux2 Klein (distilled, 4-step, no guidance)
+    register_configs(
+        sampling_param_cls=None,
+        pipeline_config_cls=Flux2KleinPipelineConfig,
+        workload_types=(WorkloadType.T2I, ),
+        hf_model_paths=[
+            "black-forest-labs/FLUX.2-klein-4B",
+            "black-forest-labs/FLUX.2-klein-9B",
+        ],
+        model_detectors=[
+            _is_flux2_klein,
+        ],
+        model_family="flux2",
+        default_preset="flux2_klein_4b",
+    )
+    # Flux2 (full, Mistral3 text encoder, embedded guidance)
+    register_configs(
+        sampling_param_cls=None,
+        pipeline_config_cls=Flux2PipelineConfig,
+        workload_types=(WorkloadType.T2I, ),
+        hf_model_paths=[
+            "black-forest-labs/FLUX.2-dev",
+        ],
+        model_detectors=[
+            _is_flux2_full,
+        ],
+        model_family="flux2",
+        default_preset="flux2_dev",
+    )
+
     # Hunyuan 1.5 (specific)
     register_configs(
         sampling_param_cls=None,
@@ -738,6 +782,18 @@ def _register_configs() -> None:
         model_family="wan",
         default_preset="fast_wan_2_2_ti2v_5b",
     )
+    register_configs(
+        sampling_param_cls=None,
+        pipeline_config_cls=LucyEditDevConfig,
+        workload_types=(),
+        hf_model_paths=[
+            "decart-ai/Lucy-Edit-Dev",
+            "decart-ai/Lucy-Edit-1.1-Dev",
+        ],
+        model_detectors=[lambda path: "lucy-edit" in path.lower()],
+        model_family="wan",
+        default_preset="lucy_edit_dev",
+    )
     register_configs(
         sampling_param_cls=None,
         pipeline_config_cls=Wan2_2_T2V_A14B_Config,
@@ -840,15 +896,19 @@ def get_model_info(
     if workload_type is None:
         workload_type = WorkloadType.T2V
 
-    if os.path.exists(model_path):
-        config = verify_model_config_and_directory(model_path)
-    else:
-        config = maybe_download_model_index(model_path)
+    config_info = _get_config_info(model_path, raise_on_missing=True)
+    assert config_info is not None, "config_info must be resolved"
 
-    pipeline_name = config.get("_class_name")
     if override_pipeline_cls_name:
-        logger.info("Overriding pipeline class name from %s to %s", pipeline_name, override_pipeline_cls_name)
         pipeline_name = override_pipeline_cls_name
+        logger.info("Using override pipeline class name %s", pipeline_name)
+    else:
+        if os.path.exists(model_path):
+            config = verify_model_config_and_directory(model_path)
+        else:
+            config = maybe_download_model_index(model_path)
+
+        pipeline_name = config.get("_class_name")
 
     if pipeline_name is None:
         raise ValueError("Model config does not contain a _class_name attribute. "
@@ -857,9 +917,6 @@ def get_model_info(
     pipeline_registry = get_pipeline_registry(pipeline_type)
     pipeline_cls = pipeline_registry.resolve_pipeline_cls(pipeline_name, pipeline_type, workload_type)
 
-    config_info = _get_config_info(model_path, raise_on_missing=True)
-    assert config_info is not None, "config_info must be resolved"
-
     sampling_param_cls = config_info.sampling_param_cls or SamplingParam
 
     return ModelInfo(
@@ -920,9 +977,12 @@ def _register_presets() -> None:
         ALL_PRESETS as TURBODIFFUSION_PRESETS, )
     from fastvideo.pipelines.basic.wan.presets import (
         ALL_PRESETS as WAN_PRESETS, )
+    from fastvideo.pipelines.basic.flux_2.presets import (
+        ALL_PRESETS as FLUX2_PRESETS, )
 
     all_preset_groups = (
         COSMOS_PRESETS,
+        FLUX2_PRESETS,
         GAMECRAFT_PRESETS,
         GEN3C_PRESETS,
         HUNYUAN_PRESETS,

fastvideo/tests/api/test_lucy_edit.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from fastvideo.api.presets import get_preset
+from fastvideo.api.sampling_param import SamplingParam
+from fastvideo.configs.pipelines.wan import LucyEditDevConfig
+from fastvideo.fastvideo_args import WorkloadType
+from fastvideo.pipelines.basic.wan.lucy_edit_pipeline import LucyEditPipeline
+from fastvideo.pipelines.pipeline_registry import PipelineType, get_pipeline_registry
+from fastvideo.registry import get_default_preset, get_pipeline_config_cls_from_name
+
+
+def test_lucy_edit_registry_and_preset() -> None:
+    import fastvideo.registry  # noqa: F401
+
+    preset = get_preset("lucy_edit_dev", "wan")
+    assert preset.model_family == "wan"
+    assert preset.defaults["height"] == 480
+    assert preset.defaults["width"] == 832
+    assert preset.defaults["num_frames"] == 81
+
+    config = LucyEditDevConfig()
+    assert config.lucy_edit_task is True
+    assert config.ti2v_task is False
+    assert config.dit_config.arch_config.out_channels == 48
+    assert config.dit_config.arch_config.in_channels == 96
+    assert config.vae_config.arch_config.z_dim == 48
+    assert config.dit_config.arch_config.in_channels == config.vae_config.arch_config.z_dim * 2
+    assert get_default_preset("decart-ai/Lucy-Edit-Dev") == "lucy_edit_dev"
+    assert get_default_preset("decart-ai/Lucy-Edit-1.1-Dev") == "lucy_edit_dev"
+    assert get_pipeline_config_cls_from_name("decart-ai/Lucy-Edit-Dev") is LucyEditDevConfig
+    assert get_pipeline_config_cls_from_name("decart-ai/Lucy-Edit-1.1-Dev") is LucyEditDevConfig
+
+    sampling_param = SamplingParam.from_pretrained("decart-ai/Lucy-Edit-Dev")
+    assert sampling_param.height == 480
+    assert sampling_param.width == 832
+    assert sampling_param.num_frames == 81
+    assert sampling_param.fps == 24
+    assert sampling_param.guidance_scale == 5.0
+    assert sampling_param.negative_prompt == ""
+
+    sampling_param_1_1 = SamplingParam.from_pretrained("decart-ai/Lucy-Edit-1.1-Dev")
+    assert sampling_param_1_1.height == 480
+    assert sampling_param_1_1.width == 832
+    assert sampling_param_1_1.num_frames == 81
+    assert sampling_param_1_1.fps == 24
+    assert sampling_param_1_1.guidance_scale == 5.0
+    assert sampling_param_1_1.negative_prompt == ""
+
+    # FastVideo has no V2V workload enum today; model_index dispatches Lucy by pipeline class name.
+    registry = get_pipeline_registry(PipelineType.BASIC)
+    assert registry.resolve_pipeline_cls("LucyEditPipeline", PipelineType.BASIC, WorkloadType.T2V) is LucyEditPipeline

fastvideo/tests/attention/test_video_sparse_attention_metadata.py

@@ -0,0 +1,38 @@
+import torch
+
+from fastvideo.attention.backends.video_sparse_attn import (
+    VideoSparseAttentionImpl,
+    VideoSparseAttentionMetadataBuilder,
+)
+
+
+def _build_metadata(cache_tile_buf: bool):
+    return VideoSparseAttentionMetadataBuilder().build(
+        current_timestep=0,
+        raw_latent_shape=(4, 4, 4),
+        patch_size=(1, 1, 1),
+        VSA_sparsity=0.5,
+        device=torch.device("cpu"),
+        cache_tile_buf=cache_tile_buf,
+    )
+
+
+def test_vsa_tile_does_not_cache_training_scratch_when_disabled():
+    metadata = _build_metadata(cache_tile_buf=False)
+    impl = object.__new__(VideoSparseAttentionImpl)
+    x = torch.ones(1, 64, 2, 2)
+
+    tiled = impl.tile(x, metadata)
+
+    assert tiled.shape == x.shape
+    assert metadata.tile_buf is None
+
+
+def test_vsa_tile_caches_scratch_by_default():
+    metadata = _build_metadata(cache_tile_buf=True)
+    impl = object.__new__(VideoSparseAttentionImpl)
+    x = torch.ones(1, 64, 2, 2)
+
+    tiled = impl.tile(x, metadata)
+
+    assert metadata.tile_buf is tiled

fastvideo/tests/layers/test_linear_mlp_quant_gating.py

@@ -0,0 +1,49 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Regression coverage for PR #1390's S2-1 plumbing finding: the dormant FP4 shape-tracking path must stay
+gated off unless explicitly enabled, and the MLP quant_config=None default path must keep using ReplicatedLinear's
+unquantized fallback.
+"""
+from __future__ import annotations
+
+import torch
+
+from fastvideo.layers.linear import ReplicatedLinear, UnquantizedLinearMethod
+from fastvideo.layers.mlp import MLP
+
+
+def test_replicated_linear_shape_tracking_default_off() -> None:
+    ReplicatedLinear.reset_shape_tracking()
+    assert ReplicatedLinear.enable_shape_tracking is False
+
+    linear = ReplicatedLinear(input_size=8, output_size=4)
+    linear(torch.randn(2, 8))
+
+    assert len(ReplicatedLinear._shape_to_layer_types) == 0
+
+
+def test_replicated_linear_shape_tracking_enabled_records_unique_shapes() -> None:
+    ReplicatedLinear.reset_shape_tracking()
+    ReplicatedLinear.enable_shape_tracking = True
+    try:
+        linear = ReplicatedLinear(input_size=8, output_size=4)
+        linear(torch.randn(2, 8))
+        linear(torch.randn(3, 8))
+
+        assert len(ReplicatedLinear._shape_to_layer_types) == 2
+        for layer_types in ReplicatedLinear._shape_to_layer_types.values():
+            assert "ReplicatedLinear" in layer_types
+
+        ReplicatedLinear.reset_shape_tracking()
+        assert len(ReplicatedLinear._shape_to_layer_types) == 0
+    finally:
+        ReplicatedLinear.enable_shape_tracking = False
+
+
+def test_mlp_quant_config_none_uses_unquantized_path() -> None:
+    mlp = MLP(input_dim=8, mlp_hidden_dim=16)
+
+    assert isinstance(mlp.fc_in.quant_method, UnquantizedLinearMethod)
+    assert isinstance(mlp.fc_out.quant_method, UnquantizedLinearMethod)
+
+    output = mlp.forward(torch.randn(2, 8))
+    assert output.shape == (2, 8)

fastvideo/tests/modal/launch_l40s_job.py

@@ -0,0 +1,531 @@
+"""Launch an arbitrary FastVideo command on Modal GPUs.
+
+Examples:
+    python -m modal run fastvideo/tests/modal/launch_l40s_job.py --command "nvidia-smi" --install-extra none
+
+    python -m modal run fastvideo/tests/modal/launch_l40s_job.py \
+        --num-gpus 2 \
+        --install-extra test \
+        --command "pytest fastvideo/tests/vaes -vs"
+
+    python -m modal run fastvideo/tests/modal/launch_l40s_job.py \
+        --gpu-type H100 \
+        --num-gpus 1 \
+        --install-extra none \
+        --command "nvidia-smi"
+
+Use ``--no-wait`` with ``modal run --detach`` when the job should keep running
+after the local Modal client exits.
+"""
+
+import os
+import base64
+import shlex
+import shutil
+import subprocess
+import sys
+import time
+from collections.abc import Callable
+from typing import Any
+
+import modal
+
+app = modal.App("fastvideo-gpu-job")
+
+REPO_DIR = "/FastVideo"
+MODEL_VOLUME_NAME = os.environ.get("FASTVIDEO_MODAL_VOLUME", "hf-model-weights")
+IMAGE_VERSION = os.environ.get("IMAGE_VERSION", "latest")
+IMAGE_TAG = os.environ.get(
+    "FASTVIDEO_MODAL_IMAGE",
+    f"ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:{IMAGE_VERSION}",
+)
+SECRET_ENV_KEYS = (
+    "HF_API_KEY",
+    "HUGGINGFACE_HUB_TOKEN",
+    "HF_TOKEN",
+    "WANDB_API_KEY",
+    "WANDB_BASE_URL",
+    "WANDB_MODE",
+)
+
+print(f"Using image: {IMAGE_TAG}")
+print(f"Using Modal volume: {MODEL_VOLUME_NAME}")
+
+model_vol = modal.Volume.from_name(MODEL_VOLUME_NAME, create_if_missing=True)
+local_secrets = modal.Secret.from_dict({
+    key: os.environ[key]
+    for key in SECRET_ENV_KEYS
+    if os.environ.get(key)
+})
+
+image = (
+    modal.Image.from_registry(IMAGE_TAG, add_python="3.12")
+    .apt_install(
+        "cmake",
+        "pkg-config",
+        "build-essential",
+        "curl",
+        "git",
+        "libssl-dev",
+        "ffmpeg",
+    )
+    .run_commands("curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --default-toolchain stable")
+    .run_commands("echo 'source ~/.cargo/env' >> ~/.bashrc")
+    .env({
+        "PATH": "/root/.cargo/bin:$PATH",
+        "HF_HOME": "/root/data/.cache",
+        "TOKENIZERS_PARALLELISM": "false",
+        "FASTVIDEO_ATTENTION_BACKEND": os.environ.get("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN"),
+    })
+)
+
+COMMON_FUNCTION_KWARGS = dict(
+    image=image,
+    timeout=86400,
+    secrets=[local_secrets],
+    volumes={"/root/data": model_vol},
+)
+
+
+def _run_local_git_command(args: list[str]) -> str:
+    result = subprocess.run(
+        ["git", *args],
+        check=True,
+        capture_output=True,
+        text=True,
+    )
+    return result.stdout.strip()
+
+
+def _run_local_git_command_allow_diff(args: list[str]) -> str:
+    result = subprocess.run(
+        ["git", *args],
+        check=False,
+        capture_output=True,
+        text=True,
+    )
+    if result.returncode not in (0, 1):
+        raise RuntimeError(result.stderr.strip() or f"git {' '.join(args)} failed")
+    return result.stdout
+
+
+def _split_patch_paths(patch_paths: str) -> list[str]:
+    return [path.strip() for path in patch_paths.split(",") if path.strip()]
+
+
+def _build_local_patch(patch_paths: str) -> str:
+    paths = _split_patch_paths(patch_paths)
+    diff_args = ["diff", "--binary"]
+    if paths:
+        diff_args.extend(["--", *paths])
+    patch_parts = [_run_local_git_command_allow_diff(diff_args)]
+
+    untracked_args = ["ls-files", "--others", "--exclude-standard"]
+    if paths:
+        untracked_args.extend(["--", *paths])
+    untracked = _run_local_git_command(untracked_args).splitlines()
+    for path in untracked:
+        if not os.path.isfile(path):
+            continue
+        patch_parts.append(_run_local_git_command_allow_diff(["diff", "--binary", "--no-index", "/dev/null", path]))
+
+    patch = "\n".join(part for part in patch_parts if part.strip())
+    if not patch.strip():
+        raise RuntimeError("Requested --apply-local-patch but no local diff was found.")
+    return patch
+
+
+def _apply_local_patch(patch_b64: str) -> None:
+    if not patch_b64:
+        return
+    patch = base64.b64decode(patch_b64.encode("ascii"))
+    print("Applying local workspace patch", flush=True)
+    result = subprocess.run(
+        ["git", "apply", "--binary", "--whitespace=nowarn", "-"],
+        cwd=REPO_DIR,
+        input=patch,
+        stdout=subprocess.PIPE,
+        stderr=sys.stderr,
+        check=False,
+    )
+    if result.stdout:
+        print(result.stdout.decode("utf-8", errors="replace"), end="", flush=True)
+    if result.returncode != 0:
+        raise RuntimeError(f"Failed to apply local patch with exit code {result.returncode}")
+
+
+def _normalize_git_repo_url(git_repo: str) -> str:
+    if git_repo.startswith("git@github.com:"):
+        return "https://github.com/" + git_repo[len("git@github.com:"):]
+    if git_repo.startswith("ssh://git@github.com/"):
+        return "https://github.com/" + git_repo[len("ssh://git@github.com/"):]
+    return git_repo
+
+
+def _resolve_git_repo(git_repo: str) -> str:
+    if git_repo.strip():
+        return _normalize_git_repo_url(git_repo.strip())
+
+    env_repo = os.environ.get("BUILDKITE_REPO", "").strip()
+    if env_repo:
+        return _normalize_git_repo_url(env_repo)
+
+    discovered_repo = _run_local_git_command(["config", "--get", "remote.origin.url"])
+    if discovered_repo:
+        return _normalize_git_repo_url(discovered_repo)
+
+    raise RuntimeError("Could not resolve git repo URL. Pass --git-repo or set BUILDKITE_REPO.")
+
+
+def _resolve_git_commit(git_commit: str) -> str:
+    if git_commit.strip():
+        return git_commit.strip()
+
+    env_commit = os.environ.get("BUILDKITE_COMMIT", "").strip()
+    if env_commit:
+        return env_commit
+
+    discovered_commit = _run_local_git_command(["rev-parse", "HEAD"])
+    if discovered_commit:
+        return discovered_commit
+
+    raise RuntimeError("Could not resolve git commit. Pass --git-commit or set BUILDKITE_COMMIT.")
+
+
+def _resolve_pull_request(pr_number: str) -> str:
+    if pr_number.strip():
+        return pr_number.strip()
+    env_pr = os.environ.get("BUILDKITE_PULL_REQUEST", "").strip()
+    if env_pr:
+        return env_pr
+    return "false"
+
+
+def _run(args: list[str], cwd: str | None = None, env: dict[str, str] | None = None) -> str:
+    print("$ " + " ".join(shlex.quote(arg) for arg in args), flush=True)
+    result = subprocess.run(
+        args,
+        cwd=cwd,
+        env=env,
+        check=True,
+        stdout=subprocess.PIPE,
+        stderr=sys.stderr,
+        text=True,
+    )
+    if result.stdout:
+        print(result.stdout, end="", flush=True)
+    return result.stdout.strip()
+
+
+def _run_shell(command: str, cwd: str, env: dict[str, str]) -> None:
+    print(f"$ {command}", flush=True)
+    result = subprocess.run(
+        ["/bin/bash", "-lc", command],
+        cwd=cwd,
+        env=env,
+        stdout=sys.stdout,
+        stderr=sys.stderr,
+        check=False,
+    )
+    if result.returncode != 0:
+        raise RuntimeError(f"Command failed with exit code {result.returncode}: {command}")
+
+
+def _parse_env_vars(env_vars: str) -> dict[str, str]:
+    parsed: dict[str, str] = {}
+    for item in env_vars.split(","):
+        item = item.strip()
+        if not item:
+            continue
+        if "=" not in item:
+            raise RuntimeError(f"Invalid env var override {item!r}; expected KEY=VALUE.")
+        key, value = item.split("=", 1)
+        parsed[key.strip()] = value.strip()
+    return parsed
+
+
+def _activate_remote_python_env(env: dict[str, str]) -> dict[str, str]:
+    venv_bin = "/opt/venv/bin"
+    if os.path.isdir(venv_bin):
+        env["VIRTUAL_ENV"] = "/opt/venv"
+        env["PATH"] = venv_bin + os.pathsep + env.get("PATH", "")
+    return env
+
+
+def _clone_checkout(git_repo: str, git_commit: str, pr_number: str) -> str:
+    last_clone_error: subprocess.CalledProcessError | None = None
+    for attempt in range(1, 4):
+        shutil.rmtree(REPO_DIR, ignore_errors=True)
+        try:
+            _run(
+                [
+                    "git",
+                    "-c",
+                    "http.version=HTTP/1.1",
+                    "clone",
+                    git_repo,
+                    REPO_DIR,
+                ],
+                cwd="/",
+            )
+            break
+        except subprocess.CalledProcessError as error:
+            last_clone_error = error
+            if attempt == 3:
+                raise
+            sleep_seconds = 5 * attempt
+            print(
+                f"git clone failed on attempt {attempt}; retrying in {sleep_seconds}s",
+                flush=True,
+            )
+            time.sleep(sleep_seconds)
+    if last_clone_error is not None and not os.path.isdir(REPO_DIR):
+        raise last_clone_error
+    if pr_number and pr_number != "false":
+        _run(["git", "fetch", "--prune", "origin", f"refs/pull/{pr_number}/head"], cwd=REPO_DIR)
+        _run(["git", "checkout", "FETCH_HEAD"], cwd=REPO_DIR)
+    else:
+        _run(["git", "checkout", git_commit], cwd=REPO_DIR)
+    _run(["git", "submodule", "update", "--init", "--recursive"], cwd=REPO_DIR)
+    return _run(["git", "rev-parse", "HEAD"], cwd=REPO_DIR)
+
+
+def _install_fastvideo(install_extra: str, env: dict[str, str]) -> None:
+    install_extra = install_extra.strip()
+    if install_extra.lower() in {"", "none", "skip", "false"}:
+        return
+    package = "." if install_extra == "." else f".[{install_extra}]"
+    _run_shell(
+        "source $HOME/.local/bin/env 2>/dev/null || true; "
+        "source /opt/venv/bin/activate 2>/dev/null || true; "
+        f"uv pip install -e {shlex.quote(package)}",
+        cwd=REPO_DIR,
+        env=env,
+    )
+
+
+def _build_kernel(env: dict[str, str]) -> None:
+    _run_shell(
+        "source $HOME/.local/bin/env 2>/dev/null || true; "
+        "source /opt/venv/bin/activate 2>/dev/null || true; "
+        "./build.sh",
+        cwd=os.path.join(REPO_DIR, "fastvideo-kernel"),
+        env=env,
+    )
+
+
+def _run_gpu_job(
+    command: str,
+    git_repo: str,
+    git_commit: str,
+    pr_number: str,
+    install_extra: str,
+    build_kernel: bool,
+    env_vars: str,
+    local_patch_b64: str,
+    commit_volume: bool,
+) -> dict[str, Any]:
+    remote_env = _activate_remote_python_env(os.environ.copy())
+    remote_env.update({
+        "HF_HOME": "/root/data/.cache",
+        "TOKENIZERS_PARALLELISM": "false",
+        "FASTVIDEO_ATTENTION_BACKEND": remote_env.get("FASTVIDEO_ATTENTION_BACKEND", "FLASH_ATTN"),
+    })
+    remote_env.update(_parse_env_vars(env_vars))
+
+    print(f"Cloning repository: {git_repo}")
+    print(f"Target commit: {git_commit}")
+    if pr_number and pr_number != "false":
+        print(f"Using PR ref: {pr_number}")
+    checked_out_commit = _clone_checkout(git_repo, git_commit, pr_number)
+    print(f"Checked out commit: {checked_out_commit}")
+    _apply_local_patch(local_patch_b64)
+
+    _install_fastvideo(install_extra, remote_env)
+    if build_kernel:
+        _build_kernel(remote_env)
+
+    try:
+        _run_shell(command, cwd=REPO_DIR, env=remote_env)
+    finally:
+        if commit_volume:
+            print("Committing Modal volume", flush=True)
+            model_vol.commit()
+    return {
+        "command": command,
+        "git_repo": git_repo,
+        "git_commit": checked_out_commit,
+        "install_extra": install_extra,
+        "build_kernel": build_kernel,
+        "local_patch_applied": bool(local_patch_b64),
+        "commit_volume": commit_volume,
+    }
+
+
+@app.function(gpu="L40S:1", **COMMON_FUNCTION_KWARGS)
+def run_l40s_1(
+    command: str,
+    git_repo: str,
+    git_commit: str,
+    pr_number: str,
+    install_extra: str,
+    build_kernel: bool,
+    env_vars: str,
+    local_patch_b64: str,
+    commit_volume: bool,
+):
+    return _run_gpu_job(command, git_repo, git_commit, pr_number, install_extra, build_kernel, env_vars,
+                        local_patch_b64, commit_volume)
+
+
+@app.function(gpu="L40S:2", **COMMON_FUNCTION_KWARGS)
+def run_l40s_2(
+    command: str,
+    git_repo: str,
+    git_commit: str,
+    pr_number: str,
+    install_extra: str,
+    build_kernel: bool,
+    env_vars: str,
+    local_patch_b64: str,
+    commit_volume: bool,
+):
+    return _run_gpu_job(command, git_repo, git_commit, pr_number, install_extra, build_kernel, env_vars,
+                        local_patch_b64, commit_volume)
+
+
+@app.function(gpu="L40S:4", **COMMON_FUNCTION_KWARGS)
+def run_l40s_4(
+    command: str,
+    git_repo: str,
+    git_commit: str,
+    pr_number: str,
+    install_extra: str,
+    build_kernel: bool,
+    env_vars: str,
+    local_patch_b64: str,
+    commit_volume: bool,
+):
+    return _run_gpu_job(command, git_repo, git_commit, pr_number, install_extra, build_kernel, env_vars,
+                        local_patch_b64, commit_volume)
+
+
+@app.function(gpu="L40S:8", **COMMON_FUNCTION_KWARGS)
+def run_l40s_8(
+    command: str,
+    git_repo: str,
+    git_commit: str,
+    pr_number: str,
+    install_extra: str,
+    build_kernel: bool,
+    env_vars: str,
+    local_patch_b64: str,
+    commit_volume: bool,
+):
+    return _run_gpu_job(command, git_repo, git_commit, pr_number, install_extra, build_kernel, env_vars,
+                        local_patch_b64, commit_volume)
+
+
+@app.function(gpu="H100:1", **COMMON_FUNCTION_KWARGS)
+def run_h100_1(
+    command: str,
+    git_repo: str,
+    git_commit: str,
+    pr_number: str,
+    install_extra: str,
+    build_kernel: bool,
+    env_vars: str,
+    local_patch_b64: str,
+    commit_volume: bool,
+):
+    return _run_gpu_job(command, git_repo, git_commit, pr_number, install_extra, build_kernel, env_vars,
+                        local_patch_b64, commit_volume)
+
+
+@app.function(gpu="H100:2", **COMMON_FUNCTION_KWARGS)
+def run_h100_2(
+    command: str,
+    git_repo: str,
+    git_commit: str,
+    pr_number: str,
+    install_extra: str,
+    build_kernel: bool,
+    env_vars: str,
+    local_patch_b64: str,
+    commit_volume: bool,
+):
+    return _run_gpu_job(command, git_repo, git_commit, pr_number, install_extra, build_kernel, env_vars,
+                        local_patch_b64, commit_volume)
+
+
+def _select_runner(gpu_type: str, num_gpus: int) -> Callable[..., Any]:
+    normalized_gpu_type = gpu_type.upper()
+    runners = {
+        ("L40S", 1): run_l40s_1,
+        ("L40S", 2): run_l40s_2,
+        ("L40S", 4): run_l40s_4,
+        ("L40S", 8): run_l40s_8,
+        ("H100", 1): run_h100_1,
+        ("H100", 2): run_h100_2,
+    }
+    try:
+        return runners[(normalized_gpu_type, num_gpus)]
+    except KeyError as error:
+        supported = ", ".join(f"{gpu}:{count}" for gpu, count in sorted(runners))
+        raise RuntimeError(f"Unsupported GPU request {gpu_type}:{num_gpus}. Supported requests: {supported}.") from error
+
+
+@app.local_entrypoint()
+def main(
+    command: str = "nvidia-smi",
+    gpu_type: str = "L40S",
+    num_gpus: int = 1,
+    git_repo: str = "",
+    git_commit: str = "",
+    pr_number: str = "",
+    install_extra: str = "dev",
+    build_kernel: bool = False,
+    env_vars: str = "",
+    apply_local_patch: bool = False,
+    patch_paths: str = "",
+    wait: bool = True,
+    commit_volume: bool = False,
+):
+    normalized_gpu_type = gpu_type.upper()
+    resolved_git_repo = _resolve_git_repo(git_repo)
+    resolved_git_commit = _resolve_git_commit(git_commit)
+    resolved_pr_number = _resolve_pull_request(pr_number)
+    runner = _select_runner(normalized_gpu_type, num_gpus)
+
+    print(f"Launching {normalized_gpu_type}:{num_gpus} job")
+    print(f"Command: {command}")
+    print(f"Repo: {resolved_git_repo}")
+    print(f"Commit: {resolved_git_commit}")
+    if resolved_pr_number and resolved_pr_number != "false":
+        print(f"PR ref: {resolved_pr_number}")
+    local_patch_b64 = ""
+    if apply_local_patch:
+        patch = _build_local_patch(patch_paths)
+        local_patch_b64 = base64.b64encode(patch.encode("utf-8")).decode("ascii")
+        print(f"Local patch payload: {len(patch)} bytes")
+
+    kwargs = dict(
+        command=command,
+        git_repo=resolved_git_repo,
+        git_commit=resolved_git_commit,
+        pr_number=resolved_pr_number,
+        install_extra=install_extra,
+        build_kernel=build_kernel,
+        env_vars=env_vars,
+        local_patch_b64=local_patch_b64,
+        commit_volume=commit_volume,
+    )
+    if wait:
+        result = runner.remote(**kwargs)
+        print(f"Completed {normalized_gpu_type} job: {result}")
+        return
+
+    function_call = runner.spawn(**kwargs)
+    print(f"Spawned Modal FunctionCall: {function_call.object_id}")
+    print("Poll later with:")
+    print(f"  python -c \"import modal; print(modal.FunctionCall.from_id('{function_call.object_id}').get())\"")

fastvideo/tests/ssim/test_flux2_similarity.py

@@ -0,0 +1,114 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Latent-slice regression tests for Flux2 text-to-image variants.
+
+Flux2 currently has local parity coverage against the official/reference
+pipeline, but CI needs a small deterministic regression gate for seeded HF
+artefacts.  Pixel-space comparisons are unnecessarily brittle for this first
+slot, so the test follows the latent helper pattern used by LTX-2: generate a
+single-image latent with the production recipe, persist the generated latent,
+and compare a stable latent signature plus the full tensor against the device
+reference.
+
+The default and full-quality parameter maps intentionally carry the same
+recipe values for now.  The ``--ssim-full-quality`` flag still switches the
+reference tier through ``conftest.py``; separate full-quality recipes can be
+introduced after the initial Flux2 references have a stable CI window.
+"""
+
+from __future__ import annotations
+
+import os
+
+import pytest
+import torch
+
+from fastvideo.logger import init_logger
+from fastvideo.tests.ssim.inference_similarity_utils import (
+    resolve_inference_device_reference_folder,
+)
+from fastvideo.tests.ssim.latent_similarity_utils import (
+    run_text_to_latent_similarity_test,
+)
+
+logger = init_logger(__name__)
+
+REQUIRED_GPUS = 1
+
+device_reference_folder = resolve_inference_device_reference_folder(logger)
+
+FLUX2_MODEL_TO_PARAMS: dict[str, dict[str, object]] = {
+    "black-forest-labs/FLUX.2-klein-4B": {
+        "num_gpus": 1,
+        "model_path": "black-forest-labs/FLUX.2-klein-4B",
+        "height": 1024,
+        "width": 1024,
+        "num_frames": 1,
+        "num_inference_steps": 4,
+        "guidance_scale": 1.0,
+        "seed": 0,
+        "sp_size": 1,
+        "tp_size": 1,
+        "fps": 1,
+    },
+    "black-forest-labs/FLUX.2-klein-9B": {
+        "num_gpus": 1,
+        "model_path": "black-forest-labs/FLUX.2-klein-9B",
+        "height": 1024,
+        "width": 1024,
+        "num_frames": 1,
+        "num_inference_steps": 4,
+        "guidance_scale": 1.0,
+        "seed": 0,
+        "sp_size": 1,
+        "tp_size": 1,
+        "fps": 1,
+    },
+}
+
+FLUX2_FULL_QUALITY_MODEL_TO_PARAMS: dict[str, dict[str, object]] = {
+    model_id: dict(params)
+    for model_id, params in FLUX2_MODEL_TO_PARAMS.items()
+}
+
+TEST_PROMPTS: dict[str, str] = {
+    "black-forest-labs/FLUX.2-klein-4B": "a brushed steel espresso machine on a marble counter, morning window light",
+    "black-forest-labs/FLUX.2-klein-9B": "a brushed steel espresso machine on a marble counter, morning window light",
+}
+
+SLICE_COSINE_THRESHOLD = 0.96
+FULL_COSINE_THRESHOLD = 0.99
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available(),
+    reason="Flux2 SSIM test requires CUDA",
+)
+@pytest.mark.parametrize("attention_backend_name", ["TORCH_SDPA"])
+@pytest.mark.parametrize("model_id", list(FLUX2_MODEL_TO_PARAMS.keys()))
+def test_flux2_similarity(
+    attention_backend_name: str,
+    model_id: str,
+) -> None:
+    _ = run_text_to_latent_similarity_test(
+        logger=logger,
+        script_dir=os.path.dirname(os.path.abspath(__file__)),
+        device_reference_folder=device_reference_folder,
+        prompt=TEST_PROMPTS[model_id],
+        attention_backend_name=attention_backend_name,
+        model_id=model_id,
+        default_params_map=FLUX2_MODEL_TO_PARAMS,
+        full_quality_params_map=FLUX2_FULL_QUALITY_MODEL_TO_PARAMS,
+        slice_cosine_threshold=SLICE_COSINE_THRESHOLD,
+        full_cosine_threshold=FULL_COSINE_THRESHOLD,
+        init_kwargs_override={
+            "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": (
+                "Flux2KleinPipeline" if "klein" in model_id.lower() else "Flux2Pipeline"
+            ),
+        },
+    )

fastvideo/train/callbacks/__init__.py

@@ -8,8 +8,6 @@ from fastvideo.train.callbacks.ema import (
     EMACallback, )
 from fastvideo.train.callbacks.grad_clip import (
     GradNormClipCallback, )
-from fastvideo.train.callbacks.log_rl_samples import (
-    LogRLSamplesCallback, )
 from fastvideo.train.callbacks.validation import (
     ValidationCallback, )
 
@@ -18,6 +16,5 @@ __all__ = [
     "CallbackDict",
     "EMACallback",
     "GradNormClipCallback",
-    "LogRLSamplesCallback",
     "ValidationCallback",
 ]

fastvideo/train/callbacks/callback.py

@@ -24,7 +24,6 @@ _BUILTIN_CALLBACKS: dict[str, str] = {
     "grad_clip": "fastvideo.train.callbacks.grad_clip.GradNormClipCallback",
     "validation": "fastvideo.train.callbacks.validation.ValidationCallback",
     "ema": "fastvideo.train.callbacks.ema.EMACallback",
-    "log_rl_samples": "fastvideo.train.callbacks.log_rl_samples.LogRLSamplesCallback",
 }
 
 
@@ -64,7 +63,6 @@ class Callback:
         self,
         method: TrainingMethod,
         iteration: int = 0,
-        outputs: dict[str, Any] | None = None,
     ) -> None:
         pass
 

fastvideo/train/callbacks/ema.py

@@ -47,11 +47,9 @@ class EMACallback(Callback):
         *,
         decay: float = 0.9999,
         start_iter: int = 0,
-        update_interval: int = 1,
     ) -> None:
         self._decay = float(decay)
         self._start_iter = int(start_iter)
-        self._update_interval = max(1, int(update_interval))
         self._ema_started = False
         self.student_ema: EMA_FSDP | None = None
 
@@ -80,10 +78,9 @@ class EMACallback(Callback):
         )
         logger.info(
             "EMA callback enabled (decay=%s, "
-            "start_iter=%d, update_interval=%d).",
+            "start_iter=%d).",
             self._decay,
             self._start_iter,
-            self._update_interval,
         )
 
     def on_training_step_end(
@@ -97,8 +94,6 @@ class EMACallback(Callback):
 
         if iteration < self._start_iter:
             return
-        if (iteration - self._start_iter) % self._update_interval != 0:
-            return
         if not self._ema_started:
             logger.info(
                 "Starting EMA updates at iteration %d "

fastvideo/train/callbacks/grad_clip.py

@@ -8,7 +8,7 @@ Optionally logs per-module grad norms to the tracker.
 
 from __future__ import annotations
 
-from typing import Any, TYPE_CHECKING
+from typing import TYPE_CHECKING
 
 from fastvideo.logger import init_logger
 from fastvideo.train.callbacks.callback import Callback
@@ -41,7 +41,6 @@ class GradNormClipCallback(Callback):
         self,
         method: TrainingMethod,
         iteration: int = 0,
-        outputs: dict[str, Any] | None = None,
     ) -> None:
         max_norm = self._max_grad_norm
         if max_norm <= 0.0:

fastvideo/train/callbacks/log_rl_samples.py

@@ -1,116 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Callback to log sampled RL videos to the tracker."""
-
-from __future__ import annotations
-
-import contextlib
-import os
-import tempfile
-from typing import Any, TYPE_CHECKING
-
-import imageio
-import numpy as np
-import torch
-
-from fastvideo.logger import init_logger
-from fastvideo.train.callbacks.callback import Callback
-
-if TYPE_CHECKING:
-    from fastvideo.train.methods.base import TrainingMethod
-
-logger = init_logger(__name__)
-
-
-class LogRLSamplesCallback(Callback):
-    """Log RL-sampled videos to the experiment tracker.
-
-    Expects ``outputs`` to contain:
-
-    - ``sample_videos``: uint8 tensor (B, 3, T, H, W).
-    - ``sample_prompts``: list of prompt strings.
-
-    Configuration (YAML ``callbacks.log_rl_samples``):
-
-    .. code-block:: yaml
-
-        callbacks:
-          log_rl_samples:
-            every_steps: 5
-            max_videos: 4
-            fps: 16
-    """
-
-    def __init__(
-        self,
-        *,
-        every_steps: int = 1,
-        max_videos: int = 4,
-        fps: int = 16,
-    ) -> None:
-        self._every_steps = int(every_steps)
-        self._max_videos = int(max_videos)
-        self._fps = int(fps)
-
-    def on_before_optimizer_step(
-        self,
-        method: TrainingMethod,
-        iteration: int = 0,
-        outputs: dict[str, Any] | None = None,
-    ) -> None:
-        if outputs is None:
-            return
-        if (self._every_steps > 0 and iteration % self._every_steps != 0):
-            return
-
-        videos = outputs.get("sample_videos")
-        prompts = outputs.get("sample_prompts")
-        if videos is None:
-            return
-
-        tracker = getattr(method, "tracker", None)
-        if tracker is None:
-            return
-
-        self._log_videos(tracker, videos, prompts, iteration)
-
-    def _log_videos(
-        self,
-        tracker: Any,
-        videos: torch.Tensor,
-        prompts: list[str] | None,
-        step: int,
-    ) -> None:
-        n = min(len(videos), self._max_videos)
-        tmp_dir = tempfile.mkdtemp(prefix="rl_samples_")
-        video_logs = []
-
-        try:
-            for i in range(n):
-                # (3, T, H, W) uint8 -> (T, H, W, 3) numpy.
-                v = videos[i].permute(1, 2, 3, 0)
-                frames = v.numpy().astype(np.uint8)
-                fname = os.path.join(tmp_dir, f"sample_{step}_{i}.mp4")
-                imageio.mimsave(fname, frames, fps=self._fps)
-
-                caption = (prompts[i] if prompts and i < len(prompts) else None)
-                art = tracker.video(fname, caption=caption, fps=self._fps)
-                if art is not None:
-                    video_logs.append(art)
-
-            if video_logs:
-                tracker.log_artifacts(
-                    {"rl_sample_videos": video_logs},
-                    step,
-                )
-                logger.info(
-                    "Logged %d RL sample videos at step %d",
-                    len(video_logs),
-                    step,
-                )
-        finally:
-            # Clean up temp files.
-            for f in os.listdir(tmp_dir):
-                with contextlib.suppress(OSError):
-                    os.remove(os.path.join(tmp_dir, f))
-            with contextlib.suppress(OSError):
-                os.rmdir(tmp_dir)

fastvideo/train/callbacks/validation.py

@@ -362,7 +362,6 @@ class ValidationCallback(Callback):
 
         batch = ForwardBatch(
             **shallow_asdict(sampling_param),
-            latents=None,
             generator=self.validation_random_generator,
             n_tokens=n_tokens,
             eta=0.0,

fastvideo/train/methods/__init__.py

@@ -9,7 +9,6 @@ __all__ = [
     "KDMethod",
     "SelfForcingMethod",
     "DiffusionForcingSFTMethod",
-    "GenRLMethod",
 ]
 
 
@@ -29,7 +28,4 @@ def __getattr__(name: str) -> object:
     if name == "DiffusionForcingSFTMethod":
         from fastvideo.train.methods.fine_tuning.dfsft import DiffusionForcingSFTMethod
         return DiffusionForcingSFTMethod
-    if name == "GenRLMethod":
-        from fastvideo.train.methods.rl.genrl import GenRLMethod
-        return GenRLMethod
     raise AttributeError(name)

fastvideo/train/methods/rl/__init__.py

@@ -1,2 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-"""Reinforcement learning methods for video generation."""

fastvideo/train/methods/rl/reward/HPSv3/__init__.py

@@ -1,17 +0,0 @@
-"""Vendored runtime subset of HPSv3.
-
-Source: https://github.com/MizzenAI/HPSv3
-Commit: bd0c5fcb5f587617b0169c07222ab78d01e2f3c2
-Purpose: Runtime reward inference integration for FastVideo GenRL.
-
-This is temporary minimal vendoring for PR integration. It is expected to be
-cleaned up and normalized later.
-
-Porting rules:
-- Include only files required by the runtime import closure used by FastVideo.
-- When an upstream file is required, copy the entire file faithfully.
-- Only adjust imports as needed to make the vendored code import through package
-  paths instead of sys.path mutation or ambiguous top-level imports.
-- Do not perform style, typing, or behavioral cleanup as part of this vendoring
-  step.
-"""

fastvideo/train/methods/rl/reward/HPSv3/hpsv3/__init__.py

@@ -1,21 +0,0 @@
-"""Vendored HPSv3 runtime package.
-
-Source: https://github.com/MizzenAI/HPSv3
-Commit: bd0c5fcb5f587617b0169c07222ab78d01e2f3c2
-Purpose: Runtime reward inference integration for FastVideo GenRL.
-
-This is temporary minimal vendoring for PR integration. It is expected to be
-cleaned up and normalized later.
-
-Porting rules:
-- Include only files required by the runtime import closure used by FastVideo.
-- When an upstream file is required, copy the entire file faithfully.
-- Only adjust imports as needed to make the vendored code import through package
-  paths instead of sys.path mutation or ambiguous top-level imports.
-- Do not perform style, typing, or behavioral cleanup as part of this vendoring
-  step.
-"""
-
-from .inference import HPSv3RewardInferencer
-
-__all__ = ["HPSv3RewardInferencer"]

fastvideo/train/methods/rl/reward/HPSv3/hpsv3/config/HPSv3_7B.yaml

@@ -1,60 +0,0 @@
-# Model Configuration
-rm_head_type: "ranknet"
-lora_enable: False
-vision_lora: False
-freeze_vision_tower: False
-freeze_llm: False
-tune_merger: True                     
-model_name_or_path: "Qwen/Qwen2-VL-7B-Instruct"
-num_lora_modules: -1
-lora_r: 512
-lora_alpha: 1024
-lora_namespan_exclude: ['lm_head', 'rm_head', 'embed_tokens']
-
-# Data Configuration
-confidence_threshold: 0.95
-tied_threshold: null
-max_pixels: 200704 # 256 * 28 * 28
-min_pixels: 200704
-with_instruction: true
-
-train_json_list:
-  - example_train.json
-test_json_list:
-  - ["Valid Set 1", ["example_set_1_part1.json", "example_set_1_part2.json"]]
-  - ['Valid Set 2',["example_set_2_part1.json"]]
-
-soft_label: False
-output_dir: output_models
-use_special_tokens: true
-reward_token: "special"
-output_dim: 2
-loss_type: "uncertainty"
-
-# Training Configuration
-disable_flash_attn2: False
-per_device_train_batch_size: 2
-per_device_eval_batch_size: 8
-gradient_accumulation_steps: 4
-num_train_epochs: 10
-learning_rate: 2.0e-6
-special_token_lr: 2.0e-6
-warmup_ratio: 0.05
-lr_scheduler_type: "constant_with_warmup"
-gradient_checkpointing: True
-gradient_checkpointing_kwargs: {"use_reentrant": False}
-
-# Evaluation and Logging
-eval_strategy: "steps"
-logging_epochs: 0.01
-eval_epochs: 0.1
-save_epochs: 0.1
-report_to: tensorboard
-
-# System Configuration
-bf16: True
-torch_dtype: "bfloat16"
-deepspeed: hpsv3/config/ds_config/zero2.json
-save_only_model: True
-save_full_model: True
-dataloader_num_workers: 8

fastvideo/train/methods/rl/reward/HPSv3/hpsv3/dataset/__init__.py

@@ -1 +0,0 @@
-"""Vendored HPSv3 dataset runtime helpers."""

fastvideo/train/methods/rl/reward/HPSv3/hpsv3/dataset/data_collator_qwen.py

@@ -1,185 +0,0 @@
-import numpy as np
-import torch
-from .utils import process_vision_info
-
-INSTRUCTION = """
-You are tasked with evaluating a generated image based on Visual Quality and Text Alignment and give a overall score to estimate the human preference. Please provide a rating from 0 to 10, with 0 being the worst and 10 being the best. 
-
-**Visual Quality:**  
-Evaluate the overall visual quality of the image. The following sub-dimensions should be considered:
-- **Reasonableness:** The image should not contain any significant biological or logical errors, such as abnormal body structures or nonsensical environmental setups.
-- **Clarity:** Evaluate the sharpness and visibility of the image. The image should be clear and easy to interpret, with no blurring or indistinct areas.
-- **Detail Richness:** Consider the level of detail in textures, materials, lighting, and other visual elements (e.g., hair, clothing, shadows).
-- **Aesthetic and Creativity:** Assess the artistic aspects of the image, including the color scheme, composition, atmosphere, depth of field, and the overall creative appeal. The scene should convey a sense of harmony and balance.
-- **Safety:** The image should not contain harmful or inappropriate content, such as political, violent, or adult material. If such content is present, the image quality and satisfaction score should be the lowest possible. 
-
-**Text Alignment:**  
-Assess how well the image matches the textual prompt across the following sub-dimensions:
-- **Subject Relevance** Evaluate how accurately the subject(s) in the image (e.g., person, animal, object) align with the textual description. The subject should match the description in terms of number, appearance, and behavior.
-- **Style Relevance:** If the prompt specifies a particular artistic or stylistic style, evaluate how well the image adheres to this style.
-- **Contextual Consistency**: Assess whether the background, setting, and surrounding elements in the image logically fit the scenario described in the prompt. The environment should support and enhance the subject without contradictions.
-- **Attribute Fidelity**: Check if specific attributes mentioned in the prompt (e.g., colors, clothing, accessories, expressions, actions) are faithfully represented in the image. Minor deviations may be acceptable, but critical attributes should be preserved.
-- **Semantic Coherence**: Evaluate whether the overall meaning and intent of the prompt are captured in the image. The generated content should not introduce elements that conflict with or distort the original description.
-Textual prompt - {text_prompt}
-
-
-"""
-
-INSTRUCTION_debug = """
-{text_prompt}
-"""
-
-prompt_with_special_token = """
-Please provide the overall ratings of this image: <|Reward|>
-
-END
-"""
-
-prompt_without_special_token = """
-Please provide the overall ratings of this image: 
-"""
-
-
-class QWen2VLDataCollator:
-
-    def __init__(
-        self,
-        processor,
-        with_instruction=True,
-        max_pixels=256 * 28 * 28,  # Default max pixels
-        min_pixels=256 * 28 * 28,  # Default min pixels
-        use_special_tokens=True,
-    ):
-        self.processor = processor
-        self.with_instruction = with_instruction
-        self.max_pixels = max_pixels
-        self.min_pixels = min_pixels
-        self.use_special_tokens = use_special_tokens
-
-    def _clean_message(
-        self,
-        texts,
-        images,
-        max_pixels=256 * 28 * 28,
-        min_pixels=256 * 28 * 28,
-        with_instruction=True,
-        use_special_tokens=True,
-    ):
-        """
-        remove unnecessary keys from message(very very necessary)
-        """
-        message_list = []
-        for text, image in zip(texts, images, strict=False):
-            out_message = [{
-                "role":
-                "user",
-                "content": [
-                    {
-                        "type": "image",
-                        "image": image,
-                        "min_pixels": min_pixels,
-                        "max_pixels": max_pixels,
-                    },
-                    {
-                        "type":
-                        "text",
-                        "text": (INSTRUCTION.format(text_prompt=text) +
-                                 prompt_with_special_token if use_special_tokens else prompt_without_special_token),
-                    },
-                ],
-            }]
-
-            message_list.append(out_message)
-
-        return message_list
-
-    def _pad_sequence(self, sequences, attention_mask, max_len, padding_side="right"):
-        """
-        Pad the sequences to the maximum length.
-        """
-        assert padding_side in ["right", "left"]
-        if sequences.shape[1] >= max_len:
-            return sequences, attention_mask
-
-        pad_len = max_len - sequences.shape[1]
-        padding = (0, pad_len) if padding_side == "right" else (pad_len, 0)
-
-        sequences_padded = torch.nn.functional.pad(sequences, padding, "constant",
-                                                   self.processor.tokenizer.pad_token_id)
-        attention_mask_padded = torch.nn.functional.pad(attention_mask, padding, "constant", 0)
-
-        return sequences_padded, attention_mask_padded
-
-    def __call__(self, inputs, with_instruction=True):
-        """
-        Preprocess inputs to token sequences and return a batch
-        """
-        images_1, images_2, texts_1, texts_2 = [], [], [], []
-
-        for idx, batch in enumerate(inputs):
-            texts_1.append(batch["text_1"])
-            texts_2.append(batch["text_2"])
-            images_1.append(batch["image_1"])
-            images_2.append(batch["image_2"])
-
-        messages_batch_1 = self._clean_message(
-            texts_1,
-            images_1,
-            max_pixels=self.max_pixels,
-            min_pixels=self.min_pixels,
-            with_instruction=self.with_instruction,
-            use_special_tokens=self.use_special_tokens,
-        )
-        messages_batch_2 = self._clean_message(
-            texts_2,
-            images_2,
-            max_pixels=self.max_pixels,
-            min_pixels=self.min_pixels,
-            with_instruction=self.with_instruction,
-            use_special_tokens=self.use_special_tokens,
-        )
-        # import pdb; pdb.set_trace()
-        image_inputs_1, _ = process_vision_info(messages_batch_1)
-        image_inputs_2, _ = process_vision_info(messages_batch_2)
-        image_inputs_1 = [np.array(image_inputs_1[i]) / 255.0 for i in range(len(image_inputs_1))]
-        image_inputs_2 = [np.array(image_inputs_2[i]) / 255.0 for i in range(len(image_inputs_2))]
-        do_rescale = False
-
-        batch_1 = self.processor(
-            text=self.processor.apply_chat_template(messages_batch_1, tokenize=False, add_generation_prompt=True),
-            images=image_inputs_1,
-            videos=None,
-            padding=True,
-            return_tensors="pt",
-            images_kwargs={"do_rescale": do_rescale},
-        )
-        batch_2 = self.processor(
-            text=self.processor.apply_chat_template(messages_batch_2, tokenize=False, add_generation_prompt=True),
-            images=image_inputs_2,
-            videos=None,
-            padding=True,
-            return_tensors="pt",
-            images_kwargs={"do_rescale": do_rescale},
-        )
-
-        # pdb.set_trace()
-        max_len = max(batch_1["input_ids"].shape[1], batch_2["input_ids"].shape[1])
-        batch_1["input_ids"], batch_1["attention_mask"] = self._pad_sequence(batch_1["input_ids"],
-                                                                             batch_1["attention_mask"], max_len,
-                                                                             "right")
-        batch_2["input_ids"], batch_2["attention_mask"] = self._pad_sequence(batch_2["input_ids"],
-                                                                             batch_2["attention_mask"], max_len,
-                                                                             "right")
-
-        batch = {
-            "batch_1": batch_1,
-            "batch_2": batch_2,
-            "choice_dist": torch.stack([batch["choice_dist"] for batch in inputs]),
-            # Store original text prompts for visualization
-            "text_1": texts_1,
-            "text_2": texts_2,
-            "image_1": image_inputs_1,
-            "image_2": image_inputs_2,
-        }
-
-        return batch

fastvideo/train/methods/rl/reward/HPSv3/hpsv3/dataset/pairwise_dataset.py

@@ -1,75 +0,0 @@
-import torch
-from torch.utils.data import Dataset
-import random
-import json
-import os
-from tqdm import tqdm
-
-
-class PairwiseOriginalDataset(Dataset):
-
-    def __init__(
-        self,
-        json_list,
-        soft_label=False,
-        confidence_threshold=None,
-    ):
-        self.samples = []
-        for json_file in json_list:
-            with open(json_file) as f:
-                data = json.load(f)
-            self.samples.extend(data)
-
-        self.soft_label = soft_label
-        self.confidence_threshold = confidence_threshold
-
-        if confidence_threshold is not None:
-            new_samples = []
-            for sample in tqdm(self.samples, desc="Filtering samples according to confidence threshold"):
-                if sample.get("confidence", float("inf")) >= confidence_threshold:
-                    new_samples.append(sample)
-            self.samples = new_samples
-
-    def __len__(self):
-        return len(self.samples)
-
-    def __getitem__(self, idx):
-        while True:
-            index = idx
-            try:
-                return self.get_single_item(index)
-            except Exception as e:
-                print(f"Error processing sample at index {idx}: {e}")
-                import traceback
-                traceback.print_exc()
-                index = random.randint(0, len(self.samples) - 1)
-                if index == idx:
-                    continue
-                idx = index
-
-    def get_single_item(self, idx):
-        sample = self.samples[idx]
-        # Load image paths
-        image_1 = sample["path1"]
-        image_2 = sample["path2"]
-        assert os.path.exists(image_1) and os.path.exists(image_2), f'{image_1} or {image_2}'
-        text_1 = sample["prompt"]
-        text_2 = sample["prompt"]
-
-        # Process Label
-        if self.soft_label:
-            choice_dist = sorted(sample["choice_dist"], reverse=True)
-            assert (torch.sum(torch.tensor(choice_dist)) > 0), "Choice distribution cannot be zero."
-            label = torch.tensor(choice_dist[0]) / torch.sum(torch.tensor(choice_dist))
-        else:
-            label = torch.tensor(1).float()
-        # breakpoint()
-        return {
-            "image_1": image_1,
-            "image_2": image_2,
-            "text_1": text_1,
-            "text_2": text_2,
-            "label": label,
-            "confidence": sample.get("confidence", 1.0),
-            "choice_dist": torch.tensor(sample.get("choice_dist", [1.0, 0.0])),
-        }

fastvideo/train/methods/rl/reward/HPSv3/hpsv3/dataset/utils.py

@@ -1,414 +0,0 @@
-from __future__ import annotations
-
-## This file is modified from https://github.com/kq-chen/qwen-vl-utils/blob/main/src/qwen_vl_utils/vision_process.py
-import base64
-import logging
-import math
-import os
-import sys
-import warnings
-from functools import lru_cache
-from io import BytesIO
-
-import requests
-import torch
-import torchvision
-from packaging import version
-from PIL import Image
-from torchvision import io, transforms
-from torchvision.transforms import InterpolationMode
-
-logger = logging.getLogger(__name__)
-
-IMAGE_FACTOR = 28
-MIN_PIXELS = 4 * 28 * 28
-MAX_PIXELS = 16384 * 28 * 28
-MAX_RATIO = 200
-
-VIDEO_MIN_PIXELS = 128 * 28 * 28
-VIDEO_MAX_PIXELS = 768 * 28 * 28
-VIDEO_TOTAL_PIXELS = 24576 * 28 * 28
-FRAME_FACTOR = 2
-FPS = 2.0
-FPS_MIN_FRAMES = 4
-FPS_MAX_FRAMES = 768
-
-
-def round_by_factor(number: int, factor: int) -> int:
-    """Returns the closest integer to 'number' that is divisible by 'factor'."""
-    return round(number / factor) * factor
-
-
-def ceil_by_factor(number: int, factor: int) -> int:
-    """Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'."""
-    return math.ceil(number / factor) * factor
-
-
-def floor_by_factor(number: int, factor: int) -> int:
-    """Returns the largest integer less than or equal to 'number' that is divisible by 'factor'."""
-    return math.floor(number / factor) * factor
-
-
-def smart_resize(height: int,
-                 width: int,
-                 factor: int = IMAGE_FACTOR,
-                 min_pixels: int = MIN_PIXELS,
-                 max_pixels: int = MAX_PIXELS) -> tuple[int, int]:
-    """
-    Rescales the image so that the following conditions are met:
-
-    1. Both dimensions (height and width) are divisible by 'factor'.
-
-    2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
-
-    3. The aspect ratio of the image is maintained as closely as possible.
-    """
-    if max(height, width) / min(height, width) > MAX_RATIO:
-        raise ValueError(
-            f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}")
-    h_bar = max(factor, round_by_factor(height, factor))
-    w_bar = max(factor, round_by_factor(width, factor))
-    if h_bar * w_bar > max_pixels:
-        beta = math.sqrt((height * width) / max_pixels)
-        h_bar = floor_by_factor(height / beta, factor)
-        w_bar = floor_by_factor(width / beta, factor)
-    elif h_bar * w_bar < min_pixels:
-        beta = math.sqrt(min_pixels / (height * width))
-        h_bar = ceil_by_factor(height * beta, factor)
-        w_bar = ceil_by_factor(width * beta, factor)
-    return h_bar, w_bar
-
-
-def fetch_image(ele: dict[str, str | Image.Image], size_factor: int = IMAGE_FACTOR) -> Image.Image:
-    image = ele["image"] if "image" in ele else ele["image_url"]
-    image_obj = None
-    if isinstance(image, Image.Image | torch.Tensor):
-        image_obj = image
-    elif image.startswith("http://") or image.startswith("https://"):
-        image_obj = Image.open(requests.get(image, stream=True).raw)
-    elif image.startswith("file://"):
-        image_obj = Image.open(image[7:])
-    elif image.startswith("data:image"):
-        if "base64," in image:
-            _, base64_data = image.split("base64,", 1)
-            data = base64.b64decode(base64_data)
-            image_obj = Image.open(BytesIO(data))
-    else:
-        image_obj = Image.open(image)
-    if image_obj is None:
-        raise ValueError(f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}")
-    if isinstance(image_obj, Image.Image):
-        image = image_obj.convert("RGB")
-    ## resize
-    if "resized_height" in ele and "resized_width" in ele:
-        resized_height, resized_width = smart_resize(
-            ele["resized_height"],
-            ele["resized_width"],
-            factor=size_factor,
-        )
-    else:
-        if isinstance(image, torch.Tensor):
-            shape = image.shape
-            if len(shape) == 4:
-                if shape[1] in [1, 3]:  # Likely [B, C, H, W]
-                    height, width = shape[2], shape[3]
-                    image_mode = 'NCHW'
-                elif shape[3] in [1, 3]:  # Likely [B, H, W, C]
-                    height, width = shape[1], shape[2]
-                    image_mode = 'NHWC'
-
-            elif len(shape) == 3:
-                if shape[0] in [1, 3]:  # Likely [C, H, W]
-                    height, width = shape[1], shape[2]
-                    image_mode = 'CHW'
-                elif shape[2] in [1, 3]:  # Likely [H, W, C]
-                    height, width = shape[0], shape[1]
-                    image_mode = 'HWC'
-                else:
-                    raise ValueError(f"Cannot determine tensor image format from shape {shape}")
-            else:
-                raise ValueError(f"Unsupported tensor image shape: {shape}")
-        else:
-            width, height = image.size
-        min_pixels = ele.get("min_pixels", MIN_PIXELS)
-        max_pixels = ele.get("max_pixels", MAX_PIXELS)
-        resized_height, resized_width = smart_resize(
-            height,
-            width,
-            factor=size_factor,
-            min_pixels=min_pixels,
-            max_pixels=max_pixels,
-        )
-
-    if isinstance(image, torch.Tensor):
-        if image_mode == 'NCHW':
-            image = transforms.functional.resize(image, [resized_height, resized_width],
-                                                 interpolation=InterpolationMode.BICUBIC,
-                                                 antialias=True)
-        elif image_mode == 'NHWC':
-            image = transforms.functional.resize(image.permute(0, 3, 1, 2), [resized_height, resized_width],
-                                                 interpolation=InterpolationMode.BICUBIC,
-                                                 antialias=True)
-        elif image_mode == 'CHW':
-            image = image.unsqueeze(0)  # Add batch dimension
-            image = transforms.functional.resize(image, [resized_height, resized_width],
-                                                 interpolation=InterpolationMode.BICUBIC,
-                                                 antialias=True)
-        elif image_mode == 'HWC':
-            image = image.permute(2, 0, 1).unsqueeze(0)  # Add batch dimension and change to CHW
-            image = transforms.functional.resize(image, [resized_height, resized_width],
-                                                 interpolation=InterpolationMode.BICUBIC,
-                                                 antialias=True)
-
-    else:
-        # If the image is a PIL Image, we resize it using PIL.
-        if image.mode != "RGB":
-            image = image.convert("RGB")
-        image = image.resize((resized_width, resized_height), Image.BICUBIC)
-
-    return image
-
-
-def smart_nframes(
-    ele: dict,
-    total_frames: int,
-    video_fps: int | float,
-) -> int:
-    """calculate the number of frames for video used for model inputs.
-
-    Args:
-        ele (dict): a dict contains the configuration of video.
-            support either `fps` or `nframes`:
-                - nframes: the number of frames to extract for model inputs.
-                - fps: the fps to extract frames for model inputs.
-                    - min_frames: the minimum number of frames of the video, only used when fps is provided.
-                    - max_frames: the maximum number of frames of the video, only used when fps is provided.
-        total_frames (int): the original total number of frames of the video.
-        video_fps (int | float): the original fps of the video.
-
-    Raises:
-        ValueError: nframes should in interval [FRAME_FACTOR, total_frames].
-
-    Returns:
-        int: the number of frames for video used for model inputs.
-    """
-    assert not ("fps" in ele and "nframes" in ele), "Only accept either `fps` or `nframes`"
-    if "nframes" in ele:
-        nframes = round_by_factor(ele["nframes"], FRAME_FACTOR)
-    else:
-        fps = ele.get("fps", FPS)
-        min_frames = ceil_by_factor(ele.get("min_frames", FPS_MIN_FRAMES), FRAME_FACTOR)
-        max_frames = floor_by_factor(ele.get("max_frames", min(FPS_MAX_FRAMES, total_frames)), FRAME_FACTOR)
-        nframes = total_frames / video_fps * fps
-        nframes = min(max(nframes, min_frames), max_frames)
-        nframes = round_by_factor(nframes, FRAME_FACTOR)
-    if nframes > total_frames:
-        nframes = total_frames
-    if not (nframes >= FRAME_FACTOR and nframes <= total_frames):
-        raise ValueError(f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}.")
-    return nframes
-
-
-def _read_video_torchvision(ele: dict, ) -> torch.Tensor:
-    """read video using torchvision.io.read_video
-
-    Args:
-        ele (dict): a dict contains the configuration of video.
-        support keys:
-            - video: the path of video. support "file://", "http://", "https://" and local path.
-            - video_start: the start time of video.
-            - video_end: the end time of video.
-    Returns:
-        torch.Tensor: the video tensor with shape (T, C, H, W).
-    """
-    video_path = ele["video"]
-    if version.parse(torchvision.__version__) < version.parse("0.19.0"):
-        if "http://" in video_path or "https://" in video_path:
-            warnings.warn(
-                "torchvision < 0.19.0 does not support http/https video path, please upgrade to 0.19.0.",
-                stacklevel=2,
-            )
-        if "file://" in video_path:
-            video_path = video_path[7:]
-    video, audio, info = io.read_video(
-        video_path,
-        start_pts=ele.get("video_start", 0.0),
-        end_pts=ele.get("video_end"),
-        pts_unit="sec",
-        output_format="TCHW",
-    )
-
-    total_frames, video_fps = video.size(0), info["video_fps"]
-    # logger.info(f"torchvision:  {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s")
-    if ele['sample_type'] == 'uniform':
-        nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
-        idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist()
-    elif ele['sample_type'] == 'multi_pts':
-        frames_each_pts = 6
-        num_pts = 4
-        fps = 8
-        nframes = int(total_frames * fps // video_fps)
-        frames_idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist()
-
-        start_pt = int(frames_each_pts // 2)
-        end_pt = int(nframes - frames_each_pts // 2 - 1)
-        pts = torch.linspace(start_pt, end_pt, num_pts).round().long().tolist()
-        idx = []
-        for pt in pts:
-            idx.extend(frames_idx[pt - frames_each_pts // 2:pt + frames_each_pts // 2])
-
-    video = video[idx]
-    return video
-
-
-def is_decord_available() -> bool:
-    import importlib.util
-
-    return importlib.util.find_spec("decord") is not None
-
-
-def _read_video_decord(ele: dict, ) -> torch.Tensor:
-    """read video using decord.VideoReader
-
-    Args:
-        ele (dict): a dict contains the configuration of video.
-        support keys:
-            - video: the path of video. support "file://", "http://", "https://" and local path.
-            - video_start: the start time of video.
-            - video_end: the end time of video.
-    Returns:
-        torch.Tensor: the video tensor with shape (T, C, H, W).
-    """
-    import decord
-    video_path = ele["video"]
-    vr = decord.VideoReader(video_path)
-    # TODO: support start_pts and end_pts
-    if 'video_start' in ele or 'video_end' in ele:
-        raise NotImplementedError("not support start_pts and end_pts in decord for now.")
-    total_frames, video_fps = len(vr), vr.get_avg_fps()
-    # logger.info(f"decord:  {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s")
-    if ele['sample_type'] == 'uniform':
-        nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
-        # nframes = max(nframes, 8)
-        # import pdb; pdb.set_trace()
-        idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist()
-    elif ele['sample_type'] == 'multi_pts':
-        frames_each_pts = 6
-        num_pts = 4
-        fps = 8
-        nframes = int(total_frames * fps // video_fps)
-        frames_idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist()
-
-        start_pt = int(frames_each_pts // 2)
-        end_pt = int(nframes - frames_each_pts // 2 - 1)
-        pts = torch.linspace(start_pt, end_pt, num_pts).round().long().tolist()
-        idx = []
-        for pt in pts:
-            idx.extend(frames_idx[pt - frames_each_pts // 2:pt + frames_each_pts // 2])
-    video = vr.get_batch(idx).asnumpy()
-    video = torch.tensor(video).permute(0, 3, 1, 2)  # Convert to TCHW format
-    return video
-
-
-VIDEO_READER_BACKENDS = {
-    "decord": _read_video_decord,
-    "torchvision": _read_video_torchvision,
-}
-
-FORCE_QWENVL_VIDEO_READER = os.getenv("FORCE_QWENVL_VIDEO_READER", None)
-
-
-@lru_cache(maxsize=1)
-def get_video_reader_backend() -> str:
-    if FORCE_QWENVL_VIDEO_READER is not None:
-        video_reader_backend = FORCE_QWENVL_VIDEO_READER
-    elif is_decord_available():
-        video_reader_backend = "decord"
-    else:
-        video_reader_backend = "torchvision"
-    print(f"qwen-vl-utils using {video_reader_backend} to read video.", file=sys.stderr)
-    return video_reader_backend
-
-
-def fetch_video(ele: dict, image_factor: int = IMAGE_FACTOR) -> torch.Tensor | list[Image.Image]:
-    if isinstance(ele["video"], str):
-        video_reader_backend = get_video_reader_backend()
-        video = VIDEO_READER_BACKENDS[video_reader_backend](ele)
-        # import pdb; pdb.set_trace()
-        nframes, _, height, width = video.shape
-
-        min_pixels = ele.get("min_pixels", VIDEO_MIN_PIXELS)
-        total_pixels = ele.get("total_pixels", VIDEO_TOTAL_PIXELS)
-        max_pixels = max(min(VIDEO_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR), int(min_pixels * 1.05))
-        max_pixels = ele.get("max_pixels", max_pixels)
-        if "resized_height" in ele and "resized_width" in ele:
-            resized_height, resized_width = smart_resize(
-                ele["resized_height"],
-                ele["resized_width"],
-                factor=image_factor,
-            )
-        else:
-            resized_height, resized_width = smart_resize(
-                height,
-                width,
-                factor=image_factor,
-                min_pixels=min_pixels,
-                max_pixels=max_pixels,
-            )
-        video = transforms.functional.resize(
-            video,
-            [resized_height, resized_width],
-            interpolation=InterpolationMode.BICUBIC,
-            antialias=True,
-        ).float()
-        return video
-    else:
-        assert isinstance(ele["video"], list | tuple)
-        process_info = ele.copy()
-        process_info.pop("type", None)
-        process_info.pop("video", None)
-        images = [
-            fetch_image({
-                "image": video_element,
-                **process_info
-            }, size_factor=image_factor) for video_element in ele["video"]
-        ]
-        nframes = ceil_by_factor(len(images), FRAME_FACTOR)
-        if len(images) < nframes:
-            images.extend([images[-1]] * (nframes - len(images)))
-        return images
-
-
-def extract_vision_info(conversations: list[dict] | list[list[dict]]) -> list[dict]:
-    vision_infos = []
-    if isinstance(conversations[0], dict):
-        conversations = [conversations]
-    for conversation in conversations:
-        for message in conversation:
-            if isinstance(message["content"], list):
-                for ele in message["content"]:
-                    if ("image" in ele or "image_url" in ele or "video" in ele
-                            or ele["type"] in ("image", "image_url", "video")):
-                        vision_infos.append(ele)
-    return vision_infos
-
-
-def process_vision_info(
-    conversations: list[dict] | list[list[dict]],
-) -> tuple[list[Image.Image] | None, list[torch.Tensor | list[Image.Image]] | None]:
-    vision_infos = extract_vision_info(conversations)
-    ## Read images or videos
-    image_inputs = []
-    video_inputs = []
-    for vision_info in vision_infos:
-        if "image" in vision_info or "image_url" in vision_info:
-            image_inputs.append(fetch_image(vision_info))
-        elif "video" in vision_info:
-            video_inputs.append(fetch_video(vision_info))
-        else:
-            raise ValueError("image, image_url or video should in content.")
-    if len(image_inputs) == 0:
-        image_inputs = None
-    if len(video_inputs) == 0:
-        video_inputs = None
-    return image_inputs, video_inputs

fastvideo/train/methods/rl/reward/HPSv3/hpsv3/inference.py

@@ -1,155 +0,0 @@
-import os
-from collections.abc import Mapping
-import torch
-import huggingface_hub
-from .dataset.utils import process_vision_info
-from .dataset.data_collator_qwen import prompt_with_special_token, prompt_without_special_token, INSTRUCTION
-from .utils.parser import ModelConfig, PEFTLoraConfig, TrainingConfig, DataConfig, parse_args_with_yaml
-from .train import create_model_and_processor
-from pathlib import Path
-
-_MODEL_CONFIG_PATH = Path(__file__).parent / "config/"
-
-
-class HPSv3RewardInferencer:
-
-    def __init__(self, config_path=None, checkpoint_path=None, device='cuda', differentiable=False):
-        if config_path is None:
-            config_path = os.path.join(_MODEL_CONFIG_PATH, 'HPSv3_7B.yaml')
-
-        if checkpoint_path is None:
-            checkpoint_path = huggingface_hub.hf_hub_download("MizzenAI/HPSv3", 'HPSv3.safetensors', repo_type='model')
-
-        (data_config, training_args, model_config, peft_lora_config), config_path = (parse_args_with_yaml(
-            (DataConfig, TrainingConfig, ModelConfig, PEFTLoraConfig), config_path, is_train=False))
-        training_args.output_dir = os.path.join(training_args.output_dir, config_path.split("/")[-1].split(".")[0])
-        model, processor, peft_config = create_model_and_processor(
-            model_config=model_config,
-            peft_lora_config=peft_lora_config,
-            training_args=training_args,
-            differentiable=differentiable,
-        )
-
-        self.device = device
-        self.use_special_tokens = model_config.use_special_tokens
-
-        if checkpoint_path.endswith('.safetensors'):
-            import safetensors.torch
-            state_dict = safetensors.torch.load_file(checkpoint_path, device="cpu")
-        else:
-            state_dict = torch.load(checkpoint_path, map_location="cpu")
-
-        if "model" in state_dict:
-            state_dict = state_dict["model"]
-        model.load_state_dict(state_dict, strict=True)
-        model.eval()
-
-        self.model = model
-        self.processor = processor
-
-        self.model.to(self.device)
-        self.data_config = data_config
-
-    def _pad_sequence(self, sequences, attention_mask, max_len, padding_side='right'):
-        """
-        Pad the sequences to the maximum length.
-        """
-        assert padding_side in ['right', 'left']
-        if sequences.shape[1] >= max_len:
-            return sequences, attention_mask
-
-        pad_len = max_len - sequences.shape[1]
-        padding = (0, pad_len) if padding_side == 'right' else (pad_len, 0)
-
-        sequences_padded = torch.nn.functional.pad(sequences, padding, 'constant',
-                                                   self.processor.tokenizer.pad_token_id)
-        attention_mask_padded = torch.nn.functional.pad(attention_mask, padding, 'constant', 0)
-
-        return sequences_padded, attention_mask_padded
-
-    def _prepare_input(self, data):
-        """
-        Prepare `inputs` before feeding them to the model, converting them to tensors if they are not already and
-        handling potential state.
-        """
-        if isinstance(data, Mapping):
-            return type(data)({k: self._prepare_input(v) for k, v in data.items()})
-        elif isinstance(data, tuple | list):
-            return type(data)(self._prepare_input(v) for v in data)
-        elif isinstance(data, torch.Tensor):
-            kwargs = {"device": self.device}
-            return data.to(**kwargs)
-        return data
-
-    def _prepare_inputs(self, inputs):
-        """
-        Prepare `inputs` before feeding them to the model, converting them to tensors if they are not already and
-        handling potential state.
-        """
-        inputs = self._prepare_input(inputs)
-        if len(inputs) == 0:
-            raise ValueError
-        return inputs
-
-    def prepare_batch(self, image_paths, prompts):
-        max_pixels = 256 * 28 * 28
-        min_pixels = 256 * 28 * 28
-        message_list = []
-        for text, image in zip(prompts, image_paths, strict=False):
-            out_message = [{
-                "role":
-                "user",
-                "content": [
-                    {
-                        "type": "image",
-                        "image": image,
-                        "min_pixels": min_pixels,
-                        "max_pixels": max_pixels,
-                    },
-                    {
-                        "type":
-                        "text",
-                        "text":
-                        (INSTRUCTION.format(text_prompt=text) +
-                         prompt_with_special_token if self.use_special_tokens else prompt_without_special_token),
-                    },
-                ],
-            }]
-
-            message_list.append(out_message)
-
-        image_inputs, _ = process_vision_info(message_list)
-
-        batch = self.processor(
-            text=self.processor.apply_chat_template(message_list, tokenize=False, add_generation_prompt=True),
-            images=image_inputs,
-            padding=True,
-            return_tensors="pt",
-            videos_kwargs={"do_rescale": True},
-        )
-        batch = self._prepare_inputs(batch)
-        return batch
-
-    @torch.inference_mode()
-    def reward(self, prompts, image_paths):
-        batch = self.prepare_batch(image_paths, prompts)
-        rewards = self.model(return_dict=True, **batch)["logits"]
-
-        return rewards
-
-
-if __name__ == "__main__":
-    config_path = 'config/inference/HPSv3_7B.yaml'
-    checkpoint_path = 'checkpoints/HPSv3_7B.pth'
-    device = 'cuda'
-    dtype = torch.bfloat16
-    inferencer = HPSv3RewardInferencer(config_path, checkpoint_path, device=device)
-
-    image_paths = ["assets/example1.png", "assets/example2.png"]
-    prompts = [
-        "cute chibi anime cartoon fox, smiling wagging tail with a small cartoon heart above sticker",
-        "cute chibi anime cartoon fox, smiling wagging tail with a small cartoon heart above sticker"
-    ]
-    rewards = inferencer.reward(image_paths, prompts)
-    print(rewards[0][0].item())  # miu and sigma. we select miu as the final output
-    print(rewards[1][0].item())

fastvideo/train/methods/rl/reward/HPSv3/hpsv3/model/__init__.py

@@ -1 +0,0 @@
-"""Vendored HPSv3 reward model runtime helpers."""

fastvideo/train/methods/rl/reward/HPSv3/hpsv3/model/differentiable_image_processor.py

@@ -1,615 +0,0 @@
-# Copyright 2024 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
-#
-# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
-# and OPT implementations in this library. It has been modified from its
-# original forms to accommodate minor architectural differences compared
-# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Image processor class for Qwen2-VL.
-
-This module provides both differentiable and non-differentiable image processing methods:
-
-1. For DIFFERENTIABLE processing (torch.autograd compatible):
-   - Pass torch.Tensor to _preprocess() method
-   - Use preprocess_tensor() method directly
-   - All operations use PyTorch functions (F.interpolate, tensor operations, etc.)
-
-2. For NON-DIFFERENTIABLE processing (original functionality):
-   - Pass PIL images or numpy arrays to preprocess() method
-   - Uses PIL/transformers image processing functions and numpy operations
-
-The differentiable path supports:
-- Bilinear interpolation for resizing (instead of PIL resampling)
-- Tensor-based rescaling and normalization
-- Differentiable patch extraction and reshaping
-"""
-
-import math
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-
-from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
-from transformers.image_transforms import (
-    convert_to_rgb,
-    resize,
-    to_channel_dimension_format,
-)
-from transformers.image_utils import (
-    OPENAI_CLIP_MEAN,
-    OPENAI_CLIP_STD,
-    ChannelDimension,
-    ImageInput,
-    PILImageResampling,
-    VideoInput,
-    get_image_size,
-    infer_channel_dimension_format,
-    is_scaled_image,
-    is_valid_image,
-    make_list_of_images,
-    to_numpy_array,
-    valid_images,
-    validate_preprocess_arguments,
-)
-from transformers.utils import TensorType, is_vision_available, logging
-
-logger = logging.get_logger(__name__)
-
-if is_vision_available():
-    from PIL import Image
-
-
-def make_batched_images(images) -> list[list[ImageInput]]:
-    """
-    Accepts images in list or nested list format, and makes a list of images for preprocessing.
-
-    Args:
-        images (`Union[List[List[ImageInput]], List[ImageInput], ImageInput]`):
-            The input image.
-
-    Returns:
-        list: A list of images.
-    """
-    if isinstance(images, list | tuple) and isinstance(images[0], list | tuple) and is_valid_image(images[0][0]):
-        return [img for img_list in images for img in img_list]
-
-    elif isinstance(images, list | tuple) and is_valid_image(images[0]):
-        return images
-
-    elif is_valid_image(images):
-        return [images]
-
-    raise ValueError(f"Could not make batched images from {images}")
-
-
-# Copied from transformers.models.llava_next_video.image_processing_llava_next_video.make_batched_videos
-def make_batched_videos(videos) -> list[VideoInput]:
-    if isinstance(videos, list | tuple) and isinstance(videos[0], list | tuple) and is_valid_image(videos[0][0]):
-        return videos
-
-    elif isinstance(videos, list | tuple) and is_valid_image(videos[0]):
-        if isinstance(videos[0], Image.Image):
-            return [videos]
-        elif len(videos[0].shape) == 4:
-            return [list(video) for video in videos]
-
-    elif is_valid_image(videos) and len(videos.shape) == 4:
-        return [list(videos)]
-
-    raise ValueError(f"Could not make batched video from {videos}")
-
-
-def smart_resize(height: int,
-                 width: int,
-                 factor: int = 28,
-                 min_pixels: int = 56 * 56,
-                 max_pixels: int = 14 * 14 * 4 * 1280):
-    """Rescales the image so that the following conditions are met:
-
-    1. Both dimensions (height and width) are divisible by 'factor'.
-
-    2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
-
-    3. The aspect ratio of the image is maintained as closely as possible.
-
-    """
-    if height < factor or width < factor:
-        raise ValueError(f"height:{height} or width:{width} must be larger than factor:{factor}")
-    elif max(height, width) / min(height, width) > 200:
-        raise ValueError(
-            f"absolute aspect ratio must be smaller than 200, got {max(height, width) / min(height, width)}")
-    h_bar = round(height / factor) * factor
-    w_bar = round(width / factor) * factor
-    if h_bar * w_bar > max_pixels:
-        beta = math.sqrt((height * width) / max_pixels)
-        h_bar = math.floor(height / beta / factor) * factor
-        w_bar = math.floor(width / beta / factor) * factor
-    elif h_bar * w_bar < min_pixels:
-        beta = math.sqrt(min_pixels / (height * width))
-        h_bar = math.ceil(height * beta / factor) * factor
-        w_bar = math.ceil(width * beta / factor) * factor
-    return h_bar, w_bar
-
-
-class Qwen2VLImageProcessor(BaseImageProcessor):
-    r"""
-    Constructs a Qwen2-VL image processor that dynamically resizes images based on the original images.
-
-    Args:
-        do_resize (`bool`, *optional*, defaults to `True`):
-            Whether to resize the image's (height, width) dimensions.
-        resample (`PILImageResampling`, *optional*, defaults to `Resampling.BICUBIC`):
-            Resampling filter to use when resizing the image.
-        do_rescale (`bool`, *optional*, defaults to `True`):
-            Whether to rescale the image by the specified scale `rescale_factor`.
-        rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):
-            Scale factor to use if rescaling the image.
-        do_normalize (`bool`, *optional*, defaults to `True`):
-            Whether to normalize the image.
-        image_mean (`float` or `List[float]`, *optional*, defaults to `[0.48145466, 0.4578275, 0.40821073]`):
-            Mean to use if normalizing the image. This is a float or list of floats for each channel in the image.
-        image_std (`float` or `List[float]`, *optional*, defaults to `[0.26862954, 0.26130258, 0.27577711]`):
-            Standard deviation to use if normalizing the image. This is a float or list of floats for each channel in the image.
-        do_convert_rgb (`bool`, *optional*, defaults to `True`):
-            Whether to convert the image to RGB.
-        min_pixels (`int`, *optional*, defaults to `56 * 56`):
-            The min pixels of the image to resize the image.
-        max_pixels (`int`, *optional*, defaults to `28 * 28 * 1280`):
-            The max pixels of the image to resize the image.
-        patch_size (`int`, *optional*, defaults to 14):
-            The spatial patch size of the vision encoder.
-        temporal_patch_size (`int`, *optional*, defaults to 2):
-            The temporal patch size of the vision encoder.
-        merge_size (`int`, *optional*, defaults to 2):
-            The merge size of the vision encoder to llm encoder.
-    """
-
-    model_input_names = ["pixel_values", "image_grid_thw", "pixel_values_videos", "video_grid_thw"]
-
-    def __init__(
-        self,
-        do_resize: bool = True,
-        resample: PILImageResampling = PILImageResampling.BICUBIC,
-        do_rescale: bool = True,
-        rescale_factor: int | float = 1 / 255,
-        do_normalize: bool = True,
-        image_mean: float | list[float] | None = None,
-        image_std: float | list[float] | None = None,
-        do_convert_rgb: bool = True,
-        min_pixels: int = 56 * 56,
-        max_pixels: int = 28 * 28 * 1280,
-        patch_size: int = 14,
-        temporal_patch_size: int = 2,
-        merge_size: int = 2,
-        **kwargs,
-    ) -> None:
-        super().__init__(**kwargs)
-        self.do_resize = do_resize
-        self.resample = resample
-        self.do_rescale = do_rescale
-        self.rescale_factor = rescale_factor
-        self.do_normalize = do_normalize
-        self.image_mean = image_mean if image_mean is not None else OPENAI_CLIP_MEAN
-        self.image_std = image_std if image_std is not None else OPENAI_CLIP_STD
-        self.min_pixels = min_pixels
-        self.max_pixels = max_pixels
-        self.patch_size = patch_size
-        self.temporal_patch_size = temporal_patch_size
-        self.merge_size = merge_size
-        self.size = {"min_pixels": min_pixels, "max_pixels": max_pixels}
-        self.do_convert_rgb = do_convert_rgb
-
-    def _preprocess_differentiable(
-        self,
-        images: torch.Tensor,
-        do_resize: bool = None,
-        do_rescale: bool = None,
-        rescale_factor: float = None,
-        do_normalize: bool = None,
-        image_mean: float | list[float] | None = None,
-        image_std: float | list[float] | None = None,
-    ):
-        """
-        Differentiable version of image preprocessing using torch operations.
-        
-        Args:
-            images: torch.Tensor of shape (B, C, H, W) or (C, H, W)
-        Returns:
-            flatten_patches: torch.Tensor - flattened patches
-            grid_thw: tuple - (grid_t, grid_h, grid_w)
-        """
-        if images.dim() == 3:
-            images = images.unsqueeze(0)  # Add batch dimension
-
-        batch_size, channels, height, width = images.shape
-
-        processed_images = []
-        resized_height, resized_width = height, width
-
-        for i in range(batch_size):
-            image = images[i]  # (C, H, W)
-
-            if do_resize:
-                resized_height, resized_width = smart_resize(
-                    height,
-                    width,
-                    factor=self.patch_size * self.merge_size,
-                    min_pixels=self.min_pixels,
-                    max_pixels=self.max_pixels,
-                )
-                # Use differentiable interpolation
-                image = F.interpolate(image.unsqueeze(0),
-                                      size=(resized_height, resized_width),
-                                      mode='bilinear',
-                                      align_corners=False).squeeze(0)
-
-            if do_rescale:
-                image = image * rescale_factor
-
-            if do_normalize:
-                if isinstance(image_mean, list | tuple):
-                    mean = torch.tensor(image_mean, device=image.device, dtype=image.dtype).view(-1, 1, 1)
-                    std = torch.tensor(image_std, device=image.device, dtype=image.dtype).view(-1, 1, 1)
-                else:
-                    mean = image_mean
-                    std = image_std
-                image = (image - mean) / std
-
-            processed_images.append(image)
-
-        # Stack all processed images
-        patches = torch.stack(processed_images)  # (B, C, H, W)
-
-        # Handle temporal dimension
-        if patches.shape[0] == 1:
-            patches = patches.repeat(self.temporal_patch_size, 1, 1, 1)
-
-        # Reshape for patch extraction
-        batch_size, channel, resized_height, resized_width = patches.shape
-        grid_t = batch_size // self.temporal_patch_size
-        grid_h, grid_w = resized_height // self.patch_size, resized_width // self.patch_size
-
-        # Differentiable patch extraction and reshaping
-        patches = patches.view(
-            grid_t,
-            self.temporal_patch_size,
-            channel,
-            grid_h // self.merge_size,
-            self.merge_size,
-            self.patch_size,
-            grid_w // self.merge_size,
-            self.merge_size,
-            self.patch_size,
-        )
-        patches = patches.permute(0, 3, 6, 4, 7, 2, 1, 5, 8)
-        flatten_patches = patches.reshape(grid_t * grid_h * grid_w,
-                                          channel * self.temporal_patch_size * self.patch_size * self.patch_size)
-
-        return flatten_patches, (grid_t, grid_h, grid_w)
-
-    def _preprocess(
-        self,
-        images: ImageInput | VideoInput,
-        do_resize: bool = None,
-        resample: PILImageResampling = None,
-        do_rescale: bool = None,
-        rescale_factor: float = None,
-        do_normalize: bool = None,
-        image_mean: float | list[float] | None = None,
-        image_std: float | list[float] | None = None,
-        do_convert_rgb: bool = None,
-        data_format: ChannelDimension | None = ChannelDimension.FIRST,
-        input_data_format: str | ChannelDimension | None = None,
-    ):
-        """
-        Preprocess an image or batch of images. Copy of the `preprocess` method from `CLIPImageProcessor`.
-
-        Args:
-            images (`ImageInput`):
-                Image or batch of images to preprocess. Expects pixel values ranging from 0 to 255. If pixel values range from 0 to 1, set `do_rescale=False`.
-            vision_info (`List[Dict]`, *optional*):
-                Optional list of dictionaries containing additional information about vision inputs.
-            do_resize (`bool`, *optional*, defaults to `self.do_resize`):
-                Whether to resize the image.
-            resample (`PILImageResampling`, *optional*, defaults to `self.resample`):
-                Resampling filter to use if resizing the image. This can be one of the `PILImageResampling` enums.
-            do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
-                Whether to rescale the image.
-            rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
-                Scale factor to use if rescaling the image.
-            do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
-                Whether to normalize the image.
-            image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`):
-                Mean to use if normalizing the image. Can be a float or a list of floats corresponding to the number of channels in the image.
-            image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):
-                Standard deviation to use if normalizing the image. Can be a float or a list of floats corresponding to the number of channels in the image.
-            do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`):
-                Whether to convert the image to RGB.
-            data_format (`ChannelDimension`, *optional*, defaults to `ChannelDimension.FIRST`):
-                The channel dimension format for the output image. Can be one of:
-                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
-                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
-                - Unset: Use the channel dimension format of the input image.
-            input_data_format (`ChannelDimension` or `str`, *optional*):
-                The channel dimension format for the input image. Can be one of:
-                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
-                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
-                - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.   - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
-        """
-        # Check if input is already a torch tensor (differentiable path)
-        if isinstance(images, torch.Tensor):
-            return self._preprocess_differentiable(
-                images,
-                do_resize=do_resize,
-                do_rescale=do_rescale,
-                rescale_factor=rescale_factor,
-                do_normalize=do_normalize,
-                image_mean=image_mean,
-                image_std=image_std,
-            )
-
-        # Original non-differentiable path for backward compatibility
-        images = make_list_of_images(images)
-
-        if do_convert_rgb:
-            images = [convert_to_rgb(image) for image in images]
-
-        # All transformations expect numpy arrays.
-        images = [to_numpy_array(image) for image in images]
-
-        if is_scaled_image(images[0]) and do_rescale:
-            logger.warning_once(
-                "It looks like you are trying to rescale already rescaled images. If the input"
-                " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again.")
-        if input_data_format is None:
-            # We assume that all images have the same channel dimension format.
-            input_data_format = infer_channel_dimension_format(images[0])
-
-        height, width = get_image_size(images[0], channel_dim=input_data_format)
-        resized_height, resized_width = height, width
-        processed_images = []
-        for image in images:
-            if do_resize:
-                resized_height, resized_width = smart_resize(
-                    height,
-                    width,
-                    factor=self.patch_size * self.merge_size,
-                    min_pixels=self.min_pixels,
-                    max_pixels=self.max_pixels,
-                )
-                image = resize(image,
-                               size=(resized_height, resized_width),
-                               resample=resample,
-                               input_data_format=input_data_format)
-
-            if do_rescale:
-                image = self.rescale(image, scale=rescale_factor, input_data_format=input_data_format)
-
-            if do_normalize:
-                image = self.normalize(image=image, mean=image_mean, std=image_std, input_data_format=input_data_format)
-
-            image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
-            processed_images.append(image)
-
-        # NOTE: The following operations use numpy and are NOT differentiable
-        # For differentiable operations, pass torch.Tensor as input to use _preprocess_differentiable
-        patches = np.array(processed_images)
-        if data_format == ChannelDimension.LAST:
-            patches = patches.transpose(0, 3, 1, 2)
-        if patches.shape[0] == 1:
-            patches = np.tile(patches, (self.temporal_patch_size, 1, 1, 1))
-        channel = patches.shape[1]
-        grid_t = patches.shape[0] // self.temporal_patch_size
-        grid_h, grid_w = resized_height // self.patch_size, resized_width // self.patch_size
-        patches = patches.reshape(
-            grid_t,
-            self.temporal_patch_size,
-            channel,
-            grid_h // self.merge_size,
-            self.merge_size,
-            self.patch_size,
-            grid_w // self.merge_size,
-            self.merge_size,
-            self.patch_size,
-        )
-        patches = patches.transpose(0, 3, 6, 4, 7, 2, 1, 5, 8)
-        flatten_patches = patches.reshape(grid_t * grid_h * grid_w,
-                                          channel * self.temporal_patch_size * self.patch_size * self.patch_size)
-
-        return flatten_patches, (grid_t, grid_h, grid_w)
-
-    def preprocess_tensor(
-        self,
-        images: torch.Tensor,
-        do_resize: bool = None,
-        do_rescale: bool = None,
-        rescale_factor: float = None,
-        do_normalize: bool = None,
-        image_mean: float | list[float] | None = None,
-        image_std: float | list[float] | None = None,
-    ):
-        """
-        Differentiable preprocessing method for torch tensors.
-        
-        Args:
-            images: torch.Tensor of shape (B, C, H, W) or (C, H, W)
-            
-        Returns:
-            dict containing:
-                - pixel_values: torch.Tensor - processed patches
-                - image_grid_thw: torch.Tensor - grid dimensions
-        """
-        do_resize = do_resize if do_resize is not None else self.do_resize
-        do_rescale = do_rescale if do_rescale is not None else self.do_rescale
-        rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
-        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
-        image_mean = image_mean if image_mean is not None else self.image_mean
-        image_std = image_std if image_std is not None else self.image_std
-
-        patches, image_grid_thw = self._preprocess_differentiable(
-            images,
-            do_resize=do_resize,
-            do_rescale=do_rescale,
-            rescale_factor=rescale_factor,
-            do_normalize=do_normalize,
-            image_mean=image_mean,
-            image_std=image_std,
-        )
-
-        return {"pixel_values": patches, "image_grid_thw": torch.tensor(image_grid_thw, device=patches.device)}
-
-    def preprocess(
-        self,
-        images: ImageInput,
-        videos: VideoInput = None,
-        do_resize: bool = None,
-        size: dict[str, int] = None,
-        resample: PILImageResampling = None,
-        do_rescale: bool = None,
-        rescale_factor: float = None,
-        do_normalize: bool = None,
-        image_mean: float | list[float] | None = None,
-        image_std: float | list[float] | None = None,
-        do_convert_rgb: bool = None,
-        return_tensors: str | TensorType | None = None,
-        data_format: ChannelDimension | None = ChannelDimension.FIRST,
-        input_data_format: str | ChannelDimension | None = None,
-    ):
-        """
-        Args:
-            images (`ImageInput`):
-                Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If
-                passing in images with pixel values between 0 and 1, set `do_rescale=False`.
-            videos (`VideoInput`):
-                Video to preprocess. Expects a single or batch of videos with pixel values ranging from 0 to 255. If
-                passing in videos with pixel values between 0 and 1, set `do_rescale=False`.
-            do_resize (`bool`, *optional*, defaults to `self.do_resize`):
-                Whether to resize the image.
-            size (`Dict[str, int]`, *optional*, defaults to `self.size`):
-                Size of the image after resizing. Shortest edge of the image is resized to size["shortest_edge"], with
-                the longest edge resized to keep the input aspect ratio.
-            resample (`int`, *optional*, defaults to `self.resample`):
-                Resampling filter to use if resizing the image. This can be one of the enum `PILImageResampling`. Only
-                has an effect if `do_resize` is set to `True`.
-            do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
-                Whether to rescale the image.
-            rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
-                Rescale factor to rescale the image by if `do_rescale` is set to `True`.
-            do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
-                Whether to normalize the image.
-            image_mean (`float` or `List[float]`, *optional*, defaults to `self.image_mean`):
-                Image mean to use for normalization. Only has an effect if `do_normalize` is set to `True`.
-            image_std (`float` or `List[float]`, *optional*, defaults to `self.image_std`):
-                Image standard deviation to use for normalization. Only has an effect if `do_normalize` is set to
-                `True`.
-            do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`):
-                Whether to convert the image to RGB.
-            return_tensors (`str` or `TensorType`, *optional*):
-                The type of tensors to return. Can be one of:
-                - Unset: Return a list of `np.ndarray`.
-                - `TensorType.TENSORFLOW` or `'tf'`: Return a batch of type `tf.Tensor`.
-                - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
-                - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
-                - `TensorType.JAX` or `'jax'`: Return a batch of type `jax.numpy.ndarray`.
-            data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):
-                The channel dimension format for the output image. Can be one of:
-                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
-                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
-                - Unset: Use the channel dimension format of the input image.
-            input_data_format (`ChannelDimension` or `str`, *optional*):
-                The channel dimension format for the input image. If unset, the channel dimension format is inferred
-                from the input image. Can be one of:
-                - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
-                - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
-                - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
-
-        """
-        do_resize = do_resize if do_resize is not None else self.do_resize
-        size = size if size is not None else self.size
-        resample = resample if resample is not None else self.resample
-        do_rescale = do_rescale if do_rescale is not None else self.do_rescale
-        rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
-        do_normalize = do_normalize if do_normalize is not None else self.do_normalize
-        image_mean = image_mean if image_mean is not None else self.image_mean
-        image_std = image_std if image_std is not None else self.image_std
-        do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
-
-        if images is not None:
-            images = make_batched_images(images)
-        if videos is not None:
-            videos = make_batched_videos(videos)
-
-        if images is not None and not valid_images(images):
-            raise ValueError("Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
-                             "torch.Tensor, tf.Tensor or jax.ndarray.")
-
-        validate_preprocess_arguments(
-            rescale_factor=rescale_factor,
-            do_normalize=do_normalize,
-            image_mean=image_mean,
-            image_std=image_std,
-            do_resize=do_resize,
-            size=size,
-            resample=resample,
-        )
-
-        if images is not None:
-            pixel_values, vision_grid_thws = [], []
-            for image in images:
-                patches, image_grid_thw = self._preprocess(
-                    image,
-                    do_resize=do_resize,
-                    resample=resample,
-                    do_rescale=do_rescale,
-                    rescale_factor=rescale_factor,
-                    do_normalize=do_normalize,
-                    image_mean=image_mean,
-                    image_std=image_std,
-                    data_format=data_format,
-                    do_convert_rgb=do_convert_rgb,
-                    input_data_format=input_data_format,
-                )
-                pixel_values.extend(patches)
-                vision_grid_thws.append(image_grid_thw)
-            if not isinstance(pixel_values[0], torch.Tensor):
-                pixel_values = np.array(pixel_values)
-            else:
-                pixel_values = torch.stack(pixel_values)
-            vision_grid_thws = np.array(vision_grid_thws)
-            data = {"pixel_values": pixel_values, "image_grid_thw": vision_grid_thws}
-
-        if videos is not None:
-            pixel_values, vision_grid_thws = [], []
-            for images in videos:
-                patches, video_grid_thw = self._preprocess(
-                    images,
-                    do_resize=do_resize,
-                    resample=resample,
-                    do_rescale=do_rescale,
-                    rescale_factor=rescale_factor,
-                    do_normalize=do_normalize,
-                    image_mean=image_mean,
-                    image_std=image_std,
-                    data_format=data_format,
-                    do_convert_rgb=do_convert_rgb,
-                    input_data_format=input_data_format,
-                )
-                pixel_values.extend(patches)
-                vision_grid_thws.append(video_grid_thw)
-            pixel_values = np.array(pixel_values)
-            vision_grid_thws = np.array(vision_grid_thws)
-            data = {"pixel_values_videos": pixel_values, "video_grid_thw": vision_grid_thws}
-
-        return BatchFeature(data=data, tensor_type=return_tensors)

fastvideo/train/methods/rl/reward/HPSv3/hpsv3/model/qwen2vl_trainer.py

@@ -1,817 +0,0 @@
-import os
-import math
-import matplotlib.pyplot as plt
-
-import safetensors
-import numpy as np
-import torch
-import torch.nn as nn
-import datasets
-from torch.utils.data import Dataset, DataLoader
-from peft import PeftModel
-from transformers import Qwen2VLForConditionalGeneration
-from transformers.modeling_utils import PreTrainedModel
-from transformers.trainer import TrainerCallback
-from transformers.trainer import (
-    is_sagemaker_mp_enabled,
-    is_peft_available,
-    is_datasets_available,
-    WEIGHTS_NAME,
-    TRAINING_ARGS_NAME,
-    SAFE_WEIGHTS_NAME,
-    PREFIX_CHECKPOINT_DIR,
-    logger,
-)
-
-from transformers.trainer_pt_utils import nested_detach
-from trl import RewardTrainer
-from ..utils.training_utils import get_peft_state_non_lora_maybe_zero_3
-
-
-class Qwen2VLRewardModelBT(Qwen2VLForConditionalGeneration):
-
-    def __init__(
-        self,
-        config,
-        output_dim=4,
-        reward_token="last",
-        special_token_ids=None,
-        rm_head_type="default",
-        rm_head_kwargs=None,
-    ):
-        super().__init__(config)
-        # pdb.set_trace()
-        self.output_dim = output_dim
-        if rm_head_type == "default":
-            self.rm_head = nn.Linear(config.hidden_size, output_dim, bias=False)
-        elif rm_head_type == "ranknet":
-            if rm_head_kwargs is not None:
-                for layer in range(rm_head_kwargs.get("num_layers", 3)):
-                    if layer == 0:
-                        self.rm_head = nn.Sequential(
-                            nn.Linear(config.hidden_size, rm_head_kwargs["hidden_size"]),
-                            nn.ReLU(),
-                            nn.Dropout(rm_head_kwargs.get("dropout", 0.1)),
-                        )
-                    elif layer < rm_head_kwargs.get("num_layers", 3) - 1:
-                        self.rm_head.add_module(
-                            f"layer_{layer}",
-                            nn.Sequential(
-                                nn.Linear(rm_head_kwargs["hidden_size"], rm_head_kwargs["hidden_size"]),
-                                nn.ReLU(),
-                                nn.Dropout(rm_head_kwargs.get("dropout", 0.1)),
-                            ),
-                        )
-                    else:
-                        self.rm_head.add_module(
-                            "output_layer",
-                            nn.Linear(rm_head_kwargs["hidden_size"], output_dim, bias=rm_head_kwargs.get("bias",
-                                                                                                         False)),
-                        )
-
-            else:
-                self.rm_head = nn.Sequential(
-                    nn.Linear(config.hidden_size, 1024),
-                    nn.ReLU(),
-                    nn.Dropout(0.05),
-                    nn.Linear(1024, 16),
-                    nn.ReLU(),
-                    nn.Linear(16, output_dim),
-                )
-
-        self.rm_head.to(torch.float32)
-        self.reward_token = reward_token
-
-        self.special_token_ids = special_token_ids
-        if self.special_token_ids is not None:
-            self.reward_token = "special"
-
-    def forward(
-        self,
-        input_ids: torch.LongTensor = None,
-        attention_mask: torch.Tensor | None = None,
-        position_ids: torch.LongTensor | None = None,
-        past_key_values: list[torch.FloatTensor] | None = None,
-        inputs_embeds: torch.FloatTensor | None = None,
-        labels: torch.LongTensor | None = None,
-        use_cache: bool | None = None,
-        output_attentions: bool | None = None,
-        output_hidden_states: bool | None = None,
-        return_dict: bool | None = None,
-        pixel_values: torch.Tensor | None = None,
-        pixel_values_videos: torch.FloatTensor | None = None,
-        image_grid_thw: torch.LongTensor | None = None,
-        video_grid_thw: torch.LongTensor | None = None,
-        rope_deltas: torch.LongTensor | None = None,
-    ):
-        ## modified from the origin class Qwen2VLForConditionalGeneration
-        output_attentions = (output_attentions if output_attentions is not None else self.config.output_attentions)
-        output_hidden_states = (output_hidden_states
-                                if output_hidden_states is not None else self.config.output_hidden_states)
-        return_dict = (return_dict if return_dict is not None else self.config.use_return_dict)
-        # pdb.set_trace()
-        if inputs_embeds is None:
-            inputs_embeds = self.model.embed_tokens(input_ids)
-            if pixel_values is not None:
-                pixel_values = pixel_values.type(self.visual.get_dtype())
-                image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)
-                image_mask = ((input_ids == self.config.image_token_id).unsqueeze(-1).expand_as(inputs_embeds))
-                image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
-                inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
-
-            if pixel_values_videos is not None:
-                pixel_values_videos = pixel_values_videos.type(self.visual.get_dtype())
-                video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw)
-                video_mask = ((input_ids == self.config.video_token_id).unsqueeze(-1).expand_as(inputs_embeds))
-                video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
-                inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)
-
-            if attention_mask is not None:
-                attention_mask = attention_mask.to(inputs_embeds.device)
-
-        outputs = self.model(
-            input_ids=None,
-            position_ids=position_ids,
-            attention_mask=attention_mask,
-            past_key_values=past_key_values,
-            inputs_embeds=inputs_embeds,
-            use_cache=use_cache,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-        )
-
-        hidden_states = outputs[0]  # [B, L, D]
-        with torch.autocast(device_type='cuda', dtype=torch.float32):
-            logits = self.rm_head(hidden_states)  # [B, L, N]
-
-        batch_size = input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0]
-
-        ## get sequence length
-        if self.config.pad_token_id is None and batch_size != 1:
-            raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
-        if self.config.pad_token_id is None:
-            sequence_lengths = -1
-        else:
-            if input_ids is not None:
-                # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
-                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1)
-                sequence_lengths = sequence_lengths % input_ids.shape[-1]
-                sequence_lengths = sequence_lengths.to(logits.device)
-            else:
-                sequence_lengths = -1
-
-        ## get the last token's logits
-        if self.reward_token == "last":
-            pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
-        elif self.reward_token == "mean":
-            ## get the mean of all valid tokens' logits
-            valid_lengths = torch.clamp(sequence_lengths, min=0, max=logits.size(1) - 1)
-            pooled_logits = torch.stack([logits[i, :valid_lengths[i]].mean(dim=0) for i in range(batch_size)])
-        elif self.reward_token == "special":
-            # special_token_ids = self.tokenizer.convert_tokens_to_ids(self.special_tokens)
-            # create a mask for special tokens
-            special_token_mask = torch.zeros_like(input_ids, dtype=torch.bool)
-            for special_token_id in self.special_token_ids:
-                special_token_mask = special_token_mask | (input_ids == special_token_id)
-            pooled_logits = logits[special_token_mask, ...]
-            pooled_logits = pooled_logits.view(batch_size, 1, -1)  # [B, 3, N] assert 3 attributes
-            pooled_logits = pooled_logits.view(batch_size, -1)
-
-            # pdb.set_trace()
-        else:
-            raise ValueError("Invalid reward_token")
-
-        return {"logits": pooled_logits}
-
-
-def _convert_A_B_to_chosen_rejected(
-    rewards_A,
-    rewards_B,
-    tied_threshold=None,
-    choice_dist=None,
-):
-    """
-    Inputs:
-        rewards_A: [B, 1]
-        rewards_B: [B, 1]
-    Outputs:
-        rewards_chosen: [B, 1]
-        rewards_rejected: [B, 1]
-        nontied_mask: [B, 1] (preference labels that is not tied)
-    """
-    chosen_label = torch.ones_like(rewards_A, dtype=torch.int64).to(rewards_A.device)  # [B, 1]
-
-    rewards_chosen = rewards_A
-    rewards_rejected = rewards_B
-
-    if tied_threshold is None:
-        nontied_mask = torch.ones_like(chosen_label, dtype=torch.float32).to(rewards_A.device)
-    else:
-        nontied_mask = (torch.abs((choice_dist[:, 0] - choice_dist[:, 1]) / torch.sum(choice_dist, dim=-1))
-                        > tied_threshold)
-        print(nontied_mask)
-    return (
-        rewards_chosen,
-        rewards_rejected,
-        nontied_mask,
-    )
-
-
-class PartialEmbeddingUpdateCallback(TrainerCallback):
-    """
-    Callback to update the embedding of special tokens
-    Only the special tokens are updated, the rest of the embeddings are kept fixed
-    """
-
-    def __init__(self, special_token_ids):
-        super().__init__()
-        self.special_token_ids = special_token_ids
-        self.orig_embeds_params = None
-
-    def on_train_begin(self, args, state, control, **kwargs):
-        model = kwargs.get("model")
-        self.orig_embeds_params = model.get_input_embeddings().weight.clone().detach()
-
-    def on_step_end(self, args, state, control, **kwargs):
-        # pdb.set_trace()
-        model = kwargs.get("model")
-        tokenizer = kwargs.get("tokenizer")
-
-        index_no_updates = torch.ones((len(tokenizer), ), dtype=torch.bool)
-        index_no_updates[self.special_token_ids] = False
-        with torch.no_grad():
-            model.get_input_embeddings().weight[index_no_updates] = (self.orig_embeds_params[index_no_updates])
-
-
-class VLMRewardTrainer(RewardTrainer):
-
-    def __init__(self,
-                 loss_type="regular",
-                 loss_hyperparameters=None,
-                 tied_threshold=None,
-                 visualization_steps=500,
-                 max_viz_samples=4,
-                 *args,
-                 **kwargs):
-        super().__init__(*args, **kwargs)
-        self.loss_type = loss_type
-        self.tied_threshold = tied_threshold
-        self.rewards_chosen_accumulated = []
-        self.rewards_rejected_accumulated = []
-        self.loss_hyperparameters = loss_hyperparameters if loss_hyperparameters is not None else {}
-        self.visualization_steps = visualization_steps
-        self.max_viz_samples = max_viz_samples
-
-    def get_eval_dataloader(self, eval_dataset: str | Dataset | None = None) -> DataLoader:
-        """
-        Returns the evaluation [`~torch.utils.data.DataLoader`].
-
-        Subclass and override this method if you want to inject some custom behavior.
-
-        Args:
-            eval_dataset (`str` or `torch.utils.data.Dataset`, *optional*):
-                If a `str`, will use `self.eval_dataset[eval_dataset]` as the evaluation dataset. If a `Dataset`, will override `self.eval_dataset` and must implement `__len__`. If it is a [`~datasets.Dataset`], columns not accepted by the `model.forward()` method are automatically removed.
-        """
-        if eval_dataset is None and self.eval_dataset is None:
-            raise ValueError("Trainer: evaluation requires an eval_dataset.")
-
-        # If we have persistent workers, don't do a fork bomb especially as eval datasets
-        # don't change during training
-        dataloader_key = eval_dataset if isinstance(eval_dataset, str) else "eval"
-        if (hasattr(self, "_eval_dataloaders") and dataloader_key in self._eval_dataloaders
-                and self.args.dataloader_persistent_workers):
-            return self.accelerator.prepare(self._eval_dataloaders[dataloader_key])
-
-        eval_dataset = (self.eval_dataset[eval_dataset] if isinstance(eval_dataset, str) else
-                        eval_dataset if eval_dataset is not None else self.eval_dataset)
-
-        data_collator = self.data_collator
-
-        if is_datasets_available() and isinstance(eval_dataset, datasets.Dataset):
-            eval_dataset = self._remove_unused_columns(eval_dataset, description="evaluation")
-        else:
-            data_collator = self._get_collator_with_removed_columns(data_collator, description="evaluation")
-
-        dataloader_params = {
-            "batch_size": self.args.eval_batch_size,
-            "collate_fn": data_collator,
-            "num_workers": self.args.dataloader_num_workers,
-            "pin_memory": self.args.dataloader_pin_memory,
-            "persistent_workers": self.args.dataloader_persistent_workers,
-        }
-
-        if not isinstance(eval_dataset, torch.utils.data.IterableDataset):
-            dataloader_params["sampler"] = self._get_eval_sampler(eval_dataset)
-            dataloader_params["drop_last"] = self.args.dataloader_drop_last
-            dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor
-
-        # accelerator.free_memory() will destroy the references, so
-        # we need to store the non-prepared version
-        eval_dataloader = DataLoader(eval_dataset, **dataloader_params)
-        if self.args.dataloader_persistent_workers:
-            if hasattr(self, "_eval_dataloaders"):
-                self._eval_dataloaders[dataloader_key] = eval_dataloader
-            else:
-                self._eval_dataloaders = {dataloader_key: eval_dataloader}
-
-        return self.accelerator.prepare(eval_dataloader)
-
-    def create_optimizer(self):
-        """
-        Setup the optimizer.
-        We provide a reasonable default that works well. If you want to use something else, you can pass a tuple in the
-        Trainer's init through `optimizers`, or subclass and override this method in a subclass.
-        """
-        if is_sagemaker_mp_enabled():
-            return super().create_optimizer()
-
-        opt_model = self.model
-
-        if self.optimizer is None:
-            decay_parameters = self.get_decay_parameter_names(opt_model)
-            decay_parameters = [name for name in decay_parameters if "bias" not in name]
-            lr_mapper = {}
-            visual_parameters = []
-            merger_parameters = []
-            rm_head_parameters = []
-
-            if self.args.vision_lr is not None:
-                lr_mapper["visual"] = self.args.vision_lr
-                visual_parameters = [
-                    name for name, _ in opt_model.named_parameters() if "visual" in name and "merger" not in name
-                ]
-            if self.args.merger_lr is not None:
-                lr_mapper["merger"] = self.args.merger_lr
-                merger_parameters = [name for name, _ in opt_model.named_parameters() if "merger" in name]
-
-            if self.args.rm_head_lr is not None:
-                lr_mapper["rm_head"] = self.args.rm_head_lr
-                rm_head_parameters = [name for name, _ in opt_model.named_parameters() if "rm_head" in name]
-
-            if len(lr_mapper) > 0:
-                special_lr_parameters = merger_parameters + visual_parameters + rm_head_parameters
-
-                optimizer_grouped_parameters = [
-                    {
-                        "params": [
-                            p for n, p in opt_model.named_parameters()
-                            if (n in decay_parameters and n not in special_lr_parameters and p.requires_grad)
-                        ],
-                        "weight_decay":
-                        self.args.weight_decay,
-                    },
-                    {
-                        "params": [
-                            p for n, p in opt_model.named_parameters()
-                            if (n not in decay_parameters and n not in special_lr_parameters and p.requires_grad)
-                        ],
-                        "weight_decay":
-                        0.0,
-                    },
-                ]
-
-                if visual_parameters:
-                    optimizer_grouped_parameters.extend([
-                        {
-                            "params": [
-                                p for n, p in opt_model.named_parameters()
-                                if (n in decay_parameters and n in visual_parameters and p.requires_grad)
-                            ],
-                            "weight_decay":
-                            self.args.weight_decay,
-                            "lr":
-                            self.args.vision_lr,
-                        },
-                        {
-                            "params": [
-                                p for n, p in opt_model.named_parameters()
-                                if (n not in decay_parameters and n in visual_parameters and p.requires_grad)
-                            ],
-                            "weight_decay":
-                            0.0,
-                            "lr":
-                            self.args.vision_lr,
-                        },
-                    ])
-
-                if merger_parameters:
-                    optimizer_grouped_parameters.extend([
-                        {
-                            "params": [
-                                p for n, p in opt_model.named_parameters()
-                                if (n in decay_parameters and n in merger_parameters and p.requires_grad)
-                            ],
-                            "weight_decay":
-                            self.args.weight_decay,
-                            "lr":
-                            self.args.merger_lr,
-                        },
-                        {
-                            "params": [
-                                p for n, p in opt_model.named_parameters()
-                                if (n not in decay_parameters and n in merger_parameters and p.requires_grad)
-                            ],
-                            "weight_decay":
-                            0.0,
-                            "lr":
-                            self.args.merger_lr,
-                        },
-                    ])
-
-                if rm_head_parameters:
-                    optimizer_grouped_parameters.extend([
-                        {
-                            "params": [
-                                p for n, p in opt_model.named_parameters()
-                                if (n in decay_parameters and n in rm_head_parameters and p.requires_grad)
-                            ],
-                            "weight_decay":
-                            self.args.weight_decay,
-                            "lr":
-                            self.args.rm_head_lr,
-                        },
-                        {
-                            "params": [
-                                p for n, p in opt_model.named_parameters()
-                                if (n not in decay_parameters and n in rm_head_parameters and p.requires_grad)
-                            ],
-                            "weight_decay":
-                            0.0,
-                            "lr":
-                            self.args.rm_head_lr,
-                        },
-                    ])
-
-            else:
-                optimizer_grouped_parameters = [
-                    {
-                        "params":
-                        [p for n, p in opt_model.named_parameters() if (n in decay_parameters and p.requires_grad)],
-                        "weight_decay":
-                        self.args.weight_decay,
-                    },
-                    {
-                        "params":
-                        [p for n, p in opt_model.named_parameters() if (n not in decay_parameters and p.requires_grad)],
-                        "weight_decay":
-                        0.0,
-                    },
-                ]
-
-            if self.model.special_token_ids:
-                special_token_embeddings = opt_model.get_input_embeddings().weight
-
-                special_token_embeddings.requires_grad = True
-
-                optimizer_grouped_parameters.extend([
-                    {
-                        # "params": [p for n, p in opt_model.get_input_embeddings().named_parameters() if (p.requires_grad)],
-                        "params": [special_token_embeddings],
-                        "lr": self.args.special_token_lr,
-                        "weight_decay": 0.0,
-                    },
-                ])
-
-            optimizer_cls, optimizer_kwargs = self.get_optimizer_cls_and_kwargs(self.args, opt_model)
-
-            self.optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)
-
-        return self.optimizer
-
-    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
-        rewards_A = model(return_dict=True, **inputs["batch_1"])["logits"]
-        rewards_B = model(return_dict=True, **inputs["batch_2"])["logits"]
-
-        # Log to TensorBoard for visualization
-        if (hasattr(self.state, 'global_step') and self.state.global_step % self.visualization_steps == 0
-                and self.state.global_step > 0):
-            # Pass the original inputs which should contain the text prompts
-            self._log_training_visualization(inputs, rewards_A, rewards_B)
-
-        # calculate loss, optionally modulate with margin
-        # get chosen and rejected rewards from the chosen label
-        (
-            rewards_chosen,
-            rewards_rejected,
-            nontied_mask,
-        ) = _convert_A_B_to_chosen_rejected(
-            rewards_A,
-            rewards_B,
-            tied_threshold=self.tied_threshold,
-            choice_dist=inputs["choice_dist"],
-        )
-
-        loss_dict = {}
-
-        if self.loss_type == "bt":
-            # Bradley-Terry model
-            loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected)
-            out_mask = nontied_mask
-            loss = loss * out_mask
-            loss = loss.mean()
-        elif self.loss_type == "likelihood_displacement":
-            # Bradley-Terry model
-            loss = -nn.functional.logsigmoid(rewards_chosen - self.loss_hyperparameters['tau'] * rewards_rejected)
-            out_mask = nontied_mask
-            loss = loss * out_mask
-            loss = loss.mean()
-
-        elif self.loss_type == "constant_margin":
-            # Bradley-Terry model with constant margin
-            loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected - 0.57)
-            out_mask = nontied_mask
-            loss = loss * out_mask
-            loss = loss.mean()
-        elif self.loss_type == "btt":
-            # Bradley-Terry-With-Ties model
-            k = 5.0
-            log_k = math.log(k)
-            log_k2_sub_1 = math.log(k**2 - 1)
-            bt_loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected - log_k)
-            same_loss = (-nn.functional.logsigmoid(rewards_chosen - rewards_rejected - log_k) -
-                         nn.functional.logsigmoid(rewards_rejected - rewards_chosen - log_k) - log_k2_sub_1)
-            loss = bt_loss * nontied_mask.float() + same_loss * (1 - nontied_mask.float())
-            out_mask = torch.ones_like(nontied_mask, dtype=torch.float32).to(rewards_A.device)  # [B, 1]
-            loss = loss * out_mask
-
-            loss = loss.mean()
-        elif self.loss_type == "hpsv2":
-            device = rewards_A.device
-            rewards = torch.nn.functional.softmax(torch.cat([rewards_A, rewards_B], dim=-1), dim=-1)
-            text_0_logits, text_1_logits = rewards[:, 0], rewards[:, 1]
-            label_0, label_1 = torch.ones_like(text_0_logits), torch.zeros_like(text_0_logits)
-
-            text_logits = torch.stack([text_0_logits, text_1_logits], dim=-1)
-            text_0_labels = torch.zeros(text_logits.shape[0], device=device, dtype=torch.long)
-            text_1_labels = text_0_labels + 1
-
-            text_0_loss = torch.nn.functional.cross_entropy(text_logits, text_0_labels, reduction="none")
-            text_1_loss = torch.nn.functional.cross_entropy(text_logits, text_1_labels, reduction="none")
-
-            loss = label_0 * text_0_loss + label_1 * text_1_loss
-
-            # absolute_example_weight = 1 / num_per_prompt
-            # denominator = absolute_example_weight.sum()
-            # weight_per_example = absolute_example_weight / denominator
-            # text_loss *= weight_per_example
-            loss = loss.sum()
-        elif self.loss_type == "uncertainty":
-            batch_size = rewards_A.shape[0]
-            mean_chosen = rewards_A[:, 0]
-            mean_rejected = rewards_B[:, 0]
-            sigma_chosen = torch.exp(rewards_A[:, 1])
-            sigma_rejected = torch.exp(rewards_B[:, 1])
-
-            mean_z = mean_chosen - mean_rejected
-            sigma_z = torch.sqrt(sigma_chosen**2 + sigma_rejected**2)
-
-            z_samples = torch.randn(batch_size, 1000).to(sigma_z.device).to(
-                torch.float16) * sigma_z.unsqueeze(1).repeat(1, 1000) + mean_z.unsqueeze(1).repeat(1, 1000)
-            loss = -torch.nn.functional.logsigmoid(z_samples).mean()
-        else:
-            raise NotImplementedError(f"Loss type {self.loss_type} not implemented.")
-
-        loss_dict.update({"loss": loss.item()})
-
-        if return_outputs:
-            ## return rewards_A/B instead of chosen/rejected
-            ## easier to calculate metrics for multi-attribute
-            return loss, {
-                "rewards_A": rewards_A,
-                "rewards_B": rewards_B,
-            }
-        return loss
-
-    def prediction_step(
-        self,
-        model,
-        inputs,
-        prediction_loss_only,
-        ignore_keys=None,
-    ):
-        model.eval()
-        inputs = self._prepare_inputs(inputs)
-        if ignore_keys is None:
-            if hasattr(self.model, "config"):
-                ignore_keys = getattr(self.model.config, "keys_to_ignore_at_inference", [])
-            else:
-                ignore_keys = []
-
-        with torch.no_grad():
-            loss, logits_dict = self.compute_loss(model, inputs, return_outputs=True)
-
-        if prediction_loss_only:
-            return (loss, None, None)
-        loss = loss.detach()
-        logits = tuple(v for k, v in logits_dict.items() if k not in ignore_keys)
-        logits = nested_detach(logits)
-        if self.loss_type != "uncertainty":
-            logits = torch.cat(logits, dim=1)  # [B, 2]
-        else:
-            logits = torch.cat([p[:, [0]] for p in logits], dim=1)
-
-        labels = torch.ones((logits.shape[0], 1)).to(logits.device)
-
-        return loss, logits, labels
-
-    def _log_training_visualization(self, inputs, rewards_A, rewards_B):
-        """Log training samples and predictions to TensorBoard"""
-        try:
-            # Get tensorboard writer from trainer
-            writer = None
-            if hasattr(self, 'log_metrics') and hasattr(self.args,
-                                                        'report_to') and 'tensorboard' in self.args.report_to:
-                # Try to get the writer from the logger
-                from torch.utils.tensorboard import SummaryWriter
-                if not hasattr(self, '_tb_writer'):
-                    self._tb_writer = SummaryWriter(log_dir=self.args.logging_dir)
-                writer = self._tb_writer
-
-            if writer is None:
-                return
-
-            step = self.state.global_step
-            batch_size = min(len(rewards_A), self.max_viz_samples)
-
-            # Log scalar metrics
-            for i in range(batch_size):
-                score_A = rewards_A[i].float().detach().cpu().numpy()
-                score_B = rewards_B[i].float().detach().cpu().numpy()
-
-                # Convert to float for logging
-                score_A_val = float(score_A.mean()) if score_A.ndim > 0 else float(score_A)
-                score_B_val = float(score_B.mean()) if score_B.ndim > 0 else float(score_B)
-                score_diff = score_A_val - score_B_val
-
-                writer.add_scalar(f'train_viz/sample_{i}/score_A', score_A_val, step)
-                writer.add_scalar(f'train_viz/sample_{i}/score_B', score_B_val, step)
-                writer.add_scalar(f'train_viz/sample_{i}/score_diff', score_diff, step)
-
-                try:
-                    # Get image data from inputs
-                    image_A = inputs['image_1'][i] if 'image_1' in inputs else None
-                    image_B = inputs['image_2'][i] if 'image_2' in inputs else None
-
-                    # Get prompt text from the original batch (now properly stored)
-                    prompt_A = inputs.get('text_1', ['Unknown prompt'])[i] if 'text_1' in inputs else 'Unknown prompt'
-
-                    fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(12, 8))
-                    fig.text(0.05,
-                             0.05,
-                             f'Prompt:\n{prompt_A[:200]}{"..." if len(prompt_A) > 200 else ""}',
-                             ha='left',
-                             va='bottom',
-                             fontsize=8,
-                             wrap=True,
-                             bbox=dict(boxstyle="round,pad=0.3", facecolor="lightblue", alpha=0.7))
-                    img_A_np = np.array(image_A)
-                    if img_A_np.ndim == 3 and img_A_np.shape[0] == 3:  # CHW format
-                        img_A_np = np.transpose(img_A_np, (1, 2, 0))
-                    img_A_np = np.clip(img_A_np, 0, 1)  # Ensure values are in [0,1]
-                    axes[0].imshow(img_A_np)
-                    axes[0].set_title(f'Image A - Score: {score_A_val:.3f}')
-                    axes[0].axis('off')
-
-                    img_B_np = np.array(image_B)
-                    if img_B_np.ndim == 3 and img_B_np.shape[0] == 3:  # CHW format
-                        img_B_np = np.transpose(img_B_np, (1, 2, 0))
-                    img_B_np = np.clip(img_B_np, 0, 1)  # Ensure values are in [0,1]
-                    axes[1].imshow(img_B_np)
-
-                    axes[1].set_title(f'Image B - Score: {score_B_val:.3f}')
-                    axes[1].axis('off')
-
-                    # Add prediction info
-                    winner = "A" if score_diff > 0 else "B"
-                    plt.suptitle(
-                        f'Step {step} - Sample {i} | Predicted Winner: Image {winner} | Diff: {score_diff:.3f}',
-                        fontsize=14)
-                    plt.tight_layout()
-
-                    # Log figure to tensorboard
-                    writer.add_figure(f'train_viz/sample_{i}_comparison', fig, step)
-                    plt.close(fig)
-                except Exception as viz_error:
-                    print(f"Warning: Could not extract images for visualization: {viz_error}")
-                    continue
-
-            # Log aggregate statistics
-            all_scores_A = rewards_A.float().detach().cpu().numpy()
-            all_scores_B = rewards_B.float().detach().cpu().numpy()
-
-            writer.add_histogram('train_viz/all_scores_A', all_scores_A, step)
-            writer.add_histogram('train_viz/all_scores_B', all_scores_B, step)
-            writer.add_scalar('train_viz/mean_score_A', float(all_scores_A.mean()), step)
-            writer.add_scalar('train_viz/mean_score_B', float(all_scores_B.mean()), step)
-            writer.add_scalar('train_viz/mean_score_diff', float((all_scores_A - all_scores_B).mean()), step)
-
-        except Exception as e:
-            print(f"Error in training visualization: {e}")
-
-    def _save_checkpoint(self, model, trial, metrics=None):
-
-        if isinstance(self.model, PeftModel):
-            checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}"
-
-            if self.hp_search_backend is None and trial is None:
-                self.store_flos()
-
-            run_dir = self._get_output_dir(trial=trial)
-            output_dir = os.path.join(run_dir, checkpoint_folder)
-            os.makedirs(output_dir, exist_ok=True)
-
-            # TODO: Just Temp
-            self.save_model(output_dir, _internal_call=True)
-            # pdb.set_trace()
-
-            if not self.args.save_full_model:
-                non_lora_weights = get_peft_state_non_lora_maybe_zero_3(self.model.named_parameters(),
-                                                                        require_grad_only=True)
-                torch.save(
-                    non_lora_weights,
-                    os.path.join(output_dir, "non_lora_state_dict.pth"),
-                )
-            # safetensors.torch.save(non_lora_weights, os.path.join(output_dir, "non_lora_model.safetensors"))
-
-            if not self.args.save_only_model:
-                # Save optimizer and scheduler
-                self._save_optimizer_and_scheduler(output_dir)
-                # Save RNG state
-                self._save_rng_state(output_dir)
-
-        else:
-            super(RewardTrainer, self)._save_checkpoint(model, trial, metrics)
-
-    def _save(self, output_dir: str | None = None, state_dict=None):
-        # If we are executing this function, we are the process zero, so we don't check for that.
-        output_dir = output_dir if output_dir is not None else self.args.output_dir
-        os.makedirs(output_dir, exist_ok=True)
-        logger.info(f"Saving model checkpoint to {output_dir}")
-        # pdb.set_trace()
-
-        supported_classes = ((PreTrainedModel, ) if not is_peft_available() else (PreTrainedModel, PeftModel))
-        # Save a trained model and configuration using `save_pretrained()`.
-        # They can then be reloaded using `from_pretrained()`
-        if not isinstance(self.model, supported_classes):
-            if state_dict is None:
-                state_dict = self.model.state_dict()
-
-            if isinstance(self.accelerator.unwrap_model(self.model), supported_classes):
-                self.accelerator.unwrap_model(self.model).save_pretrained(
-                    output_dir,
-                    state_dict=state_dict,
-                    safe_serialization=self.args.save_safetensors,
-                )
-            else:
-                logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.")
-                if self.args.save_safetensors:
-                    safetensors.torch.save_file(
-                        state_dict,
-                        os.path.join(output_dir, SAFE_WEIGHTS_NAME),
-                        metadata={"format": "pt"},
-                    )
-                else:
-                    torch.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME))
-        else:
-            if not self.args.save_full_model:
-                state_dict = {k: v for k, v in state_dict.items() if "wte" not in k}
-                self.model.save_pretrained(
-                    output_dir,
-                    state_dict=state_dict,
-                    safe_serialization=self.args.save_safetensors,
-                )
-            else:
-                torch.save(state_dict, os.path.join(output_dir, "model.pth"))
-
-        if self.tokenizer is not None:
-            os.makedirs(os.path.join(output_dir, "tokenizer"), exist_ok=True)
-            self.tokenizer.save_pretrained(os.path.join(output_dir, "tokenizer"))
-
-        # Good practice: save your training arguments together with the trained model
-        torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))
-        # pdb.set_trace()
-
-
-def compute_multi_attr_accuracy(eval_pred, metainfo_idxs=None) -> dict[str, float]:
-    predictions, labels = eval_pred
-    metrics = {}
-
-    pred_curr = predictions
-    label_curr = labels.squeeze(1)
-    total_count = np.sum(label_curr != 0)
-
-    rewards_chosen = pred_curr[:, 0]
-    rewards_rejected = pred_curr[:, 1]
-
-    rewards_chosen_avg = np.sum(rewards_chosen) / total_count
-    rewards_rejected_avg = np.sum(rewards_rejected) / total_count
-
-    accuracy = np.sum(rewards_chosen > rewards_rejected) / total_count
-
-    metrics.update({
-        "Acc": accuracy,
-        "R_chosen_avg": rewards_chosen_avg,
-        "R_rejected_avg": rewards_rejected_avg,
-    })
-    return metrics

fastvideo/train/methods/rl/reward/HPSv3/hpsv3/train.py

@@ -1,286 +0,0 @@
-import json
-import os
-import fire
-from dataclasses import asdict
-from functools import partial
-import torch
-from .model.qwen2vl_trainer import (
-    Qwen2VLRewardModelBT,
-    VLMRewardTrainer,
-    compute_multi_attr_accuracy,
-    PartialEmbeddingUpdateCallback,
-)
-from .dataset.pairwise_dataset import PairwiseOriginalDataset
-from .dataset.data_collator_qwen import QWen2VLDataCollator
-from .utils.parser import ModelConfig, PEFTLoraConfig, TrainingConfig, DataConfig
-from .utils.training_utils import load_model_from_checkpoint, find_target_linear_names
-from .utils.parser import parse_args_with_yaml
-from transformers import AutoProcessor
-from peft import LoraConfig, get_peft_model
-from trl import get_kbit_device_map, get_quantization_config
-from .model.differentiable_image_processor import Qwen2VLImageProcessor
-try:
-    import flash_attn
-except ImportError:
-    flash_attn = None
-    print("Flash Attention is not installed. Falling to SDPA.")
-
-
-def create_model_and_processor(
-    model_config,
-    peft_lora_config,
-    training_args,
-    cache_dir=None,
-    differentiable=False,
-):
-    # create model
-    torch_dtype = (model_config.torch_dtype if model_config.torch_dtype in ["auto", None] else getattr(
-        torch, model_config.torch_dtype))
-    quantization_config = get_quantization_config(model_config)
-    model_kwargs = dict(revision=model_config.model_revision,
-                        device_map=get_kbit_device_map() if quantization_config is not None else None,
-                        quantization_config=quantization_config,
-                        use_cache=False)
-
-    # create processor and set padding
-
-    processor = AutoProcessor.from_pretrained(model_config.model_name_or_path,
-                                              padding_side="right",
-                                              cache_dir=cache_dir)
-
-    if differentiable:
-        processor.image_processor = Qwen2VLImageProcessor()
-
-    special_token_ids = None
-    if model_config.use_special_tokens:
-        special_tokens = ["<|Reward|>"]
-        processor.tokenizer.add_special_tokens({"additional_special_tokens": special_tokens})
-        special_token_ids = processor.tokenizer.convert_tokens_to_ids(special_tokens)
-
-    model = Qwen2VLRewardModelBT.from_pretrained(
-        model_config.model_name_or_path,
-        output_dim=model_config.output_dim,
-        reward_token=model_config.reward_token,
-        special_token_ids=special_token_ids,
-        torch_dtype=torch_dtype,
-        attn_implementation=("flash_attention_2"
-                             if not training_args.disable_flash_attn2 and flash_attn is not None else "sdpa"),
-        cache_dir=cache_dir,
-        rm_head_type=model_config.rm_head_type,
-        rm_head_kwargs=model_config.rm_head_kwargs,
-        **model_kwargs,
-    )
-
-    if model_config.use_special_tokens:
-        model.resize_token_embeddings(len(processor.tokenizer))
-
-    if training_args.bf16:
-        model.to(torch.bfloat16)
-    if training_args.fp16:
-        model.to(torch.float16)
-
-    model.rm_head.to(torch.float32)
-
-    # create lora and peft model
-    if peft_lora_config.lora_enable:
-        target_modules = find_target_linear_names(
-            model,
-            num_lora_modules=peft_lora_config.num_lora_modules,
-            lora_namespan_exclude=peft_lora_config.lora_namespan_exclude,
-        )
-        peft_config = LoraConfig(
-            target_modules=target_modules,
-            r=peft_lora_config.lora_r,
-            lora_alpha=peft_lora_config.lora_alpha,
-            lora_dropout=peft_lora_config.lora_dropout,
-            task_type=peft_lora_config.lora_task_type,
-            use_rslora=peft_lora_config.use_rslora,
-            bias="none",
-            modules_to_save=peft_lora_config.lora_modules_to_save,
-        )
-        model = get_peft_model(model, peft_config)
-    else:
-        peft_config = None
-
-    model.config.tokenizer_padding_side = processor.tokenizer.padding_side
-    model.config.pad_token_id = processor.tokenizer.pad_token_id
-
-    return model, processor, peft_config
-
-
-def save_configs_to_json(data_config, training_args, model_config, peft_lora_config):
-    """
-    Save all configurations to a JSON file.
-    """
-    config_dict = {
-        "data_config": asdict(data_config),
-        "training_args": asdict(training_args),
-        "model_config": asdict(model_config),
-        "peft_lora_config": asdict(peft_lora_config),
-    }
-    # del information about local device
-    del config_dict["training_args"]["local_rank"]
-    del config_dict["training_args"]["_n_gpu"]
-
-    save_path = os.path.join(training_args.output_dir, "model_config.json")
-
-    os.makedirs(training_args.output_dir, exist_ok=True)
-    print(training_args.output_dir)
-
-    with open(save_path, "w") as f:
-        json.dump(config_dict, f, indent=4)
-
-
-def set_requires_grad(parameters, requires_grad):
-    for p in parameters:
-        p.requires_grad = requires_grad
-
-
-def train(config, local_rank=0, debug=False):
-
-    ## ===> Step 1: Parse arguments
-    (data_config, training_args, model_config, peft_lora_config), config_path = (parse_args_with_yaml(
-        (DataConfig, TrainingConfig, ModelConfig, PEFTLoraConfig), config, is_train=True))
-    training_args.output_dir = os.path.join(training_args.output_dir, config.split("/")[-1].split(".")[0])
-    training_args.logging_dir = training_args.output_dir
-    # check valid (lora config)
-    assert not (peft_lora_config.lora_enable and model_config.freeze_llm
-                ), "When using LoRA, the LLM should not be frozen. If you want to freeze the LLM, please disable LoRA."
-    if not peft_lora_config.lora_enable:
-        assert (not peft_lora_config.vision_lora
-                ), "Error: model_config.lora_enable is not enabled, but model_config.vision_lora is enabled."
-    else:
-        if peft_lora_config.lora_namespan_exclude is None:
-            peft_lora_config.lora_namespan_exclude = []
-        if not peft_lora_config.vision_lora:
-            peft_lora_config.lora_namespan_exclude += ["visual"]
-
-    ## ===> Step 2: Load model and configure
-    model, processor, peft_config = create_model_and_processor(
-        model_config=model_config,
-        peft_lora_config=peft_lora_config,
-        training_args=training_args,
-    )
-
-    ## load model
-    if training_args.load_from_pretrained is not None:
-        model, checkpoint_step = load_model_from_checkpoint(
-            model,
-            training_args.load_from_pretrained,
-            training_args.load_from_pretrained_step,
-        )
-    model.train()
-
-    if peft_lora_config.lora_enable:
-        model_to_configure = model.model
-    else:
-        model_to_configure = model
-        # set requires_grad for LLM
-        set_requires_grad(model_to_configure.model.parameters(), not model_config.freeze_llm)
-        set_requires_grad(model_to_configure.model.embed_tokens.parameters(), False)
-    if not peft_lora_config.vision_lora:
-        # set requires_grad for visual encoder and merger
-        set_requires_grad(model_to_configure.visual.parameters(), not model_config.freeze_vision_tower)
-        set_requires_grad(model_to_configure.visual.merger.parameters(), model_config.tune_merger)
-
-    if model_config.trainable_visual_layers:  # This is inverse order to index of model.visual.blocks, set -1 to unfreeze all layers
-        assert model_config.trainable_visual_layers <= len(
-            model_to_configure.visual.blocks
-        ), "trainable_visual_layers should be less than or equal to the number of visual blocks"
-        freeze_layer_num = len(
-            model_to_configure.visual.blocks
-        ) - model_config.trainable_visual_layers if model_config.trainable_visual_layers > 0 else 0
-        for index, layer in enumerate(model_to_configure.visual.blocks):
-            if index < freeze_layer_num:
-                set_requires_grad(layer.parameters(), False)
-            else:
-                set_requires_grad(layer.parameters(), True)
-
-    # set requires_grad for regression head
-    set_requires_grad(model_to_configure.rm_head.parameters(), True)
-
-    ## ===> Step 3: Load Dataset and configure
-    train_dataset = PairwiseOriginalDataset(
-        data_config.train_json_list,
-        data_config.soft_label,
-        data_config.confidence_threshold,
-    )
-    test_set_dict = {}
-    for item in data_config.test_json_list:
-        test_set_dict[item[0]] = PairwiseOriginalDataset(
-            item[1],
-            data_config.soft_label,
-            data_config.confidence_threshold,
-        )
-
-    print(f"===> Selected {len(train_dataset)} samples for training.")
-    for key, value in test_set_dict.items():
-        print(f"===> Selected {len(value)} samples for {key} testing.")
-
-    num_gpu = int(os.environ.get("WORLD_SIZE", 1))
-    data_collator = QWen2VLDataCollator(
-        processor,
-        max_pixels=data_config.max_pixels,
-        min_pixels=data_config.min_pixels,
-        with_instruction=data_config.with_instruction,
-        use_special_tokens=model_config.use_special_tokens,
-    )
-    compute_metrics = partial(compute_multi_attr_accuracy)
-
-    actual_batch_size = (training_args.per_device_train_batch_size * training_args.gradient_accumulation_steps *
-                         num_gpu)
-    total_steps = (training_args.num_train_epochs * len(train_dataset) // actual_batch_size)
-    if training_args.save_epochs is not None:
-        training_args.save_steps = round(training_args.save_epochs * len(train_dataset) / actual_batch_size)
-    if training_args.eval_epochs is not None:
-        training_args.eval_steps = round(training_args.eval_epochs * len(train_dataset) / actual_batch_size)
-    if training_args.logging_epochs is not None:
-        training_args.logging_steps = round(training_args.logging_epochs * len(train_dataset) / actual_batch_size)
-
-    if training_args.local_rank == -1 or training_args.local_rank == 0:
-        print(f"===> Using {num_gpu} GPUs.")
-        print(f"===> Total Batch Size: {actual_batch_size}")
-        print(f"===> Training Epochs: {training_args.num_train_epochs}")
-        print(f"===> Total Steps: {total_steps}")
-        print(f"===> Save Steps: {training_args.save_steps}")
-        print(f"===> Eval Steps: {training_args.eval_steps}")
-        print(f"===> Logging Steps: {training_args.logging_steps}")
-
-    ## ===> Step 4: Save configs for re-check
-    if training_args.local_rank == -1 or training_args.local_rank == 0:
-        save_configs_to_json(data_config, training_args, model_config, peft_lora_config)
-
-    print(train_dataset)
-    ## ===> Step 5: Start Training!
-
-    special_token_ids = model.special_token_ids
-    callbacks = []
-    if special_token_ids is not None:
-        callbacks.append(PartialEmbeddingUpdateCallback(special_token_ids))
-
-    trainer = VLMRewardTrainer(
-        model=model,
-        compute_metrics=compute_metrics,
-        data_collator=data_collator,
-        args=training_args,
-        train_dataset=train_dataset,
-        eval_dataset=(test_set_dict if training_args.conduct_eval else None),
-        peft_config=peft_config,
-        callbacks=callbacks,
-        loss_type=model_config.loss_type,
-        loss_hyperparameters=model_config.loss_hyperparameters,
-        tokenizer=processor.tokenizer,
-        tied_threshold=data_config.tied_threshold,
-        visualization_steps=training_args.visualization_steps,
-        max_viz_samples=training_args.max_viz_samples,
-    )
-    trainer.train()
-
-    if training_args.local_rank == -1 or training_args.local_rank == 0:
-        model_state_dict = model.state_dict()
-        torch.save(model_state_dict, os.path.join(training_args.output_dir, "final_model.pth"))
-        model.config.save_pretrained(training_args.output_dir)
-
-
-if __name__ == "__main__":
-    fire.Fire(train)

fastvideo/train/methods/rl/reward/HPSv3/hpsv3/utils/__init__.py

@@ -1 +0,0 @@
-"""Vendored HPSv3 runtime utilities."""

fastvideo/train/methods/rl/reward/HPSv3/hpsv3/utils/parser.py

@@ -1,142 +0,0 @@
-from typing import Any, Literal
-from omegaconf import OmegaConf
-from transformers import HfArgumentParser
-from dataclasses import dataclass, field
-from transformers import TrainingArguments
-
-
-@dataclass
-class DataConfig:
-    train_json_list: list[str] = field(default_factory=lambda: ["/path/to/dataset/meta_data.json"])
-    val_json_list: list[str] = field(default_factory=lambda: ["/path/to/dataset/meta_data.json"])
-    test_json_list: list[str] = field(default_factory=lambda: ["/path/to/dataset/meta_data.json"])
-    soft_label: bool = False
-    confidence_threshold: float | None = None
-    max_pixels: int | None = 256 * 28 * 28  # Default max pixels
-    min_pixels: int | None = 256 * 28 * 28
-    with_instruction: bool = True
-    tied_threshold: float | None = None
-
-
-@dataclass
-class TrainingConfig(TrainingArguments):
-    max_grad_norm: float | None = 1.0
-    dataset_num_proc: int | None = None
-    center_rewards_coefficient: float | None = None
-    disable_flash_attn2: bool = field(default=False)
-    disable_dropout: bool = field(default=False)
-
-    vision_lr: float | None = None
-    merger_lr: float | None = None
-    rm_head_lr: float | None = None
-    special_token_lr: float | None = None
-
-    conduct_eval: bool | None = True
-    load_from_pretrained: str = None
-    load_from_pretrained_step: int = None
-    logging_epochs: float | None = None
-    eval_epochs: float | None = None
-    save_epochs: float | None = None
-    remove_unused_columns: bool | None = False
-
-    save_full_model: bool | None = False
-
-    # Visualization parameters
-    visualization_steps: int | None = 100
-    max_viz_samples: int | None = 4
-
-
-@dataclass
-class PEFTLoraConfig:
-    lora_enable: bool = False
-    vision_lora: bool = False
-    lora_r: int = 16
-    lora_alpha: int = 32
-    lora_dropout: float = 0.05
-    lora_target_modules: list[str] | None = None
-    lora_namespan_exclude: list[str] | None = None
-    lora_modules_to_save: list[str] | None = None
-    lora_task_type: str = "CAUSAL_LM"
-    use_rslora: bool = False
-    num_lora_modules: int = -1
-
-    def __post_init__(self):
-        if (isinstance(self.lora_target_modules, list) and len(self.lora_target_modules) == 1):
-            self.lora_target_modules = self.lora_target_modules[0]
-
-        if (isinstance(self.lora_namespan_exclude, list) and len(self.lora_namespan_exclude) == 1):
-            self.lora_namespan_exclude = self.lora_namespan_exclude[0]
-
-
-@dataclass
-class ModelConfig:
-    model_name_or_path: str | None = None
-    model_revision: str = "main"
-    rm_head_type: str = "default"
-    rm_head_kwargs: dict | None = None
-    output_dim: int = 1
-
-    use_special_tokens: bool = False
-
-    freeze_vision_tower: bool = field(default=False)
-    freeze_llm: bool = field(default=False)
-    tune_merger: bool = field(default=False)
-    trainable_visual_layers: int | None = -1
-
-    torch_dtype: Literal["auto", "bfloat16", "float16", "float32"] | None = None
-    trust_remote_code: bool = False
-    attn_implementation: str | None = None
-    load_in_8bit: bool = False
-    load_in_4bit: bool = False
-    bnb_4bit_quant_type: Literal["fp4", "nf4"] = "nf4"
-    use_bnb_nested_quant: bool = False
-    reward_token: Literal["last", "mean", "special"] = "last"
-    loss_type: Literal["bt", "reg", "btt", "margin", "constant_margin", "scaled"] = ("regular")
-    loss_hyperparameters: dict = field(default_factory=lambda: {})
-    checkpoint_path: str | None = None
-
-    def __post_init__(self):
-        if self.load_in_8bit and self.load_in_4bit:
-            raise ValueError("You can't use 8 bit and 4 bit precision at the same time")
-
-        # if isinstance(self.lora_target_modules, list) and len(self.lora_target_modules) == 1:
-        #     self.lora_target_modules = self.lora_target_modules[0]
-
-        # if isinstance(self.lora_namespan_exclude, list) and len(self.lora_namespan_exclude) == 1:
-        #     self.lora_namespan_exclude = self.lora_namespan_exclude[0]
-
-
-########## Functions for get trainable modules' parameters ##########
-
-
-def parse_args_with_yaml(
-    dataclass_types: tuple[type, ...],
-    config_path: str = None,
-    allow_extra_keys: bool = True,
-    is_train: bool = True,
-) -> tuple[Any, ...]:
-    """
-    Parse arguments using HfArgumentParser with OmegaConf for YAML support.
-    
-    Args:
-        dataclass_types: Tuple of dataclass types for HfArgumentParser
-        args: Optional arguments (if None, will read from sys.argv)
-        allow_extra_keys: Whether to allow extra keys in config
-    
-    Returns:
-        Tuple of parsed dataclass instances
-    """
-    # Read arguments from command line or provided args
-    # Load YAML config and merge with command line overrides
-    args = OmegaConf.to_container(OmegaConf.load(config_path))
-    if not is_train:
-        args.pop('deepspeed', None)
-
-    # Parse with HfArgumentParser
-    parser = HfArgumentParser(dataclass_types)
-    return parser.parse_dict(args, allow_extra_keys=allow_extra_keys), config_path
-
-
-if __name__ == "__main__":
-    data_config, training_args, model_config, peft_lora_config = parse_args_with_yaml(
-        (DataConfig, TrainingConfig, ModelConfig, PEFTLoraConfig))

fastvideo/train/methods/rl/reward/HPSv3/hpsv3/utils/training_utils.py

@@ -1,144 +0,0 @@
-import torch
-import os
-import glob
-import safetensors
-
-
-def maybe_zero_3(param, ignore_status=False, name=None):
-    from deepspeed import zero
-    from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus
-
-    if hasattr(param, "ds_id"):
-        if param.ds_status == ZeroParamStatus.NOT_AVAILABLE and not ignore_status:
-            print(f"Parameter {name} is not available in ZeRO-3, please check the ZeRO-3 status.")
-        with zero.GatheredParameters([param]):
-            param = param.data.detach().cpu().clone()
-    else:
-        param = param.detach().cpu().clone()
-    return param
-
-
-# Borrowed from peft.utils.get_peft_model_state_dict
-def get_peft_state_maybe_zero_3(named_params, bias):
-    if bias == "none":
-        to_return = {k: t for k, t in named_params if "lora_" in k}
-    elif bias == "all":
-        to_return = {k: t for k, t in named_params if "lora_" in k or "bias" in k}
-    elif bias == "lora_only":
-        to_return = {}
-        maybe_lora_bias = {}
-        lora_bias_names = set()
-        for k, t in named_params:
-            if "lora_" in k:
-                to_return[k] = t
-                bias_name = k.split("lora_")[0] + "bias"
-                lora_bias_names.add(bias_name)
-            elif "bias" in k:
-                maybe_lora_bias[k] = t
-        for k, t in maybe_lora_bias:
-            if bias_name in lora_bias_names:
-                to_return[bias_name] = t
-    else:
-        raise NotImplementedError
-    to_return = {k: maybe_zero_3(v, ignore_status=True) for k, v in to_return.items()}
-    return to_return
-
-
-def get_peft_state_non_lora_maybe_zero_3(named_params, require_grad_only=True):
-    to_return = {k: t for k, t in named_params if "lora_" not in k}
-    if require_grad_only:
-        to_return = {k: t for k, t in to_return.items() if t.requires_grad}
-    to_return = {k: maybe_zero_3(v, ignore_status=True).cpu() for k, v in to_return.items()}
-    return to_return
-
-
-def _insert_adapter_name_into_state_dict(state_dict: dict[str, torch.Tensor], adapter_name: str,
-                                         parameter_prefix: str) -> dict[str, torch.Tensor]:
-    """Utility function to remap the state_dict keys to fit the PEFT model by inserting the adapter name."""
-    peft_model_state_dict = {}
-    for key, val in state_dict.items():
-        if parameter_prefix in key:
-            suffix = key.split(parameter_prefix)[1]
-            if "." in suffix:
-                suffix_to_replace = ".".join(suffix.split(".")[1:])
-                key = key.replace(suffix_to_replace, f"{adapter_name}.{suffix_to_replace}")
-            else:
-                key = f"{key}.{adapter_name}"
-            peft_model_state_dict[key] = val
-        else:
-            peft_model_state_dict[key] = val
-    return peft_model_state_dict
-
-
-def save_video(tensor, path):
-    from torchvision.io import write_video
-
-    tensor = tensor * 255.0
-    tensor = tensor.permute(0, 2, 3, 1)
-    tensor = tensor.clamp(0, 255).byte()
-    write_video(path, tensor, 4, video_codec="h264")
-
-
-def load_model_from_checkpoint(model, checkpoint_dir, checkpoint_step):
-    checkpoint_paths = glob.glob(os.path.join(checkpoint_dir, "checkpoint-*"))
-    checkpoint_paths.sort(key=lambda x: int(x.split("-")[-1]), reverse=True)
-
-    if checkpoint_step is None or checkpoint_step == -1:
-        # get the latest checkpoint
-        checkpoint_path = checkpoint_paths[0]
-        print(f"===> Checkpoint step is not provided, using the latest checkpoint: {checkpoint_path}")
-    else:
-        checkpoint_path = os.path.join(checkpoint_dir, f"checkpoint-{checkpoint_step}")
-        if checkpoint_path not in checkpoint_paths:
-            checkpoint_path = checkpoint_paths[0]
-            print(f"===> Checkpoint step {checkpoint_step} not found, using the latest checkpoint: {checkpoint_path}")
-        else:
-            print(f"===> Checkpoint step {checkpoint_step} found, using the specified checkpoint: {checkpoint_path}")
-
-    checkpoint_step = checkpoint_path.split("checkpoint-")[-1].split("/")[0]
-
-    full_ckpt = os.path.join(checkpoint_path, "model.pth")
-    lora_ckpt = os.path.join(checkpoint_path, "adapter_model.safetensors")
-    non_lora_ckpt = os.path.join(checkpoint_path, "non_lora_state_dict.pth")
-    if os.path.exists(full_ckpt):
-        model_state_dict = torch.load(full_ckpt, map_location="cpu")
-        model.load_state_dict(model_state_dict)
-    else:
-        lora_state_dict = safetensors.torch.load_file(lora_ckpt)
-        non_lora_state_dict = torch.load(non_lora_ckpt, map_location="cpu")
-
-        lora_state_dict = _insert_adapter_name_into_state_dict(lora_state_dict,
-                                                               adapter_name="default",
-                                                               parameter_prefix="lora_")
-
-        model_state_dict = model.state_dict()
-        model_state_dict.update(non_lora_state_dict)
-        model_state_dict.update(lora_state_dict)
-        model.load_state_dict(model_state_dict)
-
-    return model, checkpoint_step
-
-
-def find_target_linear_names(model, num_lora_modules=-1, lora_namespan_exclude=None, verbose=False):
-    """
-    Find the target linear modules for LoRA.
-    """
-    linear_cls = torch.nn.Linear
-    embedding_cls = torch.nn.Embedding
-    if lora_namespan_exclude is None:
-        lora_namespan_exclude = []
-    lora_module_names = []
-
-    for name, module in model.named_modules():
-        if any(ex_keyword in name for ex_keyword in lora_namespan_exclude):
-            # print(f"Excluding module: {name}")
-            continue
-
-        if isinstance(module, linear_cls | embedding_cls):
-            lora_module_names.append(name)
-
-    if num_lora_modules > 0:
-        lora_module_names = lora_module_names[-num_lora_modules:]
-    if verbose:
-        print(f"Found {len(lora_module_names)} lora modules: {lora_module_names}")
-    return lora_module_names

fastvideo/train/methods/rl/reward/VideoAlign/__init__.py

@@ -1,17 +0,0 @@
-"""Vendored runtime subset of VideoAlign.
-
-Source: https://github.com/KlingAIResearch/VideoAlign
-Commit: 219ab9db64c045e5181a2202d11f686439351292
-Purpose: Runtime reward inference integration for FastVideo GenRL.
-
-This is temporary minimal vendoring for PR integration. It is expected to be
-cleaned up and normalized later.
-
-Porting rules:
-- Include only files required by the runtime import closure used by FastVideo.
-- When an upstream file is required, copy the entire file faithfully.
-- Only adjust imports as needed to make the vendored code import through package
-  paths instead of sys.path mutation or ambiguous top-level imports.
-- Do not perform style, typing, or behavioral cleanup as part of this vendoring
-  step.
-"""

fastvideo/train/methods/rl/reward/VideoAlign/checkpoints/README.md

@@ -1,8 +0,0 @@
-Please download our checkpoints from [Huggingface](https://huggingface.co/KwaiVGI/VideoReward) and put it in `./checkpoints/`.
-
-```bash
-cd checkpoints
-git lfs install
-git clone https://huggingface.co/KwaiVGI/VideoReward
-cd ..
-```

fastvideo/train/methods/rl/reward/VideoAlign/data.py

@@ -1,280 +0,0 @@
-from dataclasses import dataclass
-
-import torch
-from .prompt_template import build_prompt
-# from qwen_vl_utils import process_vision_info
-from .vision_process import process_vision_info
-
-
-@dataclass
-class DataConfig:
-    meta_data: str = "/path/to/dataset/meta_data.csv"
-    data_dir: str = "/path/to/dataset"
-    meta_data_test: str = None
-    max_frame_pixels: int = 240 * 320
-    num_frames: float = None
-    fps: float = 2.0
-    p_shuffle_frames: float = 0.0
-    p_color_jitter: float = 0.0
-    eval_dim: str | list[str] = "VQ"
-    prompt_template_type: str = "none"
-    add_noise: bool = False
-    sample_type: str = "uniform"
-    use_tied_data: bool = True
-
-
-def convert_GSB_csv_to_reward_data(example,
-                                   data_dir,
-                                   eval_dims=None,
-                                   max_pixels=448 * 448,
-                                   fps=2.0,
-                                   num_frames=None,
-                                   prompt_template_type="none",
-                                   sample_type="uniform"):
-    """
-    Convert Good/Same/Bad csv data to reward data.
-
-    Args:
-        example (dict): A dataframe containing the GSB csv data.
-        data_dir (str): The directory path to the video files.
-        eval_dim (str): The dimension to evaluate ("VQ"/"MQ"/"TA").
-        max_pixels (int): The maximum number of pixels allowed for videos.
-        num_frames (float): Number of frames.
-        prompt_template_type (str): The type of prompt template to use ("none"/"simple"/"video_score").
-
-    Returns:
-        dict: A dictionary containing the reward data.
-    """
-    if eval_dims is None:
-        eval_dims = ["VQ"]
-
-    A_data = [{
-        "role":
-        "user",
-        "content": [
-            {
-                "type": "video",
-                "video": f"file://{data_dir}/{example['path_A']}",
-                "max_pixels": max_pixels,
-                "fps": fps if num_frames is None else None,
-                "nframes": min(num_frames, example["num_frames_A"]) if num_frames is not None else None,
-                "sample_type": sample_type,
-            },
-            {
-                "type": "text",
-                "text": build_prompt(example["prompt"], eval_dims, prompt_template_type)
-            },
-        ],
-    }]
-    B_data = [{
-        "role":
-        "user",
-        "content": [
-            {
-                "type": "video",
-                "video": f"file://{data_dir}/{example['path_B']}",
-                "max_pixels": max_pixels,
-                "fps": fps if num_frames is None else None,
-                "nframes": min(num_frames, example["num_frames_B"]) if num_frames is not None else None,
-                "sample_type": sample_type,
-            },
-            {
-                "type": "text",
-                "text": build_prompt(example["prompt"], eval_dims, prompt_template_type)
-            },
-        ],
-    }]
-
-    chosen_labels = []
-    A_scores = []
-    B_scores = []
-
-    for eval_dim in eval_dims:
-        ### chosen_label: 1 if A is chosen, -1 if B is chosen, 0 if tied.
-        ### 22 if invalid. ooaaeeaa o.O
-        try:
-            if example[f"{eval_dim}"] is not None:
-                if example[f"{eval_dim}"] == "A":
-                    chosen_label = 1
-                elif example[f"{eval_dim}"] == "B":
-                    chosen_label = -1
-                elif example[f"{eval_dim}"] == "same":
-                    chosen_label = 0
-                elif example[f"{eval_dim}"] == "invalid":
-                    chosen_label = 22
-                else:
-                    chosen_label = 22
-            else:
-                chosen_label = 22
-        except Exception:
-            chosen_label = 22
-
-        chosen_labels.append(chosen_label)
-        if f"MOS_A_{eval_dim}" in example and f"MOS_B_{eval_dim}" in example:
-            try:
-                A_score = example[f"MOS_A_{eval_dim}"] if example[f"MOS_A_{eval_dim}"] is not None else 0.0
-                B_score = example[f"MOS_B_{eval_dim}"] if example[f"MOS_B_{eval_dim}"] is not None else 0.0
-            except Exception:
-                A_score = 0.0
-                B_score = 0.0
-            A_scores.append(A_score)
-            B_scores.append(B_score)
-        else:
-            A_scores.append(0.0)
-            B_scores.append(0.0)
-
-    chosen_labels = torch.tensor(chosen_labels, dtype=torch.long)
-    A_scores = torch.tensor(A_scores, dtype=torch.float)
-    B_scores = torch.tensor(B_scores, dtype=torch.float)
-    metainfo_idx = None
-    if 'metainfo_idx' in example:
-        metainfo_idx = example['metainfo_idx']
-
-    return {
-        "A_data": A_data,
-        "B_data": B_data,
-        "A_scores": A_scores,
-        "B_scores": B_scores,
-        "chosen_label": chosen_labels,
-        "metainfo_idx": metainfo_idx,
-    }
-
-
-class QWen2VLDataCollator:
-
-    def __init__(self, processor, add_noise=False, p_shuffle_frames=0.0, p_color_jitter=0.0):
-        self.processor = processor
-        self.add_noise = add_noise
-        self.set_noise_step = None
-
-        self.p_shuffle_frames = p_shuffle_frames
-        self.p_color_jitter = p_color_jitter
-
-        self.noise_adder = None
-
-    def _clean_message(self, message):
-        """
-        remove unnecessary keys from message(very very necessary)
-        """
-        message_content = message[0]["content"][0]
-        out_message = [{
-            "role":
-            "user",
-            "content": [
-                {
-                    "type": "video",
-                    "video": message_content["video"],
-                    "max_pixels": message_content["max_pixels"],
-                    "fps": message_content.get("fps", None),
-                    "nframes": message_content.get("nframes", None),
-                    "sample_type": message_content.get("sample_type", "uniform"),
-                },
-                {
-                    "type": "text",
-                    "text": message[0]["content"][1]["text"]
-                },
-            ],
-        }]
-
-        if out_message[0]["content"][0]["fps"] is None:
-            out_message[0]["content"][0].pop("fps")
-        if out_message[0]["content"][0]["nframes"] is None:
-            out_message[0]["content"][0].pop("nframes")
-
-        return out_message
-
-    def _pad_sequence(self, sequences, attention_mask, max_len, padding_side='right'):
-        """
-        Pad the sequences to the maximum length.
-        """
-        assert padding_side in ['right', 'left']
-        if sequences.shape[1] >= max_len:
-            return sequences, attention_mask
-
-        pad_len = max_len - sequences.shape[1]
-        padding = (0, pad_len) if padding_side == 'right' else (pad_len, 0)
-
-        sequences_padded = torch.nn.functional.pad(sequences, padding, 'constant',
-                                                   self.processor.tokenizer.pad_token_id)
-        attention_mask_padded = torch.nn.functional.pad(attention_mask, padding, 'constant', 0)
-
-        return sequences_padded, attention_mask_padded
-
-    def __call__(self, features, enable_noise=True):
-        """
-        Preprocess inputs to token sequences and return a batch
-        """
-        # try:
-        features_A = []
-        features_B = []
-        # check if we have a margin. If we do, we need to batch it as well
-        # has_margin = "margin" in features[0]
-        has_idx = "metainfo_idx" in features[0] and features[0]["metainfo_idx"] is not None
-
-        for idx, feature in enumerate(features):
-            features_A.append(self._clean_message(feature["A_data"]))
-            features_B.append(self._clean_message(feature["B_data"]))
-
-        # import pdb; pdb.set_trace()
-        image_inputs_A, video_inputs_A = process_vision_info(features_A)
-        image_inputs_B, video_inputs_B = process_vision_info(features_B)
-
-        video_inputs_A = [video_inputs_A[i].float() / 255.0 for i in range(len(video_inputs_A))]
-        video_inputs_B = [video_inputs_B[i].float() / 255.0 for i in range(len(video_inputs_B))]
-        do_rescale = False
-        # print(f"{video_inputs_A[0].shape}, {video_inputs_B[0].shape}")
-
-        # if not enable_noise:
-        #     print("Not training, no noise added.")
-        batch_A = self.processor(
-            text=self.processor.apply_chat_template(features_A, tokenize=False, add_generation_prompt=True),
-            images=image_inputs_A,
-            videos=video_inputs_A,
-            padding=True,
-            return_tensors="pt",
-            videos_kwargs={"do_rescale": do_rescale},
-        )
-        batch_B = self.processor(
-            text=self.processor.apply_chat_template(features_B, tokenize=False, add_generation_prompt=True),
-            images=image_inputs_B,
-            videos=video_inputs_B,
-            padding=True,
-            return_tensors="pt",
-            videos_kwargs={"do_rescale": do_rescale},
-        )
-
-        # pdb.set_trace()
-        max_len = max(batch_A["input_ids"].shape[1], batch_B["input_ids"].shape[1])
-        batch_A["input_ids"], batch_A["attention_mask"] = self._pad_sequence(batch_A["input_ids"],
-                                                                             batch_A["attention_mask"], max_len,
-                                                                             "right")
-        batch_B["input_ids"], batch_B["attention_mask"] = self._pad_sequence(batch_B["input_ids"],
-                                                                             batch_B["attention_mask"], max_len,
-                                                                             "right")
-        # print(f"Batch A: {batch_A['input_ids'].shape}, Batch B: {batch_B['input_ids'].shape}")
-
-        chosen_label = torch.stack([torch.tensor(feature["chosen_label"]) for feature in features])
-
-        A_scores = torch.stack([torch.tensor(feature["A_scores"]) for feature in features])
-        B_scores = torch.stack([torch.tensor(feature["B_scores"]) for feature in features])
-
-        batch = {
-            "A": batch_A,
-            "B": batch_B,
-            "return_loss": True,
-            "chosen_label": chosen_label,
-            "A_scores": A_scores,
-            "B_scores": B_scores,
-        }
-
-        if has_idx:
-            metainfo_idx = torch.stack([torch.tensor(feature["metainfo_idx"]) for feature in features])
-            batch["metainfo_idx"] = metainfo_idx
-
-        # pdb.set_trace()
-        return batch
-
-        # except Exception as e:
-        #     print(f"Error processing batch: {e} in reading.")
-        #     # get next batch
-        #     return None

fastvideo/train/methods/rl/reward/VideoAlign/inference.py

@@ -1,237 +0,0 @@
-import json
-import os
-from collections.abc import Mapping
-
-import torch
-from .vision_process import process_vision_info
-
-from .data import DataConfig
-from .utils import ModelConfig, PEFTLoraConfig, TrainingConfig
-from .utils import load_model_from_checkpoint
-from .train_reward import create_model_and_processor
-from .prompt_template import build_prompt
-
-
-def load_configs_from_json(config_path):
-    with open(config_path) as f:
-        config_dict = json.load(f)
-
-    # del config_dict["training_args"]["_n_gpu"]
-    del config_dict["data_config"]["meta_data"]
-    del config_dict["data_config"]["data_dir"]
-
-    return config_dict["data_config"], None, config_dict["model_config"], config_dict["peft_lora_config"], \
-           config_dict.get("inference_config", None)
-
-
-class VideoVLMRewardInference:
-
-    def __init__(self, load_from_pretrained, load_from_pretrained_step=-1, device='cuda', dtype=torch.bfloat16):
-        config_path = os.path.join(load_from_pretrained, "model_config.json")
-        data_config, _, model_config, peft_lora_config, inference_config = load_configs_from_json(config_path)
-        data_config = DataConfig(**data_config)
-        model_config = ModelConfig(**model_config)
-        peft_lora_config = PEFTLoraConfig(**peft_lora_config)
-
-        training_args = TrainingConfig(
-            load_from_pretrained=load_from_pretrained,
-            load_from_pretrained_step=load_from_pretrained_step,
-            gradient_checkpointing=False,
-            disable_flash_attn2=False,
-            bf16=dtype == torch.bfloat16,
-            fp16=dtype == torch.float16,
-            output_dir="",
-        )
-
-        model, processor, peft_config = create_model_and_processor(
-            model_config=model_config,
-            peft_lora_config=peft_lora_config,
-            training_args=training_args,
-        )
-
-        self.device = device
-
-        model, checkpoint_step = load_model_from_checkpoint(model, load_from_pretrained, load_from_pretrained_step)
-        model.eval()
-
-        self.model = model
-        self.processor = processor
-
-        self.model.to(self.device)
-
-        self.data_config = data_config
-
-        self.inference_config = inference_config
-
-    def _norm(self, reward):
-        if self.inference_config is None:
-            return reward
-        else:
-            reward['VQ'] = (reward['VQ'] - self.inference_config['VQ_mean']) / self.inference_config['VQ_std']
-            reward['MQ'] = (reward['MQ'] - self.inference_config['MQ_mean']) / self.inference_config['MQ_std']
-            reward['TA'] = (reward['TA'] - self.inference_config['TA_mean']) / self.inference_config['TA_std']
-            return reward
-
-    def _pad_sequence(self, sequences, attention_mask, max_len, padding_side='right'):
-        """
-        Pad the sequences to the maximum length.
-        """
-        assert padding_side in ['right', 'left']
-        if sequences.shape[1] >= max_len:
-            return sequences, attention_mask
-
-        pad_len = max_len - sequences.shape[1]
-        padding = (0, pad_len) if padding_side == 'right' else (pad_len, 0)
-
-        sequences_padded = torch.nn.functional.pad(sequences, padding, 'constant',
-                                                   self.processor.tokenizer.pad_token_id)
-        attention_mask_padded = torch.nn.functional.pad(attention_mask, padding, 'constant', 0)
-
-        return sequences_padded, attention_mask_padded
-
-    def _prepare_input(self, data):
-        """
-        Prepare `inputs` before feeding them to the model, converting them to tensors if they are not already and
-        handling potential state.
-        """
-        if isinstance(data, Mapping):
-            return type(data)({k: self._prepare_input(v) for k, v in data.items()})
-        elif isinstance(data, tuple | list):
-            return type(data)(self._prepare_input(v) for v in data)
-        elif isinstance(data, torch.Tensor):
-            kwargs = {"device": self.device}
-            ## TODO: Maybe need to add dtype
-            # if self.is_deepspeed_enabled and (torch.is_floating_point(data) or torch.is_complex(data)):
-            #     # NLP models inputs are int/uint and those get adjusted to the right dtype of the
-            #     # embedding. Other models such as wav2vec2's inputs are already float and thus
-            #     # may need special handling to match the dtypes of the model
-            #     kwargs.update({"dtype": self.accelerator.state.deepspeed_plugin.hf_ds_config.dtype()})
-            return data.to(**kwargs)
-        return data
-
-    def _prepare_inputs(self, inputs):
-        """
-        Prepare `inputs` before feeding them to the model, converting them to tensors if they are not already and
-        handling potential state.
-        """
-        inputs = self._prepare_input(inputs)
-        if len(inputs) == 0:
-            raise ValueError
-        return inputs
-
-    def prepare_batch(
-        self,
-        video_paths,
-        prompts,
-        fps=None,
-        num_frames=None,
-        max_pixels=None,
-    ):
-        fps = self.data_config.fps if fps is None else fps
-        num_frames = self.data_config.num_frames if num_frames is None else num_frames
-        max_pixels = self.data_config.max_frame_pixels if max_pixels is None else max_pixels
-
-        if num_frames is None:
-            chat_data = [[
-                {
-                    "role":
-                    "user",
-                    "content": [
-                        {
-                            "type": "video",
-                            "video": f"file://{video_path}",
-                            "max_pixels": max_pixels,
-                            "fps": fps,
-                            "sample_type": self.data_config.sample_type,
-                        },
-                        {
-                            "type": "text",
-                            "text": build_prompt(prompt, self.data_config.eval_dim,
-                                                 self.data_config.prompt_template_type)
-                        },
-                    ],
-                },
-            ] for video_path, prompt in zip(video_paths, prompts, strict=False)]
-        else:
-            chat_data = [[
-                {
-                    "role":
-                    "user",
-                    "content": [
-                        {
-                            "type": "video",
-                            "video": f"file://{video_path}",
-                            "max_pixels": max_pixels,
-                            "nframes": num_frames,
-                            "sample_type": self.data_config.sample_type,
-                        },
-                        {
-                            "type": "text",
-                            "text": build_prompt(prompt, self.data_config.eval_dim,
-                                                 self.data_config.prompt_template_type)
-                        },
-                    ],
-                },
-            ] for video_path, prompt in zip(video_paths, prompts, strict=False)]
-        image_inputs, video_inputs = process_vision_info(chat_data)
-
-        batch = self.processor(
-            text=self.processor.apply_chat_template(chat_data, tokenize=False, add_generation_prompt=True),
-            images=image_inputs,
-            videos=video_inputs,
-            padding=True,
-            return_tensors="pt",
-            videos_kwargs={"do_rescale": True},
-        )
-        batch = self._prepare_inputs(batch)
-        return batch
-
-    def reward(self, video_paths, prompts, fps=None, num_frames=None, max_pixels=None, use_norm=True):
-        """
-        Inputs:
-            video_paths: List[str], B paths of the videos.
-            prompts: List[str], B prompts for the videos.
-            eval_dims: List[str], N evaluation dimensions.
-            fps: float, sample rate of the videos. If None, use the default value in the config.
-            num_frames: int, number of frames of the videos. If None, use the default value in the config.
-            max_pixels: int, maximum pixels of the videos. If None, use the default value in the config.
-            use_norm: bool, whether to rescale the output rewards
-        Outputs:
-            Rewards: List[dict], N + 1 rewards of the B videos.
-        """
-        assert fps is None or num_frames is None, "fps and num_frames cannot be set at the same time."
-
-        batch = self.prepare_batch(video_paths, prompts, fps, num_frames, max_pixels)
-        rewards = self.model(return_dict=True, **batch)["logits"]
-
-        rewards = [{'VQ': reward[0].item(), 'MQ': reward[1].item(), 'TA': reward[2].item()} for reward in rewards]
-        for i in range(len(rewards)):
-            if use_norm:
-                rewards[i] = self._norm(rewards[i])
-            rewards[i]['Overall'] = rewards[i]['VQ'] + rewards[i]['MQ'] + rewards[i]['TA']
-
-        return rewards
-
-
-if __name__ == "__main__":
-    load_from_pretrained = "./checkpoints"
-    device = "cuda:0"
-    dtype = torch.bfloat16
-
-    inferencer = VideoVLMRewardInference(load_from_pretrained, device=device, dtype=dtype)
-
-    video_paths = [
-        "datasets/train/videos/example_1_A.mp4",
-        "datasets/train/videos/example_1_B.mp4",
-        "datasets/train/videos/example_2_A.mp4",
-    ]
-
-    prompts = [
-        "The camera remains still, a girl with braided hair and wearing a pink dress approached the chair in the room and sat on it, the background is a cozy bedroom, warm indoor lighting.",
-        "The camera remains still, a girl with braided hair and wearing a pink dress approached the chair in the room and sat on it, the background is a cozy bedroom, warm indoor lighting.",
-        "The camera follows a young explorer through an abandoned urban building at night, exploring hidden corridors and forgotten spaces, with a mix of light and shadow creating a mysterious atmosphere.",
-    ]
-
-    with torch.no_grad():
-        rewards = inferencer.reward(video_paths, prompts, use_norm=True)
-        print(rewards)