wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-01-25 13:50:15 +08:00
Code Issues Actions 15 Packages Projects Releases Wiki Activity

Update __init__.py

Browse Source 
Add strength of model
This commit is contained in:
Silversith 2023-04-19 12:08:45 +02:00
parent 2ba4d8eaed
commit 332eb856af
1 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
custom_nodes/facerestore/__init__.py
View File

@ -1,4 +1,6 @@
import os import os
import sys
import model_management import model_management
import torch import torch
import comfy.utils import comfy.utils
@ -101,6 +103,7 @@ class FaceRestoreWithModel:
def INPUT_TYPES(s): def INPUT_TYPES(s):
return {"required": { "upscale_model": ("UPSCALE_MODEL",), return {"required": { "upscale_model": ("UPSCALE_MODEL",),
"image": ("IMAGE",), "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"],) "facedetection": (["retinaface_resnet50", "retinaface_mobile0.25", "YOLOv5l", "YOLOv5n"],)
}} }}
@ -113,7 +116,7 @@ class FaceRestoreWithModel:
def __init__(self): def __init__(self):
self.face_helper = None self.face_helper = None
def restore_face(self, upscale_model, image, facedetection): def restore_face(self, upscale_model, image, strength_model, facedetection):
counter = 0 counter = 0
restored_img_list = [] restored_img_list = []
for image_itm in image: for image_itm in image:
@ -145,7 +148,7 @@ class FaceRestoreWithModel:
try: try:
with torch.no_grad(): with torch.no_grad():
#output = upscale_model(cropped_face_t, w=strength, adain=True)[0] #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)) restored_face = tensor2img(output, rgb2bgr=True, min_max=(-1, 1))
del output del output
torch.cuda.empty_cache() torch.cuda.empty_cache()