| | |
| |
|
| | 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 make_image_condition(image, image_mask=None): |
| | image = np.array(image.convert("RGB")).astype(np.float32) / 255.0 |
| | if image_mask is not None: |
| | image_mask = np.array(image_mask.convert("L")) |
| | assert ( |
| | image.shape[0:1] == image_mask.shape[0:1] |
| | ), "image and image_mask must have the same image size" |
| | image[image_mask < 128] = -1.0 |
| | image = np.expand_dims(image, 0).transpose(0, 3, 1, 2) |
| | image = torch.from_numpy(image) |
| | return image |
| |
|
| |
|
| | 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 |
| |
|
| |
|
| | if __name__ == "__main__": |
| | model_name = "p_sd15_inpaint" |
| | original_image_folder = "./control_images/" |
| | 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) |
| |
|
| | original_image_filenames = [ |
| | "pexels-sound-on-3760767_512x512.png", |
| | "vermeer_512x512.png", |
| | "bird_512x512.png", |
| | ] |
| |
|
| | inpaint_image_conditions = [ |
| | make_image_condition( |
| | Image.open(f"{original_image_folder}{fn}"), |
| | Image.open(f"{original_image_folder}mask_512x512.png"), |
| | ) |
| | for fn in original_image_filenames |
| | ] |
| |
|
| | for i, control in enumerate(inpaint_image_conditions): |
| | 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") |
| |
|
