ArtiAgent - DefectFill/src/flux/util.py

import os
from dataclasses import dataclass

import torch
from einops import rearrange
from huggingface_hub import hf_hub_download
# from imwatermark import WatermarkEncoder
from safetensors.torch import load_file as load_sft

from flux.model import Flux, FluxParams
from flux.modules.autoencoder import AutoEncoder, AutoEncoderParams
from flux.modules.conditioner import HFEmbedder
from transformers import (CLIPTextModel, CLIPTokenizer, T5EncoderModel,
                          T5Tokenizer, BitsAndBytesConfig)  # <--- Added BitsAndBytesConfig


@dataclass
class ModelSpec:
    params: FluxParams
    ae_params: AutoEncoderParams
    ckpt_path: str | None
    ae_path: str | None
    repo_id: str | None
    repo_flow: str | None
    repo_ae: str | None

configs = {
    "flux-dev": ModelSpec(
        repo_id="black-forest-labs/FLUX.1-dev",
        repo_flow="flux1-dev.safetensors",
        repo_ae=None,
        ckpt_path=os.getenv("FLUX_DEV"),
        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_path=os.getenv("AE"),
        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-fill-dev": ModelSpec(
        repo_id="black-forest-labs/FLUX.1-Fill-dev",
        repo_flow="flux1-fill-dev.safetensors",
        repo_ae="ae.safetensors",
        ckpt_path=os.getenv("FLUX_FILL_DEV"),
        params=FluxParams(
            in_channels=64,
            out_channels=384,
            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_path=os.getenv("AE"),
        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-kontext-dev": ModelSpec(
        repo_id="black-forest-labs/FLUX.1-Kontext-dev",
        repo_flow="flux1-kontext-dev.safetensors",
        repo_ae="ae.safetensors",
        ckpt_path=os.getenv("FLUX_FILL_DEV"),
        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_path=os.getenv("AE"),
        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="flux1-schnell.safetensors",
        repo_ae="black-forest-labs/FLUX.1-schnell",
        ckpt_path=os.getenv("FLUX_SCHNELL"),
        params=FluxParams(
            in_channels=64, # ArtiAgent custom input dimension logic handled in load_flow_model
            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=10000.0,
            qkv_bias=True,
            guidance_embed=False,
        ),
        ae_path="ae.safetensors",
        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,
        ),
    ),
}


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 _replace_linear_with_4bit(module, compute_dtype=torch.bfloat16):
    """Recursively replace all nn.Linear with bitsandbytes 4-bit layers"""
    import bitsandbytes as bnb
    for name, child in module.named_children():
        if name == "img_in":
            continue  # Skip img_in to preserve ArtiAgent's custom shape handling
        if isinstance(child, torch.nn.Linear):
            has_bias = child.bias is not None
            new_layer = bnb.nn.Linear4bit(
                child.in_features,
                child.out_features,
                bias=has_bias,
                compute_dtype=compute_dtype,
                compress_statistics=True,
                quant_type="nf4",
            )
            new_layer.weight = bnb.nn.Params4bit(
                child.weight.data,
                requires_grad=False,
                quant_type="nf4",
            )
            if has_bias:
                new_layer.bias = torch.nn.Parameter(child.bias.data)
            setattr(module, name, new_layer)
        else:
            _replace_linear_with_4bit(child, compute_dtype)

def load_flow_model(name: str, device: str | torch.device = "cuda", hf_download: bool = True):
    # Loading Flux
    print("Init model")
    
    ckpt_path = configs[name].ckpt_path
    if (
        ckpt_path is None
        and configs[name].repo_id is not None
        and configs[name].repo_flow is not None
        and hf_download
    ):
        ckpt_path = hf_hub_download(configs[name].repo_id, configs[name].repo_flow)

    # Initialize model directly on CPU or target device (avoids meta-tensor shape replacement)
    target_device = torch.device(device)
    model = Flux(configs[name].params).to(dtype=torch.bfloat16)

    if ckpt_path is not None:
        print("Loading checkpoint")
        # load_sft doesn't support torch.device
        sd = load_sft(ckpt_path, device="cpu")

        # --- ADD THIS LINE TO STRIP FP8 / COMFYUI KEY PREFIXES ---
        sd = {k.replace("model.diffusion_model.", ""): v for k, v in sd.items()}
        # ---------------------------------------------------------

        # --- FIX: HANDLE EXPANDED IMG_IN (384 channels vs 64 channels) ---
        img_in_weight = sd.pop("img_in.weight", None)
        img_in_bias = sd.pop("img_in.bias", None)

        # Load all standard layers safely
        missing, unexpected = model.load_state_dict(sd, strict=False, assign=True)
        print_load_warning(missing, unexpected)

        # Copy base 64 channels into ArtiAgent's expanded 384-channel input layer
        # In src/flux/util.py inside load_flow_model():

        if img_in_weight is not None:
            with torch.no_grad():
                w = img_in_weight.to(device=device, dtype=torch.bfloat16)
        
                # Check if model.img_in weight expects 384 channels while checkpoint has 64
                if model.img_in.weight.shape[1] != w.shape[1]:
                    # Slice model.img_in.weight to match the 64-channel input tensor
                    model.img_in.weight = torch.nn.Parameter(model.img_in.weight[:, :w.shape[1]])
            
                model.img_in.weight.copy_(w)

        if img_in_bias is not None and getattr(model.img_in, "bias", None) is not None:
            with torch.no_grad():
                b = img_in_bias.to(device=device, dtype=torch.bfloat16)
                model.img_in.bias.copy_(b)

        # Quantize all Linear layers to NF4 on CPU before moving to GPU
        print("Quantizing model to NF4 (this may take a minute)...")
        _replace_linear_with_4bit(model, compute_dtype=torch.bfloat16)
        print("NF4 quantization complete.")

    # Move model to target CUDA device
    model = model.to(target_device)
    return model


def load_t5(device: str | torch.device = "cuda", max_length: int = 512) -> HFEmbedder:
    # Force T5 onto CPU; sampling.py already moves the encoded txt tensor to GPU
    return HFEmbedder(
        "google/t5-v1_1-xxl", 
        max_length=max_length, 
        is_clip=False, 
        torch_dtype=torch.bfloat16,
        device_map="cpu"
    )


def load_clip(device: str | torch.device = "cuda") -> HFEmbedder:
    # Keep on CPU; sampling.py moves vec to GPU after encoding
    return HFEmbedder("openai/clip-vit-large-patch14", max_length=77, is_clip=True, torch_dtype=torch.bfloat16)


def load_ae(name: str, device: str | torch.device = "cuda", hf_download: bool = True) -> AutoEncoder:
    ckpt_path = configs[name].ae_path

    # If ckpt_path is just a filename and doesn't exist locally, download it
    if ckpt_path is not None and not os.path.exists(ckpt_path) and hf_download:
        repo_id = configs[name].repo_ae or configs[name].repo_id
        ckpt_path = hf_hub_download(repo_id, ckpt_path)
    elif ckpt_path is None and configs[name].repo_id is not None and hf_download:
        repo_id = configs[name].repo_ae or configs[name].repo_id
        ckpt_path = hf_hub_download(repo_id, "ae.safetensors")

    # Loading the autoencoder
    print("Init AE")

    # Initialize directly on CPU to avoid meta-tensor initialization issues
    ae = AutoEncoder(configs[name].ae_params)

    if ckpt_path is not None:
        sd = load_sft(ckpt_path, device=str(device))
        missing, unexpected = ae.load_state_dict(sd, strict=False, assign=True)
        print_load_warning(missing, unexpected)

    ae = ae.to(device)
    return ae


# class WatermarkEmbedder:
#     def __init__(self, watermark):
#         self.watermark = watermark
#         self.num_bits = len(WATERMARK_BITS)
#         self.encoder = WatermarkEncoder()
#         self.encoder.set_watermark("bits", self.watermark)

#     def __call__(self, image: torch.Tensor) -> torch.Tensor:
#         """
#         Adds a predefined watermark to the input image

#         Args:
#             image: ([N,] B, RGB, H, W) in range [-1, 1]

#         Returns:
#             same as input but watermarked
#         """
#         image = 0.5 * image + 0.5
#         squeeze = len(image.shape) == 4
#         if squeeze:
#             image = image[None, ...]
#         n = image.shape[0]
#         image_np = rearrange((255 * image).detach().cpu(), "n b c h w -> (n b) h w c").numpy()[:, :, :, ::-1]
#         # torch (b, c, h, w) in [0, 1] -> numpy (b, h, w, c) [0, 255]
#         # watermarking libary expects input as cv2 BGR format
#         for k in range(image_np.shape[0]):
#             image_np[k] = self.encoder.encode(image_np[k], "dwtDct")
#         image = torch.from_numpy(rearrange(image_np[:, :, :, ::-1], "(n b) h w c -> n b c h w", n=n)).to(
#             image.device
#         )
#         image = torch.clamp(image / 255, min=0.0, max=1.0)
#         if squeeze:
#             image = image[0]
#         image = 2 * image - 1
#         return image


# # A fixed 48-bit message that was chosen at random
# WATERMARK_MESSAGE = 0b001010101111111010000111100111001111010100101110
# # bin(x)[2:] gives bits of x as str, use int to convert them to 0/1
# WATERMARK_BITS = [int(bit) for bit in bin(WATERMARK_MESSAGE)[2:]]
# embed_watermark = WatermarkEmbedder(WATERMARK_BITS)