|
|
|
|
|
|
|
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': |
|
control = torch.flip(control, dims=[1]) |
|
|
|
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') |
|
|
