| |
|
|
| import einops |
| import numpy as np |
| import torch |
| import sys |
| import os |
| import yaml |
|
|
| 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, |
| |
| ).images[0] |
| return image |
|
|
| def control_images(control_image_folder, model_name): |
| with open('./control_images.yaml', 'r') as f: |
| d = yaml.safe_load(f) |
| filenames = d[model_name] |
| return [Image.open(f'{control_image_folder}/{fn}').convert("RGB") for fn in filenames] |
|
|
| if __name__ == '__main__': |
| model_name = sys.argv[1] |
| control_image_folder = './control_images/converted/' |
| output_image_folder = './output_images/diffusers/' |
| os.makedirs(output_image_folder, exist_ok=True) |
|
|
| model_id = f"lllyasviel/control_v11{model_name}" |
|
|
| controlnet = ControlNetModel.from_pretrained(model_id) |
| if model_name == 'p_sd15s2_lineart_anime': |
| base_model_id = 'Linaqruf/anything-v3.0' |
| base_model_revision = None |
| else: |
| base_model_id = "runwayml/stable-diffusion-v1-5" |
| base_model_revision = 'non-ema' |
|
|
| pipe = StableDiffusionControlNetPipeline.from_pretrained(base_model_id, |
| revision=base_model_revision, |
| controlnet=controlnet, |
| safety_checker=None).to("cuda") |
| pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config) |
|
|
| for i, control_image in enumerate(control_images(control_image_folder, model_name)): |
| 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() |
|
|
| |
| |
|
|
| 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_{model_name}_{i}_{seed}.png') |
|
|
