first commit
This commit is contained in:
Tran Xen
@@ -0,0 +1,226 @@
|
||||
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 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
|
||||
importlib.reload(swapper)
|
||||
importlib.reload(faceswaplab_logging)
|
||||
importlib.reload(faceswaplab_globals)
|
||||
importlib.reload(imgutils)
|
||||
importlib.reload(upscaling)
|
||||
importlib.reload(faceswaplab_settings)
|
||||
importlib.reload(models_utils)
|
||||
importlib.reload(faceswaplab_unit_ui)
|
||||
importlib.reload(faceswaplab_api)
|
||||
|
||||
import os
|
||||
from dataclasses import fields
|
||||
from pprint import pformat
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import dill as pickle
|
||||
import gradio as gr
|
||||
import modules.scripts as scripts
|
||||
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.shared import opts
|
||||
from PIL import Image
|
||||
|
||||
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 import enhance_image
|
||||
from scripts.faceswaplab_ui.faceswaplab_unit_settings import FaceSwapUnitSettings
|
||||
|
||||
|
||||
EXTENSION_PATH=os.path.join("extensions","sd-webui-faceswaplab")
|
||||
|
||||
|
||||
# Register the tab, done here to prevent it from being added twice
|
||||
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) :
|
||||
return opts.data.get("faceswaplab_units_count", 3)
|
||||
|
||||
@property
|
||||
def upscaled_swapper_in_generated(self) :
|
||||
return opts.data.get("faceswaplab_upscaled_swapper", False)
|
||||
|
||||
@property
|
||||
def upscaled_swapper_in_source(self) :
|
||||
return opts.data.get("faceswaplab_upscaled_swapper_in_source", False)
|
||||
|
||||
@property
|
||||
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) :
|
||||
return opts.data.get("faceswaplab_keep_original", False)
|
||||
|
||||
@property
|
||||
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) :
|
||||
return [u for u in self.units if u.swap_in_source and u.enable]
|
||||
|
||||
def title(self):
|
||||
return f"faceswaplab"
|
||||
|
||||
def show(self, is_img2img):
|
||||
return scripts.AlwaysVisible
|
||||
|
||||
|
||||
def ui(self, is_img2img):
|
||||
with gr.Accordion(f"FaceSwapLab {VERSION_FLAG}", open=False):
|
||||
components = []
|
||||
for i in range(1, self.units_count + 1):
|
||||
components += faceswaplab_unit_ui.faceswap_unit_ui(is_img2img, i)
|
||||
upscaler = faceswaplab_tab.upscaler_ui()
|
||||
# If the order is modified, the before_process should be changed accordingly.
|
||||
return components + upscaler
|
||||
|
||||
# def make_script_first(self,p: StableDiffusionProcessing) :
|
||||
# FIXME : not really useful, will only impact postprocessing (kept for further testing)
|
||||
# runner : scripts.ScriptRunner = p.scripts
|
||||
# alwayson = runner.alwayson_scripts
|
||||
# alwayson.pop(alwayson.index(self))
|
||||
# alwayson.insert(0, self)
|
||||
# print("Running in ", alwayson.index(self), "position")
|
||||
# logger.info("Running scripts : %s", pformat(runner.alwayson_scripts))
|
||||
|
||||
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
|
||||
|
||||
# self.make_first_script(p)
|
||||
|
||||
self.units: List[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)
|
||||
len_conf: int = len(fields(FaceSwapUnitSettings))
|
||||
shift: int = self.units_count * len_conf
|
||||
self.postprocess_options = PostProcessingOptions(
|
||||
*components[shift : shift + len(fields(PostProcessingOptions))]
|
||||
)
|
||||
logger.debug("%s", pformat(self.postprocess_options))
|
||||
|
||||
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 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)
|
||||
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]
|
||||
|
||||
|
||||
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:]
|
||||
|
||||
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 :
|
||||
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)
|
||||
|
||||
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")
|
||||
|
||||
try :
|
||||
if self.postprocess_options is not None:
|
||||
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 :
|
||||
logger.error("swp image is None")
|
||||
else :
|
||||
keep_original=True
|
||||
|
||||
|
||||
# Generate grid :
|
||||
if opts.return_grid and len(images) > 1:
|
||||
# FIXME :Use sd method, not that if blended is not active, the result will be a bit messy.
|
||||
grid = imgutils.create_square_image(images)
|
||||
text = processed.infotexts[0]
|
||||
infotexts.insert(0, text)
|
||||
if opts.enable_pnginfo:
|
||||
grid.info["parameters"] = text
|
||||
images.insert(0, grid)
|
||||
|
||||
if keep_original:
|
||||
# If we want to keep original images, we add all existing (including grid this time)
|
||||
images += processed.images
|
||||
infotexts += processed.infotexts
|
||||
|
||||
processed.images = images
|
||||
processed.infotexts = infotexts
|
||||
@@ -0,0 +1,70 @@
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
from fastapi import FastAPI, Body
|
||||
from fastapi.exceptions import HTTPException
|
||||
from modules.api.models import *
|
||||
from modules.api import api
|
||||
from scripts.faceswaplab_api.faceswaplab_api_types import FaceSwapUnit, FaceSwapRequest, FaceSwapResponse
|
||||
from scripts.faceswaplab_globals import VERSION_FLAG
|
||||
import gradio as gr
|
||||
from typing import List, Optional
|
||||
from scripts.faceswaplab_swapping import swapper
|
||||
from scripts.faceswaplab_utils.faceswaplab_logging import save_img_debug
|
||||
from scripts.faceswaplab_ui.faceswaplab_unit_settings import FaceSwapUnitSettings
|
||||
from scripts.faceswaplab_utils.imgutils import (pil_to_cv2,check_against_nsfw, base64_to_pil)
|
||||
from scripts.faceswaplab_utils.models_utils import get_current_model
|
||||
from modules.shared import opts
|
||||
|
||||
|
||||
def encode_to_base64(image):
|
||||
if type(image) is str:
|
||||
return image
|
||||
elif type(image) is Image.Image:
|
||||
return api.encode_pil_to_base64(image)
|
||||
elif type(image) is np.ndarray:
|
||||
return encode_np_to_base64(image)
|
||||
else:
|
||||
return ""
|
||||
|
||||
def encode_np_to_base64(image):
|
||||
pil = Image.fromarray(image)
|
||||
return api.encode_pil_to_base64(pil)
|
||||
|
||||
|
||||
def faceswaplab_api(_: gr.Blocks, app: FastAPI):
|
||||
@app.get("/faceswaplab/version", tags=["faceswaplab"], description="Get faceswaplab version")
|
||||
async def version():
|
||||
return {"version": VERSION_FLAG}
|
||||
|
||||
# use post as we consider the method non idempotent (which is debatable)
|
||||
@app.post("/faceswaplab/swap_face", tags=["faceswaplab"], description="Swap a face in an image using units")
|
||||
async def swap_face(request : FaceSwapRequest) -> FaceSwapResponse:
|
||||
units : List[FaceSwapUnitSettings]= []
|
||||
src_image : Optional[Image.Image] = base64_to_pil(request.image)
|
||||
response = FaceSwapResponse(images = [], infos=[])
|
||||
if src_image is not None :
|
||||
for u in request.units:
|
||||
units.append(
|
||||
FaceSwapUnitSettings(source_img=base64_to_pil(u.source_img),
|
||||
source_face = u.source_face,
|
||||
_batch_files = u.get_batch_images(),
|
||||
blend_faces= u.blend_faces,
|
||||
enable = True,
|
||||
same_gender = u.same_gender,
|
||||
check_similarity=u.check_similarity,
|
||||
_compute_similarity=u.compute_similarity,
|
||||
min_ref_sim= u.min_ref_sim,
|
||||
min_sim= u.min_sim,
|
||||
_faces_index = ",".join([str(i) for i in (u.faces_index)]),
|
||||
swap_in_generated=True,
|
||||
swap_in_source=False
|
||||
)
|
||||
)
|
||||
|
||||
swapped_images = swapper.process_images_units(get_current_model(), images=[(src_image,None)], units=units, upscaled_swapper=opts.data.get("faceswaplab_upscaled_swapper", False))
|
||||
for img, info in swapped_images:
|
||||
response.images.append(encode_to_base64(img))
|
||||
response.infos.append(info)
|
||||
|
||||
return response
|
||||
|
||||
@@ -0,0 +1,73 @@
|
||||
from scripts.faceswaplab_swapping import swapper
|
||||
import numpy as np
|
||||
from typing import List, Tuple
|
||||
import dill as pickle
|
||||
import gradio as gr
|
||||
from insightface.app.common import Face
|
||||
from PIL import Image
|
||||
from scripts.faceswaplab_utils.imgutils import (pil_to_cv2,check_against_nsfw, base64_to_pil)
|
||||
from scripts.faceswaplab_utils.faceswaplab_logging import logger
|
||||
from pydantic import BaseModel, Field
|
||||
from scripts.faceswaplab_postprocessing.postprocessing_options import InpaintingWhen
|
||||
|
||||
|
||||
class FaceSwapUnit(BaseModel) :
|
||||
|
||||
# The image given in reference
|
||||
source_img: str = Field(description='base64 reference image', examples=["data:image/jpeg;base64,/9j/4AAQSkZJRgABAQECWAJYAAD...."], default=None)
|
||||
# The checkpoint file
|
||||
source_face : str = Field(description='face checkpoint (from models/faceswaplab/faces)',examples=["my_face.pkl"], default=None)
|
||||
# base64 batch source images
|
||||
batch_images: Tuple[str] = Field(description='list of base64 batch source images',examples=["data:image/jpeg;base64,/9j/4AAQSkZJRgABAQECWAJYAAD....", "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQECWAJYAAD...."], default=None)
|
||||
|
||||
# Will blend faces if True
|
||||
blend_faces: bool = Field(description='Will blend faces if True', default=True)
|
||||
|
||||
# Use same gender filtering
|
||||
same_gender: bool = Field(description='Use same gender filtering', default=True)
|
||||
|
||||
# If True, discard images with low similarity
|
||||
check_similarity : bool = Field(description='If True, discard images with low similarity', default=False)
|
||||
# if True will compute similarity and add it to the image info
|
||||
compute_similarity : bool = Field(description='If True will compute similarity and add it to the image info', default=False)
|
||||
|
||||
# Minimum similarity against the used face (reference, batch or checkpoint)
|
||||
min_sim: float = Field(description='Minimum similarity against the used face (reference, batch or checkpoint)', default=0.0)
|
||||
# Minimum similarity against the reference (reference or checkpoint if checkpoint is given)
|
||||
min_ref_sim: float = Field(description='Minimum similarity against the reference (reference or checkpoint if checkpoint is given)', default=0.0)
|
||||
|
||||
# The face index to use for swapping
|
||||
faces_index: Tuple[int] = Field(description='The face index to use for swapping, list of face numbers starting from 0', default=(0,))
|
||||
|
||||
def get_batch_images(self) -> List[Image.Image] :
|
||||
images = []
|
||||
if self.batch_images :
|
||||
for img in self.batch_images :
|
||||
images.append(base64_to_pil(img))
|
||||
return images
|
||||
|
||||
class PostProcessingOptions (BaseModel):
|
||||
face_restorer_name: str = Field(description='face restorer name', default=None)
|
||||
restorer_visibility: float = Field(description='face restorer visibility', default=1, le=1, ge=0)
|
||||
codeformer_weight: float = Field(description='face restorer codeformer weight', default=1, le=1, ge=0)
|
||||
|
||||
upscaler_name: str = Field(description='upscaler name', default=None)
|
||||
scale: float = Field(description='upscaling scale', default=1, le=10, ge=0)
|
||||
upscale_visibility: float = Field(description='upscaler visibility', default=1, le=1, ge=0)
|
||||
|
||||
inpainting_denoising_strengh : float = Field(description='Inpainting denoising strenght', default=0, lt=1, ge=0)
|
||||
inpainting_prompt : str = Field(description='Inpainting denoising strenght',examples=["Portrait of a [gender]"], default="Portrait of a [gender]")
|
||||
inpainting_negative_prompt : str = Field(description='Inpainting denoising strenght',examples=["Deformed, blurry, bad anatomy, disfigured, poorly drawn face, mutation"], default="")
|
||||
inpainting_steps : int = Field(description='Inpainting steps',examples=["Portrait of a [gender]"], ge=1, le=150, default=20)
|
||||
inpainting_sampler : str = Field(description='Inpainting sampler',examples=["Euler"], default="Euler")
|
||||
inpainting_when : InpaintingWhen = Field(description='When inpainting happens', examples=[e.value for e in InpaintingWhen.__members__.values()], default=InpaintingWhen.NEVER)
|
||||
|
||||
|
||||
class FaceSwapRequest(BaseModel) :
|
||||
image : str = Field(description='base64 reference image', examples=["data:image/jpeg;base64,/9j/4AAQSkZJRgABAQECWAJYAAD...."], default=None)
|
||||
units : List[FaceSwapUnit]
|
||||
postprocessing : PostProcessingOptions
|
||||
|
||||
class FaceSwapResponse(BaseModel) :
|
||||
images : List[str] = Field(description='base64 swapped image',default=None)
|
||||
infos : List[str]
|
||||
@@ -0,0 +1,11 @@
|
||||
from scripts.faceswaplab_utils.faceswaplab_logging import logger
|
||||
import os
|
||||
|
||||
MODELS_DIR = os.path.abspath(os.path.join("models","faceswaplab"))
|
||||
ANALYZER_DIR = os.path.abspath(os.path.join(MODELS_DIR, "analysers"))
|
||||
FACE_PARSER_DIR = os.path.abspath(os.path.join(MODELS_DIR, "parser"))
|
||||
|
||||
VERSION_FLAG = "v1.1.0"
|
||||
EXTENSION_PATH=os.path.join("extensions","sd-webui-faceswaplab")
|
||||
NSFW_SCORE = 0.7
|
||||
|
||||
@@ -0,0 +1,76 @@
|
||||
from modules.face_restoration import FaceRestoration
|
||||
from modules.upscaler import UpscalerData
|
||||
from scripts.faceswaplab_utils.faceswaplab_logging import logger
|
||||
from PIL import Image
|
||||
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 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 :
|
||||
return img
|
||||
|
||||
try :
|
||||
logger.info(
|
||||
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) :
|
||||
pass
|
||||
|
||||
logger.info("send faces to image to image")
|
||||
img = img.copy()
|
||||
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 = 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}
|
||||
current_model_checkpoint = shared.opts.sd_model_checkpoint
|
||||
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" :
|
||||
# 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 img2img to face : %s", e)
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
raise e
|
||||
@@ -0,0 +1,27 @@
|
||||
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.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 :
|
||||
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)
|
||||
|
||||
result_image = restore_face(result_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
|
||||
@@ -0,0 +1,43 @@
|
||||
from modules.face_restoration import FaceRestoration
|
||||
from modules.upscaler import UpscalerData
|
||||
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 = ""
|
||||
restorer_visibility: float = 0.5
|
||||
codeformer_weight: float = 1
|
||||
|
||||
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"
|
||||
|
||||
@property
|
||||
def upscaler(self) -> UpscalerData:
|
||||
for upscaler in shared.sd_upscalers:
|
||||
if upscaler.name == self.upscaler_name:
|
||||
return upscaler
|
||||
return None
|
||||
|
||||
@property
|
||||
def face_restorer(self) -> FaceRestoration:
|
||||
for face_restorer in shared.face_restorers:
|
||||
if face_restorer.name() == self.face_restorer_name:
|
||||
return face_restorer
|
||||
return None
|
||||
@@ -0,0 +1,42 @@
|
||||
|
||||
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 :
|
||||
if pp_options.upscaler is not None and pp_options.upscaler.name != "None":
|
||||
original_image = image.copy()
|
||||
logger.info(
|
||||
"Upscale with %s scale = %s",
|
||||
pp_options.upscaler.name,
|
||||
pp_options.scale,
|
||||
)
|
||||
result_image = pp_options.upscaler.scaler.upscale(
|
||||
image, pp_options.scale, pp_options.upscaler.data_path
|
||||
)
|
||||
if pp_options.scale == 1:
|
||||
result_image = Image.blend(
|
||||
original_image, result_image, pp_options.upscale_visibility
|
||||
)
|
||||
return result_image
|
||||
return 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 :
|
||||
numpy_image = pp_options.face_restorer.restore(numpy_image)
|
||||
|
||||
restored_image = Image.fromarray(numpy_image)
|
||||
result_image = Image.blend(
|
||||
original_image, restored_image, pp_options.restorer_visibility
|
||||
)
|
||||
return result_image
|
||||
return image
|
||||
@@ -0,0 +1,53 @@
|
||||
from scripts.faceswaplab_utils.models_utils import get_models
|
||||
from modules import script_callbacks, shared
|
||||
import gradio as gr
|
||||
|
||||
def on_ui_settings():
|
||||
section = ('faceswaplab', "FaceSwapLab")
|
||||
models = get_models()
|
||||
shared.opts.add_option("faceswaplab_model", shared.OptionInfo(
|
||||
models[0] if len(models) > 0 else "None", "FaceSwapLab FaceSwap Model", gr.Dropdown, {"interactive": True, "choices" : models}, section=section))
|
||||
shared.opts.add_option("faceswaplab_keep_original", shared.OptionInfo(
|
||||
False, "keep original image before swapping", gr.Checkbox, {"interactive": True}, section=section))
|
||||
shared.opts.add_option("faceswaplab_units_count", shared.OptionInfo(
|
||||
3, "Max faces units (requires restart)", gr.Slider, {"minimum": 1, "maximum": 10, "step": 1}, section=section))
|
||||
|
||||
shared.opts.add_option("faceswaplab_detection_threshold", shared.OptionInfo(
|
||||
0.5, "Detection threshold ", gr.Slider, {"minimum": 0.1, "maximum": 0.99, "step": 0.001}, section=section))
|
||||
|
||||
|
||||
|
||||
shared.opts.add_option("faceswaplab_pp_default_face_restorer", shared.OptionInfo(
|
||||
None, "UI Default post processing face restorer (requires restart)", gr.Dropdown, {"interactive": True, "choices" : ["None"] + [x.name() for x in shared.face_restorers]}, section=section))
|
||||
shared.opts.add_option("faceswaplab_pp_default_face_restorer_visibility", shared.OptionInfo(
|
||||
1, "UI Default post processing face restorer visibility (requires restart)", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.001}, section=section))
|
||||
shared.opts.add_option("faceswaplab_pp_default_face_restorer_weight", shared.OptionInfo(
|
||||
1, "UI Default post processing face restorer weight (requires restart)", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.001}, section=section))
|
||||
shared.opts.add_option("faceswaplab_pp_default_upscaler", shared.OptionInfo(
|
||||
None, "UI Default post processing upscaler (requires restart)", gr.Dropdown, {"interactive": True, "choices" : [upscaler.name for upscaler in shared.sd_upscalers]}, section=section))
|
||||
shared.opts.add_option("faceswaplab_pp_default_upscaler_visibility", shared.OptionInfo(
|
||||
1, "UI Default post processing upscaler visibility(requires restart)", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.001}, section=section))
|
||||
|
||||
|
||||
shared.opts.add_option("faceswaplab_upscaled_swapper", shared.OptionInfo(
|
||||
False, "Upscaled swapper. Applied only to the swapped faces. Apply transformations before merging with the original image.", gr.Checkbox, {"interactive": True}, section=section))
|
||||
shared.opts.add_option("faceswaplab_upscaled_swapper_upscaler", shared.OptionInfo(
|
||||
None, "Upscaled swapper upscaler (Recommanded : LDSR but slow)", gr.Dropdown, {"interactive": True, "choices" : [upscaler.name for upscaler in shared.sd_upscalers]}, section=section))
|
||||
shared.opts.add_option("faceswaplab_upscaled_swapper_sharpen", shared.OptionInfo(
|
||||
False, "Upscaled swapper sharpen", gr.Checkbox, {"interactive": True}, section=section))
|
||||
shared.opts.add_option("faceswaplab_upscaled_swapper_fixcolor", shared.OptionInfo(
|
||||
False, "Upscaled swapper color correction", gr.Checkbox, {"interactive": True}, section=section))
|
||||
shared.opts.add_option("faceswaplab_upscaled_improved_mask", shared.OptionInfo(
|
||||
True, "Use improved segmented mask (use pastenet to mask only the face)", gr.Checkbox, {"interactive": True}, section=section))
|
||||
shared.opts.add_option("faceswaplab_upscaled_swapper_face_restorer", shared.OptionInfo(
|
||||
None, "Upscaled swapper face restorer", gr.Dropdown, {"interactive": True, "choices" : ["None"] + [x.name() for x in shared.face_restorers]}, section=section))
|
||||
shared.opts.add_option("faceswaplab_upscaled_swapper_face_restorer_visibility", shared.OptionInfo(
|
||||
1, "Upscaled swapper face restorer visibility", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.001}, section=section))
|
||||
shared.opts.add_option("faceswaplab_upscaled_swapper_face_restorer_weight", shared.OptionInfo(
|
||||
1, "Upscaled swapper face restorer weight (codeformer)", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.001}, section=section))
|
||||
shared.opts.add_option("faceswaplab_upscaled_swapper_fthresh", shared.OptionInfo(
|
||||
10, "Upscaled swapper fthresh (diff sensitivity) 10 = default behaviour. Low impact.", gr.Slider, {"minimum": 5, "maximum": 250, "step": 1}, section=section))
|
||||
shared.opts.add_option("faceswaplab_upscaled_swapper_erosion", shared.OptionInfo(
|
||||
1, "Upscaled swapper mask erosion factor, 1 = default behaviour. The larger it is, the more blur is applied around the face. Too large and the facial change is no longer visible.", gr.Slider, {"minimum": 0, "maximum": 10, "step": 0.001}, section=section))
|
||||
|
||||
script_callbacks.on_ui_settings(on_ui_settings)
|
||||
@@ -0,0 +1,85 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
from torchvision.transforms.functional import normalize
|
||||
from scripts.faceswaplab_swapping.parsing import init_parsing_model
|
||||
from functools import lru_cache
|
||||
from typing import Union, List
|
||||
from torch import device as torch_device
|
||||
|
||||
@lru_cache
|
||||
def get_parsing_model(device: torch_device) -> torch.nn.Module:
|
||||
"""
|
||||
Returns an instance of the parsing model.
|
||||
The returned model is cached for faster subsequent access.
|
||||
|
||||
Args:
|
||||
device: The torch device to use for computations.
|
||||
|
||||
Returns:
|
||||
The parsing model.
|
||||
"""
|
||||
return init_parsing_model(device=device)
|
||||
|
||||
def convert_image_to_tensor(images: Union[np.ndarray, List[np.ndarray]], convert_bgr_to_rgb: bool = True, use_float32: bool = True) -> Union[torch.Tensor, List[torch.Tensor]]:
|
||||
"""
|
||||
Converts an image or a list of images to PyTorch tensor.
|
||||
|
||||
Args:
|
||||
images: An image or a list of images in numpy.ndarray format.
|
||||
convert_bgr_to_rgb: A boolean flag indicating if the conversion from BGR to RGB should be performed.
|
||||
use_float32: A boolean flag indicating if the tensor should be converted to float32.
|
||||
|
||||
Returns:
|
||||
PyTorch tensor or a list of PyTorch tensors.
|
||||
"""
|
||||
def _convert_single_image_to_tensor(image: np.ndarray, convert_bgr_to_rgb: bool, use_float32: bool) -> torch.Tensor:
|
||||
if image.shape[2] == 3 and convert_bgr_to_rgb:
|
||||
if image.dtype == 'float64':
|
||||
image = image.astype('float32')
|
||||
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
||||
image_tensor = torch.from_numpy(image.transpose(2, 0, 1))
|
||||
if use_float32:
|
||||
image_tensor = image_tensor.float()
|
||||
return image_tensor
|
||||
|
||||
if isinstance(images, list):
|
||||
return [_convert_single_image_to_tensor(image, convert_bgr_to_rgb, use_float32) for image in images]
|
||||
else:
|
||||
return _convert_single_image_to_tensor(images, convert_bgr_to_rgb, use_float32)
|
||||
|
||||
def generate_face_mask(face_image: np.ndarray, device: torch.device) -> np.ndarray:
|
||||
"""
|
||||
Generates a face mask given a face image.
|
||||
|
||||
Args:
|
||||
face_image: The face image in numpy.ndarray format.
|
||||
device: The torch device to use for computations.
|
||||
|
||||
Returns:
|
||||
The face mask as a numpy.ndarray.
|
||||
"""
|
||||
# Resize the face image for the model
|
||||
resized_face_image = cv2.resize(face_image, (512, 512), interpolation=cv2.INTER_LINEAR)
|
||||
|
||||
# Preprocess the image
|
||||
face_input = convert_image_to_tensor((resized_face_image.astype('float32') / 255.0), convert_bgr_to_rgb=True, use_float32=True)
|
||||
normalize(face_input, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
|
||||
assert isinstance(face_input,torch.Tensor)
|
||||
face_input = torch.unsqueeze(face_input, 0).to(device)
|
||||
|
||||
# Pass the image through the model
|
||||
with torch.no_grad():
|
||||
model_output = get_parsing_model(device)(face_input)[0]
|
||||
model_output = model_output.argmax(dim=1).squeeze().cpu().numpy()
|
||||
|
||||
# Generate the mask from the model output
|
||||
parse_mask = np.zeros(model_output.shape)
|
||||
MASK_COLOR_MAP = [0, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 255, 0, 255, 0, 0, 0]
|
||||
for idx, color in enumerate(MASK_COLOR_MAP):
|
||||
parse_mask[model_output == idx] = color
|
||||
|
||||
# Resize the mask to match the original image
|
||||
face_mask = cv2.resize(parse_mask, (face_image.shape[1], face_image.shape[0]))
|
||||
|
||||
return face_mask
|
||||
@@ -0,0 +1,81 @@
|
||||
"""
|
||||
Code from codeformer https://github.com/sczhou/CodeFormer
|
||||
|
||||
S-Lab License 1.0
|
||||
|
||||
Copyright 2022 S-Lab
|
||||
|
||||
Redistribution and use for non-commercial purpose in source and
|
||||
binary forms, with or without modification, are permitted provided
|
||||
that the following conditions are met:
|
||||
|
||||
1. Redistributions of source code must retain the above copyright
|
||||
notice, this list of conditions and the following disclaimer.
|
||||
|
||||
2. Redistributions in binary form must reproduce the above copyright
|
||||
notice, this list of conditions and the following disclaimer in
|
||||
the documentation and/or other materials provided with the
|
||||
distribution.
|
||||
|
||||
3. Neither the name of the copyright holder nor the names of its
|
||||
contributors may be used to endorse or promote products derived
|
||||
from this software without specific prior written permission.
|
||||
|
||||
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
|
||||
"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
|
||||
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
|
||||
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
|
||||
HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
|
||||
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
|
||||
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
|
||||
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
|
||||
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
|
||||
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
|
||||
In the event that redistribution and/or use for commercial purpose in
|
||||
source or binary forms, with or without modification is required,
|
||||
please contact the contributor(s) of the work.
|
||||
"""
|
||||
|
||||
|
||||
import torch
|
||||
import cv2
|
||||
import os
|
||||
import torch
|
||||
from torch.hub import download_url_to_file, get_dir
|
||||
from .parsenet import ParseNet
|
||||
from urllib.parse import urlparse
|
||||
from scripts.faceswaplab_globals import FACE_PARSER_DIR
|
||||
|
||||
ROOT_DIR = FACE_PARSER_DIR
|
||||
|
||||
def load_file_from_url(url, model_dir=None, progress=True, file_name=None):
|
||||
"""Ref:https://github.com/1adrianb/face-alignment/blob/master/face_alignment/utils.py
|
||||
"""
|
||||
if model_dir is None:
|
||||
hub_dir = get_dir()
|
||||
model_dir = os.path.join(hub_dir, 'checkpoints')
|
||||
|
||||
os.makedirs(os.path.join(ROOT_DIR, model_dir), exist_ok=True)
|
||||
|
||||
parts = urlparse(url)
|
||||
filename = os.path.basename(parts.path)
|
||||
if file_name is not None:
|
||||
filename = file_name
|
||||
cached_file = os.path.abspath(os.path.join(ROOT_DIR, model_dir, filename))
|
||||
if not os.path.exists(cached_file):
|
||||
print(f'Downloading: "{url}" to {cached_file}\n')
|
||||
download_url_to_file(url, cached_file, hash_prefix=None, progress=progress)
|
||||
return cached_file
|
||||
|
||||
|
||||
def init_parsing_model(device='cuda'):
|
||||
model = ParseNet(in_size=512, out_size=512, parsing_ch=19)
|
||||
model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/parsing_parsenet.pth'
|
||||
model_path = load_file_from_url(url=model_url, model_dir='weights/facelib', progress=True, file_name=None)
|
||||
load_net = torch.load(model_path, map_location=lambda storage, loc: storage)
|
||||
model.load_state_dict(load_net, strict=True)
|
||||
model.eval()
|
||||
model = model.to(device)
|
||||
return model
|
||||
@@ -0,0 +1,680 @@
|
||||
"""
|
||||
Code from codeformer https://github.com/sczhou/CodeFormer
|
||||
|
||||
S-Lab License 1.0
|
||||
|
||||
Copyright 2022 S-Lab
|
||||
|
||||
Redistribution and use for non-commercial purpose in source and
|
||||
binary forms, with or without modification, are permitted provided
|
||||
that the following conditions are met:
|
||||
|
||||
1. Redistributions of source code must retain the above copyright
|
||||
notice, this list of conditions and the following disclaimer.
|
||||
|
||||
2. Redistributions in binary form must reproduce the above copyright
|
||||
notice, this list of conditions and the following disclaimer in
|
||||
the documentation and/or other materials provided with the
|
||||
distribution.
|
||||
|
||||
3. Neither the name of the copyright holder nor the names of its
|
||||
contributors may be used to endorse or promote products derived
|
||||
from this software without specific prior written permission.
|
||||
|
||||
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
|
||||
"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
|
||||
LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
|
||||
A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
|
||||
HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
|
||||
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
|
||||
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
|
||||
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
|
||||
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
|
||||
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
||||
|
||||
In the event that redistribution and/or use for commercial purpose in
|
||||
source or binary forms, with or without modification is required,
|
||||
please contact the contributor(s) of the work.
|
||||
|
||||
Modified from https://github.com/chaofengc/PSFRGAN
|
||||
|
||||
PSFR-GAN (c) by Chaofeng Chen
|
||||
|
||||
PSFR-GAN is licensed under a
|
||||
Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International License.
|
||||
|
||||
You should have received a copy of the license along with this
|
||||
work. If not, see <http://creativecommons.org/licenses/by-nc-sa/4.0/>.
|
||||
|
||||
Attribution-NonCommercial-ShareAlike 4.0 International
|
||||
|
||||
=======================================================================
|
||||
|
||||
Creative Commons Corporation ("Creative Commons") is not a law firm and
|
||||
does not provide legal services or legal advice. Distribution of
|
||||
Creative Commons public licenses does not create a lawyer-client or
|
||||
other relationship. Creative Commons makes its licenses and related
|
||||
information available on an "as-is" basis. Creative Commons gives no
|
||||
warranties regarding its licenses, any material licensed under their
|
||||
terms and conditions, or any related information. Creative Commons
|
||||
disclaims all liability for damages resulting from their use to the
|
||||
fullest extent possible.
|
||||
|
||||
Using Creative Commons Public Licenses
|
||||
|
||||
Creative Commons public licenses provide a standard set of terms and
|
||||
conditions that creators and other rights holders may use to share
|
||||
original works of authorship and other material subject to copyright
|
||||
and certain other rights specified in the public license below. The
|
||||
following considerations are for informational purposes only, are not
|
||||
exhaustive, and do not form part of our licenses.
|
||||
|
||||
Considerations for licensors: Our public licenses are
|
||||
intended for use by those authorized to give the public
|
||||
permission to use material in ways otherwise restricted by
|
||||
copyright and certain other rights. Our licenses are
|
||||
irrevocable. Licensors should read and understand the terms
|
||||
and conditions of the license they choose before applying it.
|
||||
Licensors should also secure all rights necessary before
|
||||
applying our licenses so that the public can reuse the
|
||||
material as expected. Licensors should clearly mark any
|
||||
material not subject to the license. This includes other CC-
|
||||
licensed material, or material used under an exception or
|
||||
limitation to copyright. More considerations for licensors:
|
||||
wiki.creativecommons.org/Considerations_for_licensors
|
||||
|
||||
Considerations for the public: By using one of our public
|
||||
licenses, a licensor grants the public permission to use the
|
||||
licensed material under specified terms and conditions. If
|
||||
the licensor's permission is not necessary for any reason--for
|
||||
example, because of any applicable exception or limitation to
|
||||
copyright--then that use is not regulated by the license. Our
|
||||
licenses grant only permissions under copyright and certain
|
||||
other rights that a licensor has authority to grant. Use of
|
||||
the licensed material may still be restricted for other
|
||||
reasons, including because others have copyright or other
|
||||
rights in the material. A licensor may make special requests,
|
||||
such as asking that all changes be marked or described.
|
||||
Although not required by our licenses, you are encouraged to
|
||||
respect those requests where reasonable. More_considerations
|
||||
for the public:
|
||||
wiki.creativecommons.org/Considerations_for_licensees
|
||||
|
||||
=======================================================================
|
||||
|
||||
Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International
|
||||
Public License
|
||||
|
||||
By exercising the Licensed Rights (defined below), You accept and agree
|
||||
to be bound by the terms and conditions of this Creative Commons
|
||||
Attribution-NonCommercial-ShareAlike 4.0 International Public License
|
||||
("Public License"). To the extent this Public License may be
|
||||
interpreted as a contract, You are granted the Licensed Rights in
|
||||
consideration of Your acceptance of these terms and conditions, and the
|
||||
Licensor grants You such rights in consideration of benefits the
|
||||
Licensor receives from making the Licensed Material available under
|
||||
these terms and conditions.
|
||||
|
||||
|
||||
Section 1 -- Definitions.
|
||||
|
||||
a. Adapted Material means material subject to Copyright and Similar
|
||||
Rights that is derived from or based upon the Licensed Material
|
||||
and in which the Licensed Material is translated, altered,
|
||||
arranged, transformed, or otherwise modified in a manner requiring
|
||||
permission under the Copyright and Similar Rights held by the
|
||||
Licensor. For purposes of this Public License, where the Licensed
|
||||
Material is a musical work, performance, or sound recording,
|
||||
Adapted Material is always produced where the Licensed Material is
|
||||
synched in timed relation with a moving image.
|
||||
|
||||
b. Adapter's License means the license You apply to Your Copyright
|
||||
and Similar Rights in Your contributions to Adapted Material in
|
||||
accordance with the terms and conditions of this Public License.
|
||||
|
||||
c. BY-NC-SA Compatible License means a license listed at
|
||||
creativecommons.org/compatiblelicenses, approved by Creative
|
||||
Commons as essentially the equivalent of this Public License.
|
||||
|
||||
d. Copyright and Similar Rights means copyright and/or similar rights
|
||||
closely related to copyright including, without limitation,
|
||||
performance, broadcast, sound recording, and Sui Generis Database
|
||||
Rights, without regard to how the rights are labeled or
|
||||
categorized. For purposes of this Public License, the rights
|
||||
specified in Section 2(b)(1)-(2) are not Copyright and Similar
|
||||
Rights.
|
||||
|
||||
e. Effective Technological Measures means those measures that, in the
|
||||
absence of proper authority, may not be circumvented under laws
|
||||
fulfilling obligations under Article 11 of the WIPO Copyright
|
||||
Treaty adopted on December 20, 1996, and/or similar international
|
||||
agreements.
|
||||
|
||||
f. Exceptions and Limitations means fair use, fair dealing, and/or
|
||||
any other exception or limitation to Copyright and Similar Rights
|
||||
that applies to Your use of the Licensed Material.
|
||||
|
||||
g. License Elements means the license attributes listed in the name
|
||||
of a Creative Commons Public License. The License Elements of this
|
||||
Public License are Attribution, NonCommercial, and ShareAlike.
|
||||
|
||||
h. Licensed Material means the artistic or literary work, database,
|
||||
or other material to which the Licensor applied this Public
|
||||
License.
|
||||
|
||||
i. Licensed Rights means the rights granted to You subject to the
|
||||
terms and conditions of this Public License, which are limited to
|
||||
all Copyright and Similar Rights that apply to Your use of the
|
||||
Licensed Material and that the Licensor has authority to license.
|
||||
|
||||
j. Licensor means the individual(s) or entity(ies) granting rights
|
||||
under this Public License.
|
||||
|
||||
k. NonCommercial means not primarily intended for or directed towards
|
||||
commercial advantage or monetary compensation. For purposes of
|
||||
this Public License, the exchange of the Licensed Material for
|
||||
other material subject to Copyright and Similar Rights by digital
|
||||
file-sharing or similar means is NonCommercial provided there is
|
||||
no payment of monetary compensation in connection with the
|
||||
exchange.
|
||||
|
||||
l. Share means to provide material to the public by any means or
|
||||
process that requires permission under the Licensed Rights, such
|
||||
as reproduction, public display, public performance, distribution,
|
||||
dissemination, communication, or importation, and to make material
|
||||
available to the public including in ways that members of the
|
||||
public may access the material from a place and at a time
|
||||
individually chosen by them.
|
||||
|
||||
m. Sui Generis Database Rights means rights other than copyright
|
||||
resulting from Directive 96/9/EC of the European Parliament and of
|
||||
the Council of 11 March 1996 on the legal protection of databases,
|
||||
as amended and/or succeeded, as well as other essentially
|
||||
equivalent rights anywhere in the world.
|
||||
|
||||
n. You means the individual or entity exercising the Licensed Rights
|
||||
under this Public License. Your has a corresponding meaning.
|
||||
|
||||
|
||||
Section 2 -- Scope.
|
||||
|
||||
a. License grant.
|
||||
|
||||
1. Subject to the terms and conditions of this Public License,
|
||||
the Licensor hereby grants You a worldwide, royalty-free,
|
||||
non-sublicensable, non-exclusive, irrevocable license to
|
||||
exercise the Licensed Rights in the Licensed Material to:
|
||||
|
||||
a. reproduce and Share the Licensed Material, in whole or
|
||||
in part, for NonCommercial purposes only; and
|
||||
|
||||
b. produce, reproduce, and Share Adapted Material for
|
||||
NonCommercial purposes only.
|
||||
|
||||
2. Exceptions and Limitations. For the avoidance of doubt, where
|
||||
Exceptions and Limitations apply to Your use, this Public
|
||||
License does not apply, and You do not need to comply with
|
||||
its terms and conditions.
|
||||
|
||||
3. Term. The term of this Public License is specified in Section
|
||||
6(a).
|
||||
|
||||
4. Media and formats; technical modifications allowed. The
|
||||
Licensor authorizes You to exercise the Licensed Rights in
|
||||
all media and formats whether now known or hereafter created,
|
||||
and to make technical modifications necessary to do so. The
|
||||
Licensor waives and/or agrees not to assert any right or
|
||||
authority to forbid You from making technical modifications
|
||||
necessary to exercise the Licensed Rights, including
|
||||
technical modifications necessary to circumvent Effective
|
||||
Technological Measures. For purposes of this Public License,
|
||||
simply making modifications authorized by this Section 2(a)
|
||||
(4) never produces Adapted Material.
|
||||
|
||||
5. Downstream recipients.
|
||||
|
||||
a. Offer from the Licensor -- Licensed Material. Every
|
||||
recipient of the Licensed Material automatically
|
||||
receives an offer from the Licensor to exercise the
|
||||
Licensed Rights under the terms and conditions of this
|
||||
Public License.
|
||||
|
||||
b. Additional offer from the Licensor -- Adapted Material.
|
||||
Every recipient of Adapted Material from You
|
||||
automatically receives an offer from the Licensor to
|
||||
exercise the Licensed Rights in the Adapted Material
|
||||
under the conditions of the Adapter's License You apply.
|
||||
|
||||
c. No downstream restrictions. You may not offer or impose
|
||||
any additional or different terms or conditions on, or
|
||||
apply any Effective Technological Measures to, the
|
||||
Licensed Material if doing so restricts exercise of the
|
||||
Licensed Rights by any recipient of the Licensed
|
||||
Material.
|
||||
|
||||
6. No endorsement. Nothing in this Public License constitutes or
|
||||
may be construed as permission to assert or imply that You
|
||||
are, or that Your use of the Licensed Material is, connected
|
||||
with, or sponsored, endorsed, or granted official status by,
|
||||
the Licensor or others designated to receive attribution as
|
||||
provided in Section 3(a)(1)(A)(i).
|
||||
|
||||
b. Other rights.
|
||||
|
||||
1. Moral rights, such as the right of integrity, are not
|
||||
licensed under this Public License, nor are publicity,
|
||||
privacy, and/or other similar personality rights; however, to
|
||||
the extent possible, the Licensor waives and/or agrees not to
|
||||
assert any such rights held by the Licensor to the limited
|
||||
extent necessary to allow You to exercise the Licensed
|
||||
Rights, but not otherwise.
|
||||
|
||||
2. Patent and trademark rights are not licensed under this
|
||||
Public License.
|
||||
|
||||
3. To the extent possible, the Licensor waives any right to
|
||||
collect royalties from You for the exercise of the Licensed
|
||||
Rights, whether directly or through a collecting society
|
||||
under any voluntary or waivable statutory or compulsory
|
||||
licensing scheme. In all other cases the Licensor expressly
|
||||
reserves any right to collect such royalties, including when
|
||||
the Licensed Material is used other than for NonCommercial
|
||||
purposes.
|
||||
|
||||
|
||||
Section 3 -- License Conditions.
|
||||
|
||||
Your exercise of the Licensed Rights is expressly made subject to the
|
||||
following conditions.
|
||||
|
||||
a. Attribution.
|
||||
|
||||
1. If You Share the Licensed Material (including in modified
|
||||
form), You must:
|
||||
|
||||
a. retain the following if it is supplied by the Licensor
|
||||
with the Licensed Material:
|
||||
|
||||
i. identification of the creator(s) of the Licensed
|
||||
Material and any others designated to receive
|
||||
attribution, in any reasonable manner requested by
|
||||
the Licensor (including by pseudonym if
|
||||
designated);
|
||||
|
||||
ii. a copyright notice;
|
||||
|
||||
iii. a notice that refers to this Public License;
|
||||
|
||||
iv. a notice that refers to the disclaimer of
|
||||
warranties;
|
||||
|
||||
v. a URI or hyperlink to the Licensed Material to the
|
||||
extent reasonably practicable;
|
||||
|
||||
b. indicate if You modified the Licensed Material and
|
||||
retain an indication of any previous modifications; and
|
||||
|
||||
c. indicate the Licensed Material is licensed under this
|
||||
Public License, and include the text of, or the URI or
|
||||
hyperlink to, this Public License.
|
||||
|
||||
2. You may satisfy the conditions in Section 3(a)(1) in any
|
||||
reasonable manner based on the medium, means, and context in
|
||||
which You Share the Licensed Material. For example, it may be
|
||||
reasonable to satisfy the conditions by providing a URI or
|
||||
hyperlink to a resource that includes the required
|
||||
information.
|
||||
3. If requested by the Licensor, You must remove any of the
|
||||
information required by Section 3(a)(1)(A) to the extent
|
||||
reasonably practicable.
|
||||
|
||||
b. ShareAlike.
|
||||
|
||||
In addition to the conditions in Section 3(a), if You Share
|
||||
Adapted Material You produce, the following conditions also apply.
|
||||
|
||||
1. The Adapter's License You apply must be a Creative Commons
|
||||
license with the same License Elements, this version or
|
||||
later, or a BY-NC-SA Compatible License.
|
||||
|
||||
2. You must include the text of, or the URI or hyperlink to, the
|
||||
Adapter's License You apply. You may satisfy this condition
|
||||
in any reasonable manner based on the medium, means, and
|
||||
context in which You Share Adapted Material.
|
||||
|
||||
3. You may not offer or impose any additional or different terms
|
||||
or conditions on, or apply any Effective Technological
|
||||
Measures to, Adapted Material that restrict exercise of the
|
||||
rights granted under the Adapter's License You apply.
|
||||
|
||||
|
||||
Section 4 -- Sui Generis Database Rights.
|
||||
|
||||
Where the Licensed Rights include Sui Generis Database Rights that
|
||||
apply to Your use of the Licensed Material:
|
||||
|
||||
a. for the avoidance of doubt, Section 2(a)(1) grants You the right
|
||||
to extract, reuse, reproduce, and Share all or a substantial
|
||||
portion of the contents of the database for NonCommercial purposes
|
||||
only;
|
||||
|
||||
b. if You include all or a substantial portion of the database
|
||||
contents in a database in which You have Sui Generis Database
|
||||
Rights, then the database in which You have Sui Generis Database
|
||||
Rights (but not its individual contents) is Adapted Material,
|
||||
including for purposes of Section 3(b); and
|
||||
|
||||
c. You must comply with the conditions in Section 3(a) if You Share
|
||||
all or a substantial portion of the contents of the database.
|
||||
|
||||
For the avoidance of doubt, this Section 4 supplements and does not
|
||||
replace Your obligations under this Public License where the Licensed
|
||||
Rights include other Copyright and Similar Rights.
|
||||
|
||||
|
||||
Section 5 -- Disclaimer of Warranties and Limitation of Liability.
|
||||
|
||||
a. UNLESS OTHERWISE SEPARATELY UNDERTAKEN BY THE LICENSOR, TO THE
|
||||
EXTENT POSSIBLE, THE LICENSOR OFFERS THE LICENSED MATERIAL AS-IS
|
||||
AND AS-AVAILABLE, AND MAKES NO REPRESENTATIONS OR WARRANTIES OF
|
||||
ANY KIND CONCERNING THE LICENSED MATERIAL, WHETHER EXPRESS,
|
||||
IMPLIED, STATUTORY, OR OTHER. THIS INCLUDES, WITHOUT LIMITATION,
|
||||
WARRANTIES OF TITLE, MERCHANTABILITY, FITNESS FOR A PARTICULAR
|
||||
PURPOSE, NON-INFRINGEMENT, ABSENCE OF LATENT OR OTHER DEFECTS,
|
||||
ACCURACY, OR THE PRESENCE OR ABSENCE OF ERRORS, WHETHER OR NOT
|
||||
KNOWN OR DISCOVERABLE. WHERE DISCLAIMERS OF WARRANTIES ARE NOT
|
||||
ALLOWED IN FULL OR IN PART, THIS DISCLAIMER MAY NOT APPLY TO YOU.
|
||||
|
||||
b. TO THE EXTENT POSSIBLE, IN NO EVENT WILL THE LICENSOR BE LIABLE
|
||||
TO YOU ON ANY LEGAL THEORY (INCLUDING, WITHOUT LIMITATION,
|
||||
NEGLIGENCE) OR OTHERWISE FOR ANY DIRECT, SPECIAL, INDIRECT,
|
||||
INCIDENTAL, CONSEQUENTIAL, PUNITIVE, EXEMPLARY, OR OTHER LOSSES,
|
||||
COSTS, EXPENSES, OR DAMAGES ARISING OUT OF THIS PUBLIC LICENSE OR
|
||||
USE OF THE LICENSED MATERIAL, EVEN IF THE LICENSOR HAS BEEN
|
||||
ADVISED OF THE POSSIBILITY OF SUCH LOSSES, COSTS, EXPENSES, OR
|
||||
DAMAGES. WHERE A LIMITATION OF LIABILITY IS NOT ALLOWED IN FULL OR
|
||||
IN PART, THIS LIMITATION MAY NOT APPLY TO YOU.
|
||||
|
||||
c. The disclaimer of warranties and limitation of liability provided
|
||||
above shall be interpreted in a manner that, to the extent
|
||||
possible, most closely approximates an absolute disclaimer and
|
||||
waiver of all liability.
|
||||
|
||||
|
||||
Section 6 -- Term and Termination.
|
||||
|
||||
a. This Public License applies for the term of the Copyright and
|
||||
Similar Rights licensed here. However, if You fail to comply with
|
||||
this Public License, then Your rights under this Public License
|
||||
terminate automatically.
|
||||
|
||||
b. Where Your right to use the Licensed Material has terminated under
|
||||
Section 6(a), it reinstates:
|
||||
|
||||
1. automatically as of the date the violation is cured, provided
|
||||
it is cured within 30 days of Your discovery of the
|
||||
violation; or
|
||||
|
||||
2. upon express reinstatement by the Licensor.
|
||||
|
||||
For the avoidance of doubt, this Section 6(b) does not affect any
|
||||
right the Licensor may have to seek remedies for Your violations
|
||||
of this Public License.
|
||||
|
||||
c. For the avoidance of doubt, the Licensor may also offer the
|
||||
Licensed Material under separate terms or conditions or stop
|
||||
distributing the Licensed Material at any time; however, doing so
|
||||
will not terminate this Public License.
|
||||
|
||||
d. Sections 1, 5, 6, 7, and 8 survive termination of this Public
|
||||
License.
|
||||
|
||||
|
||||
Section 7 -- Other Terms and Conditions.
|
||||
|
||||
a. The Licensor shall not be bound by any additional or different
|
||||
terms or conditions communicated by You unless expressly agreed.
|
||||
|
||||
b. Any arrangements, understandings, or agreements regarding the
|
||||
Licensed Material not stated herein are separate from and
|
||||
independent of the terms and conditions of this Public License.
|
||||
|
||||
|
||||
Section 8 -- Interpretation.
|
||||
|
||||
a. For the avoidance of doubt, this Public License does not, and
|
||||
shall not be interpreted to, reduce, limit, restrict, or impose
|
||||
conditions on any use of the Licensed Material that could lawfully
|
||||
be made without permission under this Public License.
|
||||
|
||||
b. To the extent possible, if any provision of this Public License is
|
||||
deemed unenforceable, it shall be automatically reformed to the
|
||||
minimum extent necessary to make it enforceable. If the provision
|
||||
cannot be reformed, it shall be severed from this Public License
|
||||
without affecting the enforceability of the remaining terms and
|
||||
conditions.
|
||||
|
||||
c. No term or condition of this Public License will be waived and no
|
||||
failure to comply consented to unless expressly agreed to by the
|
||||
Licensor.
|
||||
|
||||
d. Nothing in this Public License constitutes or may be interpreted
|
||||
as a limitation upon, or waiver of, any privileges and immunities
|
||||
that apply to the Licensor or You, including from the legal
|
||||
processes of any jurisdiction or authority.
|
||||
|
||||
=======================================================================
|
||||
|
||||
Creative Commons is not a party to its public
|
||||
licenses. Notwithstanding, Creative Commons may elect to apply one of
|
||||
its public licenses to material it publishes and in those instances
|
||||
will be considered the “Licensor.” The text of the Creative Commons
|
||||
public licenses is dedicated to the public domain under the CC0 Public
|
||||
Domain Dedication. Except for the limited purpose of indicating that
|
||||
material is shared under a Creative Commons public license or as
|
||||
otherwise permitted by the Creative Commons policies published at
|
||||
creativecommons.org/policies, Creative Commons does not authorize the
|
||||
use of the trademark "Creative Commons" or any other trademark or logo
|
||||
of Creative Commons without its prior written consent including,
|
||||
without limitation, in connection with any unauthorized modifications
|
||||
to any of its public licenses or any other arrangements,
|
||||
understandings, or agreements concerning use of licensed material. For
|
||||
the avoidance of doubt, this paragraph does not form part of the
|
||||
public licenses.
|
||||
|
||||
Creative Commons may be contacted at creativecommons.org.
|
||||
|
||||
"""
|
||||
import numpy as np
|
||||
import torch.nn as nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
|
||||
class NormLayer(nn.Module):
|
||||
"""Normalization Layers.
|
||||
|
||||
Args:
|
||||
channels: input channels, for batch norm and instance norm.
|
||||
input_size: input shape without batch size, for layer norm.
|
||||
"""
|
||||
|
||||
def __init__(self, channels, normalize_shape=None, norm_type='bn'):
|
||||
super(NormLayer, self).__init__()
|
||||
norm_type = norm_type.lower()
|
||||
self.norm_type = norm_type
|
||||
if norm_type == 'bn':
|
||||
self.norm = nn.BatchNorm2d(channels, affine=True)
|
||||
elif norm_type == 'in':
|
||||
self.norm = nn.InstanceNorm2d(channels, affine=False)
|
||||
elif norm_type == 'gn':
|
||||
self.norm = nn.GroupNorm(32, channels, affine=True)
|
||||
elif norm_type == 'pixel':
|
||||
self.norm = lambda x: F.normalize(x, p=2, dim=1)
|
||||
elif norm_type == 'layer':
|
||||
self.norm = nn.LayerNorm(normalize_shape)
|
||||
elif norm_type == 'none':
|
||||
self.norm = lambda x: x * 1.0
|
||||
else:
|
||||
assert 1 == 0, f'Norm type {norm_type} not support.'
|
||||
|
||||
def forward(self, x, ref=None):
|
||||
if self.norm_type == 'spade':
|
||||
return self.norm(x, ref)
|
||||
else:
|
||||
return self.norm(x)
|
||||
|
||||
|
||||
class ReluLayer(nn.Module):
|
||||
"""Relu Layer.
|
||||
|
||||
Args:
|
||||
relu type: type of relu layer, candidates are
|
||||
- ReLU
|
||||
- LeakyReLU: default relu slope 0.2
|
||||
- PRelu
|
||||
- SELU
|
||||
- none: direct pass
|
||||
"""
|
||||
|
||||
def __init__(self, channels, relu_type='relu'):
|
||||
super(ReluLayer, self).__init__()
|
||||
relu_type = relu_type.lower()
|
||||
if relu_type == 'relu':
|
||||
self.func = nn.ReLU(True)
|
||||
elif relu_type == 'leakyrelu':
|
||||
self.func = nn.LeakyReLU(0.2, inplace=True)
|
||||
elif relu_type == 'prelu':
|
||||
self.func = nn.PReLU(channels)
|
||||
elif relu_type == 'selu':
|
||||
self.func = nn.SELU(True)
|
||||
elif relu_type == 'none':
|
||||
self.func = lambda x: x * 1.0
|
||||
else:
|
||||
assert 1 == 0, f'Relu type {relu_type} not support.'
|
||||
|
||||
def forward(self, x):
|
||||
return self.func(x)
|
||||
|
||||
|
||||
class ConvLayer(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size=3,
|
||||
scale='none',
|
||||
norm_type='none',
|
||||
relu_type='none',
|
||||
use_pad=True,
|
||||
bias=True):
|
||||
super(ConvLayer, self).__init__()
|
||||
self.use_pad = use_pad
|
||||
self.norm_type = norm_type
|
||||
if norm_type in ['bn']:
|
||||
bias = False
|
||||
|
||||
stride = 2 if scale == 'down' else 1
|
||||
|
||||
self.scale_func = lambda x: x
|
||||
if scale == 'up':
|
||||
self.scale_func = lambda x: nn.functional.interpolate(x, scale_factor=2, mode='nearest')
|
||||
|
||||
self.reflection_pad = nn.ReflectionPad2d(int(np.ceil((kernel_size - 1.) / 2)))
|
||||
self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride, bias=bias)
|
||||
|
||||
self.relu = ReluLayer(out_channels, relu_type)
|
||||
self.norm = NormLayer(out_channels, norm_type=norm_type)
|
||||
|
||||
def forward(self, x):
|
||||
out = self.scale_func(x)
|
||||
if self.use_pad:
|
||||
out = self.reflection_pad(out)
|
||||
out = self.conv2d(out)
|
||||
out = self.norm(out)
|
||||
out = self.relu(out)
|
||||
return out
|
||||
|
||||
|
||||
class ResidualBlock(nn.Module):
|
||||
"""
|
||||
Residual block recommended in: http://torch.ch/blog/2016/02/04/resnets.html
|
||||
"""
|
||||
|
||||
def __init__(self, c_in, c_out, relu_type='prelu', norm_type='bn', scale='none'):
|
||||
super(ResidualBlock, self).__init__()
|
||||
|
||||
if scale == 'none' and c_in == c_out:
|
||||
self.shortcut_func = lambda x: x
|
||||
else:
|
||||
self.shortcut_func = ConvLayer(c_in, c_out, 3, scale)
|
||||
|
||||
scale_config_dict = {'down': ['none', 'down'], 'up': ['up', 'none'], 'none': ['none', 'none']}
|
||||
scale_conf = scale_config_dict[scale]
|
||||
|
||||
self.conv1 = ConvLayer(c_in, c_out, 3, scale_conf[0], norm_type=norm_type, relu_type=relu_type)
|
||||
self.conv2 = ConvLayer(c_out, c_out, 3, scale_conf[1], norm_type=norm_type, relu_type='none')
|
||||
|
||||
def forward(self, x):
|
||||
identity = self.shortcut_func(x)
|
||||
|
||||
res = self.conv1(x)
|
||||
res = self.conv2(res)
|
||||
return identity + res
|
||||
|
||||
|
||||
class ParseNet(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
in_size=128,
|
||||
out_size=128,
|
||||
min_feat_size=32,
|
||||
base_ch=64,
|
||||
parsing_ch=19,
|
||||
res_depth=10,
|
||||
relu_type='LeakyReLU',
|
||||
norm_type='bn',
|
||||
ch_range=[32, 256]):
|
||||
super().__init__()
|
||||
self.res_depth = res_depth
|
||||
act_args = {'norm_type': norm_type, 'relu_type': relu_type}
|
||||
min_ch, max_ch = ch_range
|
||||
|
||||
ch_clip = lambda x: max(min_ch, min(x, max_ch)) # noqa: E731
|
||||
min_feat_size = min(in_size, min_feat_size)
|
||||
|
||||
down_steps = int(np.log2(in_size // min_feat_size))
|
||||
up_steps = int(np.log2(out_size // min_feat_size))
|
||||
|
||||
# =============== define encoder-body-decoder ====================
|
||||
self.encoder = []
|
||||
self.encoder.append(ConvLayer(3, base_ch, 3, 1))
|
||||
head_ch = base_ch
|
||||
for i in range(down_steps):
|
||||
cin, cout = ch_clip(head_ch), ch_clip(head_ch * 2)
|
||||
self.encoder.append(ResidualBlock(cin, cout, scale='down', **act_args))
|
||||
head_ch = head_ch * 2
|
||||
|
||||
self.body = []
|
||||
for i in range(res_depth):
|
||||
self.body.append(ResidualBlock(ch_clip(head_ch), ch_clip(head_ch), **act_args))
|
||||
|
||||
self.decoder = []
|
||||
for i in range(up_steps):
|
||||
cin, cout = ch_clip(head_ch), ch_clip(head_ch // 2)
|
||||
self.decoder.append(ResidualBlock(cin, cout, scale='up', **act_args))
|
||||
head_ch = head_ch // 2
|
||||
|
||||
self.encoder = nn.Sequential(*self.encoder)
|
||||
self.body = nn.Sequential(*self.body)
|
||||
self.decoder = nn.Sequential(*self.decoder)
|
||||
self.out_img_conv = ConvLayer(ch_clip(head_ch), 3)
|
||||
self.out_mask_conv = ConvLayer(ch_clip(head_ch), parsing_ch)
|
||||
|
||||
def forward(self, x):
|
||||
feat = self.encoder(x)
|
||||
x = feat + self.body(feat)
|
||||
x = self.decoder(x)
|
||||
out_img = self.out_img_conv(x)
|
||||
out_mask = self.out_mask_conv(x)
|
||||
return out_mask, out_img
|
||||
@@ -0,0 +1,428 @@
|
||||
import copy
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, List, Set, Tuple, Optional, Union
|
||||
|
||||
import cv2
|
||||
import insightface
|
||||
import numpy as np
|
||||
from insightface.app.common import Face
|
||||
|
||||
from PIL import Image
|
||||
from sklearn.metrics.pairwise import cosine_similarity
|
||||
|
||||
from scripts.faceswaplab_swapping import upscaled_inswapper
|
||||
from scripts.faceswaplab_utils.imgutils import cv2_to_pil, pil_to_cv2, check_against_nsfw
|
||||
from scripts.faceswaplab_utils.faceswaplab_logging import logger, save_img_debug
|
||||
from scripts import faceswaplab_globals
|
||||
from modules.shared import opts
|
||||
from functools import lru_cache
|
||||
from scripts.faceswaplab_ui.faceswaplab_unit_settings import FaceSwapUnitSettings
|
||||
|
||||
providers = ["CPUExecutionProvider"]
|
||||
|
||||
|
||||
def cosine_similarity_face(face1, face2) -> float:
|
||||
"""
|
||||
Calculates the cosine similarity between two face embeddings.
|
||||
|
||||
Args:
|
||||
face1 (Face): The first face object containing an embedding.
|
||||
face2 (Face): The second face object containing an embedding.
|
||||
|
||||
Returns:
|
||||
float: The cosine similarity between the face embeddings.
|
||||
|
||||
Note:
|
||||
The cosine similarity ranges from 0 to 1, where 1 indicates identical embeddings and 0 indicates completely
|
||||
dissimilar embeddings. In this implementation, the similarity is clamped to a minimum value of 0 to ensure a
|
||||
non-negative similarity score.
|
||||
"""
|
||||
# Reshape the face embeddings to have a shape of (1, -1)
|
||||
vec1 = face1.embedding.reshape(1, -1)
|
||||
vec2 = face2.embedding.reshape(1, -1)
|
||||
|
||||
# Calculate the cosine similarity between the reshaped embeddings
|
||||
similarity = cosine_similarity(vec1, vec2)
|
||||
|
||||
# Return the maximum of 0 and the calculated similarity as the final similarity score
|
||||
return max(0, similarity[0, 0])
|
||||
|
||||
def compare_faces(img1: Image.Image, img2: Image.Image) -> float:
|
||||
"""
|
||||
Compares the similarity between two faces extracted from images using cosine similarity.
|
||||
|
||||
Args:
|
||||
img1: The first image containing a face.
|
||||
img2: The second image containing a face.
|
||||
|
||||
Returns:
|
||||
A float value representing the similarity between the two faces (0 to 1).
|
||||
Returns -1 if one or both of the images do not contain any faces.
|
||||
"""
|
||||
|
||||
# Extract faces from the images
|
||||
face1 = get_or_default(get_faces(pil_to_cv2(img1)), 0, None)
|
||||
face2 = get_or_default(get_faces(pil_to_cv2(img2)), 0, None)
|
||||
|
||||
# Check if both faces are detected
|
||||
if face1 is not None and face2 is not None:
|
||||
# Calculate the cosine similarity between the faces
|
||||
return cosine_similarity_face(face1, face2)
|
||||
|
||||
# Return -1 if one or both of the images do not contain any faces
|
||||
return -1
|
||||
|
||||
|
||||
class FaceModelException(Exception):
|
||||
"""Exception raised when an error is encountered in the face model."""
|
||||
def __init__(self, message: str) -> None:
|
||||
"""
|
||||
Args:
|
||||
message: A string containing the error description.
|
||||
"""
|
||||
self.message = message
|
||||
super().__init__(self.message)
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def getAnalysisModel():
|
||||
"""
|
||||
Retrieves the analysis model for face analysis.
|
||||
|
||||
Returns:
|
||||
insightface.app.FaceAnalysis: The analysis model for face analysis.
|
||||
"""
|
||||
try :
|
||||
if not os.path.exists(faceswaplab_globals.ANALYZER_DIR):
|
||||
os.makedirs(faceswaplab_globals.ANALYZER_DIR)
|
||||
|
||||
logger.info("Load analysis model, will take some time.")
|
||||
# Initialize the analysis model with the specified name and providers
|
||||
return insightface.app.FaceAnalysis(
|
||||
name="buffalo_l", providers=providers, root=faceswaplab_globals.ANALYZER_DIR
|
||||
)
|
||||
except Exception as e :
|
||||
logger.error("Loading of swapping model failed, please check the requirements (On Windows, download and install Visual Studio. During the install, make sure to include the Python and C++ packages.)")
|
||||
raise FaceModelException("Loading of swapping model failed")
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def getFaceSwapModel(model_path: str):
|
||||
"""
|
||||
Retrieves the face swap model and initializes it if necessary.
|
||||
|
||||
Args:
|
||||
model_path (str): Path to the face swap model.
|
||||
|
||||
Returns:
|
||||
insightface.model_zoo.FaceModel: The face swap model.
|
||||
"""
|
||||
try :
|
||||
# Initializes the face swap model using the specified model path.
|
||||
return upscaled_inswapper.UpscaledINSwapper(insightface.model_zoo.get_model(model_path, providers=providers))
|
||||
except Exception as e :
|
||||
logger.error("Loading of swapping model failed, please check the requirements (On Windows, download and install Visual Studio. During the install, make sure to include the Python and C++ packages.)")
|
||||
|
||||
|
||||
def get_faces(img_data: np.ndarray, det_size=(640, 640), det_thresh : Optional[int]=None, sort_by_face_size = False) -> List[Face]:
|
||||
"""
|
||||
Detects and retrieves faces from an image using an analysis model.
|
||||
|
||||
Args:
|
||||
img_data (np.ndarray): The image data as a NumPy array.
|
||||
det_size (tuple): The desired detection size (width, height). Defaults to (640, 640).
|
||||
sort_by_face_size (bool) : Will sort the faces by their size from larger to smaller face
|
||||
|
||||
Returns:
|
||||
list: A list of detected faces, sorted by their x-coordinate of the bounding box.
|
||||
"""
|
||||
|
||||
if det_thresh is None :
|
||||
det_thresh = opts.data.get("faceswaplab_detection_threshold", 0.5)
|
||||
|
||||
# Create a deep copy of the analysis model (otherwise det_size is attached to the analysis model and can't be changed)
|
||||
face_analyser = copy.deepcopy(getAnalysisModel())
|
||||
|
||||
# Prepare the analysis model for face detection with the specified detection size
|
||||
face_analyser.prepare(ctx_id=0, det_thresh=det_thresh, det_size=det_size)
|
||||
|
||||
# Get the detected faces from the image using the analysis model
|
||||
face = face_analyser.get(img_data)
|
||||
|
||||
# If no faces are detected and the detection size is larger than 320x320,
|
||||
# recursively call the function with a smaller detection size
|
||||
if len(face) == 0 and det_size[0] > 320 and det_size[1] > 320:
|
||||
det_size_half = (det_size[0] // 2, det_size[1] // 2)
|
||||
return get_faces(img_data, det_size=det_size_half, det_thresh=det_thresh)
|
||||
|
||||
try:
|
||||
if sort_by_face_size :
|
||||
return sorted(face, reverse=True, key=lambda x: (x.bbox[2] - x.bbox[0]) * (x.bbox[3] - x.bbox[1]))
|
||||
|
||||
# Sort the detected faces based on their x-coordinate of the bounding box
|
||||
return sorted(face, key=lambda x: x.bbox[0])
|
||||
except Exception as e:
|
||||
return []
|
||||
|
||||
|
||||
|
||||
@dataclass
|
||||
class ImageResult:
|
||||
"""
|
||||
Represents the result of an image swap operation
|
||||
"""
|
||||
|
||||
image: Image.Image
|
||||
"""
|
||||
The image object with the swapped face
|
||||
"""
|
||||
|
||||
similarity: Dict[int, float]
|
||||
"""
|
||||
A dictionary mapping face indices to their similarity scores.
|
||||
The similarity scores are represented as floating-point values between 0 and 1.
|
||||
"""
|
||||
|
||||
ref_similarity: Dict[int, float]
|
||||
"""
|
||||
A dictionary mapping face indices to their similarity scores compared to a reference image.
|
||||
The similarity scores are represented as floating-point values between 0 and 1.
|
||||
"""
|
||||
|
||||
|
||||
def get_or_default(l, index, default):
|
||||
"""
|
||||
Retrieve the value at the specified index from the given list.
|
||||
If the index is out of bounds, return the default value instead.
|
||||
|
||||
Args:
|
||||
l (list): The input list.
|
||||
index (int): The index to retrieve the value from.
|
||||
default: The default value to return if the index is out of bounds.
|
||||
|
||||
Returns:
|
||||
The value at the specified index if it exists, otherwise the default value.
|
||||
"""
|
||||
return l[index] if index < len(l) else default
|
||||
|
||||
|
||||
def get_faces_from_img_files(files):
|
||||
"""
|
||||
Extracts faces from a list of image files.
|
||||
|
||||
Args:
|
||||
files (list): A list of file objects representing image files.
|
||||
|
||||
Returns:
|
||||
list: A list of detected faces.
|
||||
|
||||
"""
|
||||
|
||||
faces = []
|
||||
|
||||
if len(files) > 0:
|
||||
for file in files:
|
||||
img = Image.open(file.name) # Open the image file
|
||||
face = get_or_default(get_faces(pil_to_cv2(img)), 0, None) # Extract faces from the image
|
||||
if face is not None:
|
||||
faces.append(face) # Add the detected face to the list of faces
|
||||
|
||||
return faces
|
||||
|
||||
def blend_faces(faces: List[Face]) -> Face:
|
||||
"""
|
||||
Blends the embeddings of multiple faces into a single face.
|
||||
|
||||
Args:
|
||||
faces (List[Face]): List of Face objects.
|
||||
|
||||
Returns:
|
||||
Face: The blended Face object with the averaged embedding.
|
||||
Returns None if the input list is empty.
|
||||
|
||||
Raises:
|
||||
ValueError: If the embeddings have different shapes.
|
||||
|
||||
"""
|
||||
embeddings = [face.embedding for face in faces]
|
||||
|
||||
if len(embeddings) > 0:
|
||||
embedding_shape = embeddings[0].shape
|
||||
|
||||
# Check if all embeddings have the same shape
|
||||
for embedding in embeddings:
|
||||
if embedding.shape != embedding_shape:
|
||||
raise ValueError("embedding shape mismatch")
|
||||
|
||||
# Compute the mean of all embeddings
|
||||
blended_embedding = np.mean(embeddings, axis=0)
|
||||
|
||||
# Create a new Face object using the properties of the first face in the list
|
||||
# Assign the blended embedding to the blended Face object
|
||||
blended = Face(embedding=blended_embedding, gender=faces[0].gender, age=faces[0].age)
|
||||
|
||||
assert not np.array_equal(blended.embedding,faces[0].embedding) if len(faces) > 1 else True, "If len(faces)>0, the blended embedding should not be the same than the first image"
|
||||
|
||||
return blended
|
||||
|
||||
# Return None if the input list is empty
|
||||
return None
|
||||
|
||||
|
||||
def swap_face(
|
||||
reference_face: np.ndarray,
|
||||
source_face: np.ndarray,
|
||||
target_img: Image.Image,
|
||||
model: str,
|
||||
faces_index: Set[int] = {0},
|
||||
same_gender=True,
|
||||
upscaled_swapper = False,
|
||||
compute_similarity = True,
|
||||
sort_by_face_size = False
|
||||
) -> ImageResult:
|
||||
"""
|
||||
Swaps faces in the target image with the source face.
|
||||
|
||||
Args:
|
||||
reference_face (np.ndarray): The reference face used for similarity comparison.
|
||||
source_face (np.ndarray): The source face to be swapped.
|
||||
target_img (Image.Image): The target image to swap faces in.
|
||||
model (str): Path to the face swap model.
|
||||
faces_index (Set[int], optional): Set of indices specifying which faces to swap. Defaults to {0}.
|
||||
same_gender (bool, optional): If True, only swap faces with the same gender as the source face. Defaults to True.
|
||||
|
||||
Returns:
|
||||
ImageResult: An object containing the swapped image and similarity scores.
|
||||
|
||||
"""
|
||||
return_result = ImageResult(target_img, {}, {})
|
||||
try :
|
||||
target_img = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR)
|
||||
gender = source_face["gender"]
|
||||
logger.info("Source Gender %s", gender)
|
||||
if source_face is not None:
|
||||
result = target_img
|
||||
model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), model)
|
||||
face_swapper = getFaceSwapModel(model_path)
|
||||
target_faces = get_faces(target_img, sort_by_face_size=sort_by_face_size)
|
||||
logger.info("Target faces count : %s", len(target_faces))
|
||||
|
||||
if same_gender:
|
||||
target_faces = [x for x in target_faces if x["gender"] == gender]
|
||||
logger.info("Target Gender Matches count %s", len(target_faces))
|
||||
|
||||
for i, swapped_face in enumerate(target_faces):
|
||||
logger.info(f"swap face {i}")
|
||||
if i in faces_index:
|
||||
result = face_swapper.get(result, swapped_face, source_face, upscale = upscaled_swapper)
|
||||
|
||||
result_image = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB))
|
||||
return_result.image = result_image
|
||||
|
||||
|
||||
if compute_similarity :
|
||||
try:
|
||||
result_faces = get_faces(
|
||||
cv2.cvtColor(np.array(result_image), cv2.COLOR_RGB2BGR), sort_by_face_size=sort_by_face_size
|
||||
)
|
||||
if same_gender:
|
||||
result_faces = [x for x in result_faces if x["gender"] == gender]
|
||||
|
||||
for i, swapped_face in enumerate(result_faces):
|
||||
logger.info(f"compare face {i}")
|
||||
if i in faces_index and i < len(target_faces):
|
||||
return_result.similarity[i] = cosine_similarity_face(
|
||||
source_face, swapped_face
|
||||
)
|
||||
return_result.ref_similarity[i] = cosine_similarity_face(
|
||||
reference_face, swapped_face
|
||||
)
|
||||
|
||||
logger.info(f"similarity {return_result.similarity}")
|
||||
logger.info(f"ref similarity {return_result.ref_similarity}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error("Similarity processing failed %s", e)
|
||||
raise e
|
||||
except Exception as e :
|
||||
logger.error("Conversion failed %s", e)
|
||||
raise e
|
||||
return return_result
|
||||
|
||||
|
||||
def process_image_unit(model, unit : FaceSwapUnitSettings, image: Image.Image, info = None, upscaled_swapper = False, force_blend = False) -> List:
|
||||
"""Process one image and return a List of (image, info) (one if blended, many if not).
|
||||
|
||||
Args:
|
||||
unit : the current unit
|
||||
image : the image where to apply swapping
|
||||
info : The info
|
||||
|
||||
Returns:
|
||||
List of tuple of (image, info) where image is the image where swapping has been applied and info is the image info with similarity infos.
|
||||
"""
|
||||
|
||||
results = []
|
||||
if unit.enable :
|
||||
if check_against_nsfw(image) :
|
||||
return [(image, info)]
|
||||
if not unit.blend_faces and not force_blend :
|
||||
src_faces = unit.faces
|
||||
logger.info(f"will generate {len(src_faces)} images")
|
||||
else :
|
||||
logger.info("blend all faces together")
|
||||
src_faces = [unit.blended_faces]
|
||||
assert(not np.array_equal(unit.reference_face.embedding,src_faces[0].embedding) if len(unit.faces)>1 else True), "Reference face cannot be the same as blended"
|
||||
|
||||
|
||||
for i,src_face in enumerate(src_faces):
|
||||
logger.info(f"Process face {i}")
|
||||
if unit.reference_face is not None :
|
||||
reference_face = unit.reference_face
|
||||
else :
|
||||
logger.info("Use source face as reference face")
|
||||
reference_face = src_face
|
||||
|
||||
save_img_debug(image, "Before swap")
|
||||
result: ImageResult = swap_face(
|
||||
reference_face,
|
||||
src_face,
|
||||
image,
|
||||
faces_index=unit.faces_index,
|
||||
model=model,
|
||||
same_gender=unit.same_gender,
|
||||
upscaled_swapper=upscaled_swapper,
|
||||
compute_similarity=unit.compute_similarity,
|
||||
sort_by_face_size=unit.sort_by_size
|
||||
)
|
||||
save_img_debug(result.image, "After swap")
|
||||
|
||||
if result.image is None :
|
||||
logger.error("Result image is None")
|
||||
if (not unit.check_similarity) or result.similarity and all([result.similarity.values()!=0]+[x >= unit.min_sim for x in result.similarity.values()]) and all([result.ref_similarity.values()!=0]+[x >= unit.min_ref_sim for x in result.ref_similarity.values()]):
|
||||
results.append((result.image, f"{info}, similarity = {result.similarity}, ref_similarity = {result.ref_similarity}"))
|
||||
else:
|
||||
logger.warning(
|
||||
f"skip, similarity to low, sim = {result.similarity} (target {unit.min_sim}) ref sim = {result.ref_similarity} (target = {unit.min_ref_sim})"
|
||||
)
|
||||
logger.debug("process_image_unit : Unit produced %s results", len(results))
|
||||
return results
|
||||
|
||||
def process_images_units(model, units : List[FaceSwapUnitSettings], images: List[Tuple[Optional[Image.Image], Optional[str]]], upscaled_swapper = False, force_blend = False) -> Union[List,None]:
|
||||
if len(units) == 0 :
|
||||
logger.info("Finished processing image, return %s images", len(images))
|
||||
return None
|
||||
|
||||
logger.debug("%s more units", len(units))
|
||||
|
||||
processed_images = []
|
||||
for i,(image, info) in enumerate(images) :
|
||||
logger.debug("Processing image %s", i)
|
||||
swapped = process_image_unit(model,units[0],image, info, upscaled_swapper, force_blend)
|
||||
logger.debug("Image %s -> %s images", i, len(swapped))
|
||||
nexts = process_images_units(model,units[1:],swapped, upscaled_swapper,force_blend)
|
||||
if nexts :
|
||||
processed_images.extend(nexts)
|
||||
else :
|
||||
processed_images.extend(swapped)
|
||||
|
||||
return processed_images
|
||||
@@ -0,0 +1,188 @@
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import onnx
|
||||
import onnxruntime
|
||||
from insightface.model_zoo.inswapper import INSwapper
|
||||
from insightface.utils import face_align
|
||||
from modules import codeformer_model, processing, scripts, shared
|
||||
from modules.face_restoration import FaceRestoration
|
||||
from modules.shared import cmd_opts, opts, state
|
||||
from modules.upscaler import UpscalerData
|
||||
from onnx import numpy_helper
|
||||
from PIL import Image
|
||||
|
||||
from scripts.faceswaplab_utils.faceswaplab_logging import logger
|
||||
from scripts.faceswaplab_postprocessing import upscaling
|
||||
from scripts.faceswaplab_postprocessing.postprocessing_options import \
|
||||
PostProcessingOptions
|
||||
from scripts.faceswaplab_swapping.facemask import generate_face_mask
|
||||
from scripts.faceswaplab_utils.imgutils import cv2_to_pil, pil_to_cv2
|
||||
|
||||
|
||||
def get_upscaler() -> UpscalerData:
|
||||
for upscaler in shared.sd_upscalers:
|
||||
if upscaler.name == opts.data.get("faceswaplab_upscaled_swapper_upscaler", "LDSR"):
|
||||
return upscaler
|
||||
return None
|
||||
|
||||
def merge_images_with_mask(image1, image2, mask):
|
||||
if image1.shape != image2.shape or image1.shape[:2] != mask.shape:
|
||||
raise ValueError("Img should have the same shape")
|
||||
mask = mask.astype(np.uint8)
|
||||
masked_region = cv2.bitwise_and(image2, image2, mask=mask)
|
||||
inverse_mask = cv2.bitwise_not(mask)
|
||||
empty_region = cv2.bitwise_and(image1, image1, mask=inverse_mask)
|
||||
merged_image = cv2.add(empty_region, masked_region)
|
||||
return merged_image
|
||||
|
||||
def erode_mask(mask, kernel_size=3, iterations=1):
|
||||
kernel = np.ones((kernel_size, kernel_size), np.uint8)
|
||||
eroded_mask = cv2.erode(mask, kernel, iterations=iterations)
|
||||
return eroded_mask
|
||||
|
||||
def apply_gaussian_blur(mask, kernel_size=(5, 5), sigma_x=0):
|
||||
blurred_mask = cv2.GaussianBlur(mask, kernel_size, sigma_x)
|
||||
return blurred_mask
|
||||
|
||||
def dilate_mask(mask, kernel_size=5, iterations=1):
|
||||
kernel = np.ones((kernel_size, kernel_size), np.uint8)
|
||||
dilated_mask = cv2.dilate(mask, kernel, iterations=iterations)
|
||||
return dilated_mask
|
||||
|
||||
def get_face_mask(aimg,bgr_fake):
|
||||
mask1 = generate_face_mask(aimg, device = shared.device)
|
||||
mask2 = generate_face_mask(bgr_fake, device = shared.device)
|
||||
mask = dilate_mask(cv2.bitwise_or(mask1,mask2))
|
||||
return mask
|
||||
|
||||
|
||||
class UpscaledINSwapper():
|
||||
def __init__(self, inswapper : INSwapper):
|
||||
self.__dict__.update(inswapper.__dict__)
|
||||
|
||||
def forward(self, img, latent):
|
||||
img = (img - self.input_mean) / self.input_std
|
||||
pred = self.session.run(self.output_names, {self.input_names[0]: img, self.input_names[1]: latent})[0]
|
||||
return pred
|
||||
|
||||
def super_resolution(self,img, k = 2) :
|
||||
pil_img = cv2_to_pil(img)
|
||||
options = PostProcessingOptions(
|
||||
upscaler_name=opts.data.get('faceswaplab_upscaled_swapper_upscaler', 'LDSR'),
|
||||
upscale_visibility=1,
|
||||
scale=k,
|
||||
face_restorer_name=opts.data.get('faceswaplab_upscaled_swapper_face_restorer', ""),
|
||||
codeformer_weight= opts.data.get('faceswaplab_upscaled_swapper_face_restorer_weight', 1),
|
||||
restorer_visibility=opts.data.get('faceswaplab_upscaled_swapper_face_restorer_visibility', 1))
|
||||
upscaled = upscaling.upscale_img(pil_img, options)
|
||||
upscaled = upscaling.restore_face(upscaled, options)
|
||||
return pil_to_cv2(upscaled)
|
||||
|
||||
def get(self, img, target_face, source_face, paste_back=True, upscale = True):
|
||||
aimg, M = face_align.norm_crop2(img, target_face.kps, self.input_size[0])
|
||||
blob = cv2.dnn.blobFromImage(aimg, 1.0 / self.input_std, self.input_size,
|
||||
(self.input_mean, self.input_mean, self.input_mean), swapRB=True)
|
||||
latent = source_face.normed_embedding.reshape((1,-1))
|
||||
latent = np.dot(latent, self.emap)
|
||||
latent /= np.linalg.norm(latent)
|
||||
pred = self.session.run(self.output_names, {self.input_names[0]: blob, self.input_names[1]: latent})[0]
|
||||
#print(latent.shape, latent.dtype, pred.shape)
|
||||
img_fake = pred.transpose((0,2,3,1))[0]
|
||||
bgr_fake = np.clip(255 * img_fake, 0, 255).astype(np.uint8)[:,:,::-1]
|
||||
|
||||
try :
|
||||
if not paste_back:
|
||||
return bgr_fake, M
|
||||
else:
|
||||
target_img = img
|
||||
|
||||
def compute_diff(bgr_fake,aimg) :
|
||||
fake_diff = bgr_fake.astype(np.float32) - aimg.astype(np.float32)
|
||||
fake_diff = np.abs(fake_diff).mean(axis=2)
|
||||
fake_diff[:2,:] = 0
|
||||
fake_diff[-2:,:] = 0
|
||||
fake_diff[:,:2] = 0
|
||||
fake_diff[:,-2:] = 0
|
||||
return fake_diff
|
||||
|
||||
if upscale :
|
||||
|
||||
print("*"*80)
|
||||
print(f"Upscaled inswapper using {opts.data.get('faceswaplab_upscaled_swapper_upscaler', 'LDSR')}")
|
||||
print("*"*80)
|
||||
|
||||
k = 4
|
||||
aimg, M = face_align.norm_crop2(img, target_face.kps, self.input_size[0]*k)
|
||||
|
||||
# upscale and restore face :
|
||||
bgr_fake = self.super_resolution(bgr_fake, k)
|
||||
|
||||
if opts.data.get("faceswaplab_upscaled_improved_mask", True) :
|
||||
mask = get_face_mask(aimg,bgr_fake)
|
||||
bgr_fake = merge_images_with_mask(aimg, bgr_fake,mask)
|
||||
|
||||
# compute fake_diff before sharpen and color correction (better result)
|
||||
fake_diff = compute_diff(bgr_fake, aimg)
|
||||
|
||||
if opts.data.get("faceswaplab_upscaled_swapper_sharpen", True) :
|
||||
print("sharpen")
|
||||
# Add sharpness
|
||||
blurred = cv2.GaussianBlur(bgr_fake, (0, 0), 3)
|
||||
bgr_fake = cv2.addWeighted(bgr_fake, 1.5, blurred, -0.5, 0)
|
||||
|
||||
# Apply color corrections
|
||||
if opts.data.get("faceswaplab_upscaled_swapper_fixcolor", True) :
|
||||
print("color correction")
|
||||
correction = processing.setup_color_correction(cv2_to_pil(aimg))
|
||||
bgr_fake_pil = processing.apply_color_correction(correction, cv2_to_pil(bgr_fake))
|
||||
bgr_fake = pil_to_cv2(bgr_fake_pil)
|
||||
|
||||
|
||||
else :
|
||||
fake_diff = compute_diff(bgr_fake, aimg)
|
||||
|
||||
IM = cv2.invertAffineTransform(M)
|
||||
|
||||
img_white = np.full((aimg.shape[0],aimg.shape[1]), 255, dtype=np.float32)
|
||||
bgr_fake = cv2.warpAffine(bgr_fake, IM, (target_img.shape[1], target_img.shape[0]), borderValue=0.0)
|
||||
img_white = cv2.warpAffine(img_white, IM, (target_img.shape[1], target_img.shape[0]), borderValue=0.0)
|
||||
fake_diff = cv2.warpAffine(fake_diff, IM, (target_img.shape[1], target_img.shape[0]), borderValue=0.0)
|
||||
img_white[img_white>20] = 255
|
||||
fthresh = opts.data.get('faceswaplab_upscaled_swapper_fthresh', 10)
|
||||
print("fthresh", fthresh)
|
||||
fake_diff[fake_diff<fthresh] = 0
|
||||
fake_diff[fake_diff>=fthresh] = 255
|
||||
img_mask = img_white
|
||||
mask_h_inds, mask_w_inds = np.where(img_mask==255)
|
||||
mask_h = np.max(mask_h_inds) - np.min(mask_h_inds)
|
||||
mask_w = np.max(mask_w_inds) - np.min(mask_w_inds)
|
||||
mask_size = int(np.sqrt(mask_h*mask_w))
|
||||
erosion_factor = opts.data.get('faceswaplab_upscaled_swapper_erosion', 1)
|
||||
k = max(int(mask_size//10*erosion_factor), int(10*erosion_factor))
|
||||
|
||||
kernel = np.ones((k,k),np.uint8)
|
||||
img_mask = cv2.erode(img_mask,kernel,iterations = 1)
|
||||
kernel = np.ones((2,2),np.uint8)
|
||||
fake_diff = cv2.dilate(fake_diff,kernel,iterations = 1)
|
||||
k = max(int(mask_size//20*erosion_factor), int(5*erosion_factor))
|
||||
|
||||
|
||||
kernel_size = (k, k)
|
||||
blur_size = tuple(2*i+1 for i in kernel_size)
|
||||
img_mask = cv2.GaussianBlur(img_mask, blur_size, 0)
|
||||
k = int(5*erosion_factor)
|
||||
kernel_size = (k, k)
|
||||
blur_size = tuple(2*i+1 for i in kernel_size)
|
||||
fake_diff = cv2.GaussianBlur(fake_diff, blur_size, 0)
|
||||
img_mask /= 255
|
||||
fake_diff /= 255
|
||||
|
||||
img_mask = np.reshape(img_mask, [img_mask.shape[0],img_mask.shape[1],1])
|
||||
fake_merged = img_mask * bgr_fake + (1-img_mask) * target_img.astype(np.float32)
|
||||
fake_merged = fake_merged.astype(np.uint8)
|
||||
return fake_merged
|
||||
except Exception as e :
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
raise e
|
||||
@@ -0,0 +1,326 @@
|
||||
import os
|
||||
import tempfile
|
||||
from pprint import pformat, pprint
|
||||
|
||||
import dill as pickle
|
||||
import gradio as gr
|
||||
import modules.scripts as scripts
|
||||
import numpy as np
|
||||
import onnx
|
||||
import pandas as pd
|
||||
from scripts.faceswaplab_ui.faceswaplab_unit_ui import faceswap_unit_ui
|
||||
from scripts.faceswaplab_ui.faceswaplab_upscaler_ui import upscaler_ui
|
||||
from insightface.app.common import Face
|
||||
from modules import script_callbacks, scripts
|
||||
from PIL import Image
|
||||
from modules.shared import opts
|
||||
|
||||
from scripts.faceswaplab_utils import imgutils
|
||||
from scripts.faceswaplab_utils.imgutils import pil_to_cv2
|
||||
from scripts.faceswaplab_utils.models_utils import get_models
|
||||
from scripts.faceswaplab_utils.faceswaplab_logging import logger
|
||||
import scripts.faceswaplab_swapping.swapper as swapper
|
||||
from scripts.faceswaplab_postprocessing.postprocessing_options import PostProcessingOptions
|
||||
from scripts.faceswaplab_postprocessing.postprocessing import enhance_image
|
||||
from dataclasses import fields
|
||||
from typing import List
|
||||
from scripts.faceswaplab_ui.faceswaplab_unit_settings import FaceSwapUnitSettings
|
||||
from scripts.faceswaplab_utils.models_utils import get_current_model
|
||||
|
||||
def compare(img1, img2):
|
||||
if img1 is not None and img2 is not None:
|
||||
return swapper.compare_faces(img1, img2)
|
||||
|
||||
return "You need 2 images to compare"
|
||||
|
||||
|
||||
|
||||
def extract_faces(files, extract_path, face_restorer_name, face_restorer_visibility, codeformer_weight,upscaler_name,upscaler_scale, upscaler_visibility,inpainting_denoising_strengh, inpainting_prompt, inpainting_negative_prompt, inpainting_steps, inpainting_sampler,inpainting_when):
|
||||
if not extract_path :
|
||||
tempfile.mkdtemp()
|
||||
if files is not None:
|
||||
images = []
|
||||
for file in files :
|
||||
img = Image.open(file.name).convert("RGB")
|
||||
faces = swapper.get_faces(pil_to_cv2(img))
|
||||
if faces:
|
||||
face_images = []
|
||||
for face in faces:
|
||||
bbox = face.bbox.astype(int)
|
||||
x_min, y_min, x_max, y_max = bbox
|
||||
face_image = img.crop((x_min, y_min, x_max, y_max))
|
||||
if face_restorer_name or face_restorer_visibility:
|
||||
scale = 1 if face_image.width > 512 else 512//face_image.width
|
||||
face_image = enhance_image(face_image, PostProcessingOptions(face_restorer_name=face_restorer_name,
|
||||
restorer_visibility=face_restorer_visibility,
|
||||
codeformer_weight= codeformer_weight,
|
||||
upscaler_name=upscaler_name,
|
||||
upscale_visibility=upscaler_visibility,
|
||||
scale=scale,
|
||||
inpainting_denoising_strengh=inpainting_denoising_strengh,
|
||||
inpainting_prompt=inpainting_prompt,
|
||||
inpainting_steps=inpainting_steps,
|
||||
inpainting_negative_prompt=inpainting_negative_prompt,
|
||||
inpainting_when=inpainting_when,
|
||||
inpainting_sampler=inpainting_sampler))
|
||||
path = tempfile.NamedTemporaryFile(delete=False,suffix=".png",dir=extract_path).name
|
||||
face_image.save(path)
|
||||
face_images.append(path)
|
||||
images+= face_images
|
||||
return images
|
||||
return None
|
||||
|
||||
def analyse_faces(image, det_threshold = 0.5) :
|
||||
try :
|
||||
faces = swapper.get_faces(imgutils.pil_to_cv2(image), det_thresh=det_threshold)
|
||||
result = ""
|
||||
for i,face in enumerate(faces) :
|
||||
result+= f"\nFace {i} \n" + "="*40 +"\n"
|
||||
result+= pformat(face) + "\n"
|
||||
result+= "="*40
|
||||
return result
|
||||
|
||||
except Exception as e :
|
||||
logger.error("Analysis Failed : %s", e)
|
||||
return "Analysis Failed"
|
||||
|
||||
def build_face_checkpoint_and_save(batch_files, name):
|
||||
"""
|
||||
Builds a face checkpoint, swaps faces, and saves the result to a file.
|
||||
|
||||
Args:
|
||||
batch_files (list): List of image file paths.
|
||||
name (str): Name of the face checkpoint
|
||||
|
||||
Returns:
|
||||
PIL.Image.Image or None: Resulting swapped face image if successful, otherwise None.
|
||||
"""
|
||||
batch_files = batch_files or []
|
||||
print("Build", name, [x.name for x in batch_files])
|
||||
faces = swapper.get_faces_from_img_files(batch_files)
|
||||
blended_face = swapper.blend_faces(faces)
|
||||
preview_path = os.path.join(
|
||||
scripts.basedir(), "extensions", "sd-webui-faceswaplab", "references"
|
||||
)
|
||||
faces_path = os.path.join(scripts.basedir(), "models", "faceswaplab","faces")
|
||||
if not os.path.exists(faces_path):
|
||||
os.makedirs(faces_path)
|
||||
|
||||
target_img = None
|
||||
if blended_face:
|
||||
if blended_face["gender"] == 0:
|
||||
target_img = Image.open(os.path.join(preview_path, "woman.png"))
|
||||
else:
|
||||
target_img = Image.open(os.path.join(preview_path, "man.png"))
|
||||
|
||||
if name == "":
|
||||
name = "default_name"
|
||||
pprint(blended_face)
|
||||
result = swapper.swap_face(blended_face, blended_face, target_img, get_models()[0])
|
||||
result_image = enhance_image(result.image, PostProcessingOptions(face_restorer_name="CodeFormer", restorer_visibility=1))
|
||||
|
||||
file_path = os.path.join(faces_path, f"{name}.pkl")
|
||||
file_number = 1
|
||||
while os.path.exists(file_path):
|
||||
file_path = os.path.join(faces_path, f"{name}_{file_number}.pkl")
|
||||
file_number += 1
|
||||
result_image.save(file_path+".png")
|
||||
with open(file_path, "wb") as file:
|
||||
pickle.dump({"embedding" :blended_face.embedding, "gender" :blended_face.gender, "age" :blended_face.age},file)
|
||||
try :
|
||||
with open(file_path, "rb") as file:
|
||||
data = Face(pickle.load(file))
|
||||
print(data)
|
||||
except Exception as e :
|
||||
print(e)
|
||||
return result_image
|
||||
|
||||
print("No face found")
|
||||
|
||||
return target_img
|
||||
|
||||
def explore_onnx_faceswap_model(model_path):
|
||||
data = {
|
||||
'Node Name': [],
|
||||
'Op Type': [],
|
||||
'Inputs': [],
|
||||
'Outputs': [],
|
||||
'Attributes': []
|
||||
}
|
||||
if model_path:
|
||||
model = onnx.load(model_path)
|
||||
for node in model.graph.node:
|
||||
data['Node Name'].append(pformat(node.name))
|
||||
data['Op Type'].append(pformat(node.op_type))
|
||||
data['Inputs'].append(pformat(node.input))
|
||||
data['Outputs'].append(pformat(node.output))
|
||||
attributes = []
|
||||
for attr in node.attribute:
|
||||
attr_name = attr.name
|
||||
attr_value = attr.t
|
||||
attributes.append("{} = {}".format(pformat(attr_name), pformat(attr_value)))
|
||||
data['Attributes'].append(attributes)
|
||||
|
||||
df = pd.DataFrame(data)
|
||||
return df
|
||||
|
||||
def batch_process(files, save_path, *components):
|
||||
try :
|
||||
if save_path is not None:
|
||||
os.makedirs(save_path, exist_ok=True)
|
||||
|
||||
units_count = opts.data.get("faceswaplab_units_count", 3)
|
||||
units: List[FaceSwapUnitSettings] = []
|
||||
|
||||
#Parse and convert units flat components into FaceSwapUnitSettings
|
||||
for i in range(0, units_count):
|
||||
units += [FaceSwapUnitSettings.get_unit_configuration(i, components)]
|
||||
|
||||
for i, u in enumerate(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)
|
||||
len_conf: int = len(fields(FaceSwapUnitSettings))
|
||||
shift: int = units_count * len_conf
|
||||
postprocess_options = PostProcessingOptions(
|
||||
*components[shift : shift + len(fields(PostProcessingOptions))]
|
||||
)
|
||||
logger.debug("%s", pformat(postprocess_options))
|
||||
|
||||
units = [u for u in units if u.enable]
|
||||
if files is not None:
|
||||
images = []
|
||||
for file in files :
|
||||
current_images = []
|
||||
src_image = Image.open(file.name).convert("RGB")
|
||||
swapped_images = swapper.process_images_units(get_current_model(), images=[(src_image,None)], units=units, upscaled_swapper=opts.data.get("faceswaplab_upscaled_swapper", False))
|
||||
if len(swapped_images) > 0:
|
||||
current_images+= [img for img,info in swapped_images]
|
||||
|
||||
logger.info("%s images generated", len(current_images))
|
||||
for i, img in enumerate(current_images) :
|
||||
current_images[i] = enhance_image(img,postprocess_options)
|
||||
|
||||
for img in current_images :
|
||||
path = tempfile.NamedTemporaryFile(delete=False,suffix=".png",dir=save_path).name
|
||||
img.save(path)
|
||||
|
||||
images += current_images
|
||||
return images
|
||||
except Exception as e:
|
||||
logger.error("Batch Process error : %s",e)
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
return None
|
||||
|
||||
|
||||
def tools_ui():
|
||||
models = get_models()
|
||||
with gr.Tab("Tools"):
|
||||
with gr.Tab("Build"):
|
||||
gr.Markdown(
|
||||
"""Build a face based on a batch list of images. Will blend the resulting face and store the checkpoint in the faceswaplab/faces directory.""")
|
||||
with gr.Row():
|
||||
batch_files = gr.components.File(
|
||||
type="file",
|
||||
file_count="multiple",
|
||||
label="Batch Sources Images",
|
||||
optional=True,
|
||||
elem_id="faceswaplab_build_batch_files"
|
||||
)
|
||||
preview = gr.components.Image(type="pil", label="Preview", interactive=False, elem_id="faceswaplab_build_preview_face")
|
||||
name = gr.Textbox(
|
||||
value="Face",
|
||||
placeholder="Name of the character",
|
||||
label="Name of the character",
|
||||
elem_id="faceswaplab_build_character_name"
|
||||
)
|
||||
generate_checkpoint_btn = gr.Button("Save",elem_id="faceswaplab_build_save_btn")
|
||||
with gr.Tab("Compare"):
|
||||
gr.Markdown(
|
||||
"""Give a similarity score between two images (only first face is compared).""")
|
||||
|
||||
with gr.Row():
|
||||
img1 = gr.components.Image(type="pil",
|
||||
label="Face 1",
|
||||
elem_id="faceswaplab_compare_face1"
|
||||
)
|
||||
img2 = gr.components.Image(type="pil",
|
||||
label="Face 2",
|
||||
elem_id="faceswaplab_compare_face2"
|
||||
)
|
||||
compare_btn = gr.Button("Compare",elem_id="faceswaplab_compare_btn")
|
||||
compare_result_text = gr.Textbox(
|
||||
interactive=False, label="Similarity", value="0", elem_id="faceswaplab_compare_result"
|
||||
)
|
||||
with gr.Tab("Extract"):
|
||||
gr.Markdown(
|
||||
"""Extract all faces from a batch of images. Will apply enhancement in the tools enhancement tab.""")
|
||||
with gr.Row():
|
||||
extracted_source_files = gr.components.File(
|
||||
type="file",
|
||||
file_count="multiple",
|
||||
label="Batch Sources Images",
|
||||
optional=True,
|
||||
elem_id="faceswaplab_extract_batch_images"
|
||||
)
|
||||
extracted_faces = gr.Gallery(
|
||||
label="Extracted faces", show_label=False,
|
||||
elem_id="faceswaplab_extract_results"
|
||||
).style(columns=[2], rows=[2])
|
||||
extract_save_path = gr.Textbox(label="Destination Directory", value="", elem_id="faceswaplab_extract_destination")
|
||||
extract_btn = gr.Button("Extract", elem_id="faceswaplab_extract_btn")
|
||||
with gr.Tab("Explore Model"):
|
||||
model = gr.Dropdown(
|
||||
choices=models,
|
||||
label="Model not found, please download one and reload automatic 1111",
|
||||
elem_id="faceswaplab_explore_model"
|
||||
)
|
||||
explore_btn = gr.Button("Explore", elem_id="faceswaplab_explore_btn")
|
||||
explore_result_text = gr.Dataframe(
|
||||
interactive=False, label="Explored",
|
||||
elem_id="faceswaplab_explore_result"
|
||||
)
|
||||
with gr.Tab("Analyse Face"):
|
||||
img_to_analyse = gr.components.Image(type="pil", label="Face", elem_id="faceswaplab_analyse_face")
|
||||
analyse_det_threshold = gr.Slider(0.1, 1, 0.5, step=0.01, label="Detection threshold", elem_id="faceswaplab_analyse_det_threshold")
|
||||
analyse_btn = gr.Button("Analyse", elem_id="faceswaplab_analyse_btn")
|
||||
analyse_results = gr.Textbox(label="Results", interactive=False, value="", elem_id="faceswaplab_analyse_results")
|
||||
|
||||
with gr.Tab("Batch Process"):
|
||||
with gr.Tab("Source Images"):
|
||||
gr.Markdown(
|
||||
"""Batch process images. Will apply enhancement in the tools enhancement tab.""")
|
||||
with gr.Row():
|
||||
batch_source_files = gr.components.File(
|
||||
type="file",
|
||||
file_count="multiple",
|
||||
label="Batch Sources Images",
|
||||
optional=True,
|
||||
elem_id="faceswaplab_batch_images"
|
||||
)
|
||||
batch_results = gr.Gallery(
|
||||
label="Batch result", show_label=False,
|
||||
elem_id="faceswaplab_batch_results"
|
||||
).style(columns=[2], rows=[2])
|
||||
batch_save_path = gr.Textbox(label="Destination Directory", value="outputs/faceswap/", elem_id="faceswaplab_batch_destination")
|
||||
batch_save_btn= gr.Button("Process & Save", elem_id="faceswaplab_extract_btn")
|
||||
unit_components = []
|
||||
for i in range(1,opts.data.get("faceswaplab_units_count", 3)+1):
|
||||
unit_components += faceswap_unit_ui(False, i, id_prefix="faceswaplab_tab")
|
||||
|
||||
upscale_options = upscaler_ui()
|
||||
|
||||
explore_btn.click(explore_onnx_faceswap_model, inputs=[model], outputs=[explore_result_text])
|
||||
compare_btn.click(compare, inputs=[img1, img2], outputs=[compare_result_text])
|
||||
generate_checkpoint_btn.click(build_face_checkpoint_and_save, inputs=[batch_files, name], outputs=[preview])
|
||||
extract_btn.click(extract_faces, inputs=[extracted_source_files, extract_save_path]+upscale_options, outputs=[extracted_faces])
|
||||
analyse_btn.click(analyse_faces, inputs=[img_to_analyse,analyse_det_threshold], outputs=[analyse_results])
|
||||
batch_save_btn.click(batch_process, inputs=[batch_source_files, batch_save_path]+unit_components+upscale_options, outputs=[batch_results])
|
||||
|
||||
def on_ui_tabs() :
|
||||
with gr.Blocks(analytics_enabled=False) as ui_faceswap:
|
||||
tools_ui()
|
||||
return [(ui_faceswap, "FaceSwapLab", "faceswaplab_tab")]
|
||||
|
||||
@@ -0,0 +1,152 @@
|
||||
from scripts.faceswaplab_swapping import swapper
|
||||
import numpy as np
|
||||
import base64
|
||||
import io
|
||||
from dataclasses import dataclass, fields
|
||||
from typing import List, Union
|
||||
import dill as pickle
|
||||
import gradio as gr
|
||||
from insightface.app.common import Face
|
||||
from PIL import Image
|
||||
from scripts.faceswaplab_utils.imgutils import (pil_to_cv2,check_against_nsfw)
|
||||
from scripts.faceswaplab_utils.faceswaplab_logging import logger
|
||||
|
||||
@dataclass
|
||||
class FaceSwapUnitSettings:
|
||||
|
||||
# ORDER of parameters is IMPORTANT. It should match the result of faceswap_unit_ui
|
||||
|
||||
# The image given in reference
|
||||
source_img: Union[Image.Image, str]
|
||||
# The checkpoint file
|
||||
source_face : str
|
||||
# The batch source images
|
||||
_batch_files: Union[gr.components.File,List[Image.Image]]
|
||||
# Will blend faces if True
|
||||
blend_faces: bool
|
||||
# Enable this unit
|
||||
enable: bool
|
||||
# Use same gender filtering
|
||||
same_gender: bool
|
||||
# Sort faces by their size (from larger to smaller)
|
||||
sort_by_size : bool
|
||||
# If True, discard images with low similarity
|
||||
check_similarity : bool
|
||||
# if True will compute similarity and add it to the image info
|
||||
_compute_similarity :bool
|
||||
|
||||
# Minimum similarity against the used face (reference, batch or checkpoint)
|
||||
min_sim: float
|
||||
# Minimum similarity against the reference (reference or checkpoint if checkpoint is given)
|
||||
min_ref_sim: float
|
||||
# The face index to use for swapping
|
||||
_faces_index: str
|
||||
# The face index to get image from source
|
||||
reference_face_index : int
|
||||
|
||||
# Swap in the source image in img2img (before processing)
|
||||
swap_in_source: bool
|
||||
# Swap in the generated image in img2img (always on for txt2img)
|
||||
swap_in_generated: bool
|
||||
|
||||
@staticmethod
|
||||
def get_unit_configuration(unit: int, components):
|
||||
fields_count = len(fields(FaceSwapUnitSettings))
|
||||
return FaceSwapUnitSettings(
|
||||
*components[unit * fields_count : unit * fields_count + fields_count]
|
||||
)
|
||||
|
||||
@property
|
||||
def faces_index(self):
|
||||
"""
|
||||
Convert _faces_index from str to int
|
||||
"""
|
||||
faces_index = {
|
||||
int(x) for x in self._faces_index.strip(",").split(",") if x.isnumeric()
|
||||
}
|
||||
if len(faces_index) == 0:
|
||||
return {0}
|
||||
|
||||
logger.debug("FACES INDEX : %s", faces_index)
|
||||
|
||||
return faces_index
|
||||
|
||||
@property
|
||||
def compute_similarity(self) :
|
||||
return self._compute_similarity or self.check_similarity
|
||||
|
||||
@property
|
||||
def batch_files(self):
|
||||
"""
|
||||
Return empty array instead of None for batch files
|
||||
"""
|
||||
return self._batch_files or []
|
||||
|
||||
@property
|
||||
def reference_face(self) :
|
||||
"""
|
||||
Extract reference face (only once and store it for the rest of processing).
|
||||
Reference face is the checkpoint or the source image or the first image in the batch in that order.
|
||||
"""
|
||||
if not hasattr(self,"_reference_face") :
|
||||
if self.source_face and self.source_face != "None" :
|
||||
with open(self.source_face, "rb") as file:
|
||||
try :
|
||||
logger.info(f"loading pickle {file.name}")
|
||||
face = Face(pickle.load(file))
|
||||
self._reference_face = face
|
||||
except Exception as e :
|
||||
logger.error("Failed to load checkpoint : %s", e)
|
||||
elif self.source_img is not None :
|
||||
if isinstance(self.source_img, str): # source_img is a base64 string
|
||||
if 'base64,' in self.source_img: # check if the base64 string has a data URL scheme
|
||||
base64_data = self.source_img.split('base64,')[-1]
|
||||
img_bytes = base64.b64decode(base64_data)
|
||||
else:
|
||||
# if no data URL scheme, just decode
|
||||
img_bytes = base64.b64decode(self.source_img)
|
||||
self.source_img = Image.open(io.BytesIO(img_bytes))
|
||||
source_img = pil_to_cv2(self.source_img)
|
||||
self._reference_face = swapper.get_or_default(swapper.get_faces(source_img), self.reference_face_index, None)
|
||||
if self._reference_face is None :
|
||||
logger.error("Face not found in reference image")
|
||||
else :
|
||||
self._reference_face = None
|
||||
|
||||
if self._reference_face is None :
|
||||
logger.error("You need at least one reference face")
|
||||
|
||||
return self._reference_face
|
||||
|
||||
@property
|
||||
def faces(self) :
|
||||
"""_summary_
|
||||
Extract all faces (including reference face) to provide an array of faces
|
||||
Only processed once.
|
||||
"""
|
||||
if self.batch_files is not None and not hasattr(self,"_faces") :
|
||||
self._faces = [self.reference_face] if self.reference_face is not None else []
|
||||
for file in self.batch_files :
|
||||
if isinstance(file, Image.Image) :
|
||||
img = file
|
||||
else :
|
||||
img = Image.open(file.name)
|
||||
|
||||
face = swapper.get_or_default(swapper.get_faces(pil_to_cv2(img)), 0, None)
|
||||
if face is not None :
|
||||
self._faces.append(face)
|
||||
return self._faces
|
||||
|
||||
@property
|
||||
def blended_faces(self):
|
||||
"""
|
||||
Blend the faces using the mean of all embeddings
|
||||
"""
|
||||
if not hasattr(self,"_blended_faces") :
|
||||
self._blended_faces = swapper.blend_faces(self.faces)
|
||||
assert(all([not np.array_equal(self._blended_faces.embedding, face.embedding) for face in self.faces]) if len(self.faces) > 1 else True), "Blended faces cannot be the same as one of the face if len(face)>0"
|
||||
assert(not np.array_equal(self._blended_faces.embedding,self.reference_face.embedding) if len(self.faces) > 1 else True), "Blended faces cannot be the same as reference face if len(face)>0"
|
||||
|
||||
return self._blended_faces
|
||||
|
||||
|
||||
@@ -0,0 +1,111 @@
|
||||
from scripts.faceswaplab_utils.models_utils import get_face_checkpoints
|
||||
import gradio as gr
|
||||
|
||||
def faceswap_unit_ui(is_img2img, unit_num=1, id_prefix="faceswaplab"):
|
||||
with gr.Tab(f"Face {unit_num}"):
|
||||
with gr.Column():
|
||||
gr.Markdown(
|
||||
"""Reference is an image. First face will be extracted.
|
||||
First face of batches sources will be extracted and used as input (or blended if blend is activated).""")
|
||||
with gr.Row():
|
||||
img = gr.components.Image(type="pil", label="Reference", elem_id=f"{id_prefix}_face{unit_num}_reference_image")
|
||||
batch_files = gr.components.File(
|
||||
type="file",
|
||||
file_count="multiple",
|
||||
label="Batch Sources Images",
|
||||
optional=True,
|
||||
elem_id=f"{id_prefix}_face{unit_num}_batch_source_face_files"
|
||||
)
|
||||
gr.Markdown(
|
||||
"""Face checkpoint built with the checkpoint builder in tools. Will overwrite reference image.""")
|
||||
with gr.Row() :
|
||||
|
||||
face = gr.Dropdown(
|
||||
choices=get_face_checkpoints(),
|
||||
label="Face Checkpoint (precedence over reference face)",
|
||||
elem_id=f"{id_prefix}_face{unit_num}_face_checkpoint"
|
||||
)
|
||||
refresh = gr.Button(value='↻', variant='tool', elem_id=f"{id_prefix}_face{unit_num}_refresh_checkpoints")
|
||||
def refresh_fn(selected):
|
||||
return gr.Dropdown.update(value=selected, choices=get_face_checkpoints())
|
||||
refresh.click(fn=refresh_fn,inputs=face, outputs=face)
|
||||
|
||||
with gr.Row():
|
||||
enable = gr.Checkbox(False, placeholder="enable", label="Enable", elem_id=f"{id_prefix}_face{unit_num}_enable")
|
||||
blend_faces = gr.Checkbox(
|
||||
True, placeholder="Blend Faces", label="Blend Faces ((Source|Checkpoint)+References = 1)",
|
||||
elem_id=f"{id_prefix}_face{unit_num}_blend_faces",
|
||||
interactive=True
|
||||
)
|
||||
gr.Markdown("""Discard images with low similarity or no faces :""")
|
||||
with gr.Row():
|
||||
check_similarity = gr.Checkbox(False, placeholder="discard", label="Check similarity",
|
||||
elem_id=f"{id_prefix}_face{unit_num}_check_similarity")
|
||||
compute_similarity = gr.Checkbox(False, label="Compute similarity",
|
||||
elem_id=f"{id_prefix}_face{unit_num}_compute_similarity")
|
||||
min_sim = gr.Slider(0, 1, 0, step=0.01, label="Min similarity",
|
||||
elem_id=f"{id_prefix}_face{unit_num}_min_similarity")
|
||||
min_ref_sim = gr.Slider(
|
||||
0, 1, 0, step=0.01, label="Min reference similarity",
|
||||
elem_id=f"{id_prefix}_face{unit_num}_min_ref_similarity"
|
||||
)
|
||||
|
||||
gr.Markdown("""Select the face to be swapped, you can sort by size or use the same gender as the desired face:""")
|
||||
with gr.Row():
|
||||
same_gender = gr.Checkbox(
|
||||
False, placeholder="Same Gender", label="Same Gender",
|
||||
elem_id=f"{id_prefix}_face{unit_num}_same_gender"
|
||||
)
|
||||
sort_by_size = gr.Checkbox(
|
||||
False, placeholder="Sort by size", label="Sort by size (larger>smaller)",
|
||||
elem_id=f"{id_prefix}_face{unit_num}_sort_by_size"
|
||||
)
|
||||
target_faces_index = gr.Textbox(
|
||||
value="0",
|
||||
placeholder="Which face to swap (comma separated), start from 0 (by gender if same_gender is enabled)",
|
||||
label="Target face : Comma separated face number(s)",
|
||||
elem_id=f"{id_prefix}_face{unit_num}_target_faces_index"
|
||||
)
|
||||
gr.Markdown("""The following will only affect reference face image (and is not affected by sort by size) :""")
|
||||
reference_faces_index = gr.Number(
|
||||
value=0,
|
||||
precision=0,
|
||||
minimum=0,
|
||||
placeholder="Which face to get from reference image start from 0",
|
||||
label="Reference source face : start from 0",
|
||||
elem_id=f"{id_prefix}_face{unit_num}_reference_face_index"
|
||||
)
|
||||
gr.Markdown("""Configure swapping. Swapping can occure before img2img, after or both :""", visible=is_img2img)
|
||||
swap_in_source = gr.Checkbox(
|
||||
False,
|
||||
placeholder="Swap face in source image",
|
||||
label="Swap in source image (blended face)",
|
||||
visible=is_img2img,
|
||||
elem_id=f"{id_prefix}_face{unit_num}_swap_in_source"
|
||||
)
|
||||
swap_in_generated = gr.Checkbox(
|
||||
True,
|
||||
placeholder="Swap face in generated image",
|
||||
label="Swap in generated image",
|
||||
visible=is_img2img,
|
||||
elem_id=f"{id_prefix}_face{unit_num}_swap_in_generated"
|
||||
)
|
||||
# If changed, you need to change FaceSwapUnitSettings accordingly
|
||||
# ORDER of parameters is IMPORTANT. It should match the result of FaceSwapUnitSettings
|
||||
return [
|
||||
img,
|
||||
face,
|
||||
batch_files,
|
||||
blend_faces,
|
||||
enable,
|
||||
same_gender,
|
||||
sort_by_size,
|
||||
check_similarity,
|
||||
compute_similarity,
|
||||
min_sim,
|
||||
min_ref_sim,
|
||||
target_faces_index,
|
||||
reference_faces_index,
|
||||
swap_in_source,
|
||||
swap_in_generated,
|
||||
]
|
||||
@@ -0,0 +1,80 @@
|
||||
import gradio as gr
|
||||
import modules
|
||||
from modules import shared, sd_models
|
||||
from modules.shared import cmd_opts, opts, state
|
||||
|
||||
import scripts.faceswaplab_postprocessing.upscaling as upscaling
|
||||
from scripts.faceswaplab_utils.faceswaplab_logging import logger
|
||||
|
||||
def upscaler_ui():
|
||||
with gr.Tab(f"Post-Processing"):
|
||||
gr.Markdown(
|
||||
"""Upscaling is performed on the whole image. Upscaling happens before face restoration.""")
|
||||
with gr.Row():
|
||||
face_restorer_name = gr.Radio(
|
||||
label="Restore Face",
|
||||
choices=["None"] + [x.name() for x in shared.face_restorers],
|
||||
value=lambda : opts.data.get("faceswaplab_pp_default_face_restorer", shared.face_restorers[0].name()),
|
||||
type="value",
|
||||
elem_id="faceswaplab_pp_face_restorer"
|
||||
)
|
||||
with gr.Column():
|
||||
face_restorer_visibility = gr.Slider(
|
||||
0, 1, value=lambda:opts.data.get("faceswaplab_pp_default_face_restorer_visibility", 1), step=0.001, label="Restore visibility",
|
||||
elem_id="faceswaplab_pp_face_restorer_visibility"
|
||||
)
|
||||
codeformer_weight = gr.Slider(
|
||||
0, 1, value=lambda:opts.data.get("faceswaplab_pp_default_face_restorer_weight", 1), step=0.001, label="codeformer weight",
|
||||
elem_id="faceswaplab_pp_face_restorer_weight"
|
||||
)
|
||||
upscaler_name = gr.Dropdown(
|
||||
choices=[upscaler.name for upscaler in shared.sd_upscalers],
|
||||
value= lambda:opts.data.get("faceswaplab_pp_default_upscaler","None"),
|
||||
label="Upscaler",
|
||||
elem_id="faceswaplab_pp_upscaler"
|
||||
)
|
||||
upscaler_scale = gr.Slider(1, 8, 1, step=0.1, label="Upscaler scale", elem_id="faceswaplab_pp_upscaler_scale")
|
||||
upscaler_visibility = gr.Slider(
|
||||
0, 1, value=lambda:opts.data.get("faceswaplab_pp_default_upscaler_visibility", 1), step=0.1, label="Upscaler visibility (if scale = 1)",
|
||||
elem_id="faceswaplab_pp_upscaler_visibility"
|
||||
)
|
||||
with gr.Accordion(f"Post Inpainting", open=True):
|
||||
gr.Markdown(
|
||||
"""Inpainting sends image to inpainting with a mask on face (once for each faces).""")
|
||||
inpainting_when = gr.Dropdown(
|
||||
elem_id="faceswaplab_pp_inpainting_when", choices = [e.value for e in upscaling.InpaintingWhen.__members__.values()],value=[upscaling.InpaintingWhen.BEFORE_RESTORE_FACE.value], label="Enable/When")
|
||||
inpainting_denoising_strength = gr.Slider(
|
||||
0, 1, 0, step=0.01, elem_id="faceswaplab_pp_inpainting_denoising_strength", label="Denoising strenght (will send face to img2img after processing)"
|
||||
)
|
||||
|
||||
inpainting_denoising_prompt = gr.Textbox("Portrait of a [gender]",elem_id="faceswaplab_pp_inpainting_denoising_prompt", label="Inpainting prompt use [gender] instead of men or woman")
|
||||
inpainting_denoising_negative_prompt = gr.Textbox("", elem_id="faceswaplab_pp_inpainting_denoising_neg_prompt", label="Inpainting negative prompt use [gender] instead of men or woman")
|
||||
with gr.Row():
|
||||
samplers_names = [s.name for s in modules.sd_samplers.all_samplers]
|
||||
inpainting_sampler = gr.Dropdown(
|
||||
choices=samplers_names,
|
||||
value=[samplers_names[0]],
|
||||
label="Inpainting Sampler",
|
||||
elem_id="faceswaplab_pp_inpainting_sampler"
|
||||
)
|
||||
inpainting_denoising_steps = gr.Slider(
|
||||
1, 150, 20, step=1, label="Inpainting steps",
|
||||
elem_id="faceswaplab_pp_inpainting_steps"
|
||||
)
|
||||
|
||||
inpaiting_model = gr.Dropdown(choices=["Current"]+sd_models.checkpoint_tiles(), default="Current", label="sd model (experimental)", elem_id="faceswaplab_pp_inpainting_sd_model")
|
||||
return [
|
||||
face_restorer_name,
|
||||
face_restorer_visibility,
|
||||
codeformer_weight,
|
||||
upscaler_name,
|
||||
upscaler_scale,
|
||||
upscaler_visibility,
|
||||
inpainting_denoising_strength,
|
||||
inpainting_denoising_prompt,
|
||||
inpainting_denoising_negative_prompt,
|
||||
inpainting_denoising_steps,
|
||||
inpainting_sampler,
|
||||
inpainting_when,
|
||||
inpaiting_model
|
||||
]
|
||||
@@ -0,0 +1,53 @@
|
||||
import logging
|
||||
import copy
|
||||
import sys
|
||||
from modules import shared
|
||||
from PIL import Image
|
||||
|
||||
class ColoredFormatter(logging.Formatter):
|
||||
COLORS = {
|
||||
"DEBUG": "\033[0;36m", # CYAN
|
||||
"INFO": "\033[0;32m", # GREEN
|
||||
"WARNING": "\033[0;33m", # YELLOW
|
||||
"ERROR": "\033[0;31m", # RED
|
||||
"CRITICAL": "\033[0;37;41m", # WHITE ON RED
|
||||
"RESET": "\033[0m", # RESET COLOR
|
||||
}
|
||||
|
||||
def format(self, record):
|
||||
colored_record = copy.copy(record)
|
||||
levelname = colored_record.levelname
|
||||
seq = self.COLORS.get(levelname, self.COLORS["RESET"])
|
||||
colored_record.levelname = f"{seq}{levelname}{self.COLORS['RESET']}"
|
||||
return super().format(colored_record)
|
||||
|
||||
|
||||
# Create a new logger
|
||||
logger = logging.getLogger("FaceSwapLab")
|
||||
logger.propagate = False
|
||||
|
||||
# Add handler if we don't have one.
|
||||
if not logger.handlers:
|
||||
handler = logging.StreamHandler(sys.stdout)
|
||||
handler.setFormatter(
|
||||
ColoredFormatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
|
||||
)
|
||||
logger.addHandler(handler)
|
||||
|
||||
# Configure logger
|
||||
loglevel_string = getattr(shared.cmd_opts, "faceswaplab_loglevel", "INFO")
|
||||
loglevel = getattr(logging, loglevel_string.upper(), "INFO")
|
||||
logger.setLevel(loglevel)
|
||||
|
||||
import tempfile
|
||||
if logger.getEffectiveLevel() <= logging.DEBUG :
|
||||
DEBUG_DIR = tempfile.mkdtemp()
|
||||
|
||||
def save_img_debug(img: Image.Image, message: str, *opts):
|
||||
if logger.getEffectiveLevel() <= logging.DEBUG:
|
||||
with tempfile.NamedTemporaryFile(dir=DEBUG_DIR, delete=False, suffix=".png") as temp_file:
|
||||
img_path = temp_file.name
|
||||
img.save(img_path)
|
||||
|
||||
message_with_link = f"{message}\nImage: file://{img_path}"
|
||||
logger.debug(message_with_link, *opts)
|
||||
@@ -0,0 +1,185 @@
|
||||
import io
|
||||
from typing import Optional
|
||||
from PIL import Image, ImageChops, ImageOps,ImageFilter
|
||||
import cv2
|
||||
import numpy as np
|
||||
from math import isqrt, ceil
|
||||
import torch
|
||||
from ifnude import detect
|
||||
from scripts.faceswaplab_globals import NSFW_SCORE
|
||||
from modules import processing
|
||||
import base64
|
||||
|
||||
def check_against_nsfw(img):
|
||||
shapes = []
|
||||
chunks = detect(img)
|
||||
for chunk in chunks:
|
||||
shapes.append(chunk["score"] > NSFW_SCORE)
|
||||
return any(shapes)
|
||||
|
||||
def pil_to_cv2(pil_img):
|
||||
return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
|
||||
|
||||
|
||||
def cv2_to_pil(cv2_img):
|
||||
return Image.fromarray(cv2.cvtColor(cv2_img, cv2.COLOR_BGR2RGB))
|
||||
|
||||
def torch_to_pil(images):
|
||||
"""
|
||||
Convert a numpy image or a batch of images to a PIL image.
|
||||
"""
|
||||
images = images.cpu().permute(0, 2, 3, 1).numpy()
|
||||
if images.ndim == 3:
|
||||
images = images[None, ...]
|
||||
images = (images * 255).round().astype("uint8")
|
||||
pil_images = [Image.fromarray(image) for image in images]
|
||||
return pil_images
|
||||
|
||||
|
||||
def pil_to_torch(pil_images):
|
||||
"""
|
||||
Convert a PIL image or a list of PIL images to a torch tensor or a batch of torch tensors.
|
||||
"""
|
||||
if isinstance(pil_images, list):
|
||||
numpy_images = [np.array(image) for image in pil_images]
|
||||
torch_images = torch.from_numpy(np.stack(numpy_images)).permute(0, 3, 1, 2)
|
||||
return torch_images
|
||||
|
||||
numpy_image = np.array(pil_images)
|
||||
torch_image = torch.from_numpy(numpy_image).permute(2, 0, 1)
|
||||
return torch_image
|
||||
|
||||
from collections import Counter
|
||||
def create_square_image(image_list):
|
||||
"""
|
||||
Creates a square image by combining multiple images in a grid pattern.
|
||||
|
||||
Args:
|
||||
image_list (list): List of PIL Image objects to be combined.
|
||||
|
||||
Returns:
|
||||
PIL Image object: The resulting square image.
|
||||
None: If the image_list is empty or contains only one image.
|
||||
"""
|
||||
|
||||
# Count the occurrences of each image size in the image_list
|
||||
size_counter = Counter(image.size for image in image_list)
|
||||
|
||||
# Get the most common image size (size with the highest count)
|
||||
common_size = size_counter.most_common(1)[0][0]
|
||||
|
||||
# Filter the image_list to include only images with the common size
|
||||
image_list = [image for image in image_list if image.size == common_size]
|
||||
|
||||
# Get the dimensions (width and height) of the common size
|
||||
size = common_size
|
||||
|
||||
# If there are more than one image in the image_list
|
||||
if len(image_list) > 1:
|
||||
num_images = len(image_list)
|
||||
|
||||
# Calculate the number of rows and columns for the grid
|
||||
rows = isqrt(num_images)
|
||||
cols = ceil(num_images / rows)
|
||||
|
||||
# Calculate the size of the square image
|
||||
square_size = (cols * size[0], rows * size[1])
|
||||
|
||||
# Create a new RGB image with the square size
|
||||
square_image = Image.new("RGB", square_size)
|
||||
|
||||
# Paste each image onto the square image at the appropriate position
|
||||
for i, image in enumerate(image_list):
|
||||
row = i // cols
|
||||
col = i % cols
|
||||
|
||||
square_image.paste(image, (col * size[0], row * size[1]))
|
||||
|
||||
# Return the resulting square image
|
||||
return square_image
|
||||
|
||||
# Return None if there are no images or only one image in the image_list
|
||||
return None
|
||||
|
||||
def create_mask(image, box_coords):
|
||||
width, height = image.size
|
||||
mask = Image.new("L", (width, height), 255)
|
||||
x1, y1, x2, y2 = box_coords
|
||||
for x in range(width):
|
||||
for y in range(height):
|
||||
if x1 <= x <= x2 and y1 <= y <= y2:
|
||||
mask.putpixel((x, y), 255)
|
||||
else:
|
||||
mask.putpixel((x, y), 0)
|
||||
return mask
|
||||
|
||||
def apply_mask(img : Image.Image,p : processing.StableDiffusionProcessing, batch_index : int) -> Image.Image :
|
||||
"""
|
||||
Apply mask overlay and color correction to an image if enabled
|
||||
|
||||
Args:
|
||||
img: PIL Image objects.
|
||||
p : The processing object
|
||||
batch_index : the batch index
|
||||
|
||||
Returns:
|
||||
PIL Image object
|
||||
"""
|
||||
if isinstance(p, processing.StableDiffusionProcessingImg2Img) :
|
||||
if p.inpaint_full_res :
|
||||
overlays = p.overlay_images
|
||||
if overlays is None or batch_index >= len(overlays):
|
||||
return img
|
||||
overlay : Image.Image = overlays[batch_index]
|
||||
overlay = overlay.resize((img.size), resample= Image.Resampling.LANCZOS)
|
||||
img = img.copy()
|
||||
img.paste(overlay, (0, 0), overlay)
|
||||
return img
|
||||
|
||||
img = processing.apply_overlay(img, p.paste_to, batch_index, p.overlay_images)
|
||||
if p.color_corrections is not None and batch_index < len(p.color_corrections):
|
||||
img = processing.apply_color_correction(p.color_corrections[batch_index], img)
|
||||
return img
|
||||
|
||||
|
||||
|
||||
def prepare_mask(
|
||||
mask: Image.Image, p: processing.StableDiffusionProcessing
|
||||
) -> Image.Image:
|
||||
"""
|
||||
Prepare an image mask for the inpainting process. (This comes from controlnet)
|
||||
|
||||
This function takes as input a PIL Image object and an instance of the
|
||||
StableDiffusionProcessing class, and performs the following steps to prepare the mask:
|
||||
|
||||
1. Convert the mask to grayscale (mode "L").
|
||||
2. If the 'inpainting_mask_invert' attribute of the processing instance is True,
|
||||
invert the mask colors.
|
||||
3. If the 'mask_blur' attribute of the processing instance is greater than 0,
|
||||
apply a Gaussian blur to the mask with a radius equal to 'mask_blur'.
|
||||
|
||||
Args:
|
||||
mask (Image.Image): The input mask as a PIL Image object.
|
||||
p (processing.StableDiffusionProcessing): An instance of the StableDiffusionProcessing class
|
||||
containing the processing parameters.
|
||||
|
||||
Returns:
|
||||
mask (Image.Image): The prepared mask as a PIL Image object.
|
||||
"""
|
||||
mask = mask.convert("L")
|
||||
#FIXME : Properly fix blur
|
||||
# if getattr(p, "mask_blur", 0) > 0:
|
||||
# mask = mask.filter(ImageFilter.GaussianBlur(p.mask_blur))
|
||||
return mask
|
||||
|
||||
def base64_to_pil(base64str : Optional[str]) -> Optional[Image.Image] :
|
||||
if base64str is None :
|
||||
return None
|
||||
if 'base64,' in base64str: # check if the base64 string has a data URL scheme
|
||||
base64_data = base64str.split('base64,')[-1]
|
||||
img_bytes = base64.b64decode(base64_data)
|
||||
else:
|
||||
# if no data URL scheme, just decode
|
||||
img_bytes = base64.b64decode(base64str)
|
||||
return Image.open(io.BytesIO(img_bytes))
|
||||
|
||||
@@ -0,0 +1,55 @@
|
||||
|
||||
import glob
|
||||
import os
|
||||
import modules.scripts as scripts
|
||||
from modules import scripts
|
||||
from scripts.faceswaplab_globals import EXTENSION_PATH
|
||||
from modules.shared import opts
|
||||
from scripts.faceswaplab_utils.faceswaplab_logging import logger
|
||||
|
||||
def get_models():
|
||||
"""
|
||||
Retrieve a list of swap model files.
|
||||
|
||||
This function searches for model files in the specified directories and returns a list of file paths.
|
||||
The supported file extensions are ".onnx".
|
||||
|
||||
Returns:
|
||||
A list of file paths of the model files.
|
||||
"""
|
||||
models_path = os.path.join(scripts.basedir(), EXTENSION_PATH, "models", "*")
|
||||
models = glob.glob(models_path)
|
||||
|
||||
# Add an additional models directory and find files in it
|
||||
models_path = os.path.join(scripts.basedir(), "models", "faceswaplab", "*")
|
||||
models += glob.glob(models_path)
|
||||
|
||||
# Filter the list to include only files with the supported extensions
|
||||
models = [x for x in models if x.endswith(".onnx")]
|
||||
|
||||
return models
|
||||
|
||||
def get_current_model() -> str :
|
||||
model = opts.data.get("faceswaplab_model", None)
|
||||
if model is None :
|
||||
models = get_models()
|
||||
model = models[0] if len(models) else None
|
||||
logger.info("Try to use model : %s", model)
|
||||
if not os.path.isfile(model):
|
||||
logger.error("The model %s cannot be found or loaded", model)
|
||||
raise FileNotFoundError("No faceswap model found. Please add it to the faceswaplab directory.")
|
||||
return model
|
||||
|
||||
def get_face_checkpoints():
|
||||
"""
|
||||
Retrieve a list of face checkpoint paths.
|
||||
|
||||
This function searches for face files with the extension ".pkl" in the specified directory and returns a list
|
||||
containing the paths of those files.
|
||||
|
||||
Returns:
|
||||
list: A list of face paths, including the string "None" as the first element.
|
||||
"""
|
||||
faces_path = os.path.join(scripts.basedir(), "models", "faceswaplab", "faces", "*.pkl")
|
||||
faces = glob.glob(faces_path)
|
||||
return ["None"] + faces
|
||||
Reference in New Issue
Block a user