Spaces:
Running
on
Zero
Running
on
Zero
| 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 | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--name", type=str,required=True, default=None) | |
| parser.add_argument("--name_2", type=str,required=False, default=None) | |
| parser.add_argument("--dpm", type=str,required=True, default="sd") | |
| parser.add_argument("--resolution", type=int, default=1024) | |
| parser.add_argument("--seed", type=int, default=42) | |
| parser.add_argument("--embedding_learning_rate", type=float, default=1e-4) | |
| parser.add_argument("--max_emb_train_steps", type=int, default=200) | |
| parser.add_argument("--diffusion_model_learning_rate", type=float, default=5e-5) | |
| parser.add_argument("--max_diffusion_train_steps", type=int, default=200) | |
| parser.add_argument("--train_batch_size", type=int, default=1) | |
| parser.add_argument("--gradient_accumulation_steps", type=int, default=1) | |
| parser.add_argument("--num_tokens", type=int, default=1) | |
| parser.add_argument("--load_trained", default=False, action="store_true" ) | |
| parser.add_argument("--num_sampling_steps", type=int, default=50) | |
| parser.add_argument("--guidance_scale", type=float, default = 3 ) | |
| parser.add_argument("--strength", type=float, default=0.8) | |
| parser.add_argument("--train_full_lora", default=False, action="store_true" ) | |
| parser.add_argument("--lora_rank", type=int, default=4) | |
| parser.add_argument("--lora_alpha", type=int, default=4) | |
| parser.add_argument("--prompt_auxin_list", nargs="+", type=str, default = None) | |
| parser.add_argument("--prompt_auxin_idx_list", nargs="+", type=int, default = None) | |
| # general editing configs | |
| parser.add_argument("--load_edited_mask", default=False, action="store_true") | |
| parser.add_argument("--load_edited_processed_mask", default=False, action="store_true") | |
| parser.add_argument("--edge_thickness", type=int, default=20) | |
| parser.add_argument("--num_imgs", type=int, default = 1 ) | |
| parser.add_argument('--active_mask_list', nargs="+", type=int) | |
| parser.add_argument("--tgt_index", type=int, default=None) | |
| # recon | |
| parser.add_argument("--recon", default=False, action="store_true" ) | |
| parser.add_argument("--recon_an_item", default=False, action="store_true" ) | |
| parser.add_argument("--recon_prompt", type=str, default=None) | |
| # text-based editing | |
| parser.add_argument("--text", default=False, action="store_true") | |
| parser.add_argument("--tgt_prompt", type=str, default=None) | |
| # image-based editing | |
| parser.add_argument("--image", default=False, action="store_true" ) | |
| parser.add_argument("--src_index", type=int, default=None) | |
| parser.add_argument("--tgt_name", type=str, default=None) | |
| # mask-based move | |
| parser.add_argument("--move_resize", default=False, action="store_true" ) | |
| parser.add_argument('--tgt_indices_list', nargs="+", type=int) | |
| parser.add_argument("--delta_x_list", nargs="+", type=int) | |
| parser.add_argument("--delta_y_list", nargs="+", type=int) | |
| parser.add_argument("--priority_list", nargs="+", type=int) | |
| parser.add_argument("--force_mask_remain", type=int, default=None) | |
| parser.add_argument("--resize_list", nargs="+", type=float) | |
| # remove | |
| parser.add_argument("--remove", default=False, action="store_true" ) | |
| parser.add_argument("--load_edited_removemask", default=False, action="store_true") | |
| args = parser.parse_args() | |
| torch.cuda.manual_seed_all(args.seed) | |
| torch.manual_seed(args.seed) | |
| base_input_folder = "." | |
| base_output_folder = "." | |
| input_folder = os.path.join(base_input_folder, args.name) | |
| mask_list, mask_label_list = load_mask(input_folder) | |
| assert mask_list[0].shape[0] == args.resolution, "Segmentation should be done on size {}".format(args.resolution) | |
| try: | |
| image_gt = load_image(os.path.join(input_folder, "img_{}.png".format(args.resolution) ), size = args.resolution) | |
| except: | |
| image_gt = load_image(os.path.join(input_folder, "img_{}.jpg".format(args.resolution) ), size = args.resolution) | |
| if args.image: | |
| input_folder_2 = os.path.join(base_input_folder, args.name_2) | |
| mask_list_2, mask_label_list_2 = load_mask(input_folder_2) | |
| assert mask_list_2[0].shape[0] == args.resolution, "Segmentation should be done on size {}".format(args.resolution) | |
| try: | |
| image_gt_2 = load_image(os.path.join(input_folder_2, "img_{}.png".format(args.resolution) ), size = args.resolution) | |
| except: | |
| image_gt_2 = load_image(os.path.join(input_folder_2, "img_{}.jpg".format(args.resolution) ), size = args.resolution) | |
| output_dir = os.path.join(base_output_folder, args.name + "_" + args.name_2) | |
| os.makedirs(output_dir, exist_ok = True) | |
| else: | |
| output_dir = os.path.join(base_output_folder, args.name) | |
| os.makedirs(output_dir, exist_ok = True) | |
| if args.dpm == "sd": | |
| if args.image: | |
| pipe = DEditSDPipeline(mask_list, mask_label_list, mask_list_2, mask_label_list_2, resolution = args.resolution, num_tokens = args.num_tokens) | |
| else: | |
| pipe = DEditSDPipeline(mask_list, mask_label_list, resolution = args.resolution, num_tokens = args.num_tokens) | |
| elif args.dpm == "sdxl": | |
| if args.image: | |
| pipe = DEditSDXLPipeline(mask_list, mask_label_list, mask_list_2, mask_label_list_2, resolution = args.resolution, num_tokens = args.num_tokens) | |
| else: | |
| pipe = DEditSDXLPipeline(mask_list, mask_label_list, resolution = args.resolution, num_tokens = args.num_tokens) | |
| else: | |
| raise NotImplementedError | |
| set_string_list = pipe.set_string_list | |
| if args.prompt_auxin_list is not None: | |
| for auxin_idx, auxin_prompt in zip(args.prompt_auxin_idx_list, args.prompt_auxin_list): | |
| set_string_list[auxin_idx] = auxin_prompt.replace("*", set_string_list[auxin_idx] ) | |
| print(set_string_list) | |
| if args.image: | |
| set_string_list_2 = pipe.set_string_list_2 | |
| print(set_string_list_2) | |
| if args.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 args.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=args.lora_rank, | |
| lora_alpha=args.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 args.dpm == "sdxl": | |
| pipe.text_encoder_2.load_state_dict(text_encoder2_state_dict) | |
| else: | |
| if args.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 = args.gradient_accumulation_steps, | |
| embedding_learning_rate = args.embedding_learning_rate, | |
| max_emb_train_steps = args.max_emb_train_steps, | |
| train_batch_size = args.train_batch_size, | |
| ) | |
| pipe.train_model_2imgs( | |
| image_gt, | |
| image_gt_2, | |
| set_string_list, | |
| set_string_list_2, | |
| gradient_accumulation_steps = args.gradient_accumulation_steps, | |
| max_diffusion_train_steps = args.max_diffusion_train_steps, | |
| diffusion_model_learning_rate = args.diffusion_model_learning_rate , | |
| train_batch_size =args.train_batch_size, | |
| train_full_lora = args.train_full_lora, | |
| lora_rank = args.lora_rank, | |
| lora_alpha = args.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 = args.gradient_accumulation_steps, | |
| embedding_learning_rate = args.embedding_learning_rate, | |
| max_emb_train_steps = args.max_emb_train_steps, | |
| train_batch_size = args.train_batch_size, | |
| ) | |
| pipe.train_model( | |
| image_gt, | |
| set_string_list, | |
| gradient_accumulation_steps = args.gradient_accumulation_steps, | |
| max_diffusion_train_steps = args.max_diffusion_train_steps, | |
| diffusion_model_learning_rate = args.diffusion_model_learning_rate , | |
| train_batch_size = args.train_batch_size, | |
| train_full_lora = args.train_full_lora, | |
| lora_rank = args.lora_rank, | |
| lora_alpha = args.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 args.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 args.recon: | |
| output_dir = os.path.join(output_dir, "recon") | |
| os.makedirs(output_dir, exist_ok = True) | |
| if args.recon_an_item: | |
| mask_list = [torch.from_numpy(np.ones_like(mask_list[0].numpy()))] | |
| tgt_string = set_string_list[args.tgt_index] | |
| tgt_string = args.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 = args.guidance_scale, | |
| num_sampling_steps = args.num_sampling_steps, | |
| seed = args.seed, | |
| num_imgs = args.num_imgs, | |
| set_string_list = set_string_list, | |
| mask_list = mask_list | |
| ) | |
| if args.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[args.tgt_index] = args.tgt_prompt | |
| mask_active = torch.zeros_like(mask_list[0]) | |
| mask_active = mask_union_torch(mask_active, mask_list[args.tgt_index]) | |
| if args.active_mask_list is not None: | |
| for midx in args.active_mask_list: | |
| mask_active = mask_union_torch(mask_active, mask_list[midx]) | |
| if args.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 = args.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 = args.guidance_scale, | |
| strength = args.strength, | |
| num_imgs = args.num_imgs, | |
| mask_hard= mask_hard, | |
| mask_soft = mask_soft, | |
| mask_list = mask_list, | |
| seed = args.seed, | |
| num_sampling_steps = args.num_sampling_steps | |
| ) | |
| if args.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 args.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 args.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, args.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 args.active_mask_list is not None: | |
| for midx in args.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 = args.edge_thickness) | |
| mask_hard = mask_substract_torch(mask_hard, mask_soft) | |
| pipe.inference_with_mask( | |
| save_path, | |
| orig_image = image_gt, | |
| guidance_scale = args.guidance_scale, | |
| strength = args.strength, | |
| num_imgs = args.num_imgs, | |
| mask_hard= mask_hard, | |
| mask_soft = mask_soft, | |
| mask_list = mask_list_processed, | |
| seed = args.seed, | |
| num_sampling_steps = args.num_sampling_steps | |
| ) | |
| if args.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 (args.tgt_name, args.src_index, args.tgt_index): | |
| if args.tgt_name == args.name: | |
| set_string_list_tgt = set_string_list | |
| set_string_list_src = set_string_list_2 | |
| image_tgt = image_gt | |
| if args.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 args.tgt_name == args.name_2: | |
| set_string_list_tgt = set_string_list_2 | |
| set_string_list_src = set_string_list | |
| image_tgt = image_gt_2 | |
| if args.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[args.tgt_index] = set_string_list_src[args.src_index] | |
| mask_active = mask_list_tgt[args.tgt_index] | |
| mask_frozen = (1-mask_active.float()).to(mask_active.device) | |
| mask_soft = create_outer_edge_mask_torch(mask_active.cpu(), edge_thickness = args.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 = args.guidance_scale, | |
| num_sampling_steps = args.num_sampling_steps, | |
| mask_hard = mask_hard.cuda(), | |
| mask_soft = mask_soft.cuda(), | |
| num_imgs = args.num_imgs, | |
| orig_image = image_tgt, | |
| strength = args.strength, | |
| ) | |
| if args.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 args.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 args.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, | |
| args.tgt_indices_list, | |
| args.delta_x_list, | |
| args.delta_y_list, args.priority_list, | |
| force_mask_remain = args.force_mask_remain, | |
| resize_list = args.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 = args.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 args.active_mask_list is not None: | |
| for midx in args.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 = args.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 = args.strength, | |
| orig_image = image_gt, | |
| guidance_scale = args.guidance_scale, | |
| num_sampling_steps = args.num_sampling_steps, | |
| num_imgs = args.num_imgs, | |
| mask_hard= mask_hard, | |
| mask_soft = mask_soft, | |
| mask_list = mask_list_processed, | |
| seed = args.seed | |
| ) | |