[chore] remove force_upcast from autoencoder flux2

[sayakpaul]

As discussed internally.

I have checked if the existing checkpoints can load successfully. Additionally, here's a script that one can run using HF Jobs that verifies that checkpoints can load successfully with changes, without any regression.

Script

import argparse
import inspect
import os
import subprocess
import tempfile
from pathlib import Path


CHECKPOINTS = (
    "black-forest-labs/FLUX.2-dev",
    "black-forest-labs/FLUX.2-klein-4B",
    "black-forest-labs/FLUX.2-klein-9B",
)


def check_checkpoints(phase, output_dir, device, dtype):
    import torch

    import diffusers
    from diffusers import AutoencoderKLFlux2

    torch.set_num_threads(1)
    torch.use_deterministic_algorithms(True)
    device = torch.device(device)
    dtype = getattr(torch, dtype)
    if device.type == "cuda":
        assert torch.cuda.is_available(), "CUDA is unavailable"
        if dtype == torch.bfloat16:
            assert torch.cuda.is_bf16_supported(), "GPU does not support bfloat16"
        print(f"GPU: {torch.cuda.get_device_name()}; CUDA: {torch.version.cuda}", flush=True)
    print(f"{phase}: device={device}; dtype={dtype}", flush=True)
    has_force_upcast = "force_upcast" in inspect.signature(AutoencoderKLFlux2.__init__).parameters
    assert has_force_upcast == (phase == "baseline"), (phase, has_force_upcast)
    print(f"{phase}: Diffusers {diffusers.__version__} from {diffusers.__file__}", flush=True)

    for checkpoint in CHECKPOINTS:
        print(f"{phase}: checking {checkpoint}", flush=True)
        config = AutoencoderKLFlux2.load_config(checkpoint, subfolder="vae")
        assert config["force_upcast"] is True
        model, loading_info = AutoencoderKLFlux2.from_pretrained(
            checkpoint, subfolder="vae", torch_dtype=dtype, output_loading_info=True
        )
        for key in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs"):
            assert not loading_info[key], (checkpoint, key, loading_info[key])

        model.to(device).eval()
        assert model.dtype == dtype
        sample = torch.randn(
            1, 3, 64, 64, generator=torch.Generator(device=device).manual_seed(0), device=device, dtype=dtype
        )
        with torch.inference_mode():
            latents = model.encode(sample).latent_dist.mode()
            decoded = model.decode(latents).sample
        assert latents.shape == (1, 32, 8, 8), latents.shape
        assert decoded.shape == sample.shape, decoded.shape
        assert latents.dtype == decoded.dtype == dtype
        assert torch.isfinite(latents).all()
        assert torch.isfinite(decoded).all()

        saved_path = output_dir / checkpoint.split("/")[-1]
        if phase == "baseline":
            saved_path.mkdir()
            torch.save({"latents": latents, "decoded": decoded}, saved_path / "outputs.pt")
            model.save_pretrained(saved_path / "checkpoint")
        else:
            expected = torch.load(saved_path / "outputs.pt", map_location=device, weights_only=True)
            torch.testing.assert_close(latents, expected["latents"], rtol=0, atol=0)
            torch.testing.assert_close(decoded, expected["decoded"], rtol=0, atol=0)
            reloaded = (
                AutoencoderKLFlux2.from_pretrained(saved_path / "checkpoint", torch_dtype=dtype).to(device).eval()
            )
            with torch.inference_mode():
                round_trip = reloaded(sample, sample_posterior=False).sample
            torch.testing.assert_close(round_trip, expected["decoded"], rtol=0, atol=0)

        print(f"PASS: {phase}: {checkpoint}", flush=True)


def compare_revisions(args):
    env = dict(os.environ)
    env.pop("PYTHONPATH", None)
    env["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
    with tempfile.TemporaryDirectory() as directory:
        python = str(Path(directory) / "venv" / "bin" / "python")
        subprocess.run(["uv", "venv", "--python", "3.12", str(Path(directory) / "venv")], env=env, check=True)
        subprocess.run(
            [
                "uv",
                "pip",
                "install",
                "--python",
                python,
                "torch>=2.6",
                *(
                    ["--torch-backend", "auto"]
                    if args.device == "cuda"
                    else ["--index-url", "https://download.pytorch.org/whl/cpu"]
                ),
            ],
            env=env,
            check=True,
        )
        for phase, revision in (("baseline", args.base_ref), ("candidate", args.branch_ref)):
            print(f"Installing Diffusers {revision} for {phase}", flush=True)
            dependencies = (
                ["accelerate>=0.31.0"] if phase == "baseline" else ["--no-deps", "--reinstall-package", "diffusers"]
            )
            subprocess.run(
                [
                    "uv",
                    "pip",
                    "install",
                    "--python",
                    python,
                    *dependencies,
                    f"diffusers @ https://github.com/huggingface/diffusers/archive/{revision}.tar.gz",
                ],
                env=env,
                check=True,
            )
            subprocess.run(
                [
                    python,
                    str(Path(__file__).resolve()),
                    "--phase",
                    phase,
                    "--output-dir",
                    directory,
                    "--device",
                    args.device,
                    "--dtype",
                    args.dtype,
                ],
                env=env,
                check=True,
            )
    print("PASS: All three checkpoints load and run with identical outputs on main and the branch.", flush=True)


def main():
    parser = argparse.ArgumentParser(description="Compare FLUX.2 VAE checkpoints on Diffusers main and a branch.")
    parser.add_argument("--submit", action="store_true", help="Submit the comparison to HF Jobs.")
    parser.add_argument("--base-ref", default="main")
    parser.add_argument("--branch-ref", default="autoencoderkl-flux2-remove-upcast")
    parser.add_argument("--flavor", default="cpu-basic")
    parser.add_argument("--timeout", default="30m")
    parser.add_argument("--device", choices=("cpu", "cuda"), default="cpu")
    parser.add_argument("--dtype", choices=("float32", "bfloat16"), default="float32")
    parser.add_argument("--phase", choices=("baseline", "candidate"), help=argparse.SUPPRESS)
    parser.add_argument("--output-dir", type=Path, help=argparse.SUPPRESS)
    args = parser.parse_args()

    if args.submit:
        from huggingface_hub import HfApi, get_token

        token = get_token()
        if token is None:
            raise ValueError("Authenticate with `hf auth login` or set HF_TOKEN before submitting a job.")
        job = HfApi(token=token).run_uv_job(
            str(Path(__file__).resolve()),
            script_args=[
                "--base-ref",
                args.base_ref,
                "--branch-ref",
                args.branch_ref,
                "--device",
                args.device,
                "--dtype",
                args.dtype,
            ],
            python="3.12",
            flavor=args.flavor,
            timeout=args.timeout,
            secrets={"HF_TOKEN": token},
            name=f"flux2-vae-main-vs-remove-upcast-{args.device}-{args.dtype}",
        )
        print(f"Job: {job.url}", flush=True)
        print(f"Logs: hf jobs logs {job.id}", flush=True)
        print(f"Status: hf jobs inspect {job.id}", flush=True)
    elif args.phase:
        if args.output_dir is None:
            parser.error("--output-dir is required with --phase.")
        check_checkpoints(args.phase, args.output_dir, args.device, args.dtype)
    else:
        compare_revisions(args)


if __name__ == "__main__":
    main()

Run script like so:

python check_flux2_vae_hf_jobs.py \
  --submit \
  --flavor l4x1 \
  --device cuda \
  --dtype bfloat16

HF Jobs link: https://huggingface.co/jobs/sayakpaul/6ac36e0b404719ba37653b95

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[yiyixuxu]

thanks!