add pre-commit hooks configuration
This commit is contained in:
@@ -6,26 +6,29 @@ import numpy as np
|
||||
from modules import shared
|
||||
from scripts.faceswaplab_utils import imgutils
|
||||
from modules import shared, processing, codeformer_model
|
||||
from modules.processing import (StableDiffusionProcessingImg2Img)
|
||||
from scripts.faceswaplab_postprocessing.postprocessing_options import PostProcessingOptions, InpaintingWhen
|
||||
from modules.processing import StableDiffusionProcessingImg2Img
|
||||
from scripts.faceswaplab_postprocessing.postprocessing_options import (
|
||||
PostProcessingOptions,
|
||||
InpaintingWhen,
|
||||
)
|
||||
from modules import sd_models
|
||||
|
||||
from scripts.faceswaplab_swapping import swapper
|
||||
|
||||
|
||||
def img2img_diffusion(img : Image.Image, pp: PostProcessingOptions) -> Image.Image :
|
||||
if pp.inpainting_denoising_strengh == 0 :
|
||||
def img2img_diffusion(img: Image.Image, pp: PostProcessingOptions) -> Image.Image:
|
||||
if pp.inpainting_denoising_strengh == 0:
|
||||
return img
|
||||
|
||||
try :
|
||||
try:
|
||||
logger.info(
|
||||
f"""Inpainting face
|
||||
f"""Inpainting face
|
||||
Sampler : {pp.inpainting_sampler}
|
||||
inpainting_denoising_strength : {pp.inpainting_denoising_strengh}
|
||||
inpainting_steps : {pp.inpainting_steps}
|
||||
"""
|
||||
)
|
||||
if not isinstance(pp.inpainting_sampler, str) :
|
||||
)
|
||||
if not isinstance(pp.inpainting_sampler, str):
|
||||
pass
|
||||
|
||||
logger.info("send faces to image to image")
|
||||
@@ -33,44 +36,51 @@ inpainting_steps : {pp.inpainting_steps}
|
||||
faces = swapper.get_faces(imgutils.pil_to_cv2(img))
|
||||
if faces:
|
||||
for face in faces:
|
||||
bbox =face.bbox.astype(int)
|
||||
bbox = face.bbox.astype(int)
|
||||
mask = imgutils.create_mask(img, bbox)
|
||||
prompt = pp.inpainting_prompt.replace("[gender]", "man" if face["gender"] == 1 else "woman")
|
||||
negative_prompt = pp.inpainting_negative_prompt.replace("[gender]", "man" if face["gender"] == 1 else "woman")
|
||||
prompt = pp.inpainting_prompt.replace(
|
||||
"[gender]", "man" if face["gender"] == 1 else "woman"
|
||||
)
|
||||
negative_prompt = pp.inpainting_negative_prompt.replace(
|
||||
"[gender]", "man" if face["gender"] == 1 else "woman"
|
||||
)
|
||||
logger.info("Denoising prompt : %s", prompt)
|
||||
logger.info("Denoising strenght : %s", pp.inpainting_denoising_strengh)
|
||||
|
||||
i2i_kwargs = {"sampler_name" :pp.inpainting_sampler,
|
||||
"do_not_save_samples":True,
|
||||
"steps" :pp.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" :pp.inpainting_denoising_strengh}
|
||||
logger.info("Denoising strenght : %s", pp.inpainting_denoising_strengh)
|
||||
|
||||
i2i_kwargs = {
|
||||
"sampler_name": pp.inpainting_sampler,
|
||||
"do_not_save_samples": True,
|
||||
"steps": pp.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": pp.inpainting_denoising_strengh,
|
||||
}
|
||||
current_model_checkpoint = shared.opts.sd_model_checkpoint
|
||||
if pp.inpainting_model and pp.inpainting_model != "Current" :
|
||||
if pp.inpainting_model and pp.inpainting_model != "Current":
|
||||
# Change checkpoint
|
||||
shared.opts.sd_model_checkpoint = pp.inpainting_model
|
||||
sd_models.select_checkpoint
|
||||
sd_models.load_model()
|
||||
i2i_p = StableDiffusionProcessingImg2Img([img], **i2i_kwargs)
|
||||
i2i_processed = processing.process_images(i2i_p)
|
||||
if pp.inpainting_model and pp.inpainting_model != "Current" :
|
||||
if pp.inpainting_model and pp.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 :
|
||||
if len(images) > 0:
|
||||
img = images[0]
|
||||
return img
|
||||
except Exception as e :
|
||||
except Exception as e:
|
||||
logger.error("Failed to apply img2img to face : %s", e)
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
raise e
|
||||
|
||||
@@ -1,27 +1,30 @@
|
||||
from modules.face_restoration import FaceRestoration
|
||||
from scripts.faceswaplab_utils.faceswaplab_logging import logger
|
||||
from PIL import Image
|
||||
from scripts.faceswaplab_postprocessing.postprocessing_options import PostProcessingOptions, InpaintingWhen
|
||||
from scripts.faceswaplab_postprocessing.postprocessing_options import (
|
||||
PostProcessingOptions,
|
||||
InpaintingWhen,
|
||||
)
|
||||
from scripts.faceswaplab_postprocessing.i2i_pp import img2img_diffusion
|
||||
from scripts.faceswaplab_postprocessing.upscaling import upscale_img, restore_face
|
||||
|
||||
|
||||
def enhance_image(image: Image.Image, pp_options: PostProcessingOptions) -> Image.Image:
|
||||
result_image = image
|
||||
try :
|
||||
if pp_options.inpainting_when == InpaintingWhen.BEFORE_UPSCALING.value :
|
||||
try:
|
||||
if pp_options.inpainting_when == InpaintingWhen.BEFORE_UPSCALING.value:
|
||||
result_image = img2img_diffusion(image, pp_options)
|
||||
result_image = upscale_img(result_image, pp_options)
|
||||
|
||||
if pp_options.inpainting_when == InpaintingWhen.BEFORE_RESTORE_FACE.value :
|
||||
result_image = img2img_diffusion(image,pp_options)
|
||||
if pp_options.inpainting_when == InpaintingWhen.BEFORE_RESTORE_FACE.value:
|
||||
result_image = img2img_diffusion(image, pp_options)
|
||||
|
||||
result_image = restore_face(result_image, pp_options)
|
||||
|
||||
if pp_options.inpainting_when == InpaintingWhen.AFTER_ALL.value :
|
||||
result_image = img2img_diffusion(image,pp_options)
|
||||
|
||||
if pp_options.inpainting_when == InpaintingWhen.AFTER_ALL.value:
|
||||
result_image = img2img_diffusion(image, pp_options)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Failed to upscale %s", e)
|
||||
|
||||
return result_image
|
||||
return result_image
|
||||
|
||||
@@ -4,12 +4,14 @@ from dataclasses import dataclass
|
||||
from modules import shared
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class InpaintingWhen(Enum):
|
||||
NEVER = "Never"
|
||||
BEFORE_UPSCALING = "Before Upscaling/all"
|
||||
BEFORE_RESTORE_FACE = "After Upscaling/Before Restore Face"
|
||||
AFTER_ALL = "After All"
|
||||
|
||||
|
||||
@dataclass
|
||||
class PostProcessingOptions:
|
||||
face_restorer_name: str = ""
|
||||
@@ -19,15 +21,15 @@ class PostProcessingOptions:
|
||||
upscaler_name: str = ""
|
||||
scale: int = 1
|
||||
upscale_visibility: float = 0.5
|
||||
|
||||
inpainting_denoising_strengh : float = 0
|
||||
inpainting_prompt : str = ""
|
||||
inpainting_negative_prompt : str = ""
|
||||
inpainting_steps : int = 20
|
||||
inpainting_sampler : str = "Euler"
|
||||
inpainting_when : InpaintingWhen = InpaintingWhen.BEFORE_UPSCALING
|
||||
inpainting_model : str = "Current"
|
||||
|
||||
|
||||
inpainting_denoising_strengh: float = 0
|
||||
inpainting_prompt: str = ""
|
||||
inpainting_negative_prompt: str = ""
|
||||
inpainting_steps: int = 20
|
||||
inpainting_sampler: str = "Euler"
|
||||
inpainting_when: InpaintingWhen = InpaintingWhen.BEFORE_UPSCALING
|
||||
inpainting_model: str = "Current"
|
||||
|
||||
@property
|
||||
def upscaler(self) -> UpscalerData:
|
||||
for upscaler in shared.sd_upscalers:
|
||||
@@ -40,4 +42,4 @@ class PostProcessingOptions:
|
||||
for face_restorer in shared.face_restorers:
|
||||
if face_restorer.name() == self.face_restorer_name:
|
||||
return face_restorer
|
||||
return None
|
||||
return None
|
||||
|
||||
@@ -1,11 +1,14 @@
|
||||
|
||||
from scripts.faceswaplab_postprocessing.postprocessing_options import PostProcessingOptions, InpaintingWhen
|
||||
from scripts.faceswaplab_postprocessing.postprocessing_options import (
|
||||
PostProcessingOptions,
|
||||
InpaintingWhen,
|
||||
)
|
||||
from scripts.faceswaplab_utils.faceswaplab_logging import logger
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
from modules import shared, processing, codeformer_model
|
||||
|
||||
def upscale_img(image : Image.Image, pp_options :PostProcessingOptions) -> Image.Image :
|
||||
|
||||
def upscale_img(image: Image.Image, pp_options: PostProcessingOptions) -> Image.Image:
|
||||
if pp_options.upscaler is not None and pp_options.upscaler.name != "None":
|
||||
original_image = image.copy()
|
||||
logger.info(
|
||||
@@ -23,15 +26,17 @@ def upscale_img(image : Image.Image, pp_options :PostProcessingOptions) -> Image
|
||||
return result_image
|
||||
return image
|
||||
|
||||
def restore_face(image : Image.Image, pp_options : PostProcessingOptions) -> Image.Image :
|
||||
|
||||
|
||||
def restore_face(image: Image.Image, pp_options: PostProcessingOptions) -> Image.Image:
|
||||
if pp_options.face_restorer is not None:
|
||||
original_image = image.copy()
|
||||
logger.info("Restore face with %s", pp_options.face_restorer.name())
|
||||
numpy_image = np.array(image)
|
||||
if pp_options.face_restorer_name == "CodeFormer" :
|
||||
numpy_image = codeformer_model.codeformer.restore(numpy_image, w=pp_options.codeformer_weight)
|
||||
else :
|
||||
if pp_options.face_restorer_name == "CodeFormer":
|
||||
numpy_image = codeformer_model.codeformer.restore(
|
||||
numpy_image, w=pp_options.codeformer_weight
|
||||
)
|
||||
else:
|
||||
numpy_image = pp_options.face_restorer.restore(numpy_image)
|
||||
|
||||
restored_image = Image.fromarray(numpy_image)
|
||||
@@ -39,4 +44,4 @@ def restore_face(image : Image.Image, pp_options : PostProcessingOptions) -> Ima
|
||||
original_image, restored_image, pp_options.restorer_visibility
|
||||
)
|
||||
return result_image
|
||||
return image
|
||||
return image
|
||||
|
||||
Reference in New Issue
Block a user