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, ).images[0] return image if __name__ == "__main__": output_image_root_folder = "./canny" model_id = f"../../control_sd15_canny" base_model_id = sys.argv[1] if len(sys.argv) == 2 else None canny_edged_image = load_image( "https://huggingface.co/takuma104/controlnet_dev/resolve/main/vermeer_canny_edged.png" ) if base_model_id: unet = UNet2DConditionModel.from_pretrained(base_model_id, subfolder="unet").to( "cuda" ) vae = AutoencoderKL.from_pretrained(base_model_id, subfolder="vae").to("cuda") output_types = [ base_model_id.split("/")[1] + suffix for suffix in ["_unet", "_unet_vae"] ] 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" ) elif output_type.endswith("_unet"): print(f"{base_model_id} unet only override config") pipe = StableDiffusionControlNetPipeline.from_pretrained( model_id, unet=unet ).to("cuda") elif output_type.endswith("_unet_vae"): print(f"{base_model_id} unet & vae override config") pipe = StableDiffusionControlNetPipeline.from_pretrained( model_id, unet=unet, vae=vae ).to("cuda") output_folder = f"{output_image_root_folder}/{output_type}" os.makedirs(output_folder, exist_ok=True) for seed in range(32): image = generate_image(seed=seed, control=canny_edged_image) image.save(f"{output_folder}/output_{seed:02d}.png")