-
Notifications
You must be signed in to change notification settings - Fork 92
/
i2i_pp.py
91 lines (82 loc) · 3.66 KB
/
i2i_pp.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
from scripts.faceswaplab_inpainting.faceswaplab_inpainting import InpaintingOptions
from scripts.faceswaplab_utils.faceswaplab_logging import logger
from PIL import Image
from modules import shared
from scripts.faceswaplab_utils import imgutils
from modules import shared, processing
from modules.processing import StableDiffusionProcessingImg2Img
from modules import sd_models
import traceback
from scripts.faceswaplab_swapping import swapper
from scripts.faceswaplab_utils.typing import *
from typing import *
def img2img_diffusion(
img: PILImage, options: InpaintingOptions, faces: Optional[List[Face]] = None
) -> Image.Image:
if not options or options.inpainting_denoising_strengh == 0:
logger.info("Discard inpainting denoising strength is 0 or no inpainting")
return img
try:
logger.info(
f"""Inpainting face
Sampler : {options.inpainting_sampler}
inpainting_denoising_strength : {options.inpainting_denoising_strengh}
inpainting_steps : {options.inpainting_steps}
"""
)
if not isinstance(options.inpainting_sampler, str):
options.inpainting_sampler = "Euler"
logger.info("send faces to image to image")
img = img.copy()
if not faces:
faces = swapper.get_faces(imgutils.pil_to_cv2(img))
if faces:
for face in faces:
bbox = face.bbox.astype(int)
mask = imgutils.create_mask(img, bbox)
prompt = options.inpainting_prompt.replace(
"[gender]", "man" if face["gender"] == 1 else "woman"
)
negative_prompt = options.inpainting_negative_prompt.replace(
"[gender]", "man" if face["gender"] == 1 else "woman"
)
logger.info("Denoising prompt : %s", prompt)
logger.info(
"Denoising strenght : %s", options.inpainting_denoising_strengh
)
i2i_kwargs = {
"init_images": [img],
"sampler_name": options.inpainting_sampler,
"do_not_save_samples": True,
"steps": options.inpainting_steps,
"width": img.width,
"inpainting_fill": 1,
"inpaint_full_res": True,
"height": img.height,
"mask": mask,
"prompt": prompt,
"negative_prompt": negative_prompt,
"denoising_strength": options.inpainting_denoising_strengh,
"seed": options.inpainting_seed,
}
current_model_checkpoint = shared.opts.sd_model_checkpoint
if options.inpainting_model and options.inpainting_model != "Current":
# Change checkpoint
shared.opts.sd_model_checkpoint = options.inpainting_model
sd_models.select_checkpoint
sd_models.load_model()
i2i_p = StableDiffusionProcessingImg2Img(**i2i_kwargs)
i2i_processed = processing.process_images(i2i_p)
if options.inpainting_model and options.inpainting_model != "Current":
# Restore checkpoint
shared.opts.sd_model_checkpoint = current_model_checkpoint
sd_models.select_checkpoint
sd_models.load_model()
images = i2i_processed.images
if len(images) > 0:
img = images[0]
return img
except Exception as e:
logger.error("Failed to apply inpainting to face : %s", e)
traceback.print_exc()
raise e