Created
October 22, 2023 12:10
-
-
Save donaldafeith/f5c027d7bb4ee37154d90f42519bcb39 to your computer and use it in GitHub Desktop.
This is the fix for faceswap.py for sd-webui-roop. It fixes the depreciation error.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
from gradio import components as gr_components | |
import os | |
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, | |
StableDiffusionProcessingImg2Img, | |
) | |
from modules.shared import cmd_opts, opts, state | |
from PIL import Image | |
import glob | |
from modules.face_restoration import FaceRestoration | |
from scripts.roop_logging import logger | |
from scripts.swapper import UpscaleOptions, swap_face, ImageResult | |
from scripts.roop_version import version_flag | |
import os | |
def get_models(): | |
models_path = os.path.join(scripts.basedir(), "models" + os.path.sep + "roop" + os.path.sep + "*") | |
models = glob.glob(models_path) | |
models = [x for x in models if x.endswith(".onnx") or x.endswith(".pth")] | |
return models | |
class FaceSwapScript(scripts.Script): | |
def title(self): | |
return f"roop" | |
def show(self, is_img2img): | |
return scripts.AlwaysVisible | |
def ui(self, is_img2img): | |
with gr.Accordion(f"roop {version_flag}", open=False): | |
with gr.Column(): | |
img = gr_components.Image(type="pil") | |
enable = gr.Checkbox(False, placeholder="enable", label="Enable") | |
faces_index = gr.Textbox( | |
value="0", | |
placeholder="Which face to swap (comma separated), start from 0", | |
label="Comma separated face number(s)", | |
) | |
with gr.Row(): | |
face_restorer_name = gr.Radio( | |
label="Restore Face", | |
choices=["None"] + [x.name() for x in shared.face_restorers], | |
value=shared.face_restorers[0].name(), | |
type="value", | |
) | |
face_restorer_visibility = gr.Slider( | |
0, 1, 1, step=0.1, label="Restore visibility" | |
) | |
upscaler_name = gr_components.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( | |
"You should at least have one model in models directory, please read the doc here : https://github.com/s0md3v/sd-webui-roop/" | |
) | |
model = gr_components.Dropdown( | |
choices=models, | |
label="Model not found, please download one and reload automatic 1111", | |
) | |
else: | |
model = gr_components.Dropdown( | |
choices=models, label="Model", default=models[0] | |
) | |
swap_in_source = gr.Checkbox( | |
False, | |
placeholder="Swap face in source image", | |
label="Swap in source image", | |
visible=is_img2img, | |
) | |
swap_in_generated = gr.Checkbox( | |
True, | |
placeholder="Swap face in generated image", | |
label="Swap in generated image", | |
visible=is_img2img, | |
) | |
return [ | |
img, | |
enable, | |
faces_index, | |
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: | |
if face_restorer.name() == self.face_restorer_name: | |
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, | |
img, | |
enable, | |
faces_index, | |
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 = { | |
int(x) for x in faces_index.strip(",").split(",") if x.isnumeric() | |
} | |
if len(self.faces_index) == 0: | |
self.faces_index = {0} | |
if self.enable: | |
if self.source is not None: | |
if isinstance(p, StableDiffusionProcessingImg2Img) and swap_in_source: | |
logger.info(f"roop enabled, face index %s", self.faces_index) | |
for i in range(len(p.init_images)): | |
logger.info(f"Swap in source %s", i) | |
result = swap_face( | |
self.source, | |
p.init_images[i], | |
faces_index=self.faces_index, | |
model=self.model, | |
upscale_options=self.upscale_options, | |
) | |
p.init_images[i] = result.image() | |
else: | |
logger.error(f"Please provide a source face") | |
def postprocess_batch(self, *args, **kwargs): | |
if self.enable: | |
return images | |
def postprocess_image(self, p, script_pp: scripts.PostprocessImageArgs, *args): | |
if self.enable and self.swap_in_generated: | |
if self.source is not None: | |
image: Image.Image = script_pp.image | |
result: ImageResult = swap_face( | |
self.source, | |
image, | |
faces_index=self.faces_index, | |
model=self.model, | |
upscale_options=self.upscale_options, | |
) | |
pp = scripts_postprocessing.PostprocessedImage(result.image()) | |
pp.info = {} | |
p.extra_generation_params.update(pp.info) | |
script_pp.image = pp.image |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment