import peft import torch import logging from contextlib import contextmanager from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline, QwenImageEditPlusPipeline, QwenImageTransformer2DModel from PIL import Image COLOR_GRAY = "\033[37m" COLOR_BLUE = "\033[94m" COLOR_RED = "\033[31m" COLOR_GREEN = "\033[32m" COLOR_RESET = "\033[0m" def get_logger(name: str, log_file: str): # 1. Get a logger instance logger = logging.Logger(name) # 2. Create a console handler console_handler = logging.StreamHandler() console_handler.setLevel(logging.INFO) # Set the level for console output console_formatter = logging.Formatter(f"{COLOR_GREEN}[%(asctime)s][%(name)s][%(levelname)s]{COLOR_RESET} - %(message)s") console_handler.setFormatter(console_formatter) # 3. Create a file handler file_handler = logging.FileHandler(log_file) # Specify the log file name file_handler.setLevel(logging.DEBUG) # Set the level for file output (e.g., capture more details in the file) file_formatter = logging.Formatter("[%(asctime)s][%(name)s][%(levelname)s] - %(message)s") file_handler.setFormatter(file_formatter) # 4. Add handlers to the logger logger.addHandler(console_handler) logger.addHandler(file_handler) return logger @contextmanager def use_lora_adapter( model, # 应该是 QwenImageEditPlusPipeline 类型 lora_paths: str | list[str], adapter_names: str | list[str], merge_weights: float | list[float], logger: logging.Logger, ): # normalize to list if isinstance(lora_paths, str): lora_paths = [lora_paths] if isinstance(adapter_names, str): adapter_names = [adapter_names] if isinstance(merge_weights, float): merge_weights = [merge_weights] # --- 获取原始 Base Model 的引用 --- # 确保我们始终操作最底层的 Transformer 模块 transformer = model.transformer while isinstance(transformer, peft.PeftModel): transformer = transformer.base_model.model # original_transformer = transformer # 存储一个干净的引用，以备后续清理和恢复 # ----------------------------- # 1. 彻底清理 Base Model (ENTER 阶段) # ----------------------------- # 1a. 如果当前 Pipeline 中有 PeftModel，先卸载它 if isinstance(model.transformer, peft.PeftModel): logger.info("[LoRA] Unloading existing PeftModel wrapper.") try: # unload() 会返回 Base Model，并清理 LoRA 权重 model.transformer = model.transformer.unload() except AttributeError: # 兼容不同 peft 版本 model.transformer = model.transformer.base_model.model # 1b. 【关键】从 Base Model 实例上强制清理残留的 PEFT 配置 # Base Model 现在是 QwenImageTransformer2DModel 或其基类 if hasattr(model.transformer, "peft_config"): # 尝试删除 PEFT 配置字典，这是配置残留的主要原因 del model.transformer.peft_config logger.info(f"[LoRA] Successfully removed residual peft_config from Base Model.") # 确保 transformer 变量指向当前干净的 Base Model transformer = model.transformer # ----------------------------- # 2. Load new LoRA(s) # ----------------------------- try: if torch.cuda.is_available(): mem_before = torch.cuda.memory_allocated(model.device) / 1024**2 logger.info(f"[LoRA] GPU memory before load: {mem_before:.2f} MB") # 2a. 检查配置一致性并获取配置 configs = [peft.LoraConfig.from_pretrained(p) for p in lora_paths] for i in range(1, len(configs)): assert configs[i] == configs[0], "All LoRA configs must be identical to merge." lora_config = configs[0] # 2b. 使用第一个适配器配置，将 Base Model 包装成 PeftModel # 注意：这里我们使用 adapter_names[0] 初始化第一个配置 transformer = peft.get_peft_model(transformer, lora_config, adapter_name=adapter_names[0]) logger.info(f"After get_peft_model: {type(transformer)=}") # 2c. 加载所有适配器（包括第一个，以确保权重加载） for path, name in zip(lora_paths, adapter_names): # 虽然第一个已初始化，但 load_adapter 确保权重被加载 logger.info(f"[LoRA] Loading adapter: {name=}, {path=}") transformer.load_adapter(path, adapter_name=name) # 2d. 合并权重并设置模型状态 # set_adapters 会设置 active_adapters 列表以及对应的加权系数 transformer.set_adapters(adapter_names, merge_weights) transformer.requires_grad_(False) transformer.eval() # --- swap into pipeline --- model.transformer = transformer if torch.cuda.is_available(): mem_after = torch.cuda.memory_allocated(model.device) / 1024**2 logger.info(f"[LoRA] GPU memory after load: {mem_after:.2f} MB") yield model finally: # ----------------------------- # 3. 彻底卸载 LoRA(s) (EXIT 阶段) # ----------------------------- current_transformer = model.transformer if isinstance(current_transformer, peft.PeftModel): logger.info("[LoRA] Starting LoRA unload and cleanup...") # 3a. 卸载 PeftModel 包装器和权重 try: # 这一步将清理 GPU 上的 LoRA 权重，并返回 Base Model unloaded_transformer = current_transformer.unload() model.transformer = unloaded_transformer except Exception as e: logger.error(f"[LoRA] Failed to use unload(): {e}. Manually retrieving base model.") model.transformer = current_transformer.base_model.model # 3b. 【关键】再次从 Base Model 上强制清理残留的 PEFT 配置 # 确保 Base Model 在离开上下文管理器后是完全干净的 if hasattr(model.transformer, "peft_config"): del model.transformer.peft_config logger.info(f"[LoRA] Final cleanup: Removed residual peft_config from Base Model.") torch.cuda.empty_cache() if torch.cuda.is_available(): mem_final = torch.cuda.memory_allocated(model.device) / 1024**2 logger.info(f"[LoRA] GPU memory after final unload: {mem_final:.2f} MB") if __name__ == "__main__": device = "cuda" lora_paths = [ "/temp/xr/StyQA-latency/lora_adapters/3D_Chibi/3D_Chibi_adapter", ] lora_paths2 = [ "/temp/xr/StyQA-latency/lora_adapters/pixel_level_0/pixel_level_0_adapter", "/temp/xr/StyQA-latency/lora_adapters/pixel_level_1/pixel_level_1_adapter", ] adapter_names = [ "3D_Chibi_adapter", ] adapter_names2 = [ "pixel_level_0_adapter", "pixel_level_1_adapter", ] merge_weights = [1.0] merge_weights2 = [0.5, 0.5] model = QwenImageEditPlusPipeline.from_pretrained( "Qwen/Qwen-Image-Edit-2509", torch_dtype=torch.bfloat16, ).to(device) logger = get_logger(__name__, "lora_test.log") cnt_image = "/temp/rey/eval_datasets/unsplash_subset/_ZO2obWfpFQ.jpg" ref_image = "/temp/rey/eval_datasets/omniconsistency_subset/3D_Chibi/scene/040.jpg" ref_image2 = "/temp/rey/eval_datasets/style30k_subset/s0416____0912_01_query_2_img_000050_1683438561136_035474498831515366.jpg.jpg" limit = 3 for i in range(limit): with use_lora_adapter(model, lora_paths, adapter_names, merge_weights, logger): output = model( image=Image.open(cnt_image).convert("RGB"), prompt="Modify the image to 3D Chibi style. Keep the content unchanged.", num_inference_steps=25, height=1024, width=1024, output_type="pil", return_dict=True, generator=torch.manual_seed(42), ).images[0] output.save(f"lora_test_output1_{i}.jpg") with use_lora_adapter(model, lora_paths2, adapter_names2, merge_weights2, logger): output = model( image=[Image.open(cnt_image).convert("RGB"), Image.open(ref_image2).convert("RGB")], prompt="", num_inference_steps=25, height=1024, width=1024, output_type="pil", return_dict=True, generator=torch.manual_seed(42), ).images[0] output.save(f"lora_test_output2_{i}.jpg")