Fixes and updates
This commit is contained in:
Art Gourieff
GitHub
parent
941b94cd0f
commit
abddfb0611
@ -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 = {}
|
||||
|
||||
@ -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}")
|
||||
|
||||
@ -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:
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user