add pre-commit hooks configuration
This commit is contained in:
Tran Xen
+91
-59
@@ -2,17 +2,20 @@ import importlib
|
||||
from scripts.faceswaplab_api import faceswaplab_api
|
||||
from scripts.faceswaplab_settings import faceswaplab_settings
|
||||
from scripts.faceswaplab_ui import faceswaplab_tab, faceswaplab_unit_ui
|
||||
from scripts.faceswaplab_utils.models_utils import get_current_model, get_face_checkpoints
|
||||
from scripts.faceswaplab_utils.models_utils import (
|
||||
get_current_model,
|
||||
get_face_checkpoints,
|
||||
)
|
||||
|
||||
from scripts import (faceswaplab_globals)
|
||||
from scripts import faceswaplab_globals
|
||||
from scripts.faceswaplab_swapping import swapper
|
||||
from scripts.faceswaplab_utils import faceswaplab_logging, imgutils
|
||||
from scripts.faceswaplab_utils import models_utils
|
||||
from scripts.faceswaplab_postprocessing import upscaling
|
||||
import numpy as np
|
||||
|
||||
#Reload all the modules when using "apply and restart"
|
||||
#This is mainly done for development purposes
|
||||
# Reload all the modules when using "apply and restart"
|
||||
# This is mainly done for development purposes
|
||||
importlib.reload(swapper)
|
||||
importlib.reload(faceswaplab_logging)
|
||||
importlib.reload(faceswaplab_globals)
|
||||
@@ -35,20 +38,25 @@ from modules import script_callbacks, scripts
|
||||
from insightface.app.common import Face
|
||||
from modules import scripts, shared
|
||||
from modules.images import save_image, image_grid
|
||||
from modules.processing import (Processed, StableDiffusionProcessing,
|
||||
StableDiffusionProcessingImg2Img)
|
||||
from modules.processing import (
|
||||
Processed,
|
||||
StableDiffusionProcessing,
|
||||
StableDiffusionProcessingImg2Img,
|
||||
)
|
||||
from modules.shared import opts
|
||||
from PIL import Image
|
||||
|
||||
from scripts.faceswaplab_utils.imgutils import (pil_to_cv2,check_against_nsfw)
|
||||
from scripts.faceswaplab_utils.imgutils import pil_to_cv2, check_against_nsfw
|
||||
from scripts.faceswaplab_utils.faceswaplab_logging import logger, save_img_debug
|
||||
from scripts.faceswaplab_globals import VERSION_FLAG
|
||||
from scripts.faceswaplab_postprocessing.postprocessing_options import PostProcessingOptions
|
||||
from scripts.faceswaplab_postprocessing.postprocessing_options import (
|
||||
PostProcessingOptions,
|
||||
)
|
||||
from scripts.faceswaplab_postprocessing.postprocessing import enhance_image
|
||||
from scripts.faceswaplab_ui.faceswaplab_unit_settings import FaceSwapUnitSettings
|
||||
|
||||
|
||||
EXTENSION_PATH=os.path.join("extensions","sd-webui-faceswaplab")
|
||||
EXTENSION_PATH = os.path.join("extensions", "sd-webui-faceswaplab")
|
||||
|
||||
|
||||
# Register the tab, done here to prevent it from being added twice
|
||||
@@ -56,44 +64,44 @@ script_callbacks.on_ui_tabs(faceswaplab_tab.on_ui_tabs)
|
||||
|
||||
try:
|
||||
import modules.script_callbacks as script_callbacks
|
||||
|
||||
script_callbacks.on_app_started(faceswaplab_api.faceswaplab_api)
|
||||
except:
|
||||
pass
|
||||
|
||||
|
||||
class FaceSwapScript(scripts.Script):
|
||||
|
||||
def __init__(self) -> None:
|
||||
logger.info(f"FaceSwapLab {VERSION_FLAG}")
|
||||
super().__init__()
|
||||
|
||||
@property
|
||||
def units_count(self) :
|
||||
def units_count(self):
|
||||
return opts.data.get("faceswaplab_units_count", 3)
|
||||
|
||||
|
||||
@property
|
||||
def upscaled_swapper_in_generated(self) :
|
||||
def upscaled_swapper_in_generated(self):
|
||||
return opts.data.get("faceswaplab_upscaled_swapper", False)
|
||||
|
||||
|
||||
@property
|
||||
def upscaled_swapper_in_source(self) :
|
||||
def upscaled_swapper_in_source(self):
|
||||
return opts.data.get("faceswaplab_upscaled_swapper_in_source", False)
|
||||
|
||||
|
||||
@property
|
||||
def enabled(self) -> bool :
|
||||
def enabled(self) -> bool:
|
||||
"""Return True if any unit is enabled and the state is not interupted"""
|
||||
return any([u.enable for u in self.units]) and not shared.state.interrupted
|
||||
|
||||
@property
|
||||
def keep_original_images(self) :
|
||||
def keep_original_images(self):
|
||||
return opts.data.get("faceswaplab_keep_original", False)
|
||||
|
||||
@property
|
||||
def swap_in_generated_units(self) :
|
||||
def swap_in_generated_units(self):
|
||||
return [u for u in self.units if u.swap_in_generated and u.enable]
|
||||
|
||||
@property
|
||||
def swap_in_source_units(self) :
|
||||
def swap_in_source_units(self):
|
||||
return [u for u in self.units if u.swap_in_source and u.enable]
|
||||
|
||||
def title(self):
|
||||
@@ -102,7 +110,6 @@ class FaceSwapScript(scripts.Script):
|
||||
def show(self, is_img2img):
|
||||
return scripts.AlwaysVisible
|
||||
|
||||
|
||||
def ui(self, is_img2img):
|
||||
with gr.Accordion(f"FaceSwapLab {VERSION_FLAG}", open=False):
|
||||
components = []
|
||||
@@ -121,23 +128,23 @@ class FaceSwapScript(scripts.Script):
|
||||
# print("Running in ", alwayson.index(self), "position")
|
||||
# logger.info("Running scripts : %s", pformat(runner.alwayson_scripts))
|
||||
|
||||
def read_config(self, p : StableDiffusionProcessing, *components) :
|
||||
def read_config(self, p: StableDiffusionProcessing, *components):
|
||||
# The order of processing for the components is important
|
||||
# The method first process faceswap units then postprocessing units
|
||||
# The method first process faceswap units then postprocessing units
|
||||
|
||||
# self.make_first_script(p)
|
||||
|
||||
self.units: List[FaceSwapUnitSettings] = []
|
||||
|
||||
#Parse and convert units flat components into FaceSwapUnitSettings
|
||||
|
||||
# Parse and convert units flat components into FaceSwapUnitSettings
|
||||
for i in range(0, self.units_count):
|
||||
self.units += [FaceSwapUnitSettings.get_unit_configuration(i, components)]
|
||||
|
||||
for i, u in enumerate(self.units):
|
||||
logger.debug("%s, %s", pformat(i), pformat(u))
|
||||
|
||||
#Parse the postprocessing options
|
||||
#We must first find where to start from (after face swapping units)
|
||||
# Parse the postprocessing options
|
||||
# We must first find where to start from (after face swapping units)
|
||||
len_conf: int = len(fields(FaceSwapUnitSettings))
|
||||
shift: int = self.units_count * len_conf
|
||||
self.postprocess_options = PostProcessingOptions(
|
||||
@@ -145,67 +152,92 @@ class FaceSwapScript(scripts.Script):
|
||||
)
|
||||
logger.debug("%s", pformat(self.postprocess_options))
|
||||
|
||||
if self.enabled :
|
||||
if self.enabled:
|
||||
p.do_not_save_samples = not self.keep_original_images
|
||||
|
||||
|
||||
def process(self, p: StableDiffusionProcessing, *components):
|
||||
self.read_config(p, *components)
|
||||
|
||||
#If is instance of img2img, we check if face swapping in source is required.
|
||||
# If is instance of img2img, we check if face swapping in source is required.
|
||||
if isinstance(p, StableDiffusionProcessingImg2Img):
|
||||
if self.enabled and len(self.swap_in_source_units) > 0:
|
||||
init_images : List[Tuple[Optional[Image.Image], Optional[str]]] = [(img,None) for img in p.init_images]
|
||||
new_inits = swapper.process_images_units(get_current_model(), self.swap_in_source_units,images=init_images, upscaled_swapper=self.upscaled_swapper_in_source,force_blend=True)
|
||||
init_images: List[Tuple[Optional[Image.Image], Optional[str]]] = [
|
||||
(img, None) for img in p.init_images
|
||||
]
|
||||
new_inits = swapper.process_images_units(
|
||||
get_current_model(),
|
||||
self.swap_in_source_units,
|
||||
images=init_images,
|
||||
upscaled_swapper=self.upscaled_swapper_in_source,
|
||||
force_blend=True,
|
||||
)
|
||||
logger.info(f"processed init images: {len(init_images)}")
|
||||
if new_inits is not None :
|
||||
p.init_images = [img[0] for img in new_inits]
|
||||
if new_inits is not None:
|
||||
p.init_images = [img[0] for img in new_inits]
|
||||
|
||||
|
||||
def postprocess(self, p : StableDiffusionProcessing, processed: Processed, *args):
|
||||
if self.enabled :
|
||||
def postprocess(self, p: StableDiffusionProcessing, processed: Processed, *args):
|
||||
if self.enabled:
|
||||
# Get the original images without the grid
|
||||
orig_images : List[Image.Image] = processed.images[processed.index_of_first_image:]
|
||||
orig_infotexts : List[str] = processed.infotexts[processed.index_of_first_image:]
|
||||
orig_images: List[Image.Image] = processed.images[
|
||||
processed.index_of_first_image :
|
||||
]
|
||||
orig_infotexts: List[str] = processed.infotexts[
|
||||
processed.index_of_first_image :
|
||||
]
|
||||
|
||||
keep_original = self.keep_original_images
|
||||
|
||||
# These are were images and infos of swapped images will be stored
|
||||
images = []
|
||||
infotexts = []
|
||||
if (len(self.swap_in_generated_units))>0 :
|
||||
for i,(img,info) in enumerate(zip(orig_images, orig_infotexts)):
|
||||
batch_index = i%p.batch_size
|
||||
swapped_images = swapper.process_images_units(get_current_model(), self.swap_in_generated_units, images=[(img,info)], upscaled_swapper=self.upscaled_swapper_in_generated)
|
||||
if swapped_images is None :
|
||||
if (len(self.swap_in_generated_units)) > 0:
|
||||
for i, (img, info) in enumerate(zip(orig_images, orig_infotexts)):
|
||||
batch_index = i % p.batch_size
|
||||
swapped_images = swapper.process_images_units(
|
||||
get_current_model(),
|
||||
self.swap_in_generated_units,
|
||||
images=[(img, info)],
|
||||
upscaled_swapper=self.upscaled_swapper_in_generated,
|
||||
)
|
||||
if swapped_images is None:
|
||||
continue
|
||||
|
||||
logger.info(f"{len(swapped_images)} images swapped")
|
||||
for swp_img, new_info in swapped_images :
|
||||
img = swp_img # Will only swap the last image in the batch in next units (FIXME : hard to fix properly but not really critical)
|
||||
for swp_img, new_info in swapped_images:
|
||||
img = swp_img # Will only swap the last image in the batch in next units (FIXME : hard to fix properly but not really critical)
|
||||
|
||||
if swp_img is not None :
|
||||
|
||||
save_img_debug(swp_img,"Before apply mask")
|
||||
if swp_img is not None:
|
||||
save_img_debug(swp_img, "Before apply mask")
|
||||
swp_img = imgutils.apply_mask(swp_img, p, batch_index)
|
||||
save_img_debug(swp_img,"After apply mask")
|
||||
save_img_debug(swp_img, "After apply mask")
|
||||
|
||||
try :
|
||||
try:
|
||||
if self.postprocess_options is not None:
|
||||
swp_img = enhance_image(swp_img, self.postprocess_options)
|
||||
swp_img = enhance_image(
|
||||
swp_img, self.postprocess_options
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error("Failed to upscale : %s", e)
|
||||
|
||||
logger.info("Add swp image to processed")
|
||||
images.append(swp_img)
|
||||
infotexts.append(new_info)
|
||||
if p.outpath_samples and opts.samples_save :
|
||||
save_image(swp_img, p.outpath_samples, "", p.all_seeds[batch_index], p.all_prompts[batch_index], opts.samples_format,info=new_info, p=p, suffix="-swapped")
|
||||
else :
|
||||
if p.outpath_samples and opts.samples_save:
|
||||
save_image(
|
||||
swp_img,
|
||||
p.outpath_samples,
|
||||
"",
|
||||
p.all_seeds[batch_index],
|
||||
p.all_prompts[batch_index],
|
||||
opts.samples_format,
|
||||
info=new_info,
|
||||
p=p,
|
||||
suffix="-swapped",
|
||||
)
|
||||
else:
|
||||
logger.error("swp image is None")
|
||||
else :
|
||||
keep_original=True
|
||||
|
||||
else:
|
||||
keep_original = True
|
||||
|
||||
# Generate grid :
|
||||
if opts.return_grid and len(images) > 1:
|
||||
@@ -223,4 +255,4 @@ class FaceSwapScript(scripts.Script):
|
||||
infotexts += processed.infotexts
|
||||
|
||||
processed.images = images
|
||||
processed.infotexts = infotexts
|
||||
processed.infotexts = infotexts
|
||||
|
||||
Reference in New Issue
Block a user