import os, yaml, random import torch import numpy as np from typing import Union import pickle from diffusers.pipelines.flux.pipeline_flux import calculate_shift, retrieve_timesteps from peft import LoraConfig from peft.utils import get_peft_model_state_dict, set_peft_model_state_dict from models.mmdit import CustomFluxTransformer2DModel from models.pipeline import CustomFluxPipeline from models.multiLayer_adapter import MultiLayerAdapter def save_checkpoint(transformer, multiLayer_adater, optimizer, optimizer_adapter, scheduler, scheduler_adapter, step, save_dir): import gc trans_dir = os.path.join(save_dir, "transformer") adapter_dir = os.path.join(save_dir, "adapter") os.makedirs(trans_dir, exist_ok=True) os.makedirs(adapter_dir, exist_ok=True) # Get state dicts and IMMEDIATELY move to CPU to avoid GPU memory buildup flux_transformer_lora_state_dict = get_peft_model_state_dict(transformer) flux_transformer_lora_state_dict = {k: v.detach().cpu().to(torch.float32) for k, v in flux_transformer_lora_state_dict.items()} flux_adapter_lora_state_dict = get_peft_model_state_dict(multiLayer_adater) flux_adapter_lora_state_dict = {k: v.detach().cpu().to(torch.float32) for k, v in flux_adapter_lora_state_dict.items()} CustomFluxPipeline.save_lora_weights( os.path.join(trans_dir), flux_transformer_lora_state_dict, safe_serialization=True, ) # Clear after saving del flux_transformer_lora_state_dict CustomFluxPipeline.save_lora_weights( os.path.join(adapter_dir), flux_adapter_lora_state_dict, safe_serialization=True, ) # Clear after saving del flux_adapter_lora_state_dict torch.save({"layer_pe": transformer.layer_pe.detach().cpu().to(torch.float32)}, os.path.join(save_dir, "layer_pe.pth")) torch.save(optimizer.state_dict(), os.path.join(trans_dir, "optimizer.bin")) torch.save(optimizer_adapter.state_dict(), os.path.join(adapter_dir, "optimizer.bin")) torch.save(scheduler.state_dict(), os.path.join(trans_dir, "scheduler.bin")) torch.save(scheduler_adapter.state_dict(), os.path.join(adapter_dir, "scheduler.bin")) save_path = os.path.join(save_dir, f"random_states_0.pkl") state = { "step": step, "random_state": random.getstate(), "numpy_random_seed": np.random.get_state(), "torch_manual_seed": torch.get_rng_state(), } if torch.cuda.is_available(): state["torch_cuda_manual_seed"] = torch.cuda.get_rng_state_all() # list of tensors with open(save_path, "wb") as f: pickle.dump(state, f) # Force garbage collection and clear CUDA cache gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() print(f"[INFO] Saved RNG states + step {step} to {save_path}") def load_checkpoint(transformer, multiLayer_adater, optimizer, optimizer_adapter, scheduler, scheduler_adapter, ckpt_dir, device="cuda"): trans_dir = os.path.join(ckpt_dir, "transformer") adapter_dir = os.path.join(ckpt_dir, "adapter") start_step = 0 lora_path = os.path.join(trans_dir, "pytorch_lora_weights.safetensors") lora_path_adapter = os.path.join(adapter_dir, "pytorch_lora_weights.safetensors") if os.path.exists(lora_path): lora_state_dict = CustomFluxPipeline.lora_state_dict(lora_path) stripped = {k.replace("transformer.", "", 1) if k.startswith("transformer.") else k: v for k, v in lora_state_dict.items()} result = set_peft_model_state_dict(transformer, stripped) if result.unexpected_keys: print(f"[WARN] Transformer LoRA: {len(result.unexpected_keys)} unexpected keys") print(f"[INFO] Loaded Transformer LoRA weights ({len(stripped)} keys).") if os.path.exists(lora_path_adapter): lora_state_dict = CustomFluxPipeline.lora_state_dict(lora_path_adapter) stripped = {k.replace("transformer.", "", 1) if k.startswith("transformer.") else k: v for k, v in lora_state_dict.items()} result = set_peft_model_state_dict(multiLayer_adater, stripped) if result.unexpected_keys: print(f"[WARN] Adapter LoRA: {len(result.unexpected_keys)} unexpected keys") print(f"[INFO] Loaded Adapter LoRA weights ({len(stripped)} keys).") pe_path = os.path.join(ckpt_dir, "layer_pe.pth") if os.path.exists(pe_path): layer_pe = torch.load(pe_path) missing_keys, unexpected_keys = transformer.load_state_dict(layer_pe, strict=False) opt_path = os.path.join(trans_dir, "optimizer.bin") opt_path_adapter = os.path.join(adapter_dir, "optimizer.bin") if os.path.exists(opt_path): optimizer.load_state_dict(torch.load(opt_path, map_location=device)) print("[INFO] Loaded optimizer state.") if os.path.exists(opt_path_adapter): optimizer_adapter.load_state_dict(torch.load(opt_path_adapter, map_location=device)) print("[INFO] Loaded optimizer state.") sch_path = os.path.join(trans_dir, "scheduler.bin") sch_path_adapter = os.path.join(adapter_dir, "scheduler.bin") if os.path.exists(sch_path): scheduler.load_state_dict(torch.load(sch_path, map_location=device)) print("[INFO] Loaded scheduler state.") if os.path.exists(sch_path_adapter): scheduler_adapter.load_state_dict(torch.load(sch_path_adapter, map_location=device)) print("[INFO] Loaded scheduler state.") rng_file = None for f in os.listdir(ckpt_dir): if f.startswith("random_states_") and f.endswith(".pkl"): rng_file = os.path.join(ckpt_dir, f) break if rng_file: with open(rng_file, "rb") as f: state = pickle.load(f) start_step = state.get("step", 0) if "random_state" in state: random.setstate(state["random_state"]) if "numpy_random_seed" in state: np.random.set_state(state["numpy_random_seed"]) if "torch_manual_seed" in state: torch.set_rng_state(state["torch_manual_seed"]) if "torch_cuda_manual_seed" in state and torch.cuda.is_available(): torch.cuda.set_rng_state_all(state["torch_cuda_manual_seed"]) print(f"[INFO] Resumed RNG states + step {start_step}") return start_step def load_config(path): with open(path, "r") as f: return yaml.safe_load(f) def seed_everything(seed: int): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic = True def get_input_box(layer_boxes, image_size=512): """ Quantize layer boxes to 16-pixel grid for latent space alignment. Args: layer_boxes: List of boxes in xyxy format [x0, y0, x1, y1] image_size: Image size to clamp bounds (default 512) Returns: List of quantized boxes in xyxy format """ list_layer_box = [] for layer_box in layer_boxes: min_col, min_row = layer_box[0], layer_box[1] max_col, max_row = layer_box[2], layer_box[3] # Floor for min (start of box) quantized_min_row = (min_row // 16) * 16 quantized_min_col = (min_col // 16) * 16 # Ceiling for max (end of box) - use (val + 15) // 16 * 16 for proper ceiling quantized_max_row = ((max_row + 15) // 16) * 16 quantized_max_col = ((max_col + 15) // 16) * 16 # Clamp to image bounds quantized_min_row = max(0, quantized_min_row) quantized_min_col = max(0, quantized_min_col) quantized_max_row = min(image_size, quantized_max_row) quantized_max_col = min(image_size, quantized_max_col) # Ensure minimum box size of 16 pixels (1 latent token) in each dimension # This prevents zero-size boxes that cause reshape errors if quantized_max_col <= quantized_min_col: # Expand the box, preferring to expand max if there's room if quantized_min_col + 16 <= image_size: quantized_max_col = quantized_min_col + 16 else: quantized_min_col = max(0, quantized_max_col - 16) quantized_max_col = quantized_min_col + 16 if quantized_max_row <= quantized_min_row: # Expand the box, preferring to expand max if there's room if quantized_min_row + 16 <= image_size: quantized_max_row = quantized_min_row + 16 else: quantized_min_row = max(0, quantized_max_row - 16) quantized_max_row = quantized_min_row + 16 list_layer_box.append((quantized_min_col, quantized_min_row, quantized_max_col, quantized_max_row)) return list_layer_box def set_lora_into_transformer( model: Union[CustomFluxTransformer2DModel, MultiLayerAdapter], lora_rank: int, lora_alpha: float = 1.0, lora_dropout: float = 0.1, ): target_modules = [ "to_k", "to_q", "to_v", "to_out.0", "add_k_proj", "add_q_proj", "add_v_proj", "to_add_out", ] + [f"single_transformer_blocks.{i}.proj_out" for i in range(model.config.num_single_layers)] + [f"transformer_blocks.{i}.proj_out" for i in range(model.config.num_layers)] transformer_lora_config = LoraConfig( r=lora_rank, lora_alpha=lora_alpha, lora_dropout=lora_dropout, init_lora_weights="gaussian", target_modules=target_modules, ) model.add_adapter(transformer_lora_config) return model def build_layer_mask(n_layers, H_lat, W_lat, list_layer_box): mask = torch.zeros((n_layers, 1, H_lat, W_lat), dtype=torch.float32) for i, box in enumerate(list_layer_box): if box is None: continue x1, y1, x2, y2 = box x1_t, y1_t, x2_t, y2_t = x1 // 8, y1 // 8, x2 // 8, y2 // 8 x1_t, y1_t = max(0, x1_t), max(0, y1_t) x2_t, y2_t = min(W_lat, x2_t), min(H_lat, y2_t) if x2_t > x1_t and y2_t > y1_t: mask[i, :, y1_t:y2_t, x1_t:x2_t] = 1.0 return mask def encode_target_latents(pipeline, pixel_bchw, n_layers, list_layer_box): device = pixel_bchw.device dtype = pixel_bchw.dtype vae = pipeline.vae.eval() bs, n_layers_in, C, H, W = pixel_bchw.shape assert n_layers_in == n_layers, f"The number of input layers {n_layers_in} does not match the specified number of layers {n_layers}" with torch.no_grad(): dummy_lat = vae.encode(pixel_bchw[:,0]).latent_dist.sample() _, C_lat, H_lat, W_lat = dummy_lat.shape x0 = torch.zeros((bs, n_layers, C_lat, H_lat, W_lat), device=device, dtype=dtype) with torch.no_grad(): for i in range(n_layers): pixel_i = pixel_bchw[:, i] lat = vae.encode(pixel_i).latent_dist.sample() # [1,C_lat,H_lat,W_lat] lat = (lat - vae.config.shift_factor) * vae.config.scaling_factor x0[:, i] = lat latent_ids = pipeline._prepare_latent_image_ids(H_lat, W_lat, list_layer_box, device, dtype) return x0, latent_ids def get_timesteps(pipeline, image_seq_len, num_inference_steps, device): sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) mu = calculate_shift( image_seq_len, pipeline.scheduler.config.base_image_seq_len, pipeline.scheduler.config.max_image_seq_len, pipeline.scheduler.config.base_shift, pipeline.scheduler.config.max_shift, ) timesteps, num_inference_steps = retrieve_timesteps( scheduler=pipeline.scheduler, num_inference_steps=num_inference_steps, device=device, sigmas=sigmas, mu=mu, ) return timesteps # ============================================================================ # Box utilities for Prism blended dataset # ============================================================================ def scale_box_xyxy(box, source_size: int, target_size: int): """ Scale a box from source_size to target_size. Box is already in xyxy format: [x0, y0, x1, y1]. Args: box: [x0, y0, x1, y1] in source_size coordinates source_size: Original data size (e.g., 512) target_size: Target inference size (e.g., 512) Returns: (x0, y0, x1, y1) in target_size coordinates """ scale = target_size / source_size x0, y0, x1, y1 = box x0_s = int(x0 * scale) y0_s = int(y0 * scale) x1_s = int(x1 * scale) y1_s = int(y1 * scale) # Clamp to valid range x0_s = max(0, x0_s) y0_s = max(0, y0_s) x1_s = min(target_size, x1_s) y1_s = min(target_size, y1_s) return (x0_s, y0_s, x1_s, y1_s) def quantize_box_16(box, target_size: int): """ Quantize box to 16-pixel grid for latent space alignment. Box is in xyxy format. """ x0, y0, x1, y1 = box # Quantize to 16-pixel grid x0_q = (x0 // 16) * 16 y0_q = (y0 // 16) * 16 x1_q = ((x1 + 15) // 16) * 16 y1_q = ((y1 + 15) // 16) * 16 # Clamp to image bounds x0_q = max(0, x0_q) y0_q = max(0, y0_q) x1_q = min(target_size, x1_q) y1_q = min(target_size, y1_q) return (x0_q, y0_q, x1_q, y1_q) def get_prism_layer_boxes_xyxy(layers, source_size: int, target_size: int): """ Extract and scale layer boxes from prism blended metadata. Note: Our blended dataset uses xyxy format [x0, y0, x1, y1]. Args: layers: List of layer metadata dicts with 'box' field (xyxy format) source_size: Size the data was generated at (e.g., 512) target_size: Size to run inference at (e.g., 512) Returns: List of quantized boxes in xyxy format """ boxes = [] for layer in layers: box = layer.get('box', [0, 0, source_size, source_size]) # Scale from source to target size (box is already xyxy) scaled_box = scale_box_xyxy(box, source_size, target_size) # Quantize to 16-pixel grid quantized_box = quantize_box_16(scaled_box, target_size) boxes.append(quantized_box) return boxes def xywh_to_xyxy(box): """Convert (x, y, w, h) to (x0, y0, x1, y1).""" x, y, w, h = box return (x, y, x + w, y + h) def xyxy_to_xywh(box): """Convert (x0, y0, x1, y1) to (x, y, w, h).""" x0, y0, x1, y1 = box return (x0, y0, x1 - x0, y1 - y0)