add a ldsr inswapper option in sd roop settings
Tran Xen
parent
44ccfc0521
commit
b7872cb9e4
|
|
@ -102,6 +102,10 @@ class FaceSwapUnitSettings:
|
|||
enable: bool
|
||||
# Use same gender filtering
|
||||
same_gender: bool
|
||||
|
||||
# If True, discard images with low similarity
|
||||
check_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)
|
||||
|
|
@ -345,7 +349,7 @@ def upscaler_ui():
|
|||
with gr.Accordion(f"Post Inpainting (Beta)", open=True):
|
||||
gr.Markdown(
|
||||
"""Inpainting sends image to inpainting with a mask on face (once for each faces).""")
|
||||
inpainting_when = gr.Dropdown(choices = [e.value for e in upscaling.InpaintingWhen.__members__.values()],value=[upscaling.InpaintingWhen.BEFORE_RESTORE_FACE.value], label="When")
|
||||
inpainting_when = gr.Dropdown(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, label="Denoising strenght (will send face to img2img after processing)"
|
||||
)
|
||||
|
|
@ -451,7 +455,8 @@ def on_ui_settings():
|
|||
section = ('roop', "Roop")
|
||||
shared.opts.add_option("roop_units_count", shared.OptionInfo(
|
||||
3, "Max faces units (requires restart)", gr.Slider, {"minimum": 1, "maximum": 10, "step": 1}, section=section))
|
||||
|
||||
shared.opts.add_option("roop_upscaled_swapper", shared.OptionInfo(
|
||||
False, "Upscaled swapper", gr.Checkbox, {"interactive": True}, section=section))
|
||||
script_callbacks.on_ui_settings(on_ui_settings)
|
||||
|
||||
|
||||
|
|
@ -461,6 +466,10 @@ class FaceSwapScript(scripts.Script):
|
|||
def units_count(self) :
|
||||
return opts.data.get("roop_units_count", 3)
|
||||
|
||||
@property
|
||||
def upscaled_swapper(self) :
|
||||
return opts.data.get("roop_upscaled_swapper", False)
|
||||
|
||||
def title(self):
|
||||
return f"roop"
|
||||
|
||||
|
|
@ -502,7 +511,8 @@ class FaceSwapScript(scripts.Script):
|
|||
blend_faces = gr.Checkbox(
|
||||
True, placeholder="Blend Faces", label="Blend Faces ((Source|Checkpoint)+References = 1)"
|
||||
)
|
||||
gr.Markdown("""Discard images with low similarity or no faces :""")
|
||||
gr.Markdown("""Discard images with low similarity or no faces :""")
|
||||
check_similarity = gr.Checkbox(False, placeholder="discard", label="Check similarity")
|
||||
min_sim = gr.Slider(0, 1, 0, step=0.01, label="Min similarity")
|
||||
min_ref_sim = gr.Slider(
|
||||
0, 1, 0, step=0.01, label="Min reference similarity"
|
||||
|
|
@ -532,6 +542,7 @@ class FaceSwapScript(scripts.Script):
|
|||
blend_faces,
|
||||
enable,
|
||||
same_gender,
|
||||
check_similarity,
|
||||
min_sim,
|
||||
min_ref_sim,
|
||||
faces_index,
|
||||
|
|
@ -599,7 +610,7 @@ class FaceSwapScript(scripts.Script):
|
|||
p.init_images = init_images
|
||||
|
||||
|
||||
def process_images_unit(self, unit, images, infos = None) :
|
||||
def process_images_unit(self, unit : FaceSwapUnitSettings, images, infos = None) :
|
||||
if unit.enable :
|
||||
result_images = []
|
||||
result_infos = []
|
||||
|
|
@ -625,8 +636,9 @@ class FaceSwapScript(scripts.Script):
|
|||
faces_index=unit.faces_index,
|
||||
model=self.model,
|
||||
same_gender=unit.same_gender,
|
||||
upscaled_swapper=self.upscaled_swapper
|
||||
)
|
||||
if 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()]):
|
||||
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()]):
|
||||
result_infos.append(f"{info}, similarity = {result.similarity}, ref_similarity = {result.ref_similarity}")
|
||||
result_images.append(result.image)
|
||||
else:
|
||||
|
|
@ -644,7 +656,7 @@ class FaceSwapScript(scripts.Script):
|
|||
if any([u.enable for u in self.units]):
|
||||
result_images = processed.images[:]
|
||||
result_infos = processed.infotexts[:]
|
||||
if p.batch_size > 1 and p.n_iter > 1:
|
||||
if p.batch_size > 1 or p.n_iter > 1:
|
||||
# Remove grid image if batch size is greater than 1 :
|
||||
result_images = result_images[1:]
|
||||
result_infos = result_infos[1:]
|
||||
|
|
|
|||
|
|
@ -9,9 +9,11 @@ import insightface
|
|||
import numpy as np
|
||||
import onnxruntime
|
||||
from insightface.app.common import Face
|
||||
|
||||
from PIL import Image
|
||||
from sklearn.metrics.pairwise import cosine_similarity
|
||||
|
||||
from scripts import upscaled_inswapper
|
||||
from scripts.imgutils import cv2_to_pil, pil_to_cv2
|
||||
from scripts.roop_logging import logger
|
||||
|
||||
|
|
@ -124,7 +126,7 @@ def getFaceSwapModel(model_path: str):
|
|||
try :
|
||||
CURRENT_FS_MODEL_PATH = model_path
|
||||
# Initializes the face swap model using the specified model path.
|
||||
FS_MODEL = insightface.model_zoo.get_model(model_path, providers=providers)
|
||||
FS_MODEL = 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.)")
|
||||
return FS_MODEL
|
||||
|
|
@ -274,6 +276,7 @@ def swap_face(
|
|||
model: str,
|
||||
faces_index: Set[int] = {0},
|
||||
same_gender=True,
|
||||
upscaled_swapper = False
|
||||
) -> ImageResult:
|
||||
"""
|
||||
Swaps faces in the target image with the source face.
|
||||
|
|
@ -309,7 +312,7 @@ def swap_face(
|
|||
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)
|
||||
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
|
||||
|
|
|
|||
|
|
@ -0,0 +1,119 @@
|
|||
import time
|
||||
import numpy as np
|
||||
import onnxruntime
|
||||
import cv2
|
||||
import onnx
|
||||
from onnx import numpy_helper
|
||||
from insightface.utils import face_align
|
||||
from insightface.model_zoo.inswapper import INSwapper
|
||||
from modules import scripts, shared, processing
|
||||
from modules.face_restoration import FaceRestoration
|
||||
from modules.upscaler import UpscalerData
|
||||
from PIL import Image
|
||||
from scripts.roop_logging import logger
|
||||
from scripts.imgutils import cv2_to_pil, pil_to_cv2
|
||||
|
||||
|
||||
def get_ldsr() -> UpscalerData:
|
||||
for upscaler in shared.sd_upscalers:
|
||||
if upscaler.name == "LDSR":
|
||||
return upscaler
|
||||
return None
|
||||
|
||||
|
||||
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)
|
||||
upscaled = get_ldsr().scaler.upscale(
|
||||
pil_img, k, get_ldsr().data_path
|
||||
)
|
||||
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:
|
||||
if upscale :
|
||||
print("*"*80)
|
||||
print("LDSR inswapper")
|
||||
print("*"*80)
|
||||
k = 4
|
||||
aimg, M = face_align.norm_crop2(img, target_face.kps, self.input_size[0]*k)
|
||||
bgr_fake = self.super_resolution(bgr_fake, k)
|
||||
|
||||
# Add sharpness
|
||||
blurred = cv2.GaussianBlur(bgr_fake, (0, 0), 3)
|
||||
bgr_fake = cv2.addWeighted(bgr_fake, 1.5, blurred, -0.5, 0)
|
||||
|
||||
|
||||
target_img = img
|
||||
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
|
||||
|
||||
|
||||
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 = 10
|
||||
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))
|
||||
k = max(mask_size//10, 10)
|
||||
|
||||
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(mask_size//20, 5)
|
||||
|
||||
|
||||
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 = 5
|
||||
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
|
||||
|
|
@ -115,27 +115,7 @@ def resize_bbox(bbox):
|
|||
y_max = int(y_max // 8 + 1) * 8 if y_max % 8 != 0 else y_max
|
||||
return x_min, y_min, x_max, y_max
|
||||
|
||||
def get_ldsr() -> UpscalerData:
|
||||
for upscaler in shared.sd_upscalers:
|
||||
if upscaler.name == "LDSR":
|
||||
return upscaler
|
||||
return None
|
||||
|
||||
def resize_small_image(img: Image.Image, min_resolution=512, use_ldsr = True):
|
||||
width, height = img.size
|
||||
if min(width, height) > min_resolution:
|
||||
return img
|
||||
k = float(min_resolution) / float(min(width, height))
|
||||
target_width = int(round(width * k))
|
||||
target_height = int(round(height * k))
|
||||
if not use_ldsr :
|
||||
resized_img = img.resize((target_width, target_height), resample=Image.Resampling.LANCZOS)
|
||||
else :
|
||||
logger.info("Upscale face with LDSR")
|
||||
resized_img = get_ldsr().scaler.upscale(
|
||||
img, k, get_ldsr().data_path
|
||||
)
|
||||
return resized_img
|
||||
|
||||
|
||||
def create_mask(image, box_coords):
|
||||
width, height = image.size
|
||||
|
|
@ -161,7 +141,9 @@ inpainting_denoising_strength : {inpainting_denoising_strength}
|
|||
inpainting_steps : {inpainting_steps}
|
||||
"""
|
||||
)
|
||||
|
||||
if not isinstance(inpainting_sampler, str) :
|
||||
inpainting_sampler = "Euler"
|
||||
|
||||
logger.info("send faces to image to image")
|
||||
img = img.copy()
|
||||
faces = swapper.get_faces(imgutils.pil_to_cv2(img))
|
||||
|
|
|
|||
Loading…
Reference in New Issue