models/flux/util.py

import getpass
import math
import os
from dataclasses import dataclass
from pathlib import Path

import requests
import torch
from einops import rearrange
from huggingface_hub import hf_hub_download, login
from PIL import ExifTags, Image
from safetensors.torch import load_file as load_sft

from .model import Flux, FluxLoraWrapper, FluxParams
from .modules.autoencoder import AutoEncoder, AutoEncoderParams
from .modules.conditioner import HFEmbedder
from shared.utils import files_locator as fl 
from shared.convert import convert_diffusers_to_flux

CHECKPOINTS_DIR = Path("checkpoints")

BFL_API_KEY = os.getenv("BFL_API_KEY")


def ensure_hf_auth():
    hf_token = os.environ.get("HF_TOKEN")
    if hf_token:
        print("Trying to authenticate to HuggingFace with the HF_TOKEN environment variable.")
        try:
            login(token=hf_token)
            print("Successfully authenticated with HuggingFace using HF_TOKEN")
            return True
        except Exception as e:
            print(f"Warning: Failed to authenticate with HF_TOKEN: {e}")

    if os.path.exists(os.path.expanduser("~/.cache/huggingface/token")):
        print("Already authenticated with HuggingFace")
        return True

    return False


def prompt_for_hf_auth():
    try:
        token = getpass.getpass("HF Token (hidden input): ").strip()
        if not token:
            print("No token provided. Aborting.")
            return False

        login(token=token)
        print("Successfully authenticated!")
        return True
    except KeyboardInterrupt:
        print("\nAuthentication cancelled by user.")
        return False
    except Exception as auth_e:
        print(f"Authentication failed: {auth_e}")
        print("Tip: You can also run 'huggingface-cli login' or set HF_TOKEN environment variable")
        return False


def get_checkpoint_path(repo_id: str, filename: str, env_var: str) -> Path:
    """Get the local path for a checkpoint file, downloading if necessary."""
    # if os.environ.get(env_var) is not None:
    #     local_path = os.environ[env_var]
    #     if os.path.exists(local_path):
    #         return Path(local_path)

    #     print(
    #         f"Trying to load model {repo_id}, {filename} from environment "
    #         f"variable {env_var}. But file {local_path} does not exist. "
    #         "Falling back to default location."
    #     )

    # # Create a safe directory name from repo_id
    # safe_repo_name = repo_id.replace("/", "_")
    # checkpoint_dir = CHECKPOINTS_DIR / safe_repo_name
    # checkpoint_dir.mkdir(exist_ok=True)

    # local_path = checkpoint_dir / filename

    local_path = filename
    from mmgp import offload

    if False:
        print(f"Downloading {filename} from {repo_id} to {local_path}")
        try:
            ensure_hf_auth()
            hf_hub_download(repo_id=repo_id, filename=filename, local_dir=checkpoint_dir)
        except Exception as e:
            if "gated repo" in str(e).lower() or "restricted" in str(e).lower():
                print(f"\nError: Cannot access {repo_id} -- this is a gated repository.")

                # Try one more time to authenticate
                if prompt_for_hf_auth():
                    # Retry the download after authentication
                    print("Retrying download...")
                    hf_hub_download(repo_id=repo_id, filename=filename, local_dir=checkpoint_dir)
                else:
                    print("Authentication failed or cancelled.")
                    print("You can also run 'huggingface-cli login' or set HF_TOKEN environment variable")
                    raise RuntimeError(f"Authentication required for {repo_id}")
            else:
                raise e

    return local_path


def download_onnx_models_for_trt(model_name: str, trt_transformer_precision: str = "bf16") -> str | None:
    """Download ONNX models for TRT to our checkpoints directory"""
    onnx_repo_map = {
        "flux-dev": "black-forest-labs/FLUX.1-dev-onnx",
        "flux-schnell": "black-forest-labs/FLUX.1-schnell-onnx",
        "flux-dev-canny": "black-forest-labs/FLUX.1-Canny-dev-onnx",
        "flux-dev-depth": "black-forest-labs/FLUX.1-Depth-dev-onnx",
        "flux-dev-redux": "black-forest-labs/FLUX.1-Redux-dev-onnx",
        "flux-dev-fill": "black-forest-labs/FLUX.1-Fill-dev-onnx",
        "flux-dev-kontext": "black-forest-labs/FLUX.1-Kontext-dev-onnx",
    }

    if model_name not in onnx_repo_map:
        return None  # No ONNX repository required for this model

    repo_id = onnx_repo_map[model_name]
    safe_repo_name = repo_id.replace("/", "_")
    onnx_dir = CHECKPOINTS_DIR / safe_repo_name

    # Map of module names to their ONNX file paths (using specified precision)
    onnx_file_map = {
        "clip": "clip.opt/model.onnx",
        "transformer": f"transformer.opt/{trt_transformer_precision}/model.onnx",
        "transformer_data": f"transformer.opt/{trt_transformer_precision}/backbone.onnx_data",
        "t5": "t5.opt/model.onnx",
        "t5_data": "t5.opt/backbone.onnx_data",
        "vae": "vae.opt/model.onnx",
    }

    # If all files exist locally, return the custom_onnx_paths format
    if onnx_dir.exists():
        all_files_exist = True
        custom_paths = []
        for module, onnx_file in onnx_file_map.items():
            if module.endswith("_data"):
                continue  # Skip data files
            local_path = onnx_dir / onnx_file
            if not local_path.exists():
                all_files_exist = False
                break
            custom_paths.append(f"{module}:{local_path}")

        if all_files_exist:
            print(f"ONNX models ready in {onnx_dir}")
            return ",".join(custom_paths)

    # If not all files exist, download them
    print(f"Downloading ONNX models from {repo_id} to {onnx_dir}")
    print(f"Using transformer precision: {trt_transformer_precision}")
    onnx_dir.mkdir(exist_ok=True)

    # Download all ONNX files
    for module, onnx_file in onnx_file_map.items():
        local_path = onnx_dir / onnx_file
        if local_path.exists():
            continue  # Already downloaded

        # Create parent directories
        local_path.parent.mkdir(parents=True, exist_ok=True)

        try:
            print(f"Downloading {onnx_file}")
            hf_hub_download(repo_id=repo_id, filename=onnx_file, local_dir=onnx_dir)
        except Exception as e:
            if "does not exist" in str(e).lower() or "not found" in str(e).lower():
                continue
            elif "gated repo" in str(e).lower() or "restricted" in str(e).lower():
                print(f"Cannot access {repo_id} - requires license acceptance")
                print("Please follow these steps:")
                print(f"   1. Visit: https://huggingface.co/{repo_id}")
                print("   2. Log in to your HuggingFace account")
                print("   3. Accept the license terms and conditions")
                print("   4. Then retry this command")
                raise RuntimeError(f"License acceptance required for {model_name}")
            else:
                # Re-raise other errors
                raise

    print(f"ONNX models ready in {onnx_dir}")

    # Return the custom_onnx_paths format that TRT expects: "module1:path1,module2:path2"
    # Note: Only return the actual module paths, not the data file
    custom_paths = []
    for module, onnx_file in onnx_file_map.items():
        if module.endswith("_data"):
            continue  # Skip the data file in the return paths
        full_path = onnx_dir / onnx_file
        if full_path.exists():
            custom_paths.append(f"{module}:{full_path}")

    return ",".join(custom_paths)


def check_onnx_access_for_trt(model_name: str, trt_transformer_precision: str = "bf16") -> str | None:
    """Check ONNX access and download models for TRT - returns ONNX directory path"""
    return download_onnx_models_for_trt(model_name, trt_transformer_precision)


def track_usage_via_api(name: str, n=1) -> None:
    """
    Track usage of licensed models via the BFL API for commercial licensing compliance.

    For more information on licensing BFL's models for commercial use and usage reporting,
    see the README.md or visit: https://dashboard.bfl.ai/licensing/subscriptions?showInstructions=true
    """
    assert BFL_API_KEY is not None, "BFL_API_KEY is not set"

    model_slug_map = {
        "flux-dev": "flux-1-dev",
        "flux-dev-kontext": "flux-1-kontext-dev",
        "flux-dev-fill": "flux-tools",
        "flux-dev-depth": "flux-tools",
        "flux-dev-canny": "flux-tools",
        "flux-dev-canny-lora": "flux-tools",
        "flux-dev-depth-lora": "flux-tools",
        "flux-dev-redux": "flux-tools",
    }

    if name not in model_slug_map:
        print(f"Skipping tracking usage for {name}, as it cannot be tracked. Please check the model name.")
        return

    model_slug = model_slug_map[name]
    url = f"https://api.bfl.ai/v1/licenses/models/{model_slug}/usage"
    headers = {"x-key": BFL_API_KEY, "Content-Type": "application/json"}
    payload = {"number_of_generations": n}

    response = requests.post(url, headers=headers, json=payload)
    if response.status_code != 200:
        raise Exception(f"Failed to track usage: {response.status_code} {response.text}")
    else:
        print(f"Successfully tracked usage for {name} with {n} generations")


def save_image(
    nsfw_classifier,
    name: str,
    output_name: str,
    idx: int,
    x: torch.Tensor,
    add_sampling_metadata: bool,
    prompt: str,
    nsfw_threshold: float = 0.85,
    track_usage: bool = False,
) -> int:
    fn = output_name.format(idx=idx)
    print(f"Saving {fn}")
    # bring into PIL format and save
    x = x.clamp(-1, 1)
    x = rearrange(x[0], "c h w -> h w c")
    img = Image.fromarray((127.5 * (x + 1.0)).cpu().byte().numpy())
    
    if nsfw_classifier is not None:
        nsfw_score = [x["score"] for x in nsfw_classifier(img) if x["label"] == "nsfw"][0]
    else:
        nsfw_score = nsfw_threshold - 1.0

    if nsfw_score < nsfw_threshold:
        exif_data = Image.Exif()
        if name in ["flux-dev", "flux-schnell"]:
            exif_data[ExifTags.Base.Software] = "AI generated;txt2img;flux"
        else:
            exif_data[ExifTags.Base.Software] = "AI generated;img2img;flux"
        exif_data[ExifTags.Base.Make] = "Black Forest Labs"
        exif_data[ExifTags.Base.Model] = name
        if add_sampling_metadata:
            exif_data[ExifTags.Base.ImageDescription] = prompt
        img.save(fn, exif=exif_data, quality=95, subsampling=0)
        if track_usage:
            track_usage_via_api(name, 1)
        idx += 1
    else:
        print("Your generated image may contain NSFW content.")

    return idx


@dataclass
class ModelSpec:
    params: FluxParams
    ae_params: AutoEncoderParams
    repo_id: str
    repo_flow: str
    repo_ae: str
    lora_repo_id: str | None = None
    lora_filename: str | None = None


configs = {
    "flux2-dev": ModelSpec(
        repo_id="",
        repo_flow="",
        repo_ae="ckpts/flux2_vae.safetensors",
        params=FluxParams(
            in_channels=128,
            out_channels=128,
            vec_in_dim=1,
            context_in_dim=15360,
            hidden_size=6144,
            mlp_ratio=3.0,
            single_linear1_mlp_ratio=6.0,
            single_mlp_hidden_ratio=3.0,
            double_mlp_ratio=3.0,
            double_linear1_mlp_ratio=6.0,
            num_heads=48,
            depth=8,
            depth_single_blocks=48,
            axes_dim=[32, 32, 32, 32],
            theta=2000,
            qkv_bias=False,
            guidance_embed=True,
            flux2=True,
        ),
        ae_params=AutoEncoderParams(
            resolution=1024,
            in_channels=3,
            ch=128,
            out_ch=3,
            ch_mult=[1, 2, 4, 4],
            num_res_blocks=2,
            z_channels=32,
            scale_factor=0.5,
            shift_factor=0.0,
        ),
    ),
    "flux2-klein-4b": ModelSpec(
        repo_id="",
        repo_flow="",
        repo_ae="ckpts/flux2_vae.safetensors",
        params=FluxParams(
            in_channels=128,
            out_channels=128,
            vec_in_dim=1,
            context_in_dim=7680,
            hidden_size=3072,
            mlp_ratio=3.0,
            single_linear1_mlp_ratio=6.0,
            single_mlp_hidden_ratio=3.0,
            double_mlp_ratio=3.0,
            double_linear1_mlp_ratio=6.0,
            num_heads=24,
            depth=5,
            depth_single_blocks=20,
            axes_dim=[32, 32, 32, 32],
            theta=2000,
            qkv_bias=False,
            guidance_embed=False,
            flux2=True,
        ),
        ae_params=AutoEncoderParams(
            resolution=1024,
            in_channels=3,
            ch=128,
            out_ch=3,
            ch_mult=[1, 2, 4, 4],
            num_res_blocks=2,
            z_channels=32,
            scale_factor=0.5,
            shift_factor=0.0,
        ),
    ),
    "flux2-klein-9b": ModelSpec(
        repo_id="",
        repo_flow="",
        repo_ae="ckpts/flux2_vae.safetensors",
        params=FluxParams(
            in_channels=128,
            out_channels=128,
            vec_in_dim=1,
            context_in_dim=12288,
            hidden_size=4096,
            mlp_ratio=3.0,
            single_linear1_mlp_ratio=6.0,
            single_mlp_hidden_ratio=3.0,
            double_mlp_ratio=3.0,
            double_linear1_mlp_ratio=6.0,
            num_heads=32,
            depth=8,
            depth_single_blocks=24,
            axes_dim=[32, 32, 32, 32],
            theta=2000,
            qkv_bias=False,
            guidance_embed=False,
            flux2=True,
        ),
        ae_params=AutoEncoderParams(
            resolution=1024,
            in_channels=3,
            ch=128,
            out_ch=3,
            ch_mult=[1, 2, 4, 4],
            num_res_blocks=2,
            z_channels=32,
            scale_factor=0.5,
            shift_factor=0.0,
        ),
    ),
    "pi-flux2": ModelSpec(
        repo_id="",
        repo_flow="",
        repo_ae="ckpts/flux2_vae.safetensors",
        params=FluxParams(
            in_channels=128,
            out_channels=128,
            vec_in_dim=1,
            context_in_dim=15360,
            hidden_size=6144,
            mlp_ratio=3.0,
            single_linear1_mlp_ratio=6.0,
            single_mlp_hidden_ratio=3.0,
            double_mlp_ratio=3.0,
            double_linear1_mlp_ratio=6.0,
            num_heads=48,
            depth=8,
            depth_single_blocks=48,
            axes_dim=[32, 32, 32, 32],
            theta=2000,
            qkv_bias=False,
            guidance_embed=True,
            flux2=True,
            piflow=True,
        ),
        ae_params=AutoEncoderParams(
            resolution=1024,
            in_channels=3,
            ch=128,
            out_ch=3,
            ch_mult=[1, 2, 4, 4],
            num_res_blocks=2,
            z_channels=32,
            scale_factor=0.5,
            shift_factor=0.0,
        ),
    ),
    "flux-dev": ModelSpec(
        repo_id="",
        repo_flow="",
        repo_ae="ckpts/flux_vae.safetensors",
        params=FluxParams(
            in_channels=64,
            out_channels=64,
            vec_in_dim=768,
            context_in_dim=4096,
            hidden_size=3072,
            mlp_ratio=4.0,
            num_heads=24,
            depth=19,
            depth_single_blocks=38,
            axes_dim=[16, 56, 56],
            theta=10_000,
            qkv_bias=True,
            guidance_embed=True,
        ),
        ae_params=AutoEncoderParams(
            resolution=256,
            in_channels=3,
            ch=128,
            out_ch=3,
            ch_mult=[1, 2, 4, 4],
            num_res_blocks=2,
            z_channels=16,
            scale_factor=0.3611,
            shift_factor=0.1159,
        ),
    ),
    "flux-schnell": ModelSpec(
        repo_id="black-forest-labs/FLUX.1-schnell",
        repo_flow="",
        repo_ae="ckpts/flux_vae.safetensors",
        params=FluxParams(
            in_channels=64,
            out_channels=64,
            vec_in_dim=768,
            context_in_dim=4096,
            hidden_size=3072,
            mlp_ratio=4.0,
            num_heads=24,
            depth=19,
            depth_single_blocks=38,
            axes_dim=[16, 56, 56],
            theta=10_000,
            qkv_bias=True,
            guidance_embed=False,
        ),
        ae_params=AutoEncoderParams(
            resolution=256,
            in_channels=3,
            ch=128,
            out_ch=3,
            ch_mult=[1, 2, 4, 4],
            num_res_blocks=2,
            z_channels=16,
            scale_factor=0.3611,
            shift_factor=0.1159,
        ),
    ),
    "flux-chroma": ModelSpec(
        repo_id="lodestones/Chroma1-HD",
        repo_flow="",
        repo_ae="ckpts/flux_vae.safetensors",
        params=FluxParams(
            in_channels=64,
            out_channels=64,
            vec_in_dim=768,
            context_in_dim=4096,
            hidden_size=3072,
            mlp_ratio=4.0,
            num_heads=24,
            depth=19,
            depth_single_blocks=38,
            axes_dim=[16, 56, 56],
            theta=10_000,
            qkv_bias=True,
            guidance_embed=False,
            chroma=True,
        ),
        ae_params=AutoEncoderParams(
            resolution=256,
            in_channels=3,
            ch=128,
            out_ch=3,
            ch_mult=[1, 2, 4, 4],
            num_res_blocks=2,
            z_channels=16,
            scale_factor=0.3611,
            shift_factor=0.1159,
        ),
    ),
    "flux-chroma-radiance": ModelSpec(
        repo_id="lodestones/Chroma1-Radiance",
        repo_flow="",
        repo_ae="ckpts/flux_vae.safetensors",
        params=FluxParams(
            in_channels=64,
            out_channels=3,
            vec_in_dim=768,
            context_in_dim=4096,
            hidden_size=3072,
            mlp_ratio=4.0,
            num_heads=24,
            depth=19,
            depth_single_blocks=38,
            axes_dim=[16, 56, 56],
            theta=10_000,
            qkv_bias=True,
            guidance_embed=False,
            chroma=True,
            radiance=True,
            radiance_patch_size=16,
            radiance_hidden_size=64,
            radiance_mlp_ratio=4,
            radiance_depth=4,
            radiance_max_freqs=8,
            radiance_tile_size=0,
            radiance_final_head_type="conv",
        ),
        ae_params=AutoEncoderParams(
            resolution=256,
            in_channels=3,
            ch=128,
            out_ch=3,
            ch_mult=[1, 2, 4, 4],
            num_res_blocks=2,
            z_channels=16,
            scale_factor=0.3611,
            shift_factor=0.1159,
        ),
    ),
    "flux-dev-canny": ModelSpec(
        repo_id="black-forest-labs/FLUX.1-Canny-dev",
        repo_flow="",
        repo_ae="ckpts/flux_vae.safetensors",
        params=FluxParams(
            in_channels=128,
            out_channels=64,
            vec_in_dim=768,
            context_in_dim=4096,
            hidden_size=3072,
            mlp_ratio=4.0,
            num_heads=24,
            depth=19,
            depth_single_blocks=38,
            axes_dim=[16, 56, 56],
            theta=10_000,
            qkv_bias=True,
            guidance_embed=True,
        ),
        ae_params=AutoEncoderParams(
            resolution=256,
            in_channels=3,
            ch=128,
            out_ch=3,
            ch_mult=[1, 2, 4, 4],
            num_res_blocks=2,
            z_channels=16,
            scale_factor=0.3611,
            shift_factor=0.1159,
        ),
    ),
    "flux-dev-canny-lora": ModelSpec(
        repo_id="black-forest-labs/FLUX.1-dev",
        repo_flow="",
        repo_ae="ckpts/flux_vae.safetensors",
        lora_repo_id="black-forest-labs/FLUX.1-Canny-dev-lora",
        lora_filename="flux1-canny-dev-lora.safetensors",
        params=FluxParams(
            in_channels=128,
            out_channels=64,
            vec_in_dim=768,
            context_in_dim=4096,
            hidden_size=3072,
            mlp_ratio=4.0,
            num_heads=24,
            depth=19,
            depth_single_blocks=38,
            axes_dim=[16, 56, 56],
            theta=10_000,
            qkv_bias=True,
            guidance_embed=True,
        ),
        ae_params=AutoEncoderParams(
            resolution=256,
            in_channels=3,
            ch=128,
            out_ch=3,
            ch_mult=[1, 2, 4, 4],
            num_res_blocks=2,
            z_channels=16,
            scale_factor=0.3611,
            shift_factor=0.1159,
        ),
    ),
    "flux-dev-depth": ModelSpec(
        repo_id="black-forest-labs/FLUX.1-Depth-dev",
        repo_flow="",
        repo_ae="ckpts/flux_vae.safetensors",
        params=FluxParams(
            in_channels=128,
            out_channels=64,
            vec_in_dim=768,
            context_in_dim=4096,
            hidden_size=3072,
            mlp_ratio=4.0,
            num_heads=24,
            depth=19,
            depth_single_blocks=38,
            axes_dim=[16, 56, 56],
            theta=10_000,
            qkv_bias=True,
            guidance_embed=True,
        ),
        ae_params=AutoEncoderParams(
            resolution=256,
            in_channels=3,
            ch=128,
            out_ch=3,
            ch_mult=[1, 2, 4, 4],
            num_res_blocks=2,
            z_channels=16,
            scale_factor=0.3611,
            shift_factor=0.1159,
        ),
    ),
    "flux-dev-depth-lora": ModelSpec(
        repo_id="black-forest-labs/FLUX.1-dev",
        repo_flow="",
        repo_ae="ckpts/flux_vae.safetensors",
        lora_repo_id="black-forest-labs/FLUX.1-Depth-dev-lora",
        lora_filename="flux1-depth-dev-lora.safetensors",
        params=FluxParams(
            in_channels=128,
            out_channels=64,
            vec_in_dim=768,
            context_in_dim=4096,
            hidden_size=3072,
            mlp_ratio=4.0,
            num_heads=24,
            depth=19,
            depth_single_blocks=38,
            axes_dim=[16, 56, 56],
            theta=10_000,
            qkv_bias=True,
            guidance_embed=True,
        ),
        ae_params=AutoEncoderParams(
            resolution=256,
            in_channels=3,
            ch=128,
            out_ch=3,
            ch_mult=[1, 2, 4, 4],
            num_res_blocks=2,
            z_channels=16,
            scale_factor=0.3611,
            shift_factor=0.1159,
        ),
    ),
    "flux-dev-redux": ModelSpec(
        repo_id="black-forest-labs/FLUX.1-Redux-dev",
        repo_flow="",
        repo_ae="ckpts/flux_vae.safetensors",
        params=FluxParams(
            in_channels=64,
            out_channels=64,
            vec_in_dim=768,
            context_in_dim=4096,
            hidden_size=3072,
            mlp_ratio=4.0,
            num_heads=24,
            depth=19,
            depth_single_blocks=38,
            axes_dim=[16, 56, 56],
            theta=10_000,
            qkv_bias=True,
            guidance_embed=True,
        ),
        ae_params=AutoEncoderParams(
            resolution=256,
            in_channels=3,
            ch=128,
            out_ch=3,
            ch_mult=[1, 2, 4, 4],
            num_res_blocks=2,
            z_channels=16,
            scale_factor=0.3611,
            shift_factor=0.1159,
        ),
    ),
    "flux-dev-fill": ModelSpec(
        repo_id="black-forest-labs/FLUX.1-Fill-dev",
        repo_flow="",
        repo_ae="ckpts/flux_vae.safetensors",
        params=FluxParams(
            in_channels=384,
            out_channels=64,
            vec_in_dim=768,
            context_in_dim=4096,
            hidden_size=3072,
            mlp_ratio=4.0,
            num_heads=24,
            depth=19,
            depth_single_blocks=38,
            axes_dim=[16, 56, 56],
            theta=10_000,
            qkv_bias=True,
            guidance_embed=True,
        ),
        ae_params=AutoEncoderParams(
            resolution=256,
            in_channels=3,
            ch=128,
            out_ch=3,
            ch_mult=[1, 2, 4, 4],
            num_res_blocks=2,
            z_channels=16,
            scale_factor=0.3611,
            shift_factor=0.1159,
        ),
    ),
    "flux-dev-kontext": ModelSpec(
        repo_id="black-forest-labs/FLUX.1-Kontext-dev",
        repo_flow="",
        repo_ae="ckpts/flux_vae.safetensors",
        params=FluxParams(
            in_channels=64,
            out_channels=64,
            vec_in_dim=768,
            context_in_dim=4096,
            hidden_size=3072,
            mlp_ratio=4.0,
            num_heads=24,
            depth=19,
            depth_single_blocks=38,
            axes_dim=[16, 56, 56],
            theta=10_000,
            qkv_bias=True,
            guidance_embed=True,
        ),
        ae_params=AutoEncoderParams(
            resolution=256,
            in_channels=3,
            ch=128,
            out_ch=3,
            ch_mult=[1, 2, 4, 4],
            num_res_blocks=2,
            z_channels=16,
            scale_factor=0.3611,
            shift_factor=0.1159,
        ),
    ),
    "flux-dev-uso": ModelSpec(
        repo_id="",
        repo_flow="",
        repo_ae="ckpts/flux_vae.safetensors",
        params=FluxParams(
            in_channels=64,
            out_channels=64,
            vec_in_dim=768,
            context_in_dim=4096,
            hidden_size=3072,
            mlp_ratio=4.0,
            num_heads=24,
            depth=19,
            depth_single_blocks=38,
            axes_dim=[16, 56, 56],
            theta=10_000,
            qkv_bias=True,
            guidance_embed=True,
            eso= True,
        ),
        ae_params=AutoEncoderParams(
            resolution=256,
            in_channels=3,
            ch=128,
            out_ch=3,
            ch_mult=[1, 2, 4, 4],
            num_res_blocks=2,
            z_channels=16,
            scale_factor=0.3611,
            shift_factor=0.1159,
        ),
    ),
    "flux-dev-umo": ModelSpec(
        repo_id="",
        repo_flow="",
        repo_ae="ckpts/flux_vae.safetensors",
        params=FluxParams(
            in_channels=64,
            out_channels=64,
            vec_in_dim=768,
            context_in_dim=4096,
            hidden_size=3072,
            mlp_ratio=4.0,
            num_heads=24,
            depth=19,
            depth_single_blocks=38,
            axes_dim=[16, 56, 56],
            theta=10_000,
            qkv_bias=True,
            guidance_embed=True,
            eso= True,
        ),
        ae_params=AutoEncoderParams(
            resolution=256,
            in_channels=3,
            ch=128,
            out_ch=3,
            ch_mult=[1, 2, 4, 4],
            num_res_blocks=2,
            z_channels=16,
            scale_factor=0.3611,
            shift_factor=0.1159,
        ),
    ),
    "flux-dev-kontext-dreamomni2": ModelSpec(
        repo_id="",
        repo_flow="",
        repo_ae="ckpts/flux_vae.safetensors",
        params=FluxParams(
            in_channels=64,
            out_channels=64,
            vec_in_dim=768,
            context_in_dim=4096,
            hidden_size=3072,
            mlp_ratio=4.0,
            num_heads=24,
            depth=19,
            depth_single_blocks=38,
            axes_dim=[16, 56, 56],
            theta=10_000,
            qkv_bias=True,
            guidance_embed=True,
            eso= True,
        ),
        ae_params=AutoEncoderParams(
            resolution=256,
            in_channels=3,
            ch=128,
            out_ch=3,
            ch_mult=[1, 2, 4, 4],
            num_res_blocks=2,
            z_channels=16,
            scale_factor=0.3611,
            shift_factor=0.1159,
        ),
    ),
}


PREFERED_KONTEXT_RESOLUTIONS = [
    (672, 1568),
    (688, 1504),
    (720, 1456),
    (752, 1392),
    (800, 1328),
    (832, 1248),
    (880, 1184),
    (944, 1104),
    (1024, 1024),
    (1104, 944),
    (1184, 880),
    (1248, 832),
    (1328, 800),
    (1392, 752),
    (1456, 720),
    (1504, 688),
    (1568, 672),
]


def aspect_ratio_to_height_width(aspect_ratio: str, area: int = 1024**2) -> tuple[int, int]:
    width = float(aspect_ratio.split(":")[0])
    height = float(aspect_ratio.split(":")[1])
    ratio = width / height
    width = round(math.sqrt(area * ratio))
    height = round(math.sqrt(area / ratio))
    return 16 * (width // 16), 16 * (height // 16)


def print_load_warning(missing: list[str], unexpected: list[str]) -> None:
    if len(missing) > 0 and len(unexpected) > 0:
        print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))
        print("\n" + "-" * 79 + "\n")
        print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))
    elif len(missing) > 0:
        print(f"Got {len(missing)} missing keys:\n\t" + "\n\t".join(missing))
    elif len(unexpected) > 0:
        print(f"Got {len(unexpected)} unexpected keys:\n\t" + "\n\t".join(unexpected))


def preprocess_flux2_state_dict(state_dict: dict, params: FluxParams) -> dict:
    """
    Remap Flux2 checkpoint keys to the shared Flux1-style module layout.
    - Duplicate shared modulation weights to each single/double block.
    - Drop unused guidance embeddings that are identity in Flux2.
    """
    sd = dict(state_dict)
    def pop_mod(prefix: str):
        w = sd.pop(f"{prefix}.lin.weight", None)
        b = sd.pop(f"{prefix}.lin.bias", None)
        if w is None:
            w = sd.pop(f"{prefix}.linear.weight", None)
            b = sd.pop(f"{prefix}.linear.bias", None) if b is None else b
        return w, b

    img_mod_w, img_mod_b = pop_mod("double_stream_modulation_img")
    txt_mod_w, txt_mod_b = pop_mod("double_stream_modulation_txt")
    single_mod_w, single_mod_b = pop_mod("single_stream_modulation")

    if img_mod_w is not None:
        for i in range(params.depth):
            sd[f"double_blocks.{i}.img_mod.lin.weight"] = img_mod_w
            if img_mod_b is not None:
                sd[f"double_blocks.{i}.img_mod.lin.bias"] = img_mod_b
    if txt_mod_w is not None:
        for i in range(params.depth):
            sd[f"double_blocks.{i}.txt_mod.lin.weight"] = txt_mod_w
            if txt_mod_b is not None:
                sd[f"double_blocks.{i}.txt_mod.lin.bias"] = txt_mod_b
    if single_mod_w is not None:
        for i in range(params.depth_single_blocks):
            sd[f"single_blocks.{i}.modulation.lin.weight"] = single_mod_w
            if single_mod_b is not None:
                sd[f"single_blocks.{i}.modulation.lin.bias"] = single_mod_b

    # Remove unused inputs from Flux2 that are identity for us.
    for unused in (
        "double_stream_modulation_img.lin.weight",
        "double_stream_modulation_img.lin.bias",
        "double_stream_modulation_txt.lin.weight",
        "double_stream_modulation_txt.lin.bias",
        "single_stream_modulation.lin.weight",
        "single_stream_modulation.lin.bias",
        "double_stream_modulation_img.linear.weight",
        "double_stream_modulation_img.linear.bias",
        "double_stream_modulation_txt.linear.weight",
        "double_stream_modulation_txt.linear.bias",
        "single_stream_modulation.linear.weight",
        "single_stream_modulation.linear.bias",
        "guidance_in.in_layer.weight",
        "guidance_in.in_layer.bias",
        "guidance_in.out_layer.weight",
        "guidance_in.out_layer.bias",
    ):
        sd.pop(unused, None)
    return sd


def preprocess_flux_state_dict(state_dict: dict) -> dict:
    return convert_diffusers_to_flux.convert_state_dict(state_dict)


def load_flow_model(
    name: str,
    model_filename,
    device: str | torch.device = "cuda",
    verbose: bool = True,
    preprocess_sd=preprocess_flux_state_dict,
) -> Flux:
    # Loading Flux
    config = configs[name]

    ckpt_path =  model_filename #config.repo_flow

    with torch.device("meta"):
        if config.lora_repo_id is not None and config.lora_filename is not None:
            model = FluxLoraWrapper(params=config.params).to(torch.bfloat16)
        else:
            model = Flux(config.params).to(torch.bfloat16)

    # print(f"Loading checkpoint: {ckpt_path}")
    from mmgp import offload
    offload.load_model_data(model, model_filename, preprocess_sd=preprocess_sd, writable_tensors=False)

    # # load_sft doesn't support torch.device
    # sd = load_sft(ckpt_path, device=str(device))
    # sd = optionally_expand_state_dict(model, sd)
    # missing, unexpected = model.load_state_dict(sd, strict=False, assign=True)
    # if verbose:
    #     print_load_warning(missing, unexpected)

    # if config.lora_repo_id is not None and config.lora_filename is not None:
    #     print("Loading LoRA")
    #     lora_path = str(get_checkpoint_path(config.lora_repo_id, config.lora_filename, "FLUX_LORA"))
    #     lora_sd = load_sft(lora_path, device=str(device))
    #     # loading the lora params + overwriting scale values in the norms
    #     missing, unexpected = model.load_state_dict(lora_sd, strict=False, assign=True)
    #     if verbose:
    #         print_load_warning(missing, unexpected)
    return model


def load_t5(device: str | torch.device = "cuda", text_encoder_filename = None, max_length: int = 512) -> HFEmbedder:
    # max length 64, 128, 256 and 512 should work (if your sequence is short enough)
    return HFEmbedder("",text_encoder_filename, max_length=max_length, torch_dtype=torch.bfloat16).to(device)


def load_clip(device: str | torch.device = "cuda") -> HFEmbedder:
    return HFEmbedder( fl.locate_folder("clip_vit_large_patch14"), "", max_length=77, torch_dtype=torch.bfloat16, is_clip  =True).to(device)


def load_ae(name: str, device: str | torch.device = "cuda") -> AutoEncoder:
    config = configs[name]
    ckpt_path = str(get_checkpoint_path(config.repo_id, fl.locate_file("flux_vae.safetensors"), "FLUX_AE"))

    # Loading the autoencoder
    with torch.device("meta"):
        ae = AutoEncoder(config.ae_params)

    # print(f"Loading AE checkpoint: {ckpt_path}")
    sd = load_sft(ckpt_path, device=str(device))
    missing, unexpected = ae.load_state_dict(sd, strict=False, assign=True)
    print_load_warning(missing, unexpected)
    return ae


def optionally_expand_state_dict(model: torch.nn.Module, state_dict: dict) -> dict:
    """
    Optionally expand the state dict to match the model's parameters shapes.
    """
    for name, param in model.named_parameters():
        if name in state_dict:
            if state_dict[name].shape != param.shape:
                print(
                    f"Expanding '{name}' with shape {state_dict[name].shape} to model parameter with shape {param.shape}."
                )
                # expand with zeros:
                expanded_state_dict_weight = torch.zeros_like(param, device=state_dict[name].device)
                slices = tuple(slice(0, dim) for dim in state_dict[name].shape)
                expanded_state_dict_weight[slices] = state_dict[name]
                state_dict[name] = expanded_state_dict_weight

    return state_dict


class Flux2TransformerShared:
    """
    Convenience wrapper to instantiate the shared Flux2 transformer with explicit defaults.
    """

    def __new__(cls, **kwargs):
        from diffusers.models.transformers.transformer_flux2 import Flux2Transformer2DModel

        defaults = dict(
            patch_size=1,
            in_channels=128,
            out_channels=None,
            num_layers=8,
            num_single_layers=48,
            attention_head_dim=128,
            num_attention_heads=48,
            joint_attention_dim=15360,
            timestep_guidance_channels=256,
            mlp_ratio=3.0,
            axes_dims_rope=(32, 32, 32, 32),
            rope_theta=2000,
            eps=1e-6,
        )
        defaults.update(kwargs)
        return Flux2Transformer2DModel(**defaults)