| import numpy as np | |
| import torch | |
| import sys | |
| import os | |
| from diffusers import ( | |
| StableDiffusionControlNetPipeline, | |
| AutoencoderKL, | |
| UNet2DConditionModel, | |
| ) | |
| from diffusers.utils import load_image | |
| test_prompt = "best quality, extremely detailed" | |
| test_negative_prompt = "lowres, bad anatomy, worst quality, low quality" | |
| def generate_image(seed, control): | |
| image = pipe( | |
| prompt=test_prompt, | |
| negative_prompt=test_negative_prompt, | |
| width=512, | |
| height=512, | |
| generator=torch.Generator(device="cuda").manual_seed(seed), | |
| image=control, | |
| num_inference_steps=30, | |
| ).images[0] | |
| return image | |
| if __name__ == "__main__": | |
| contrlnet_model = sys.argv[1] | |
| output_image_root_folder = f"./all_controlnet_models/{contrlnet_model}" | |
| os.makedirs(output_image_root_folder, exist_ok=True) | |
| model_id = f"../../control_sd15_{contrlnet_model}" | |
| base_model_id = sys.argv[2] if len(sys.argv) == 3 else None | |
| control_image_dict = { | |
| "canny":"control_bird_canny.png", | |
| "depth":"control_vermeer_depth.png", | |
| "hed":"control_bird_hed.png", | |
| "mlsd":"control_room_mlsd.png", | |
| "normal":"control_human_normal.png", | |
| "openpose":"control_human_openpose.png", | |
| "scribble":"control_vermeer_scribble.png", | |
| "seg":"control_room_seg.png", | |
| } | |
| control_image = load_image( | |
| f"https://huggingface.co/takuma104/controlnet_dev/resolve/main/gen_compare/control_images/converted/{control_image_dict[contrlnet_model]}" | |
| ) | |
| if base_model_id: | |
| unet = UNet2DConditionModel.from_pretrained(base_model_id, subfolder="unet").to( | |
| "cuda" | |
| ) | |
| output_types = [base_model_id.split("/")[1]] | |
| else: | |
| output_types = ["sd15"] | |
| for output_type in output_types: | |
| if output_type == "sd15": | |
| print("SD15 no override config") | |
| pipe = StableDiffusionControlNetPipeline.from_pretrained(model_id).to( | |
| "cuda" | |
| ) | |
| else: | |
| print(f"{base_model_id} unet only override config") | |
| pipe = StableDiffusionControlNetPipeline.from_pretrained( | |
| model_id, unet=unet | |
| ).to("cuda") | |
| output_folder = f"{output_image_root_folder}/{output_type}" | |
| os.makedirs(output_folder, exist_ok=True) | |
| for seed in range(16): | |
| image = generate_image(seed=seed, control=control_image) | |
| image.save(f"{output_folder}/output_{seed:02d}.png") | |