From 332eb856af150631b0ff4c622fbfd17d59f19aa0 Mon Sep 17 00:00:00 2001 From: Silversith Date: Wed, 19 Apr 2023 12:08:45 +0200 Subject: [PATCH] Update __init__.py Add strength of model --- custom_nodes/facerestore/__init__.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/custom_nodes/facerestore/__init__.py b/custom_nodes/facerestore/__init__.py index 6557306aa..17e3e0649 100644 --- a/custom_nodes/facerestore/__init__.py +++ b/custom_nodes/facerestore/__init__.py @@ -1,4 +1,6 @@ import os +import sys + import model_management import torch import comfy.utils @@ -101,6 +103,7 @@ class FaceRestoreWithModel: def INPUT_TYPES(s): return {"required": { "upscale_model": ("UPSCALE_MODEL",), "image": ("IMAGE",), + "strength_model": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}), "facedetection": (["retinaface_resnet50", "retinaface_mobile0.25", "YOLOv5l", "YOLOv5n"],) }} @@ -113,7 +116,7 @@ class FaceRestoreWithModel: def __init__(self): self.face_helper = None - def restore_face(self, upscale_model, image, facedetection): + def restore_face(self, upscale_model, image, strength_model, facedetection): counter = 0 restored_img_list = [] for image_itm in image: @@ -145,7 +148,7 @@ class FaceRestoreWithModel: try: with torch.no_grad(): #output = upscale_model(cropped_face_t, w=strength, adain=True)[0] - output = upscale_model(cropped_face_t)[0] + output = upscale_model(cropped_face_t, w=strength_model)[0] restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1)) del output torch.cuda.empty_cache()