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")