refactor: ⚡️ small local fixes

made while writting wiki
This commit is contained in:
melMass
2023-10-21 03:14:12 +02:00
parent 049983dbe2
commit bcac66508d
4 changed files with 95 additions and 5 deletions
+7
View File
@@ -0,0 +1,7 @@
class ModelNotFound(Exception):
def __init__(self, model_name, *args, **kwargs):
super().__init__(
f"The model {model_name} could not be found, make sure to download it using ComfyManager first.\nrepository: https://github.com/ltdrdata/ComfyUI-Manager",
*args,
**kwargs,
)
-1
View File
@@ -1,7 +1,6 @@
import tempfile
from pathlib import Path
import folder_paths
import numpy as np
import onnxruntime as ort
import torch
+3 -4
View File
@@ -8,7 +8,6 @@ import cv2
import folder_paths
import numpy as np
import torch
from basicsr.utils import imwrite
from comfy import model_management
from gfpgan import GFPGANer
from PIL import Image
@@ -247,16 +246,16 @@ class RestoreFace:
):
face_id = idx + 1
file = self.get_step_image_path("cropped_faces", face_id)
imwrite(cropped_face, file)
cv2.imwrite(file, cropped_face)
file = self.get_step_image_path("cropped_faces_restored", face_id)
imwrite(restored_face, file)
cv2.imwrite(file, restored_face)
file = self.get_step_image_path("cropped_faces_compare", face_id)
# save comparison image
cmp_img = np.concatenate((cropped_face, restored_face), axis=1)
imwrite(cmp_img, file)
cv2.imwrite(file, cmp_img)
__nodes__ = [RestoreFace, LoadFaceEnhanceModel]
+85
View File
@@ -22,6 +22,18 @@ from ..utils import pil2tensor, tensor2np, tensor2pil
# log.warning("cv2.ximgproc.guidedFilter not found, use opencv-contrib-python")
def gaussian_kernel(kernel_size: int, sigma_x: float, sigma_y: float, device=None):
x, y = torch.meshgrid(
torch.linspace(-1, 1, kernel_size, device=device),
torch.linspace(-1, 1, kernel_size, device=device),
indexing="ij",
)
d_x = x * x / (2.0 * sigma_x * sigma_x)
d_y = y * y / (2.0 * sigma_y * sigma_y)
g = torch.exp(-(d_x + d_y))
return g / g.sum()
class ColorCorrect:
"""Various color correction methods"""
@@ -273,6 +285,78 @@ class Blur_:
return (torch.from_numpy(image),)
class Sharpen_:
"""Sharpens an image using a Gaussian kernel."""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"image": ("IMAGE",),
"sharpen_radius": (
"INT",
{"default": 1, "min": 1, "max": 31, "step": 1},
),
"sigma_x": (
"FLOAT",
{"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1},
),
"sigma_y": (
"FLOAT",
{"default": 1.0, "min": 0.1, "max": 10.0, "step": 0.1},
),
"alpha": (
"FLOAT",
{"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.1},
),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "do_sharp"
CATEGORY = "mtb/image processing"
def do_sharp(
self,
image: torch.Tensor,
sharpen_radius: int,
sigma_x: float,
sigma_y: float,
alpha: float,
):
if sharpen_radius == 0:
return (image,)
channels = image.shape[3]
kernel_size = 2 * sharpen_radius + 1
kernel = gaussian_kernel(kernel_size, sigma_x, sigma_y) * -(alpha * 10)
# Modify center of kernel to make it a sharpening kernel
center = kernel_size // 2
kernel[center, center] = kernel[center, center] - kernel.sum() + 1.0
kernel = kernel.repeat(channels, 1, 1).unsqueeze(1)
tensor_image = image.permute(0, 3, 1, 2)
tensor_image = F.pad(
tensor_image,
(sharpen_radius, sharpen_radius, sharpen_radius, sharpen_radius),
"reflect",
)
sharpened = F.conv2d(tensor_image, kernel, padding=center, groups=channels)
# Remove padding
sharpened = sharpened[
:, :, sharpen_radius:-sharpen_radius, sharpen_radius:-sharpen_radius
]
sharpened = sharpened.permute(0, 2, 3, 1)
result = torch.clamp(sharpened, 0, 1)
return (result,)
# https://github.com/lllyasviel/AdverseCleaner/blob/main/clean.py
# def deglaze_np_img(np_img):
# y = np_img.copy()
@@ -702,4 +786,5 @@ __nodes__ = [
ImageResizeFactor,
SaveImageGrid_,
LoadImageFromUrl_,
Sharpen_,
]