|
import json |
|
from pathlib import Path |
|
|
|
import piexif |
|
import torch |
|
from PIL import Image |
|
from torch import autocast |
|
|
|
from diffusers.pipelines.stable_diffusion import StableDiffusionInpaintPipeline |
|
from diffusers.schedulers import DDIMScheduler |
|
|
|
ddim_sched = scheduler = DDIMScheduler( |
|
beta_start=0.00085, |
|
beta_end=0.012, |
|
clip_sample=False, |
|
set_alpha_to_one=False, |
|
) |
|
|
|
pipe = StableDiffusionInpaintPipeline.from_pretrained( |
|
"CompVis/stable-diffusion-v1-4", |
|
revision="fp16", |
|
torch_dtype=torch.float16, |
|
use_auth_token=True, |
|
scheduler=ddim_sched, |
|
).to("cuda") |
|
|
|
prompt = "photo of a man smiling" |
|
prompt_short = "me-smiling-10" |
|
size = 512 + int(32 * 10) |
|
|
|
strength = 0.7 |
|
num_inference_steps = 150 |
|
guidance_scale = 11.0 |
|
|
|
|
|
def make_prompt_exif(**kwargs): |
|
"""https://stackoverflow.com/a/63649983/""" |
|
exif_ifd = {piexif.ExifIFD.UserComment: f"Params: {json.dumps(kwargs)}".encode()} |
|
exif_dict = {"0th": {}, "Exif": exif_ifd, "1st": {}, "thumbnail": None, "GPS": {}} |
|
return piexif.dump(exif_dict) |
|
|
|
|
|
def load_image(image_path): |
|
return Image.open(image_path).convert("RGB").resize((size, size)) |
|
|
|
|
|
image_path = Path("inputs") / "twitter-pic.jpg" |
|
mask_path = Path("inputs") / "twitter-pic-mask.jpg" |
|
init_image = load_image(image_path) |
|
mask_image = load_image(mask_path) |
|
|
|
out_dir = Path("results") / prompt_short |
|
out_dir.mkdir(exist_ok=False) |
|
|
|
for i in range(100): |
|
with autocast("cuda"): |
|
image = pipe( |
|
prompt=prompt, |
|
strength=strength, |
|
init_image=init_image, |
|
mask_image=mask_image, |
|
num_inference_steps=num_inference_steps, |
|
guidance_scale=guidance_scale, |
|
)["sample"][0] |
|
|
|
image.save( |
|
str(out_dir / f"{prompt_short}_{i:03d}.png"), |
|
exif=make_prompt_exif( |
|
prompt=prompt, |
|
prompt_short=prompt_short, |
|
strength=strength, |
|
num_inference_steps=num_inference_steps, |
|
guidance_scale=guidance_scale, |
|
size=size, |
|
), |
|
) |