142 lines
4.3 KiB
Python
142 lines
4.3 KiB
Python
import sys
|
|
from os import path
|
|
sys.path.insert(0, path.dirname(__file__))
|
|
from .ldsrlib.LDSR import LDSR
|
|
from folder_paths import get_filename_list, get_full_path
|
|
from comfy.model_management import get_torch_device
|
|
from comfy.utils import ProgressBar
|
|
import torch
|
|
|
|
|
|
class LDSRModelLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
model_list = get_filename_list("upscale_models")
|
|
candidates = [name for name in model_list if 'last.ckpt' in name]
|
|
if len(candidates) > 0:
|
|
default_path = candidates[0]
|
|
else:
|
|
default_path = 'last.ckpt'
|
|
|
|
return {
|
|
"required": {
|
|
"model": (model_list, {'default': default_path}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("UPSCALE_MODEL",)
|
|
FUNCTION = "load"
|
|
|
|
CATEGORY = "Flowty LDSR"
|
|
|
|
def load(self, model):
|
|
model_path = get_full_path("upscale_models", model)
|
|
model = LDSR.load_model_from_path(model_path)
|
|
model['model'].cpu()
|
|
return (model, )
|
|
|
|
|
|
class LDSRUpscale:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"upscale_model": ("UPSCALE_MODEL",),
|
|
"images": ("IMAGE",),
|
|
"steps": (["25", "50", "100", "250", "500", "1000"], {"default": "100"}),
|
|
"pre_downscale": (['None', '1/2', '1/4'], {"default": "None"}),
|
|
"post_downscale": (['None', 'Original Size', '1/2', '1/4'], {"default": "None"}),
|
|
"downsample_method": (['Nearest', 'Lanczos'], {"default": "Lanczos"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("images",)
|
|
FUNCTION = "upscale"
|
|
|
|
CATEGORY = "Flowty LDSR"
|
|
|
|
def upscale(self, upscale_model, images, steps, pre_downscale="None", post_downscale="None", downsample_method="Lanczos"):
|
|
pbar = ProgressBar(int(steps))
|
|
p = {"prev": 0}
|
|
|
|
def prog(i):
|
|
i = i + 1
|
|
if i < p["prev"]:
|
|
p["prev"] = 0
|
|
pbar.update(i - p["prev"])
|
|
p["prev"] = i
|
|
|
|
ldsr = LDSR(model=upscale_model, on_progress=prog)
|
|
|
|
outputs = []
|
|
|
|
for image in images:
|
|
outputs.append(ldsr.superResolution(image, int(steps), pre_downscale, post_downscale, downsample_method))
|
|
|
|
return (torch.stack(outputs),)
|
|
|
|
|
|
class LDSRUpscaler:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
model_list = get_filename_list("upscale_models")
|
|
candidates = [name for name in model_list if 'last.ckpt' in name]
|
|
if len(candidates) > 0:
|
|
default_path = candidates[0]
|
|
else:
|
|
default_path = 'last.ckpt'
|
|
|
|
return {
|
|
"required": {
|
|
"model": (model_list, {'default': default_path}),
|
|
"images": ("IMAGE",),
|
|
"steps": (["25", "50", "100", "250", "500", "1000"], {"default": "100"}),
|
|
"pre_downscale": (['None', '1/2', '1/4'], {"default": "None"}),
|
|
"post_downscale": (['None', 'Original Size', '1/2', '1/4'], {"default": "None"}),
|
|
"downsample_method": (['Nearest', 'Lanczos'], {"default": "Lanczos"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("images",)
|
|
FUNCTION = "upscale"
|
|
|
|
CATEGORY = "Flowty LDSR"
|
|
|
|
def upscale(self, model, images, steps, pre_downscale="None", post_downscale="None", downsample_method="Lanczos"):
|
|
model_path = get_full_path("upscale_models", model)
|
|
pbar = ProgressBar(int(steps))
|
|
p = {"prev": 0}
|
|
|
|
def prog(i):
|
|
i = i + 1
|
|
if i < p["prev"]:
|
|
p["prev"] = 0
|
|
pbar.update(i - p["prev"])
|
|
p["prev"] = i
|
|
|
|
ldsr = LDSR(modelPath=model_path, torchdevice=get_torch_device(), on_progress=prog)
|
|
|
|
outputs = []
|
|
|
|
for image in images:
|
|
outputs.append(ldsr.superResolution(image, int(steps), pre_downscale, post_downscale, downsample_method))
|
|
|
|
return (torch.stack(outputs),)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"LDSRUpscaler": LDSRUpscaler,
|
|
"LDSRModelLoader": LDSRModelLoader,
|
|
"LDSRUpscale": LDSRUpscale
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"LDSRUpscaler": "LDSR Upscale (all-in-one)",
|
|
"LDSRModelLoader": "Load LDSR Model",
|
|
"LDSRUpscale": "LDSR Upscale"
|
|
}
|
|
|
|
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|