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)