import os import torch import numpy as np import argparse from peft import LoraConfig from pipeline_dedit_sdxl import DEditSDXLPipeline from pipeline_dedit_sd import DEditSDPipeline from utils import load_image, load_mask, load_mask_edit from utils_mask import process_mask_move_torch, process_mask_remove_torch, mask_union_torch, mask_substract_torch, create_outer_edge_mask_torch from utils_mask import check_mask_overlap_torch, check_cover_all_torch, visualize_mask_list, get_mask_difference_torch, save_mask_list_to_npys print("PyTorch 版本:", torch.__version__) print("CUDA 版本:", torch.version.cuda) print("CUDA 是否可用:", torch.cuda.is_available()) import subprocess # 检查 nvidia-smi 命令的输出 try: result = subprocess.run(['nvidia-smi'], stdout=subprocess.PIPE, stderr=subprocess.PIPE) print(result.stdout.decode()) except FileNotFoundError: print("nvidia-smi 命令不可用，可能未安装 NVIDIA 驱动。") device_count = torch.cuda.device_count() print("可用的 CUDA 设备数量:", device_count) for i in range(device_count): print(f"设备 {i} 名称:", torch.cuda.get_device_name(i)) import os try: ctypes.CDLL('libnvidia-ml.so') print("成功加载 libnvidia-ml.so。") except Exception as e: print("无法加载 libnvidia-ml.so，错误信息:", e) import glob libs = glob.glob('/usr/lib*/**/libnvidia-*.so*', recursive=True) print("NVIDIA 驱动库文件:", libs) print("LD_LIBRARY_PATH:", os.environ.get('LD_LIBRARY_PATH', '未设置')) try: with open('/proc/driver/nvidia/version') as f: print(f.read()) except Exception as e: print("无法读取驱动程序版本信息，错误信息:", e) print("PyTorch 安装路径:", torch.__file__) cuda_paths = [path for path in os.listdir('/usr/local') if path.startswith('cuda')] print("CUDA 安装路径:", cuda_paths) print("PyTorch CUDA 版本:", torch.version.cuda) print("系统 CUDA 版本:", os.environ.get('CUDA_VERSION', '未设置')) def run_main( name="example_tmp", name_2=None, mask_np_list=None, mask_label_list=None, image_gt=None, dpm="sd", resolution=512, seed=42, embedding_learning_rate=1e-4, max_emb_train_steps=200, diffusion_model_learning_rate=5e-5, max_diffusion_train_steps=200, train_batch_size=1, gradient_accumulation_steps=1, num_tokens=1, load_trained=False , num_sampling_steps=50, guidance_scale= 3 , strength=0.8, train_full_lora=False , lora_rank=4, lora_alpha=4, prompt_auxin_list = None, prompt_auxin_idx_list= None, load_edited_mask=False, load_edited_processed_mask=False, edge_thickness=20, num_imgs= 1 , active_mask_list = None, tgt_index=None, recon=False , recon_an_item=False, recon_prompt=None, text=False, tgt_prompt=None, image=False , src_index=None, tgt_name=None, move_resize=False , tgt_indices_list=None, delta_x_list=None, delta_y_list=None, priority_list=None, force_mask_remain=None, resize_list=None, remove=False, load_edited_removemask=False ): torch.cuda.manual_seed_all(seed) torch.manual_seed(seed) base_input_folder = "." base_output_folder = "." input_folder = os.path.join(base_input_folder, name) mask_list = [] for mask_np in mask_np_list: mask = torch.from_numpy(mask_np.astype(np.uint8)) mask_list.append(mask) #mask_list, mask_label_list = load_mask(input_folder) assert mask_list[0].shape[0] == resolution, "Segmentation should be done on size {}".format(resolution) #try: # image_gt = load_image(os.path.join(input_folder, "img_{}.png".format(resolution) ), size = resolution) #except: # image_gt = load_image(os.path.join(input_folder, "img_{}.jpg".format(resolution) ), size = resolution) if image: input_folder_2 = os.path.join(base_input_folder, name_2) mask_list_2, mask_label_list_2 = load_mask(input_folder_2) assert mask_list_2[0].shape[0] == resolution, "Segmentation should be done on size {}".format(resolution) try: image_gt_2 = load_image(os.path.join(input_folder_2, "img_{}.png".format(resolution) ), size = resolution) except: image_gt_2 = load_image(os.path.join(input_folder_2, "img_{}.jpg".format(resolution) ), size = resolution) output_dir = os.path.join(base_output_folder, name + "_" + name_2) os.makedirs(output_dir, exist_ok = True) else: output_dir = os.path.join(base_output_folder, name) os.makedirs(output_dir, exist_ok = True) if dpm == "sd": if image: pipe = DEditSDPipeline(mask_list, mask_label_list, mask_list_2, mask_label_list_2, resolution = resolution, num_tokens = num_tokens) else: pipe = DEditSDPipeline(mask_list, mask_label_list, resolution = resolution, num_tokens = num_tokens) elif dpm == "sdxl": if image: pipe = DEditSDXLPipeline(mask_list, mask_label_list, mask_list_2, mask_label_list_2, resolution = resolution, num_tokens = num_tokens) else: pipe = DEditSDXLPipeline(mask_list, mask_label_list, resolution = resolution, num_tokens = num_tokens) else: raise NotImplementedError set_string_list = pipe.set_string_list if prompt_auxin_list is not None: for auxin_idx, auxin_prompt in zip(prompt_auxin_idx_list, prompt_auxin_list): set_string_list[auxin_idx] = auxin_prompt.replace("*", set_string_list[auxin_idx] ) print(set_string_list) if image: set_string_list_2 = pipe.set_string_list_2 print(set_string_list_2) if load_trained: unet_save_path = os.path.join(output_dir, "unet.pt") unet_state_dict = torch.load(unet_save_path) text_encoder1_save_path = os.path.join(output_dir, "text_encoder1.pt") text_encoder1_state_dict = torch.load(text_encoder1_save_path) if dpm == "sdxl": text_encoder2_save_path = os.path.join(output_dir, "text_encoder2.pt") text_encoder2_state_dict = torch.load(text_encoder2_save_path) if 'lora' in ''.join(unet_state_dict.keys()): unet_lora_config = LoraConfig( r=lora_rank, lora_alpha=lora_alpha, init_lora_weights="gaussian", target_modules=["to_k", "to_q", "to_v", "to_out.0"], ) pipe.unet.add_adapter(unet_lora_config) pipe.unet.load_state_dict(unet_state_dict) pipe.text_encoder.load_state_dict(text_encoder1_state_dict) if dpm == "sdxl": pipe.text_encoder_2.load_state_dict(text_encoder2_state_dict) else: if image: pipe.mask_list = [m.cuda() for m in pipe.mask_list] pipe.mask_list_2 = [m.cuda() for m in pipe.mask_list_2] pipe.train_emb_2imgs( image_gt, image_gt_2, set_string_list, set_string_list_2, gradient_accumulation_steps = gradient_accumulation_steps, embedding_learning_rate = embedding_learning_rate, max_emb_train_steps = max_emb_train_steps, train_batch_size = train_batch_size, ) pipe.train_model_2imgs( image_gt, image_gt_2, set_string_list, set_string_list_2, gradient_accumulation_steps = gradient_accumulation_steps, max_diffusion_train_steps = max_diffusion_train_steps, diffusion_model_learning_rate = diffusion_model_learning_rate , train_batch_size =train_batch_size, train_full_lora = train_full_lora, lora_rank = lora_rank, lora_alpha = lora_alpha ) else: pipe.mask_list = [m.cuda() for m in pipe.mask_list] pipe.train_emb( image_gt, set_string_list, gradient_accumulation_steps = gradient_accumulation_steps, embedding_learning_rate = embedding_learning_rate, max_emb_train_steps = max_emb_train_steps, train_batch_size = train_batch_size, ) pipe.train_model( image_gt, set_string_list, gradient_accumulation_steps = gradient_accumulation_steps, max_diffusion_train_steps = max_diffusion_train_steps, diffusion_model_learning_rate = diffusion_model_learning_rate , train_batch_size = train_batch_size, train_full_lora = train_full_lora, lora_rank = lora_rank, lora_alpha = lora_alpha ) unet_save_path = os.path.join(output_dir, "unet.pt") torch.save(pipe.unet.state_dict(),unet_save_path ) text_encoder1_save_path = os.path.join(output_dir, "text_encoder1.pt") torch.save(pipe.text_encoder.state_dict(), text_encoder1_save_path) if dpm == "sdxl": text_encoder2_save_path = os.path.join(output_dir, "text_encoder2.pt") torch.save(pipe.text_encoder_2.state_dict(), text_encoder2_save_path ) if recon: output_dir = os.path.join(output_dir, "recon") os.makedirs(output_dir, exist_ok = True) if recon_an_item: mask_list = [torch.from_numpy(np.ones_like(mask_list[0].numpy()))] tgt_string = set_string_list[tgt_index] tgt_string = recon_prompt.replace("*", tgt_string) set_string_list = [tgt_string] print(set_string_list) save_path = os.path.join(output_dir, "out_recon.png") x_np = pipe.inference_with_mask( save_path, guidance_scale = guidance_scale, num_sampling_steps = num_sampling_steps, seed = seed, num_imgs = num_imgs, set_string_list = set_string_list, mask_list = mask_list ) if text: print("*** Text-guided editing ") output_dir = os.path.join(output_dir, "text") os.makedirs(output_dir, exist_ok = True) save_path = os.path.join(output_dir, "out_text.png") set_string_list[tgt_index] = tgt_prompt mask_active = torch.zeros_like(mask_list[0]) mask_active = mask_union_torch(mask_active, mask_list[tgt_index]) if active_mask_list is not None: for midx in active_mask_list: mask_active = mask_union_torch(mask_active, mask_list[midx]) if load_edited_mask: mask_list_edited, mask_label_list_edited = load_mask_edit(input_folder) mask_diff = get_mask_difference_torch(mask_list_edited, mask_list) mask_active = mask_union_torch(mask_active, mask_diff) mask_list = mask_list_edited save_path = os.path.join(output_dir, "out_textEdited.png") mask_hard = mask_substract_torch(torch.ones_like(mask_list[0]), mask_active) mask_soft = create_outer_edge_mask_torch(mask_active, edge_thickness = edge_thickness) mask_hard = mask_substract_torch(mask_hard, mask_soft) pipe.inference_with_mask( save_path, orig_image = image_gt, set_string_list = set_string_list, guidance_scale = guidance_scale, strength = strength, num_imgs = num_imgs, mask_hard= mask_hard, mask_soft = mask_soft, mask_list = mask_list, seed = seed, num_sampling_steps = num_sampling_steps ) if remove: output_dir = os.path.join(output_dir, "remove") save_path = os.path.join(output_dir, "out_remove.png") os.makedirs(output_dir, exist_ok = True) mask_active = torch.zeros_like(mask_list[0]) if load_edited_mask: mask_list_edited, _ = load_mask_edit(input_folder) mask_diff = get_mask_difference_torch(mask_list_edited, mask_list) mask_active = mask_union_torch(mask_active, mask_diff) mask_list = mask_list_edited if load_edited_processed_mask: # manually edit or draw masks after removing one index, then load mask_list_processed, _ = load_mask_edit(output_dir) mask_remain = get_mask_difference_torch(mask_list_processed, mask_list) else: # generate masks after removing one index, using nearest neighbor algorithm mask_list_processed, mask_remain = process_mask_remove_torch(mask_list, tgt_index) save_mask_list_to_npys(output_dir, mask_list_processed, mask_label_list, name = "mask") visualize_mask_list(mask_list_processed, os.path.join(output_dir, "seg_removed.png")) check_cover_all_torch(*mask_list_processed) mask_active = mask_union_torch(mask_active, mask_remain) if active_mask_list is not None: for midx in active_mask_list: mask_active = mask_union_torch(mask_active, mask_list[midx]) mask_hard = 1 - mask_active mask_soft = create_outer_edge_mask_torch(mask_remain, edge_thickness = edge_thickness) mask_hard = mask_substract_torch(mask_hard, mask_soft) pipe.inference_with_mask( save_path, orig_image = image_gt, guidance_scale = guidance_scale, strength = strength, num_imgs = num_imgs, mask_hard= mask_hard, mask_soft = mask_soft, mask_list = mask_list_processed, seed = seed, num_sampling_steps = num_sampling_steps ) if image: output_dir = os.path.join(output_dir, "image") save_path = os.path.join(output_dir, "out_image.png") os.makedirs(output_dir, exist_ok = True) mask_active = torch.zeros_like(mask_list[0]) if None not in (tgt_name, src_index, tgt_index): if tgt_name == name: set_string_list_tgt = set_string_list set_string_list_src = set_string_list_2 image_tgt = image_gt if load_edited_mask: mask_list_edited, _ = load_mask_edit(input_folder) mask_diff = get_mask_difference_torch(mask_list_edited, mask_list) mask_active = mask_union_torch(mask_active, mask_diff) mask_list = mask_list_edited save_path = os.path.join(output_dir, "out_imageEdited.png") mask_list_tgt = mask_list elif tgt_name == name_2: set_string_list_tgt = set_string_list_2 set_string_list_src = set_string_list image_tgt = image_gt_2 if load_edited_mask: mask_list_2_edited, _ = load_mask_edit(input_folder_2) mask_diff = get_mask_difference_torch(mask_list_2_edited, mask_list_2) mask_active = mask_union_torch(mask_active, mask_diff) mask_list_2 = mask_list_2_edited save_path = os.path.join(output_dir, "out_imageEdited.png") mask_list_tgt = mask_list_2 else: exit("tgt_name should be either name or name_2") set_string_list_tgt[tgt_index] = set_string_list_src[src_index] mask_active = mask_list_tgt[tgt_index] mask_frozen = (1-mask_active.float()).to(mask_active.device) mask_soft = create_outer_edge_mask_torch(mask_active.cpu(), edge_thickness = edge_thickness) mask_hard = mask_substract_torch(mask_frozen.cpu(), mask_soft.cpu()) mask_list_tgt = [m.cuda() for m in mask_list_tgt] pipe.inference_with_mask( save_path, set_string_list = set_string_list_tgt, mask_list = mask_list_tgt, guidance_scale = guidance_scale, num_sampling_steps = num_sampling_steps, mask_hard = mask_hard.cuda(), mask_soft = mask_soft.cuda(), num_imgs = num_imgs, orig_image = image_tgt, strength = strength, ) if move_resize: output_dir = os.path.join(output_dir, "move_resize") os.makedirs(output_dir, exist_ok = True) save_path = os.path.join(output_dir, "out_moveresize.png") mask_active = torch.zeros_like(mask_list[0]) if load_edited_mask: mask_list_edited, _ = load_mask_edit(input_folder) mask_diff = get_mask_difference_torch(mask_list_edited, mask_list) mask_active = mask_union_torch(mask_active, mask_diff) mask_list = mask_list_edited # save_path = os.path.join(output_dir, "out_moveresizeEdited.png") if load_edited_processed_mask: mask_list_processed, _ = load_mask_edit(output_dir) mask_remain = get_mask_difference_torch(mask_list_processed, mask_list) else: mask_list_processed, mask_remain = process_mask_move_torch( mask_list, tgt_indices_list, delta_x_list, delta_y_list, priority_list, force_mask_remain = force_mask_remain, resize_list = resize_list ) save_mask_list_to_npys(output_dir, mask_list_processed, mask_label_list, name = "mask") visualize_mask_list(mask_list_processed, os.path.join(output_dir, "seg_move_resize.png")) active_idxs = tgt_indices_list mask_active = mask_union_torch(mask_active, *[m for midx, m in enumerate(mask_list_processed) if midx in active_idxs]) mask_active = mask_union_torch(mask_remain, mask_active) if active_mask_list is not None: for midx in active_mask_list: mask_active = mask_union_torch(mask_active, mask_list_processed[midx]) mask_frozen =(1 - mask_active.float()) mask_soft = create_outer_edge_mask_torch(mask_active, edge_thickness = edge_thickness) mask_hard = mask_substract_torch(mask_frozen, mask_soft) check_mask_overlap_torch(mask_hard, mask_soft) pipe.inference_with_mask( save_path, strength = strength, orig_image = image_gt, guidance_scale = guidance_scale, num_sampling_steps = num_sampling_steps, num_imgs = num_imgs, mask_hard= mask_hard, mask_soft = mask_soft, mask_list = mask_list_processed, seed = seed )