Fixes and updates

scripts/faceswap.py

@@ -1,6 +1,5 @@
 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,
@@ -12,32 +11,28 @@ import glob
 from modules.face_restoration import FaceRestoration
 
 from scripts.roop_logging import logger
-from scripts.swapper import UpscaleOptions, swap_face, ImageResult
+from scripts.swapper import swap_face, ImageResult 
 from scripts.cimage import check_batch
 from scripts.roop_version import version_flag
 import os
 
 
 def get_models():
-    models_path = os.path.join(
-        scripts.basedir(), "extensions/sd-webui-roop/models/*"
-    )
-    models = glob.glob(models_path)
     models_path = os.path.join(scripts.basedir(), "models/roop/*")
-    models += glob.glob(models_path)
+    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"
+        return f"nsfw-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.Accordion(f"nsfw-roop {version_flag}", open=False):
             with gr.Column():
                 img = gr.inputs.Image(type="pil")
                 enable = gr.Checkbox(False, placeholder="enable", label="Enable")
@@ -56,19 +51,11 @@ 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(
-                        "You should at least have one model in models directory, please read the doc here : https://github.com/s0md3v/sd-webui-roop/"
+                        "You should at least have one model in models directory, please read the doc here : https://github.com/s0md3v/sd-webui-roop-nsfw/"
                     )
                     model = gr.inputs.Dropdown(
                         choices=models,
@@ -99,20 +86,10 @@ 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:
@@ -120,16 +97,6 @@ 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,
@@ -139,19 +106,13 @@ 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 = {
@@ -162,7 +123,7 @@ class FaceSwapScript(scripts.Script):
         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)
+                    logger.info(f"nsfw-roop enabled, face index %s", self.faces_index)
 
                     for i in range(len(p.init_images)):
                         logger.info(f"Swap in source %s", i)
@@ -171,7 +132,6 @@ 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:
@@ -191,7 +151,6 @@ 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 = {}

scripts/roop_version.py

@@ -1,5 +1,5 @@
-version_flag = "v0.0.1"
+version_flag = "v0.0.2"
 
 from scripts.roop_logging import logger
 
-logger.info(f"roop {version_flag}")
+logger.info(f"nswf-roop {version_flag}")

scripts/swapper.py

@@ -14,21 +14,11 @@ 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)
 
@@ -77,36 +67,6 @@ 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)
@@ -138,7 +98,6 @@ def swap_face(
     target_img: Image.Image,
     model: Union[str, None] = None,
     faces_index: Set[int] = {0},
-    upscale_options: Union[UpscaleOptions, None] = None,
 ) -> ImageResult:
     fn = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
     if model is not None:
@@ -158,8 +117,6 @@ 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)
         else: