File size: 2,339 Bytes
e75ac26 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
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")
|