Merge pull request #3 from ltdrdata/fix/resolution

fix resolution issue  and new nodes

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!
 
 <picture>

__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']

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