diff --git a/README.md b/README.md index ac29b7c..f3ecb52 100644 --- a/README.md +++ b/README.md @@ -28,6 +28,8 @@ I've created this node for experimentation, feel free to submit PRs for performa ### Workflow: ![example](workflow.png) +![example2](workflow2.png) + This is a community project from [flowt.ai](https://flowt.ai). If you like it, check us out! diff --git a/__init__.py b/__init__.py index 84178f1..9502473 100644 --- a/__init__.py +++ b/__init__.py @@ -6,16 +6,93 @@ from folder_paths import get_filename_list, get_full_path from comfy.model_management import get_torch_device from comfy.utils import ProgressBar -class LDSRUpscaler: + +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": { - "model": (get_filename_list("upscale_models"),), + "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"}), + "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 (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"}), } } @@ -29,7 +106,7 @@ class LDSRUpscaler: 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} + p = {"prev": 0} def prog(i): i = i + 1 @@ -38,7 +115,7 @@ class LDSRUpscaler: pbar.update(i - p["prev"]) p["prev"] = i - ldsr = LDSR(model_path, get_torch_device(), prog) + ldsr = LDSR(modelPath=model_path, torchdevice=get_torch_device(), on_progress=prog) outputs = [] @@ -49,11 +126,15 @@ class LDSRUpscaler: NODE_CLASS_MAPPINGS = { - "LDSRUpscaler": LDSRUpscaler + "LDSRUpscaler": LDSRUpscaler, + "LDSRModelLoader": LDSRModelLoader, + "LDSRUpscale": LDSRUpscale } NODE_DISPLAY_NAME_MAPPINGS = { - "LDSRUpscaler": "LDSR Upscale" + "LDSRUpscaler": "LDSR Upscale (all-in-one)", + "LDSRModelLoader": "Load LDSR Model", + "LDSRUpscale": "LDSR Upscale" } __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/ldsrlib/LDSR.py b/ldsrlib/LDSR.py index 2deccb7..9493b3d 100644 --- a/ldsrlib/LDSR.py +++ b/ldsrlib/LDSR.py @@ -1,3 +1,6 @@ +import PIL.Image + +import nodes from .ldm.util import instantiate_from_config from .ldm.models.diffusion.ddim import DDIMSampler from .ldm.util import ismap @@ -10,28 +13,83 @@ from omegaconf import OmegaConf import numpy as np from os import path import warnings +from comfy import model_management +import comfy +from PIL import ImageOps warnings.filterwarnings("ignore", category=UserWarning) + class LDSR(): - def __init__(self, modelPath, torchdevice=torch.device("cuda"), on_progress=None, yamlPath=path.join(path.dirname(__file__), "config.yaml")): + def __init__(self, modelPath=None, model=None, torchdevice=model_management.get_torch_device(), on_progress=None, yamlPath=path.join(path.dirname(__file__), "config.yaml")): self.modelPath = modelPath + self.model = model self.yamlPath = yamlPath - self.torchdevice=torchdevice + self.torchdevice = torchdevice self.progress_hook = on_progress if on_progress else None + @staticmethod + def normalize_image(image): + w, h = image.size - def load_model_from_config(self): - print(f"Loading model from {self.modelPath}") - pl_sd = torch.load(self.modelPath, map_location="cpu") - global_step = pl_sd["global_step"] + # ensure (min length > 128) + if h < w and h < 128: + scale_ratio = 128 / h + h = 128 + w = int(scale_ratio * w) + elif w < 128: + scale_ratio = 128 / w + w = 128 + h = int(scale_ratio * h) + + resample = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS) + image = image.resize((w, h), resample=resample) + + # ensure (multiply of 64) + w_pad = 64 - w % 64 + h_pad = 64 - h % 64 + + padded_image = Image.new("RGB", (w + w_pad, h + h_pad), color="black") + padded_image.paste(image, (0, 0)) + + return padded_image, w_pad, h_pad + + @staticmethod + def remove_padding(prev_pil, image, w_pad, h_pad): + if w_pad == 0 and h_pad == 0: + return image + + w1, h1 = prev_pil.size + h2, w2, _ = image.size() + + scale_ratio = h2 / h1 + w_pad = float.__ceil__(w_pad * scale_ratio) + h_pad = float.__ceil__(h_pad * scale_ratio) + + return image[:h2-h_pad, :w2-w_pad, :] + + + @staticmethod + def load_model_from_path(modelPath, device=model_management.get_torch_device(), yamlPath=path.join(path.dirname(__file__), "config.yaml")): + print(f"Loading model from {modelPath}") + pl_sd = torch.load(modelPath, map_location="cpu") sd = pl_sd["state_dict"] - config = OmegaConf.load(self.yamlPath) + config = OmegaConf.load(yamlPath) model = instantiate_from_config(config.model) model.load_state_dict(sd, strict=False) - model.to(self.torchdevice) + + model.to(device) model.eval() - return {"model": model} # , global_step + + return {"model": model} + + def load_model_from_config(self): + if self.model is None: + self.model = LDSR.load_model_from_path(self.modelPath, self.torchdevice) + else: + self.model['model'].to(self.torchdevice) + + return self.model def progress_callback(self, i): if self.progress_hook: @@ -237,6 +295,8 @@ class LDSR(): print(f'Downsampling from [{width_og}, {height_og}] to [{width_downsampled_pre}, {height_downsampled_pre}]') im_og = im_og.resize((width_downsampled_pre, height_downsampled_pre), Image.LANCZOS) + im_og, w_pad, h_pad = LDSR.normalize_image(im_og) + logs = self.run(model["model"], im_og, diffMode, diffusion_steps, eta) sample = logs["sample"] @@ -269,8 +329,15 @@ class LDSR(): a = a.resize((width_downsampled_post, height_downsampled_post), aliasing) elif post_downsample == 'Original Size': print(f'Downsampling from [{width}, {height}] to Original Size [{width_og}, {height_og}]') - a = a.resize((width_og, height_og), aliasing) + a = a.resize((width_og+w_pad, height_og+h_pad), aliasing) out = np.array(a).astype(np.float32) / 255.0 - return torch.from_numpy(out)[None,][0] + # Finalize + result_image = torch.from_numpy(out) + result_image = LDSR.remove_padding(im_og, result_image, w_pad, h_pad) + + model['model'].cpu() + result_image.cpu() + + return result_image diff --git a/workflow.png b/workflow.png index e1b7e73..ede2c81 100644 Binary files a/workflow.png and b/workflow.png differ diff --git a/workflow2.png b/workflow2.png new file mode 100644 index 0000000..89d36a0 Binary files /dev/null and b/workflow2.png differ