import argparse import copy import math import random from typing import Any import pdb import os import time from PIL import Image, ImageOps import torch from accelerate import Accelerator from library.device_utils import clean_memory_on_device from safetensors.torch import load_file from networks import lora_flux from library import flux_models, flux_train_utils_recraft as flux_train_utils, flux_utils, sd3_train_utils, \ strategy_base, strategy_flux, train_util from torchvision import transforms import train_network from library.utils import setup_logging from diffusers.utils import load_image import numpy as np setup_logging() import logging logger = logging.getLogger(__name__) def load_target_model( fp8_base: bool, pretrained_model_name_or_path: str, disable_mmap_load_safetensors: bool, clip_l_path: str, fp8_base_unet: bool, t5xxl_path: str, ae_path: str, weight_dtype: torch.dtype, accelerator: Accelerator ): # Determine the loading data type loading_dtype = None if fp8_base else weight_dtype # Load the main model to the accelerator's device _, model = flux_utils.load_flow_model( pretrained_model_name_or_path, # loading_dtype, torch.float8_e4m3fn, # accelerator.device, # Changed from "cpu" to accelerator.device "cpu", disable_mmap=disable_mmap_load_safetensors ) if fp8_base: # Check dtype of the model if model.dtype in {torch.float8_e4m3fnuz, torch.float8_e5m2, torch.float8_e5m2fnuz}: raise ValueError(f"Unsupported fp8 model dtype: {model.dtype}") elif model.dtype == torch.float8_e4m3fn: logger.info("Loaded fp8 FLUX model") # Load the CLIP model to the accelerator's device clip_l = flux_utils.load_clip_l( clip_l_path, weight_dtype, # accelerator.device, # Changed from "cpu" to accelerator.device "cpu", disable_mmap=disable_mmap_load_safetensors ) clip_l.eval() # Determine the loading data type for T5XXL if fp8_base and not fp8_base_unet: loading_dtype_t5xxl = None # as is else: loading_dtype_t5xxl = weight_dtype # Load the T5XXL model to the accelerator's device t5xxl = flux_utils.load_t5xxl( t5xxl_path, loading_dtype_t5xxl, # accelerator.device, # Changed from "cpu" to accelerator.device "cpu", disable_mmap=disable_mmap_load_safetensors ) t5xxl.eval() if fp8_base and not fp8_base_unet: # Check dtype of the T5XXL model if t5xxl.dtype in {torch.float8_e4m3fnuz, torch.float8_e5m2, torch.float8_e5m2fnuz}: raise ValueError(f"Unsupported fp8 model dtype: {t5xxl.dtype}") elif t5xxl.dtype == torch.float8_e4m3fn: logger.info("Loaded fp8 T5XXL model") # Load the AE model to the accelerator's device ae = flux_utils.load_ae( ae_path, weight_dtype, # accelerator.device, # Changed from "cpu" to accelerator.device "cpu", disable_mmap=disable_mmap_load_safetensors ) # # Wrap models with Accelerator for potential distributed setups # model, clip_l, t5xxl, ae = accelerator.prepare(model, clip_l, t5xxl, ae) return flux_utils.MODEL_VERSION_FLUX_V1, [clip_l, t5xxl], ae, model import torchvision.transforms as transforms class ResizeWithPadding: def __init__(self, size, fill=255): self.size = size self.fill = fill def __call__(self, img): if isinstance(img, np.ndarray): img = Image.fromarray(img) elif not isinstance(img, Image.Image): raise TypeError("Input must be a PIL Image or a NumPy array") width, height = img.size if width == height: img = img.resize((self.size, self.size), Image.LANCZOS) else: max_dim = max(width, height) new_img = Image.new("RGB", (max_dim, max_dim), (self.fill, self.fill, self.fill)) new_img.paste(img, ((max_dim - width) // 2, (max_dim - height) // 2)) img = new_img.resize((self.size, self.size), Image.LANCZOS) return img def sample(args, accelerator, vae, text_encoder, flux, output_dir, sample_images, sample_prompts): def encode_images_to_latents(vae, images): # Get image dimensions b, c, h, w = images.shape num_split = 2 if args.frame_num == 4 else 3 # Split the image into three parts img_parts = [images[:, :, :, i * w // num_split:(i + 1) * w // num_split] for i in range(num_split)] # Encode each part latents = [vae.encode(img) for img in img_parts] # Concatenate latents in the latent space to reconstruct the full image latents = torch.cat(latents, dim=-1) return latents def encode_images_to_latents2(vae, images): latents = vae.encode(images) return latents # Directly use precomputed conditions conditions = {} with torch.no_grad(): for image_path, prompt_dict in zip(sample_images, sample_prompts): prompt = prompt_dict.get("prompt", "") if prompt not in conditions: logger.info(f"Cache conditions for image: {image_path} with prompt: {prompt}") resize_transform = ResizeWithPadding(size=512, fill=255) if args.frame_num == 4 else ResizeWithPadding(size=352, fill=255) img_transforms = transforms.Compose([ resize_transform, transforms.ToTensor(), transforms.Normalize([0.5], [0.5]), ]) # Load and preprocess image image = img_transforms(np.array(load_image(image_path), dtype=np.uint8)).unsqueeze(0).to( # accelerator.device, # Move image to CUDA vae.device, dtype=vae.dtype ) latents = encode_images_to_latents2(vae, image) # Log the shape of latents logger.debug(f"Encoded latents shape for prompt '{prompt}': {latents.shape}") # Store conditions on CUDA # conditions[prompt] = latents[:,:,latents.shape[2]//2:latents.shape[2], :latents.shape[3]//2].to("cpu") conditions[prompt] = latents.to("cpu") sample_conditions = conditions if sample_conditions is not None: conditions = {k: v for k, v in sample_conditions.items()} # Already on CUDA sample_prompts_te_outputs = {} # key: prompt, value: text encoder outputs text_encoder[0].to(accelerator.device) text_encoder[1].to(accelerator.device) tokenize_strategy = strategy_flux.FluxTokenizeStrategy(512) text_encoding_strategy = strategy_flux.FluxTextEncodingStrategy(True) with accelerator.autocast(), torch.no_grad(): for prompt_dict in sample_prompts: for p in [prompt_dict.get("prompt", ""), prompt_dict.get("negative_prompt", "")]: if p not in sample_prompts_te_outputs: logger.info(f"Cache Text Encoder outputs for prompt: {p}") tokens_and_masks = tokenize_strategy.tokenize(p) sample_prompts_te_outputs[p] = text_encoding_strategy.encode_tokens( tokenize_strategy, text_encoder, tokens_and_masks, True ) logger.info(f"Generating image") save_dir = output_dir os.makedirs(save_dir, exist_ok=True) with torch.no_grad(), accelerator.autocast(): for prompt_dict in sample_prompts: sample_image_inference( args, accelerator, flux, text_encoder, vae, save_dir, prompt_dict, sample_prompts_te_outputs, None, conditions ) clean_memory_on_device(accelerator.device) def sample_image_inference( args, accelerator: Accelerator, flux: flux_models.Flux, text_encoder, ae: flux_models.AutoEncoder, save_dir, prompt_dict, sample_prompts_te_outputs, prompt_replacement, sample_images_ae_outputs ): # Extract parameters from prompt_dict sample_steps = prompt_dict.get("sample_steps", 20) width = prompt_dict.get("width", 1024) if args.frame_num == 4 else prompt_dict.get("width", 1056) height = prompt_dict.get("height", 1024) if args.frame_num == 4 else prompt_dict.get("height", 1056) scale = prompt_dict.get("scale", 1.0) seed = prompt_dict.get("seed") prompt: str = prompt_dict.get("prompt", "") if prompt_replacement is not None: prompt = prompt.replace(prompt_replacement[0], prompt_replacement[1]) if seed is not None: torch.manual_seed(seed) torch.cuda.manual_seed(seed) else: # True random sample image generation torch.seed() torch.cuda.seed() # Ensure height and width are divisible by 16 height = max(64, height - height % 16) width = max(64, width - width % 16) logger.info(f"prompt: {prompt}") logger.info(f"height: {height}") logger.info(f"width: {width}") logger.info(f"sample_steps: {sample_steps}") logger.info(f"scale: {scale}") if seed is not None: logger.info(f"seed: {seed}") # Encode prompts # Assuming that TokenizeStrategy and TextEncodingStrategy are compatible with Accelerator text_encoder_conds = [] if sample_prompts_te_outputs and prompt in sample_prompts_te_outputs: text_encoder_conds = sample_prompts_te_outputs[prompt] logger.info(f"Using cached text encoder outputs for prompt: {prompt}") if sample_images_ae_outputs and prompt in sample_images_ae_outputs: ae_outputs = sample_images_ae_outputs[prompt] else: ae_outputs = None # ae_outputs = torch.load('ae_outputs.pth', map_location='cuda:0') # text_encoder_conds = torch.load('text_encoder_conds.pth', map_location='cuda:0') l_pooled, t5_out, txt_ids, t5_attn_mask = text_encoder_conds # 打印调试信息 logger.debug( f"l_pooled shape: {l_pooled.shape}, t5_out shape: {t5_out.shape}, txt_ids shape: {txt_ids.shape}, t5_attn_mask shape: {t5_attn_mask.shape}") # 采样图像 weight_dtype = ae.dtype # TODO: give dtype as argument packed_latent_height = height // 16 packed_latent_width = width // 16 # 打印调试信息 logger.debug(f"packed_latent_height: {packed_latent_height}, packed_latent_width: {packed_latent_width}") # 准备噪声张量在 CUDA 上 noise = torch.randn( 1, packed_latent_height * packed_latent_width, 16 * 2 * 2, device=accelerator.device, dtype=weight_dtype, generator=torch.Generator(device=accelerator.device).manual_seed(seed) if seed is not None else None, ) timesteps = flux_train_utils.get_schedule(sample_steps, noise.shape[1], shift=True) # FLUX.1 dev -> shift=True img_ids = flux_utils.prepare_img_ids(1, packed_latent_height, packed_latent_width).to( accelerator.device, dtype=weight_dtype ) t5_attn_mask = t5_attn_mask.to(accelerator.device) clip_l, t5xxl = text_encoder # ae.to("cpu") clip_l.to("cpu") t5xxl.to("cpu") clean_memory_on_device(accelerator.device) flux.to("cuda") for param in flux.parameters(): param.requires_grad = False # 执行去噪 with accelerator.autocast(), torch.no_grad(): x = flux_train_utils.denoise(args, flux, noise, img_ids, t5_out, txt_ids, l_pooled, timesteps=timesteps, guidance=scale, t5_attn_mask=t5_attn_mask, ae_outputs=ae_outputs) # 打印x的形状 logger.debug(f"x shape after denoise: {x.shape}") x = x.float() x = flux_utils.unpack_latents(x, packed_latent_height, packed_latent_width) # 将潜在向量转换为图像 # clean_memory_on_device(accelerator.device) ae.to(accelerator.device) with accelerator.autocast(), torch.no_grad(): x = ae.decode(x) ae.to("cpu") clean_memory_on_device(accelerator.device) x = x.clamp(-1, 1) x = x.permute(0, 2, 3, 1) image = Image.fromarray((127.5 * (x + 1.0)).float().cpu().numpy().astype(np.uint8)[0]) # 生成唯一的文件名 ts_str = time.strftime("%Y%m%d%H%M%S", time.localtime()) seed_suffix = "" if seed is None else f"_{seed}" i: int = prompt_dict.get("enum", 0) # Ensure 'enum' exists img_filename = f"{ts_str}{seed_suffix}_{i}.png" # Added 'i' to filename for uniqueness image.save(os.path.join(save_dir, img_filename)) def setup_argparse(): parser = argparse.ArgumentParser(description="FLUX-Controlnet-Inpainting Inference Script") # Paths parser.add_argument('--base_flux_checkpoint', type=str, required=True, help='Path to BASE_FLUX_CHECKPOINT') parser.add_argument('--lora_weights_path', type=str, required=True, help='Path to LORA_WEIGHTS_PATH') parser.add_argument('--clip_l_path', type=str, required=True, help='Path to CLIP_L_PATH') parser.add_argument('--t5xxl_path', type=str, required=True, help='Path to T5XXL_PATH') parser.add_argument('--ae_path', type=str, required=True, help='Path to AE_PATH') parser.add_argument('--sample_images_file', type=str, required=True, help='Path to SAMPLE_IMAGES_FILE') parser.add_argument('--sample_prompts_file', type=str, required=True, help='Path to SAMPLE_PROMPTS_FILE') parser.add_argument('--output_dir', type=str, required=True, help='Directory to save OUTPUT_DIR') parser.add_argument('--frame_num', type=int, choices=[4, 9], required=True, help="The number of steps in the generated step diagram (choose 4 or 9)") return parser.parse_args() def main(args): accelerator = Accelerator(mixed_precision='bf16', device_placement=True) BASE_FLUX_CHECKPOINT = args.base_flux_checkpoint LORA_WEIGHTS_PATH = args.lora_weights_path CLIP_L_PATH = args.clip_l_path T5XXL_PATH = args.t5xxl_path AE_PATH = args.ae_path SAMPLE_IMAGES_FILE = args.sample_images_file SAMPLE_PROMPTS_FILE = args.sample_prompts_file OUTPUT_DIR = args.output_dir with open(SAMPLE_IMAGES_FILE, "r", encoding="utf-8") as f: image_lines = f.readlines() sample_images = [line.strip() for line in image_lines if line.strip() and not line.strip().startswith("#")] sample_prompts = train_util.load_prompts(SAMPLE_PROMPTS_FILE) # Load models onto CUDA via Accelerator _, [clip_l, t5xxl], ae, model = load_target_model( fp8_base=True, pretrained_model_name_or_path=BASE_FLUX_CHECKPOINT, disable_mmap_load_safetensors=False, clip_l_path=CLIP_L_PATH, fp8_base_unet=False, t5xxl_path=T5XXL_PATH, ae_path=AE_PATH, weight_dtype=torch.bfloat16, accelerator=accelerator ) model.eval() clip_l.eval() t5xxl.eval() ae.eval() # LoRA multiplier = 1.0 weights_sd = load_file(LORA_WEIGHTS_PATH) lora_model, _ = lora_flux.create_network_from_weights(multiplier, None, ae, [clip_l, t5xxl], model, weights_sd, True) lora_model.apply_to([clip_l, t5xxl], model) info = lora_model.load_state_dict(weights_sd, strict=True) logger.info(f"Loaded LoRA weights from {LORA_WEIGHTS_PATH}: {info}") lora_model.eval() lora_model.to("cuda") # Set text encoders text_encoder = [clip_l, t5xxl] sample(args, accelerator, vae=ae, text_encoder=text_encoder, flux=model, output_dir=OUTPUT_DIR, sample_images=sample_images, sample_prompts=sample_prompts) if __name__ == "__main__": args = setup_argparse() main(args)