File size: 3,373 Bytes
75805ca a49d82c 75805ca a49d82c 75805ca a49d82c 75805ca 5871943 75805ca 6d2cbb5 |
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 70 71 72 73 74 75 |
# Diffusers' ControlNet Implementation Subjective Evaluation
# https://github.com/takuma104/diffusers/tree/controlnet
import einops
import numpy as np
import torch
import sys
import os
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, DDIMScheduler
from PIL import Image
test_prompt = "best quality, extremely detailed"
test_negative_prompt = "lowres, bad anatomy, worst quality, low quality"
def generate_image(seed, prompt, negative_prompt, control, guess_mode=False):
latent = torch.randn((1,4,64,64), device="cpu", generator=torch.Generator(device="cpu").manual_seed(seed)).cuda()
image = pipe(
prompt=prompt,
negative_prompt=negative_prompt,
guidance_scale=4.0 if guess_mode else 9.0,
num_inference_steps=50 if guess_mode else 20,
latents=latent,
image=control,
guess_mode=guess_mode,
).images[0]
return image
if __name__ == '__main__':
model_name = sys.argv[1]
control_image_folder = '../gen_compare/control_images/converted/'
output_image_folder = './output_images/diffusers/'
os.makedirs(output_image_folder, exist_ok=True)
model_id = f'lllyasviel/sd-controlnet-{model_name}'
controlnet = ControlNetModel.from_pretrained(model_id)
pipe = StableDiffusionControlNetPipeline.from_pretrained("runwayml/stable-diffusion-v1-5",
revision="non-ema",
controlnet=controlnet,
safety_checker=None).to("cuda")
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
image_types = {'bird', 'human', 'room', 'vermeer'}
for image_type in image_types:
control_image = Image.open(f'{control_image_folder}control_{image_type}_{model_name}.png')
control = np.array(control_image)[:,:,::-1].copy()
control = torch.from_numpy(control).float().cuda() / 255.0
control = torch.stack([control for _ in range(1)], dim=0)
control = einops.rearrange(control, 'b h w c -> b c h w').clone()
if model_name == 'normal': # workaround, this should not be necessary
control = torch.flip(control, dims=[1]) # RGB -> BGR
for seed in range(4):
image = generate_image(seed=seed,
prompt=test_prompt,
negative_prompt=test_negative_prompt,
control=control)
image.save(f'{output_image_folder}output_{image_type}_{model_name}_{seed}.png')
image = generate_image(seed=seed,
prompt="",
negative_prompt="",
control=control,
guess_mode=True)
image.save(f'{output_image_folder}output_{image_type}_{model_name}_{seed}_gm.png')
image = generate_image(seed=seed,
prompt=test_prompt,
negative_prompt=test_negative_prompt,
control=control,
guess_mode=True)
image.save(f'{output_image_folder}output_{image_type}_{model_name}_{seed}_gm_wp.png')
|