| 27 | return parser.parse_args() |
| 28 | |
| 29 | def generate(args, weight_dtype): |
| 30 | random.seed(args.seed) |
| 31 | unet = JoDiffusionModel.from_pretrained(args.finetuned_model_path, torch_dtype=weight_dtype) |
| 32 | pipeline_kwargs = { |
| 33 | "pretrained_model_name_or_path": args.pretrained_model_name_or_path, |
| 34 | "unet": unet, "torch_dtype": weight_dtype, |
| 35 | } |
| 36 | if args.lightweight_label_vae: |
| 37 | from pipelines.modeling_lightweight_vae import LightweightLabelVAE |
| 38 | pipeline_kwargs["label_vae"] = LightweightLabelVAE.from_pretrained(args.pretrained_label_vae_path, torch_dtype=weight_dtype) |
| 39 | else: |
| 40 | pipeline_kwargs["label_vae"] = AutoencoderKL.from_pretrained(args.pretrained_label_vae_path, torch_dtype=weight_dtype) |
| 41 | pipeline = JoDiffusionPipeline.from_pretrained(**pipeline_kwargs) |
| 42 | pipeline.set_progress_bar_config(disable=True) |
| 43 | |
| 44 | distributed_state = PartialState() |
| 45 | pipeline = pipeline.to(distributed_state.device) |
| 46 | generator = torch.Generator(device=distributed_state.device).manual_seed(args.seed) |
| 47 | if distributed_state.is_main_process: |
| 48 | if os.path.exists(f"{args.save_dir}/image"): |
| 49 | print(f"Directory {args.save_dir}/image already exists, please doubleclick it.") |
| 50 | os.makedirs(f"{args.save_dir}/image", exist_ok=True) |
| 51 | os.makedirs(f"{args.save_dir}/label", exist_ok=True) |
| 52 | print(f"Generate args: {args}") |
| 53 | |
| 54 | # prepare dataset to args.num_images |
| 55 | assert args.dataset_name.split("_")[0] in args.finetuned_model_path |
| 56 | assert args.dataset_name.split("_")[0] in args.pretrained_label_vae_path |
| 57 | if args.dataset_name == "ade20k_semantic": |
| 58 | caption_path = "../dataset/ADE20K/annotations_caption/training.json" |
| 59 | elif args.dataset_name == "voc_semantic": |
| 60 | caption_path = "../dataset/VOC2012/ImageSets/Caption/trainaug.json" |
| 61 | elif args.dataset_name == "coco_semantic": |
| 62 | caption_path = "../dataset/COCO/annotations/captions_train2017.json" |
| 63 | else: |
| 64 | raise ValueError(f"Unknown dataset {args.dataset_name}") |
| 65 | |
| 66 | caption = json.load(open(caption_path, "r")) |
| 67 | if args.dataset_name == "coco_semantic": |
| 68 | caption_list = [c["caption"] for c in caption["annotations"]] |
| 69 | else: |
| 70 | caption_list = list(caption.values()) |
| 71 | caption_list = caption_list * math.ceil(args.num_images / len(caption_list)) |
| 72 | |
| 73 | random.shuffle(caption_list) |
| 74 | caption_list = caption_list[:args.num_images] |
| 75 | assert len(caption_list) == args.num_images |
| 76 | caption_dataset = [[idx, c] for idx, c in enumerate(caption_list)] |
| 77 | |
| 78 | with distributed_state.split_between_processes(caption_dataset) as split_caption: |
| 79 | for idx, prompt in tqdm(split_caption): |
| 80 | img_path = f"{args.save_dir}/image/{idx}.jpg" |
| 81 | lbl_path = f"{args.save_dir}/label/{idx}.png" |
| 82 | if os.path.isfile(img_path) and os.path.isfile(lbl_path): |
| 83 | continue |
| 84 | sample = pipeline( |
| 85 | mode=args.generate_mode, |
| 86 | prompt=prompt, |