|
|
|
|
|
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") |
|
|
