Recrosoftware/sd-webui-reactor
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Face Restoration fix, Upscaler back

Browse Source 
Upscaler is still not working properly
This commit is contained in:
Art Gourieff 2023-06-19 16:42:14 +07:00 committed by GitHub
parent 77e5ccc0b9
commit 3e52982eca
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 85 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

41
scripts/faceswap.py
View File

@ -1,5 +1,6 @@
import gradio as gr
import modules.scripts as scripts
from modules.upscaler import Upscaler, UpscalerData
from modules import scripts, shared, images, scripts_postprocessing
from modules.processing import (
StableDiffusionProcessing,
@ -11,7 +12,7 @@ import glob
from modules.face_restoration import FaceRestoration
from scripts.roop_logging import logger
from scripts.swapper import swap_face, ImageResult
from scripts.swapper import UpscaleOptions, swap_face, ImageResult
from scripts.cimage import check_batch
from scripts.roop_version import version_flag
import os
@ -51,7 +52,15 @@ class FaceSwapScript(scripts.Script):
face_restorer_visibility = gr.Slider(
0, 1, 1, step=0.1, label="Restore visibility"
)
upscaler_name = gr.inputs.Dropdown(
choices=[upscaler.name for upscaler in shared.sd_upscalers],
label="Upscaler",
)
upscaler_scale = gr.Slider(1, 8, 1, step=0.1, label="Upscaler scale")
upscaler_visibility = gr.Slider(
0, 1, 1, step=0.1, label="Upscaler visibility (if scale = 1)"
)
models = get_models()
if len(models) == 0:
logger.warning(
@ -86,10 +95,20 @@ class FaceSwapScript(scripts.Script):
model,
face_restorer_name,
face_restorer_visibility,
upscaler_name,
upscaler_scale,
upscaler_visibility,
swap_in_source,
swap_in_generated,
]
@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:
@ -97,6 +116,16 @@ class FaceSwapScript(scripts.Script):
return face_restorer
return None
@property
def upscale_options(self) -> UpscaleOptions:
return UpscaleOptions(
scale=self.upscaler_scale,
upscaler=self.upscaler,
face_restorer=self.face_restorer,
upscale_visibility=self.upscaler_visibility,
restorer_visibility=self.face_restorer_visibility,
)
def process(
self,
p: StableDiffusionProcessing,
@ -106,13 +135,19 @@ class FaceSwapScript(scripts.Script):
model,
face_restorer_name,
face_restorer_visibility,
upscaler_name,
upscaler_scale,
upscaler_visibility,
swap_in_source,
swap_in_generated,
):
self.source = img
self.face_restorer_name = face_restorer_name
self.upscaler_scale = upscaler_scale
self.upscaler_visibility = upscaler_visibility
self.face_restorer_visibility = face_restorer_visibility
self.enable = enable
self.upscaler_name = upscaler_name
self.swap_in_generated = swap_in_generated
self.model = model
self.faces_index = {
@ -132,6 +167,7 @@ class FaceSwapScript(scripts.Script):
p.init_images[i],
faces_index=self.faces_index,
model=self.model,
upscale_options=self.upscale_options,
)
p.init_images[i] = result.image()
else:
@ -151,6 +187,7 @@ class FaceSwapScript(scripts.Script):
image,
faces_index=self.faces_index,
model=self.model,
upscale_options=self.upscale_options,
)
pp = scripts_postprocessing.PostprocessedImage(result.image())
pp.info = {}

47
scripts/swapper.py
View File

@ -14,11 +14,21 @@ import onnxruntime
from scripts.cimage import convert_to_sd
from modules.face_restoration import FaceRestoration, restore_faces
from modules.upscaler import Upscaler, UpscalerData
from scripts.roop_logging import logger
providers = onnxruntime.get_available_providers()
@dataclass
class UpscaleOptions:
scale: int = 1
upscaler: UpscalerData = None
upscale_visibility: float = 0.5
face_restorer: FaceRestoration = None
restorer_visibility: float = 0.5
def save_image(img: Image, filename: str):
convert_to_sd(img).save(filename)
@ -67,6 +77,36 @@ def getFaceSwapModel(model_path: str):
return FS_MODEL
def upscale_image(image: Image, upscale_options: UpscaleOptions):
result_image = image
if upscale_options.upscaler is not None and upscale_options.upscaler.name != "None":
original_image = result_image.copy()
logger.info(
"Upscale with %s scale = %s",
upscale_options.upscaler.name,
upscale_options.scale,
)
result_image = upscale_options.upscaler.scaler.upscale(
image, upscale_options.scale, upscale_options.upscaler.data_path
)
if upscale_options.scale == 1:
result_image = Image.blend(
original_image, result_image, upscale_options.upscale_visibility
)
if upscale_options.face_restorer is not None:
original_image = result_image.copy()
logger.info("Restore face with %s", upscale_options.face_restorer.name())
numpy_image = np.array(result_image)
numpy_image = upscale_options.face_restorer.restore(numpy_image)
restored_image = Image.fromarray(numpy_image)
result_image = Image.blend(
original_image, restored_image, upscale_options.restorer_visibility
)
return result_image
def get_face_single(img_data: np.ndarray, face_index=0, det_size=(640, 640)):
face_analyser = copy.deepcopy(getAnalysisModel())
face_analyser.prepare(ctx_id=0, det_size=det_size)
@ -98,7 +138,9 @@ def swap_face(
target_img: Image.Image,
model: Union[str, None] = None,
faces_index: Set[int] = {0},
upscale_options: Union[UpscaleOptions, None] = None,
) -> ImageResult:
result_image = target_img
fn = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
if model is not None:
source_img = cv2.cvtColor(np.array(source_img), cv2.COLOR_RGB2BGR)
@ -117,8 +159,11 @@ def swap_face(
logger.info(f"No target face found for {face_num}")
result_image = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB))
if upscale_options is not None:
result_image = upscale_image(result_image, upscale_options)
save_image(result_image, fn.name)
# save_image(result_image, fn.name)
else:
logger.info("No source face found")
save_image(result_image, fn.name)
return ImageResult(path=fn.name)