Delete ldm_patched

ldm_patched/contrib/external.py

@@ -1,1954 +0,0 @@
-# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
-
-import torch
-
-import os
-import sys
-import json
-import hashlib
-import traceback
-import math
-import time
-import random
-
-from PIL import Image, ImageOps, ImageSequence
-from PIL.PngImagePlugin import PngInfo
-import numpy as np
-import safetensors.torch
-
-pass  # sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "ldm_patched"))
-
-
-import ldm_patched.modules.diffusers_load
-import ldm_patched.modules.samplers
-import ldm_patched.modules.sample
-import ldm_patched.modules.sd
-import ldm_patched.modules.utils
-import ldm_patched.modules.controlnet
-
-import ldm_patched.modules.clip_vision
-
-import ldm_patched.modules.model_management
-from ldm_patched.modules.args_parser import args
-
-import importlib
-
-import ldm_patched.utils.path_utils
-import ldm_patched.utils.latent_visualization
-
-def before_node_execution():
-    ldm_patched.modules.model_management.throw_exception_if_processing_interrupted()
-
-def interrupt_processing(value=True):
-    ldm_patched.modules.model_management.interrupt_current_processing(value)
-
-MAX_RESOLUTION=8192
-
-class CLIPTextEncode:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {"text": ("STRING", {"multiline": True}), "clip": ("CLIP", )}}
-    RETURN_TYPES = ("CONDITIONING",)
-    FUNCTION = "encode"
-
-    CATEGORY = "conditioning"
-
-    def encode(self, clip, text):
-        tokens = clip.tokenize(text)
-        cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)
-        return ([[cond, {"pooled_output": pooled}]], )
-
-class ConditioningCombine:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {"conditioning_1": ("CONDITIONING", ), "conditioning_2": ("CONDITIONING", )}}
-    RETURN_TYPES = ("CONDITIONING",)
-    FUNCTION = "combine"
-
-    CATEGORY = "conditioning"
-
-    def combine(self, conditioning_1, conditioning_2):
-        return (conditioning_1 + conditioning_2, )
-
-class ConditioningAverage :
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {"conditioning_to": ("CONDITIONING", ), "conditioning_from": ("CONDITIONING", ),
-                              "conditioning_to_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01})
-                             }}
-    RETURN_TYPES = ("CONDITIONING",)
-    FUNCTION = "addWeighted"
-
-    CATEGORY = "conditioning"
-
-    def addWeighted(self, conditioning_to, conditioning_from, conditioning_to_strength):
-        out = []
-
-        if len(conditioning_from) > 1:
-            print("Warning: ConditioningAverage conditioning_from contains more than 1 cond, only the first one will actually be applied to conditioning_to.")
-
-        cond_from = conditioning_from[0][0]
-        pooled_output_from = conditioning_from[0][1].get("pooled_output", None)
-
-        for i in range(len(conditioning_to)):
-            t1 = conditioning_to[i][0]
-            pooled_output_to = conditioning_to[i][1].get("pooled_output", pooled_output_from)
-            t0 = cond_from[:,:t1.shape[1]]
-            if t0.shape[1] < t1.shape[1]:
-                t0 = torch.cat([t0] + [torch.zeros((1, (t1.shape[1] - t0.shape[1]), t1.shape[2]))], dim=1)
-
-            tw = torch.mul(t1, conditioning_to_strength) + torch.mul(t0, (1.0 - conditioning_to_strength))
-            t_to = conditioning_to[i][1].copy()
-            if pooled_output_from is not None and pooled_output_to is not None:
-                t_to["pooled_output"] = torch.mul(pooled_output_to, conditioning_to_strength) + torch.mul(pooled_output_from, (1.0 - conditioning_to_strength))
-            elif pooled_output_from is not None:
-                t_to["pooled_output"] = pooled_output_from
-
-            n = [tw, t_to]
-            out.append(n)
-        return (out, )
-
-class ConditioningConcat:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {
-            "conditioning_to": ("CONDITIONING",),
-            "conditioning_from": ("CONDITIONING",),
-            }}
-    RETURN_TYPES = ("CONDITIONING",)
-    FUNCTION = "concat"
-
-    CATEGORY = "conditioning"
-
-    def concat(self, conditioning_to, conditioning_from):
-        out = []
-
-        if len(conditioning_from) > 1:
-            print("Warning: ConditioningConcat conditioning_from contains more than 1 cond, only the first one will actually be applied to conditioning_to.")
-
-        cond_from = conditioning_from[0][0]
-
-        for i in range(len(conditioning_to)):
-            t1 = conditioning_to[i][0]
-            tw = torch.cat((t1, cond_from),1)
-            n = [tw, conditioning_to[i][1].copy()]
-            out.append(n)
-
-        return (out, )
-
-class ConditioningSetArea:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {"conditioning": ("CONDITIONING", ),
-                              "width": ("INT", {"default": 64, "min": 64, "max": MAX_RESOLUTION, "step": 8}),
-                              "height": ("INT", {"default": 64, "min": 64, "max": MAX_RESOLUTION, "step": 8}),
-                              "x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
-                              "y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
-                              "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
-                             }}
-    RETURN_TYPES = ("CONDITIONING",)
-    FUNCTION = "append"
-
-    CATEGORY = "conditioning"
-
-    def append(self, conditioning, width, height, x, y, strength):
-        c = []
-        for t in conditioning:
-            n = [t[0], t[1].copy()]
-            n[1]['area'] = (height // 8, width // 8, y // 8, x // 8)
-            n[1]['strength'] = strength
-            n[1]['set_area_to_bounds'] = False
-            c.append(n)
-        return (c, )
-
-class ConditioningSetAreaPercentage:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {"conditioning": ("CONDITIONING", ),
-                              "width": ("FLOAT", {"default": 1.0, "min": 0, "max": 1.0, "step": 0.01}),
-                              "height": ("FLOAT", {"default": 1.0, "min": 0, "max": 1.0, "step": 0.01}),
-                              "x": ("FLOAT", {"default": 0, "min": 0, "max": 1.0, "step": 0.01}),
-                              "y": ("FLOAT", {"default": 0, "min": 0, "max": 1.0, "step": 0.01}),
-                              "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
-                             }}
-    RETURN_TYPES = ("CONDITIONING",)
-    FUNCTION = "append"
-
-    CATEGORY = "conditioning"
-
-    def append(self, conditioning, width, height, x, y, strength):
-        c = []
-        for t in conditioning:
-            n = [t[0], t[1].copy()]
-            n[1]['area'] = ("percentage", height, width, y, x)
-            n[1]['strength'] = strength
-            n[1]['set_area_to_bounds'] = False
-            c.append(n)
-        return (c, )
-
-class ConditioningSetMask:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {"conditioning": ("CONDITIONING", ),
-                              "mask": ("MASK", ),
-                              "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
-                              "set_cond_area": (["default", "mask bounds"],),
-                             }}
-    RETURN_TYPES = ("CONDITIONING",)
-    FUNCTION = "append"
-
-    CATEGORY = "conditioning"
-
-    def append(self, conditioning, mask, set_cond_area, strength):
-        c = []
-        set_area_to_bounds = False
-        if set_cond_area != "default":
-            set_area_to_bounds = True
-        if len(mask.shape) < 3:
-            mask = mask.unsqueeze(0)
-        for t in conditioning:
-            n = [t[0], t[1].copy()]
-            _, h, w = mask.shape
-            n[1]['mask'] = mask
-            n[1]['set_area_to_bounds'] = set_area_to_bounds
-            n[1]['mask_strength'] = strength
-            c.append(n)
-        return (c, )
-
-class ConditioningZeroOut:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {"conditioning": ("CONDITIONING", )}}
-    RETURN_TYPES = ("CONDITIONING",)
-    FUNCTION = "zero_out"
-
-    CATEGORY = "advanced/conditioning"
-
-    def zero_out(self, conditioning):
-        c = []
-        for t in conditioning:
-            d = t[1].copy()
-            if "pooled_output" in d:
-                d["pooled_output"] = torch.zeros_like(d["pooled_output"])
-            n = [torch.zeros_like(t[0]), d]
-            c.append(n)
-        return (c, )
-
-class ConditioningSetTimestepRange:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {"conditioning": ("CONDITIONING", ),
-                             "start": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
-                             "end": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001})
-                             }}
-    RETURN_TYPES = ("CONDITIONING",)
-    FUNCTION = "set_range"
-
-    CATEGORY = "advanced/conditioning"
-
-    def set_range(self, conditioning, start, end):
-        c = []
-        for t in conditioning:
-            d = t[1].copy()
-            d['start_percent'] = start
-            d['end_percent'] = end
-            n = [t[0], d]
-            c.append(n)
-        return (c, )
-
-class VAEDecode:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "samples": ("LATENT", ), "vae": ("VAE", )}}
-    RETURN_TYPES = ("IMAGE",)
-    FUNCTION = "decode"
-
-    CATEGORY = "latent"
-
-    def decode(self, vae, samples):
-        return (vae.decode(samples["samples"]), )
-
-class VAEDecodeTiled:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {"samples": ("LATENT", ), "vae": ("VAE", ),
-                             "tile_size": ("INT", {"default": 512, "min": 320, "max": 4096, "step": 64})
-                            }}
-    RETURN_TYPES = ("IMAGE",)
-    FUNCTION = "decode"
-
-    CATEGORY = "_for_testing"
-
-    def decode(self, vae, samples, tile_size):
-        return (vae.decode_tiled(samples["samples"], tile_x=tile_size // 8, tile_y=tile_size // 8, ), )
-
-class VAEEncode:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "pixels": ("IMAGE", ), "vae": ("VAE", )}}
-    RETURN_TYPES = ("LATENT",)
-    FUNCTION = "encode"
-
-    CATEGORY = "latent"
-
-    @staticmethod
-    def vae_encode_crop_pixels(pixels):
-        x = (pixels.shape[1] // 8) * 8
-        y = (pixels.shape[2] // 8) * 8
-        if pixels.shape[1] != x or pixels.shape[2] != y:
-            x_offset = (pixels.shape[1] % 8) // 2
-            y_offset = (pixels.shape[2] % 8) // 2
-            pixels = pixels[:, x_offset:x + x_offset, y_offset:y + y_offset, :]
-        return pixels
-
-    def encode(self, vae, pixels):
-        pixels = self.vae_encode_crop_pixels(pixels)
-        t = vae.encode(pixels[:,:,:,:3])
-        return ({"samples":t}, )
-
-class VAEEncodeTiled:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {"pixels": ("IMAGE", ), "vae": ("VAE", ),
-                             "tile_size": ("INT", {"default": 512, "min": 320, "max": 4096, "step": 64})
-                            }}
-    RETURN_TYPES = ("LATENT",)
-    FUNCTION = "encode"
-
-    CATEGORY = "_for_testing"
-
-    def encode(self, vae, pixels, tile_size):
-        pixels = VAEEncode.vae_encode_crop_pixels(pixels)
-        t = vae.encode_tiled(pixels[:,:,:,:3], tile_x=tile_size, tile_y=tile_size, )
-        return ({"samples":t}, )
-
-class VAEEncodeForInpaint:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "pixels": ("IMAGE", ), "vae": ("VAE", ), "mask": ("MASK", ), "grow_mask_by": ("INT", {"default": 6, "min": 0, "max": 64, "step": 1}),}}
-    RETURN_TYPES = ("LATENT",)
-    FUNCTION = "encode"
-
-    CATEGORY = "latent/inpaint"
-
-    def encode(self, vae, pixels, mask, grow_mask_by=6):
-        x = (pixels.shape[1] // 8) * 8
-        y = (pixels.shape[2] // 8) * 8
-        mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(pixels.shape[1], pixels.shape[2]), mode="bilinear")
-
-        pixels = pixels.clone()
-        if pixels.shape[1] != x or pixels.shape[2] != y:
-            x_offset = (pixels.shape[1] % 8) // 2
-            y_offset = (pixels.shape[2] % 8) // 2
-            pixels = pixels[:,x_offset:x + x_offset, y_offset:y + y_offset,:]
-            mask = mask[:,:,x_offset:x + x_offset, y_offset:y + y_offset]
-
-        #grow mask by a few pixels to keep things seamless in latent space
-        if grow_mask_by == 0:
-            mask_erosion = mask
-        else:
-            kernel_tensor = torch.ones((1, 1, grow_mask_by, grow_mask_by))
-            padding = math.ceil((grow_mask_by - 1) / 2)
-
-            mask_erosion = torch.clamp(torch.nn.functional.conv2d(mask.round(), kernel_tensor, padding=padding), 0, 1)
-
-        m = (1.0 - mask.round()).squeeze(1)
-        for i in range(3):
-            pixels[:,:,:,i] -= 0.5
-            pixels[:,:,:,i] *= m
-            pixels[:,:,:,i] += 0.5
-        t = vae.encode(pixels)
-
-        return ({"samples":t, "noise_mask": (mask_erosion[:,:,:x,:y].round())}, )
-
-
-class InpaintModelConditioning:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {"positive": ("CONDITIONING", ),
-                             "negative": ("CONDITIONING", ),
-                             "vae": ("VAE", ),
-                             "pixels": ("IMAGE", ),
-                             "mask": ("MASK", ),
-                             }}
-
-    RETURN_TYPES = ("CONDITIONING","CONDITIONING","LATENT")
-    RETURN_NAMES = ("positive", "negative", "latent")
-    FUNCTION = "encode"
-
-    CATEGORY = "conditioning/inpaint"
-
-    def encode(self, positive, negative, pixels, vae, mask):
-        x = (pixels.shape[1] // 8) * 8
-        y = (pixels.shape[2] // 8) * 8
-        mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(pixels.shape[1], pixels.shape[2]), mode="bilinear")
-
-        orig_pixels = pixels
-        pixels = orig_pixels.clone()
-        if pixels.shape[1] != x or pixels.shape[2] != y:
-            x_offset = (pixels.shape[1] % 8) // 2
-            y_offset = (pixels.shape[2] % 8) // 2
-            pixels = pixels[:,x_offset:x + x_offset, y_offset:y + y_offset,:]
-            mask = mask[:,:,x_offset:x + x_offset, y_offset:y + y_offset]
-
-        m = (1.0 - mask.round()).squeeze(1)
-        for i in range(3):
-            pixels[:,:,:,i] -= 0.5
-            pixels[:,:,:,i] *= m
-            pixels[:,:,:,i] += 0.5
-        concat_latent = vae.encode(pixels)
-        orig_latent = vae.encode(orig_pixels)
-
-        out_latent = {}
-
-        out_latent["samples"] = orig_latent
-        out_latent["noise_mask"] = mask
-
-        out = []
-        for conditioning in [positive, negative]:
-            c = []
-            for t in conditioning:
-                d = t[1].copy()
-                d["concat_latent_image"] = concat_latent
-                d["concat_mask"] = mask
-                n = [t[0], d]
-                c.append(n)
-            out.append(c)
-        return (out[0], out[1], out_latent)
-
-
-class SaveLatent:
-    def __init__(self):
-        self.output_dir = ldm_patched.utils.path_utils.get_output_directory()
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "samples": ("LATENT", ),
-                              "filename_prefix": ("STRING", {"default": "latents/ldm_patched"})},
-                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
-                }
-    RETURN_TYPES = ()
-    FUNCTION = "save"
-
-    OUTPUT_NODE = True
-
-    CATEGORY = "_for_testing"
-
-    def save(self, samples, filename_prefix="ldm_patched", prompt=None, extra_pnginfo=None):
-        full_output_folder, filename, counter, subfolder, filename_prefix = ldm_patched.utils.path_utils.get_save_image_path(filename_prefix, self.output_dir)
-
-        # support save metadata for latent sharing
-        prompt_info = ""
-        if prompt is not None:
-            prompt_info = json.dumps(prompt)
-
-        metadata = None
-        if not args.disable_server_info:
-            metadata = {"prompt": prompt_info}
-            if extra_pnginfo is not None:
-                for x in extra_pnginfo:
-                    metadata[x] = json.dumps(extra_pnginfo[x])
-
-        file = f"{filename}_{counter:05}_.latent"
-
-        results = list()
-        results.append({
-            "filename": file,
-            "subfolder": subfolder,
-            "type": "output"
-        })
-
-        file = os.path.join(full_output_folder, file)
-
-        output = {}
-        output["latent_tensor"] = samples["samples"]
-        output["latent_format_version_0"] = torch.tensor([])
-
-        ldm_patched.modules.utils.save_torch_file(output, file, metadata=metadata)
-        return { "ui": { "latents": results } }
-
-
-class LoadLatent:
-    @classmethod
-    def INPUT_TYPES(s):
-        input_dir = ldm_patched.utils.path_utils.get_input_directory()
-        files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f)) and f.endswith(".latent")]
-        return {"required": {"latent": [sorted(files), ]}, }
-
-    CATEGORY = "_for_testing"
-
-    RETURN_TYPES = ("LATENT", )
-    FUNCTION = "load"
-
-    def load(self, latent):
-        latent_path = ldm_patched.utils.path_utils.get_annotated_filepath(latent)
-        latent = safetensors.torch.load_file(latent_path, device="cpu")
-        multiplier = 1.0
-        if "latent_format_version_0" not in latent:
-            multiplier = 1.0 / 0.18215
-        samples = {"samples": latent["latent_tensor"].float() * multiplier}
-        return (samples, )
-
-    @classmethod
-    def IS_CHANGED(s, latent):
-        image_path = ldm_patched.utils.path_utils.get_annotated_filepath(latent)
-        m = hashlib.sha256()
-        with open(image_path, 'rb') as f:
-            m.update(f.read())
-        return m.digest().hex()
-
-    @classmethod
-    def VALIDATE_INPUTS(s, latent):
-        if not ldm_patched.utils.path_utils.exists_annotated_filepath(latent):
-            return "Invalid latent file: {}".format(latent)
-        return True
-
-
-class CheckpointLoader:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "config_name": (ldm_patched.utils.path_utils.get_filename_list("configs"), ),
-                              "ckpt_name": (ldm_patched.utils.path_utils.get_filename_list("checkpoints"), )}}
-    RETURN_TYPES = ("MODEL", "CLIP", "VAE")
-    FUNCTION = "load_checkpoint"
-
-    CATEGORY = "advanced/loaders"
-
-    def load_checkpoint(self, config_name, ckpt_name, output_vae=True, output_clip=True):
-        config_path = ldm_patched.utils.path_utils.get_full_path("configs", config_name)
-        ckpt_path = ldm_patched.utils.path_utils.get_full_path("checkpoints", ckpt_name)
-        return ldm_patched.modules.sd.load_checkpoint(config_path, ckpt_path, output_vae=True, output_clip=True, embedding_directory=ldm_patched.utils.path_utils.get_folder_paths("embeddings"))
-
-class CheckpointLoaderSimple:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "ckpt_name": (ldm_patched.utils.path_utils.get_filename_list("checkpoints"), ),
-                             }}
-    RETURN_TYPES = ("MODEL", "CLIP", "VAE")
-    FUNCTION = "load_checkpoint"
-
-    CATEGORY = "loaders"
-
-    def load_checkpoint(self, ckpt_name, output_vae=True, output_clip=True):
-        ckpt_path = ldm_patched.utils.path_utils.get_full_path("checkpoints", ckpt_name)
-        out = ldm_patched.modules.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, embedding_directory=ldm_patched.utils.path_utils.get_folder_paths("embeddings"))
-        return out[:3]
-
-class DiffusersLoader:
-    @classmethod
-    def INPUT_TYPES(cls):
-        paths = []
-        for search_path in ldm_patched.utils.path_utils.get_folder_paths("diffusers"):
-            if os.path.exists(search_path):
-                for root, subdir, files in os.walk(search_path, followlinks=True):
-                    if "model_index.json" in files:
-                        paths.append(os.path.relpath(root, start=search_path))
-
-        return {"required": {"model_path": (paths,), }}
-    RETURN_TYPES = ("MODEL", "CLIP", "VAE")
-    FUNCTION = "load_checkpoint"
-
-    CATEGORY = "advanced/loaders/deprecated"
-
-    def load_checkpoint(self, model_path, output_vae=True, output_clip=True):
-        for search_path in ldm_patched.utils.path_utils.get_folder_paths("diffusers"):
-            if os.path.exists(search_path):
-                path = os.path.join(search_path, model_path)
-                if os.path.exists(path):
-                    model_path = path
-                    break
-
-        return ldm_patched.modules.diffusers_load.load_diffusers(model_path, output_vae=output_vae, output_clip=output_clip, embedding_directory=ldm_patched.utils.path_utils.get_folder_paths("embeddings"))
-
-
-class unCLIPCheckpointLoader:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "ckpt_name": (ldm_patched.utils.path_utils.get_filename_list("checkpoints"), ),
-                             }}
-    RETURN_TYPES = ("MODEL", "CLIP", "VAE", "CLIP_VISION")
-    FUNCTION = "load_checkpoint"
-
-    CATEGORY = "loaders"
-
-    def load_checkpoint(self, ckpt_name, output_vae=True, output_clip=True):
-        ckpt_path = ldm_patched.utils.path_utils.get_full_path("checkpoints", ckpt_name)
-        out = ldm_patched.modules.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, output_clipvision=True, embedding_directory=ldm_patched.utils.path_utils.get_folder_paths("embeddings"))
-        return out
-
-class CLIPSetLastLayer:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "clip": ("CLIP", ),
-                              "stop_at_clip_layer": ("INT", {"default": -1, "min": -24, "max": -1, "step": 1}),
-                              }}
-    RETURN_TYPES = ("CLIP",)
-    FUNCTION = "set_last_layer"
-
-    CATEGORY = "conditioning"
-
-    def set_last_layer(self, clip, stop_at_clip_layer):
-        clip = clip.clone()
-        clip.clip_layer(stop_at_clip_layer)
-        return (clip,)
-
-class LoraLoader:
-    def __init__(self):
-        self.loaded_lora = None
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model": ("MODEL",),
-                              "clip": ("CLIP", ),
-                              "lora_name": (ldm_patched.utils.path_utils.get_filename_list("loras"), ),
-                              "strength_model": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}),
-                              "strength_clip": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}),
-                              }}
-    RETURN_TYPES = ("MODEL", "CLIP")
-    FUNCTION = "load_lora"
-
-    CATEGORY = "loaders"
-
-    def load_lora(self, model, clip, lora_name, strength_model, strength_clip):
-        if strength_model == 0 and strength_clip == 0:
-            return (model, clip)
-
-        lora_path = ldm_patched.utils.path_utils.get_full_path("loras", lora_name)
-        lora = None
-        if self.loaded_lora is not None:
-            if self.loaded_lora[0] == lora_path:
-                lora = self.loaded_lora[1]
-            else:
-                temp = self.loaded_lora
-                self.loaded_lora = None
-                del temp
-
-        if lora is None:
-            lora = ldm_patched.modules.utils.load_torch_file(lora_path, safe_load=True)
-            self.loaded_lora = (lora_path, lora)
-
-        model_lora, clip_lora = ldm_patched.modules.sd.load_lora_for_models(model, clip, lora, strength_model, strength_clip)
-        return (model_lora, clip_lora)
-
-class LoraLoaderModelOnly(LoraLoader):
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model": ("MODEL",),
-                              "lora_name": (ldm_patched.utils.path_utils.get_filename_list("loras"), ),
-                              "strength_model": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}),
-                              }}
-    RETURN_TYPES = ("MODEL",)
-    FUNCTION = "load_lora_model_only"
-
-    def load_lora_model_only(self, model, lora_name, strength_model):
-        return (self.load_lora(model, None, lora_name, strength_model, 0)[0],)
-
-class VAELoader:
-    @staticmethod
-    def vae_list():
-        vaes = ldm_patched.utils.path_utils.get_filename_list("vae")
-        approx_vaes = ldm_patched.utils.path_utils.get_filename_list("vae_approx")
-        sdxl_taesd_enc = False
-        sdxl_taesd_dec = False
-        sd1_taesd_enc = False
-        sd1_taesd_dec = False
-
-        for v in approx_vaes:
-            if v.startswith("taesd_decoder."):
-                sd1_taesd_dec = True
-            elif v.startswith("taesd_encoder."):
-                sd1_taesd_enc = True
-            elif v.startswith("taesdxl_decoder."):
-                sdxl_taesd_dec = True
-            elif v.startswith("taesdxl_encoder."):
-                sdxl_taesd_enc = True
-        if sd1_taesd_dec and sd1_taesd_enc:
-            vaes.append("taesd")
-        if sdxl_taesd_dec and sdxl_taesd_enc:
-            vaes.append("taesdxl")
-        return vaes
-
-    @staticmethod
-    def load_taesd(name):
-        sd = {}
-        approx_vaes = ldm_patched.utils.path_utils.get_filename_list("vae_approx")
-
-        encoder = next(filter(lambda a: a.startswith("{}_encoder.".format(name)), approx_vaes))
-        decoder = next(filter(lambda a: a.startswith("{}_decoder.".format(name)), approx_vaes))
-
-        enc = ldm_patched.modules.utils.load_torch_file(ldm_patched.utils.path_utils.get_full_path("vae_approx", encoder))
-        for k in enc:
-            sd["taesd_encoder.{}".format(k)] = enc[k]
-
-        dec = ldm_patched.modules.utils.load_torch_file(ldm_patched.utils.path_utils.get_full_path("vae_approx", decoder))
-        for k in dec:
-            sd["taesd_decoder.{}".format(k)] = dec[k]
-
-        if name == "taesd":
-            sd["vae_scale"] = torch.tensor(0.18215)
-        elif name == "taesdxl":
-            sd["vae_scale"] = torch.tensor(0.13025)
-        return sd
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "vae_name": (s.vae_list(), )}}
-    RETURN_TYPES = ("VAE",)
-    FUNCTION = "load_vae"
-
-    CATEGORY = "loaders"
-
-    #TODO: scale factor?
-    def load_vae(self, vae_name):
-        if vae_name in ["taesd", "taesdxl"]:
-            sd = self.load_taesd(vae_name)
-        else:
-            vae_path = ldm_patched.utils.path_utils.get_full_path("vae", vae_name)
-            sd = ldm_patched.modules.utils.load_torch_file(vae_path)
-        vae = ldm_patched.modules.sd.VAE(sd=sd)
-        return (vae,)
-
-class ControlNetLoader:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "control_net_name": (ldm_patched.utils.path_utils.get_filename_list("controlnet"), )}}
-
-    RETURN_TYPES = ("CONTROL_NET",)
-    FUNCTION = "load_controlnet"
-
-    CATEGORY = "loaders"
-
-    def load_controlnet(self, control_net_name):
-        controlnet_path = ldm_patched.utils.path_utils.get_full_path("controlnet", control_net_name)
-        controlnet = ldm_patched.modules.controlnet.load_controlnet(controlnet_path)
-        return (controlnet,)
-
-class DiffControlNetLoader:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model": ("MODEL",),
-                              "control_net_name": (ldm_patched.utils.path_utils.get_filename_list("controlnet"), )}}
-
-    RETURN_TYPES = ("CONTROL_NET",)
-    FUNCTION = "load_controlnet"
-
-    CATEGORY = "loaders"
-
-    def load_controlnet(self, model, control_net_name):
-        controlnet_path = ldm_patched.utils.path_utils.get_full_path("controlnet", control_net_name)
-        controlnet = ldm_patched.modules.controlnet.load_controlnet(controlnet_path, model)
-        return (controlnet,)
-
-
-class ControlNetApply:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {"conditioning": ("CONDITIONING", ),
-                             "control_net": ("CONTROL_NET", ),
-                             "image": ("IMAGE", ),
-                             "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01})
-                             }}
-    RETURN_TYPES = ("CONDITIONING",)
-    FUNCTION = "apply_controlnet"
-
-    CATEGORY = "conditioning"
-
-    def apply_controlnet(self, conditioning, control_net, image, strength):
-        if strength == 0:
-            return (conditioning, )
-
-        c = []
-        control_hint = image.movedim(-1,1)
-        for t in conditioning:
-            n = [t[0], t[1].copy()]
-            c_net = control_net.copy().set_cond_hint(control_hint, strength)
-            if 'control' in t[1]:
-                c_net.set_previous_controlnet(t[1]['control'])
-            n[1]['control'] = c_net
-            n[1]['control_apply_to_uncond'] = True
-            c.append(n)
-        return (c, )
-
-
-class ControlNetApplyAdvanced:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {"positive": ("CONDITIONING", ),
-                             "negative": ("CONDITIONING", ),
-                             "control_net": ("CONTROL_NET", ),
-                             "image": ("IMAGE", ),
-                             "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
-                             "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
-                             "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001})
-                             }}
-
-    RETURN_TYPES = ("CONDITIONING","CONDITIONING")
-    RETURN_NAMES = ("positive", "negative")
-    FUNCTION = "apply_controlnet"
-
-    CATEGORY = "conditioning"
-
-    def apply_controlnet(self, positive, negative, control_net, image, strength, start_percent, end_percent):
-        if strength == 0:
-            return (positive, negative)
-
-        control_hint = image.movedim(-1,1)
-        cnets = {}
-
-        out = []
-        for conditioning in [positive, negative]:
-            c = []
-            for t in conditioning:
-                d = t[1].copy()
-
-                prev_cnet = d.get('control', None)
-                if prev_cnet in cnets:
-                    c_net = cnets[prev_cnet]
-                else:
-                    c_net = control_net.copy().set_cond_hint(control_hint, strength, (start_percent, end_percent))
-                    c_net.set_previous_controlnet(prev_cnet)
-                    cnets[prev_cnet] = c_net
-
-                d['control'] = c_net
-                d['control_apply_to_uncond'] = False
-                n = [t[0], d]
-                c.append(n)
-            out.append(c)
-        return (out[0], out[1])
-
-
-class UNETLoader:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "unet_name": (ldm_patched.utils.path_utils.get_filename_list("unet"), ),
-                             }}
-    RETURN_TYPES = ("MODEL",)
-    FUNCTION = "load_unet"
-
-    CATEGORY = "advanced/loaders"
-
-    def load_unet(self, unet_name):
-        unet_path = ldm_patched.utils.path_utils.get_full_path("unet", unet_name)
-        model = ldm_patched.modules.sd.load_unet(unet_path)
-        return (model,)
-
-class CLIPLoader:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "clip_name": (ldm_patched.utils.path_utils.get_filename_list("clip"), ),
-                             }}
-    RETURN_TYPES = ("CLIP",)
-    FUNCTION = "load_clip"
-
-    CATEGORY = "advanced/loaders"
-
-    def load_clip(self, clip_name):
-        clip_path = ldm_patched.utils.path_utils.get_full_path("clip", clip_name)
-        clip = ldm_patched.modules.sd.load_clip(ckpt_paths=[clip_path], embedding_directory=ldm_patched.utils.path_utils.get_folder_paths("embeddings"))
-        return (clip,)
-
-class DualCLIPLoader:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "clip_name1": (ldm_patched.utils.path_utils.get_filename_list("clip"), ), "clip_name2": (ldm_patched.utils.path_utils.get_filename_list("clip"), ),
-                             }}
-    RETURN_TYPES = ("CLIP",)
-    FUNCTION = "load_clip"
-
-    CATEGORY = "advanced/loaders"
-
-    def load_clip(self, clip_name1, clip_name2):
-        clip_path1 = ldm_patched.utils.path_utils.get_full_path("clip", clip_name1)
-        clip_path2 = ldm_patched.utils.path_utils.get_full_path("clip", clip_name2)
-        clip = ldm_patched.modules.sd.load_clip(ckpt_paths=[clip_path1, clip_path2], embedding_directory=ldm_patched.utils.path_utils.get_folder_paths("embeddings"))
-        return (clip,)
-
-class CLIPVisionLoader:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "clip_name": (ldm_patched.utils.path_utils.get_filename_list("clip_vision"), ),
-                             }}
-    RETURN_TYPES = ("CLIP_VISION",)
-    FUNCTION = "load_clip"
-
-    CATEGORY = "loaders"
-
-    def load_clip(self, clip_name):
-        clip_path = ldm_patched.utils.path_utils.get_full_path("clip_vision", clip_name)
-        clip_vision = ldm_patched.modules.clip_vision.load(clip_path)
-        return (clip_vision,)
-
-class CLIPVisionEncode:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "clip_vision": ("CLIP_VISION",),
-                              "image": ("IMAGE",)
-                             }}
-    RETURN_TYPES = ("CLIP_VISION_OUTPUT",)
-    FUNCTION = "encode"
-
-    CATEGORY = "conditioning"
-
-    def encode(self, clip_vision, image):
-        output = clip_vision.encode_image(image)
-        return (output,)
-
-class StyleModelLoader:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "style_model_name": (ldm_patched.utils.path_utils.get_filename_list("style_models"), )}}
-
-    RETURN_TYPES = ("STYLE_MODEL",)
-    FUNCTION = "load_style_model"
-
-    CATEGORY = "loaders"
-
-    def load_style_model(self, style_model_name):
-        style_model_path = ldm_patched.utils.path_utils.get_full_path("style_models", style_model_name)
-        style_model = ldm_patched.modules.sd.load_style_model(style_model_path)
-        return (style_model,)
-
-
-class StyleModelApply:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {"conditioning": ("CONDITIONING", ),
-                             "style_model": ("STYLE_MODEL", ),
-                             "clip_vision_output": ("CLIP_VISION_OUTPUT", ),
-                             }}
-    RETURN_TYPES = ("CONDITIONING",)
-    FUNCTION = "apply_stylemodel"
-
-    CATEGORY = "conditioning/style_model"
-
-    def apply_stylemodel(self, clip_vision_output, style_model, conditioning):
-        cond = style_model.get_cond(clip_vision_output).flatten(start_dim=0, end_dim=1).unsqueeze(dim=0)
-        c = []
-        for t in conditioning:
-            n = [torch.cat((t[0], cond), dim=1), t[1].copy()]
-            c.append(n)
-        return (c, )
-
-class unCLIPConditioning:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {"conditioning": ("CONDITIONING", ),
-                             "clip_vision_output": ("CLIP_VISION_OUTPUT", ),
-                             "strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
-                             "noise_augmentation": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}),
-                             }}
-    RETURN_TYPES = ("CONDITIONING",)
-    FUNCTION = "apply_adm"
-
-    CATEGORY = "conditioning"
-
-    def apply_adm(self, conditioning, clip_vision_output, strength, noise_augmentation):
-        if strength == 0:
-            return (conditioning, )
-
-        c = []
-        for t in conditioning:
-            o = t[1].copy()
-            x = {"clip_vision_output": clip_vision_output, "strength": strength, "noise_augmentation": noise_augmentation}
-            if "unclip_conditioning" in o:
-                o["unclip_conditioning"] = o["unclip_conditioning"][:] + [x]
-            else:
-                o["unclip_conditioning"] = [x]
-            n = [t[0], o]
-            c.append(n)
-        return (c, )
-
-class GLIGENLoader:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "gligen_name": (ldm_patched.utils.path_utils.get_filename_list("gligen"), )}}
-
-    RETURN_TYPES = ("GLIGEN",)
-    FUNCTION = "load_gligen"
-
-    CATEGORY = "loaders"
-
-    def load_gligen(self, gligen_name):
-        gligen_path = ldm_patched.utils.path_utils.get_full_path("gligen", gligen_name)
-        gligen = ldm_patched.modules.sd.load_gligen(gligen_path)
-        return (gligen,)
-
-class GLIGENTextBoxApply:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {"conditioning_to": ("CONDITIONING", ),
-                              "clip": ("CLIP", ),
-                              "gligen_textbox_model": ("GLIGEN", ),
-                              "text": ("STRING", {"multiline": True}),
-                              "width": ("INT", {"default": 64, "min": 8, "max": MAX_RESOLUTION, "step": 8}),
-                              "height": ("INT", {"default": 64, "min": 8, "max": MAX_RESOLUTION, "step": 8}),
-                              "x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
-                              "y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
-                             }}
-    RETURN_TYPES = ("CONDITIONING",)
-    FUNCTION = "append"
-
-    CATEGORY = "conditioning/gligen"
-
-    def append(self, conditioning_to, clip, gligen_textbox_model, text, width, height, x, y):
-        c = []
-        cond, cond_pooled = clip.encode_from_tokens(clip.tokenize(text), return_pooled=True)
-        for t in conditioning_to:
-            n = [t[0], t[1].copy()]
-            position_params = [(cond_pooled, height // 8, width // 8, y // 8, x // 8)]
-            prev = []
-            if "gligen" in n[1]:
-                prev = n[1]['gligen'][2]
-
-            n[1]['gligen'] = ("position", gligen_textbox_model, prev + position_params)
-            c.append(n)
-        return (c, )
-
-class EmptyLatentImage:
-    def __init__(self):
-        self.device = ldm_patched.modules.model_management.intermediate_device()
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "width": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8}),
-                              "height": ("INT", {"default": 512, "min": 16, "max": MAX_RESOLUTION, "step": 8}),
-                              "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096})}}
-    RETURN_TYPES = ("LATENT",)
-    FUNCTION = "generate"
-
-    CATEGORY = "latent"
-
-    def generate(self, width, height, batch_size=1):
-        latent = torch.zeros([batch_size, 4, height // 8, width // 8], device=self.device)
-        return ({"samples":latent}, )
-
-
-class LatentFromBatch:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "samples": ("LATENT",),
-                              "batch_index": ("INT", {"default": 0, "min": 0, "max": 63}),
-                              "length": ("INT", {"default": 1, "min": 1, "max": 64}),
-                              }}
-    RETURN_TYPES = ("LATENT",)
-    FUNCTION = "frombatch"
-
-    CATEGORY = "latent/batch"
-
-    def frombatch(self, samples, batch_index, length):
-        s = samples.copy()
-        s_in = samples["samples"]
-        batch_index = min(s_in.shape[0] - 1, batch_index)
-        length = min(s_in.shape[0] - batch_index, length)
-        s["samples"] = s_in[batch_index:batch_index + length].clone()
-        if "noise_mask" in samples:
-            masks = samples["noise_mask"]
-            if masks.shape[0] == 1:
-                s["noise_mask"] = masks.clone()
-            else:
-                if masks.shape[0] < s_in.shape[0]:
-                    masks = masks.repeat(math.ceil(s_in.shape[0] / masks.shape[0]), 1, 1, 1)[:s_in.shape[0]]
-                s["noise_mask"] = masks[batch_index:batch_index + length].clone()
-        if "batch_index" not in s:
-            s["batch_index"] = [x for x in range(batch_index, batch_index+length)]
-        else:
-            s["batch_index"] = samples["batch_index"][batch_index:batch_index + length]
-        return (s,)
-    
-class RepeatLatentBatch:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "samples": ("LATENT",),
-                              "amount": ("INT", {"default": 1, "min": 1, "max": 64}),
-                              }}
-    RETURN_TYPES = ("LATENT",)
-    FUNCTION = "repeat"
-
-    CATEGORY = "latent/batch"
-
-    def repeat(self, samples, amount):
-        s = samples.copy()
-        s_in = samples["samples"]
-        
-        s["samples"] = s_in.repeat((amount, 1,1,1))
-        if "noise_mask" in samples and samples["noise_mask"].shape[0] > 1:
-            masks = samples["noise_mask"]
-            if masks.shape[0] < s_in.shape[0]:
-                masks = masks.repeat(math.ceil(s_in.shape[0] / masks.shape[0]), 1, 1, 1)[:s_in.shape[0]]
-            s["noise_mask"] = samples["noise_mask"].repeat((amount, 1,1,1))
-        if "batch_index" in s:
-            offset = max(s["batch_index"]) - min(s["batch_index"]) + 1
-            s["batch_index"] = s["batch_index"] + [x + (i * offset) for i in range(1, amount) for x in s["batch_index"]]
-        return (s,)
-
-class LatentUpscale:
-    upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "bislerp"]
-    crop_methods = ["disabled", "center"]
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "samples": ("LATENT",), "upscale_method": (s.upscale_methods,),
-                              "width": ("INT", {"default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
-                              "height": ("INT", {"default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
-                              "crop": (s.crop_methods,)}}
-    RETURN_TYPES = ("LATENT",)
-    FUNCTION = "upscale"
-
-    CATEGORY = "latent"
-
-    def upscale(self, samples, upscale_method, width, height, crop):
-        if width == 0 and height == 0:
-            s = samples
-        else:
-            s = samples.copy()
-
-            if width == 0:
-                height = max(64, height)
-                width = max(64, round(samples["samples"].shape[3] * height / samples["samples"].shape[2]))
-            elif height == 0:
-                width = max(64, width)
-                height = max(64, round(samples["samples"].shape[2] * width / samples["samples"].shape[3]))
-            else:
-                width = max(64, width)
-                height = max(64, height)
-
-            s["samples"] = ldm_patched.modules.utils.common_upscale(samples["samples"], width // 8, height // 8, upscale_method, crop)
-        return (s,)
-
-class LatentUpscaleBy:
-    upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "bislerp"]
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "samples": ("LATENT",), "upscale_method": (s.upscale_methods,),
-                              "scale_by": ("FLOAT", {"default": 1.5, "min": 0.01, "max": 8.0, "step": 0.01}),}}
-    RETURN_TYPES = ("LATENT",)
-    FUNCTION = "upscale"
-
-    CATEGORY = "latent"
-
-    def upscale(self, samples, upscale_method, scale_by):
-        s = samples.copy()
-        width = round(samples["samples"].shape[3] * scale_by)
-        height = round(samples["samples"].shape[2] * scale_by)
-        s["samples"] = ldm_patched.modules.utils.common_upscale(samples["samples"], width, height, upscale_method, "disabled")
-        return (s,)
-
-class LatentRotate:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "samples": ("LATENT",),
-                              "rotation": (["none", "90 degrees", "180 degrees", "270 degrees"],),
-                              }}
-    RETURN_TYPES = ("LATENT",)
-    FUNCTION = "rotate"
-
-    CATEGORY = "latent/transform"
-
-    def rotate(self, samples, rotation):
-        s = samples.copy()
-        rotate_by = 0
-        if rotation.startswith("90"):
-            rotate_by = 1
-        elif rotation.startswith("180"):
-            rotate_by = 2
-        elif rotation.startswith("270"):
-            rotate_by = 3
-
-        s["samples"] = torch.rot90(samples["samples"], k=rotate_by, dims=[3, 2])
-        return (s,)
-
-class LatentFlip:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "samples": ("LATENT",),
-                              "flip_method": (["x-axis: vertically", "y-axis: horizontally"],),
-                              }}
-    RETURN_TYPES = ("LATENT",)
-    FUNCTION = "flip"
-
-    CATEGORY = "latent/transform"
-
-    def flip(self, samples, flip_method):
-        s = samples.copy()
-        if flip_method.startswith("x"):
-            s["samples"] = torch.flip(samples["samples"], dims=[2])
-        elif flip_method.startswith("y"):
-            s["samples"] = torch.flip(samples["samples"], dims=[3])
-
-        return (s,)
-
-class LatentComposite:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "samples_to": ("LATENT",),
-                              "samples_from": ("LATENT",),
-                              "x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
-                              "y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
-                              "feather": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
-                              }}
-    RETURN_TYPES = ("LATENT",)
-    FUNCTION = "composite"
-
-    CATEGORY = "latent"
-
-    def composite(self, samples_to, samples_from, x, y, composite_method="normal", feather=0):
-        x =  x // 8
-        y = y // 8
-        feather = feather // 8
-        samples_out = samples_to.copy()
-        s = samples_to["samples"].clone()
-        samples_to = samples_to["samples"]
-        samples_from = samples_from["samples"]
-        if feather == 0:
-            s[:,:,y:y+samples_from.shape[2],x:x+samples_from.shape[3]] = samples_from[:,:,:samples_to.shape[2] - y, :samples_to.shape[3] - x]
-        else:
-            samples_from = samples_from[:,:,:samples_to.shape[2] - y, :samples_to.shape[3] - x]
-            mask = torch.ones_like(samples_from)
-            for t in range(feather):
-                if y != 0:
-                    mask[:,:,t:1+t,:] *= ((1.0/feather) * (t + 1))
-
-                if y + samples_from.shape[2] < samples_to.shape[2]:
-                    mask[:,:,mask.shape[2] -1 -t: mask.shape[2]-t,:] *= ((1.0/feather) * (t + 1))
-                if x != 0:
-                    mask[:,:,:,t:1+t] *= ((1.0/feather) * (t + 1))
-                if x + samples_from.shape[3] < samples_to.shape[3]:
-                    mask[:,:,:,mask.shape[3]- 1 - t: mask.shape[3]- t] *= ((1.0/feather) * (t + 1))
-            rev_mask = torch.ones_like(mask) - mask
-            s[:,:,y:y+samples_from.shape[2],x:x+samples_from.shape[3]] = samples_from[:,:,:samples_to.shape[2] - y, :samples_to.shape[3] - x] * mask + s[:,:,y:y+samples_from.shape[2],x:x+samples_from.shape[3]] * rev_mask
-        samples_out["samples"] = s
-        return (samples_out,)
-
-class LatentBlend:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {
-            "samples1": ("LATENT",),
-            "samples2": ("LATENT",),
-            "blend_factor": ("FLOAT", {
-                "default": 0.5,
-                "min": 0,
-                "max": 1,
-                "step": 0.01
-            }),
-        }}
-
-    RETURN_TYPES = ("LATENT",)
-    FUNCTION = "blend"
-
-    CATEGORY = "_for_testing"
-
-    def blend(self, samples1, samples2, blend_factor:float, blend_mode: str="normal"):
-
-        samples_out = samples1.copy()
-        samples1 = samples1["samples"]
-        samples2 = samples2["samples"]
-
-        if samples1.shape != samples2.shape:
-            samples2.permute(0, 3, 1, 2)
-            samples2 = ldm_patched.modules.utils.common_upscale(samples2, samples1.shape[3], samples1.shape[2], 'bicubic', crop='center')
-            samples2.permute(0, 2, 3, 1)
-
-        samples_blended = self.blend_mode(samples1, samples2, blend_mode)
-        samples_blended = samples1 * blend_factor + samples_blended * (1 - blend_factor)
-        samples_out["samples"] = samples_blended
-        return (samples_out,)
-
-    def blend_mode(self, img1, img2, mode):
-        if mode == "normal":
-            return img2
-        else:
-            raise ValueError(f"Unsupported blend mode: {mode}")
-
-class LatentCrop:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "samples": ("LATENT",),
-                              "width": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 8}),
-                              "height": ("INT", {"default": 512, "min": 64, "max": MAX_RESOLUTION, "step": 8}),
-                              "x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
-                              "y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
-                              }}
-    RETURN_TYPES = ("LATENT",)
-    FUNCTION = "crop"
-
-    CATEGORY = "latent/transform"
-
-    def crop(self, samples, width, height, x, y):
-        s = samples.copy()
-        samples = samples['samples']
-        x =  x // 8
-        y = y // 8
-
-        #enfonce minimum size of 64
-        if x > (samples.shape[3] - 8):
-            x = samples.shape[3] - 8
-        if y > (samples.shape[2] - 8):
-            y = samples.shape[2] - 8
-
-        new_height = height // 8
-        new_width = width // 8
-        to_x = new_width + x
-        to_y = new_height + y
-        s['samples'] = samples[:,:,y:to_y, x:to_x]
-        return (s,)
-
-class SetLatentNoiseMask:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "samples": ("LATENT",),
-                              "mask": ("MASK",),
-                              }}
-    RETURN_TYPES = ("LATENT",)
-    FUNCTION = "set_mask"
-
-    CATEGORY = "latent/inpaint"
-
-    def set_mask(self, samples, mask):
-        s = samples.copy()
-        s["noise_mask"] = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1]))
-        return (s,)
-
-def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, denoise=1.0, disable_noise=False, start_step=None, last_step=None, force_full_denoise=False):
-    latent_image = latent["samples"]
-    if disable_noise:
-        noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
-    else:
-        batch_inds = latent["batch_index"] if "batch_index" in latent else None
-        noise = ldm_patched.modules.sample.prepare_noise(latent_image, seed, batch_inds)
-
-    noise_mask = None
-    if "noise_mask" in latent:
-        noise_mask = latent["noise_mask"]
-
-    callback = ldm_patched.utils.latent_visualization.prepare_callback(model, steps)
-    disable_pbar = not ldm_patched.modules.utils.PROGRESS_BAR_ENABLED
-    samples = ldm_patched.modules.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,
-                                  denoise=denoise, disable_noise=disable_noise, start_step=start_step, last_step=last_step,
-                                  force_full_denoise=force_full_denoise, noise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed)
-    out = latent.copy()
-    out["samples"] = samples
-    return (out, )
-
-class KSampler:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"model": ("MODEL",),
-                    "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
-                    "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
-                    "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
-                    "sampler_name": (ldm_patched.modules.samplers.KSampler.SAMPLERS, ),
-                    "scheduler": (ldm_patched.modules.samplers.KSampler.SCHEDULERS, ),
-                    "positive": ("CONDITIONING", ),
-                    "negative": ("CONDITIONING", ),
-                    "latent_image": ("LATENT", ),
-                    "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
-                     }
-                }
-
-    RETURN_TYPES = ("LATENT",)
-    FUNCTION = "sample"
-
-    CATEGORY = "sampling"
-
-    def sample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0):
-        return common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise)
-
-class KSamplerAdvanced:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"model": ("MODEL",),
-                    "add_noise": (["enable", "disable"], ),
-                    "noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
-                    "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
-                    "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
-                    "sampler_name": (ldm_patched.modules.samplers.KSampler.SAMPLERS, ),
-                    "scheduler": (ldm_patched.modules.samplers.KSampler.SCHEDULERS, ),
-                    "positive": ("CONDITIONING", ),
-                    "negative": ("CONDITIONING", ),
-                    "latent_image": ("LATENT", ),
-                    "start_at_step": ("INT", {"default": 0, "min": 0, "max": 10000}),
-                    "end_at_step": ("INT", {"default": 10000, "min": 0, "max": 10000}),
-                    "return_with_leftover_noise": (["disable", "enable"], ),
-                     }
-                }
-
-    RETURN_TYPES = ("LATENT",)
-    FUNCTION = "sample"
-
-    CATEGORY = "sampling"
-
-    def sample(self, model, add_noise, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, start_at_step, end_at_step, return_with_leftover_noise, denoise=1.0):
-        force_full_denoise = True
-        if return_with_leftover_noise == "enable":
-            force_full_denoise = False
-        disable_noise = False
-        if add_noise == "disable":
-            disable_noise = True
-        return common_ksampler(model, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise)
-
-class SaveImage:
-    def __init__(self):
-        self.output_dir = ldm_patched.utils.path_utils.get_output_directory()
-        self.type = "output"
-        self.prefix_append = ""
-        self.compress_level = 4
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": 
-                    {"images": ("IMAGE", ),
-                     "filename_prefix": ("STRING", {"default": "ldm_patched"})},
-                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
-                }
-
-    RETURN_TYPES = ()
-    FUNCTION = "save_images"
-
-    OUTPUT_NODE = True
-
-    CATEGORY = "image"
-
-    def save_images(self, images, filename_prefix="ldm_patched", prompt=None, extra_pnginfo=None):
-        filename_prefix += self.prefix_append
-        full_output_folder, filename, counter, subfolder, filename_prefix = ldm_patched.utils.path_utils.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0])
-        results = list()
-        for image in images:
-            i = 255. * image.cpu().numpy()
-            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
-            metadata = None
-            if not args.disable_server_info:
-                metadata = PngInfo()
-                if prompt is not None:
-                    metadata.add_text("prompt", json.dumps(prompt))
-                if extra_pnginfo is not None:
-                    for x in extra_pnginfo:
-                        metadata.add_text(x, json.dumps(extra_pnginfo[x]))
-
-            file = f"{filename}_{counter:05}_.png"
-            img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=self.compress_level)
-            results.append({
-                "filename": file,
-                "subfolder": subfolder,
-                "type": self.type
-            })
-            counter += 1
-
-        return { "ui": { "images": results } }
-
-class PreviewImage(SaveImage):
-    def __init__(self):
-        self.output_dir = ldm_patched.utils.path_utils.get_temp_directory()
-        self.type = "temp"
-        self.prefix_append = "_temp_" + ''.join(random.choice("abcdefghijklmnopqrstupvxyz") for x in range(5))
-        self.compress_level = 1
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"images": ("IMAGE", ), },
-                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
-                }
-
-class LoadImage:
-    @classmethod
-    def INPUT_TYPES(s):
-        input_dir = ldm_patched.utils.path_utils.get_input_directory()
-        files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
-        return {"required":
-                    {"image": (sorted(files), {"image_upload": True})},
-                }
-
-    CATEGORY = "image"
-
-    RETURN_TYPES = ("IMAGE", "MASK")
-    FUNCTION = "load_image"
-    def load_image(self, image):
-        image_path = ldm_patched.utils.path_utils.get_annotated_filepath(image)
-        img = Image.open(image_path)
-        output_images = []
-        output_masks = []
-        for i in ImageSequence.Iterator(img):
-            i = ImageOps.exif_transpose(i)
-            if i.mode == 'I':
-                i = i.point(lambda i: i * (1 / 255))
-            image = i.convert("RGB")
-            image = np.array(image).astype(np.float32) / 255.0
-            image = torch.from_numpy(image)[None,]
-            if 'A' in i.getbands():
-                mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
-                mask = 1. - torch.from_numpy(mask)
-            else:
-                mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
-            output_images.append(image)
-            output_masks.append(mask.unsqueeze(0))
-
-        if len(output_images) > 1:
-            output_image = torch.cat(output_images, dim=0)
-            output_mask = torch.cat(output_masks, dim=0)
-        else:
-            output_image = output_images[0]
-            output_mask = output_masks[0]
-
-        return (output_image, output_mask)
-
-    @classmethod
-    def IS_CHANGED(s, image):
-        image_path = ldm_patched.utils.path_utils.get_annotated_filepath(image)
-        m = hashlib.sha256()
-        with open(image_path, 'rb') as f:
-            m.update(f.read())
-        return m.digest().hex()
-
-    @classmethod
-    def VALIDATE_INPUTS(s, image):
-        if not ldm_patched.utils.path_utils.exists_annotated_filepath(image):
-            return "Invalid image file: {}".format(image)
-
-        return True
-
-class LoadImageMask:
-    _color_channels = ["alpha", "red", "green", "blue"]
-    @classmethod
-    def INPUT_TYPES(s):
-        input_dir = ldm_patched.utils.path_utils.get_input_directory()
-        files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
-        return {"required":
-                    {"image": (sorted(files), {"image_upload": True}),
-                     "channel": (s._color_channels, ), }
-                }
-
-    CATEGORY = "mask"
-
-    RETURN_TYPES = ("MASK",)
-    FUNCTION = "load_image"
-    def load_image(self, image, channel):
-        image_path = ldm_patched.utils.path_utils.get_annotated_filepath(image)
-        i = Image.open(image_path)
-        i = ImageOps.exif_transpose(i)
-        if i.getbands() != ("R", "G", "B", "A"):
-            if i.mode == 'I':
-                i = i.point(lambda i: i * (1 / 255))
-            i = i.convert("RGBA")
-        mask = None
-        c = channel[0].upper()
-        if c in i.getbands():
-            mask = np.array(i.getchannel(c)).astype(np.float32) / 255.0
-            mask = torch.from_numpy(mask)
-            if c == 'A':
-                mask = 1. - mask
-        else:
-            mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
-        return (mask.unsqueeze(0),)
-
-    @classmethod
-    def IS_CHANGED(s, image, channel):
-        image_path = ldm_patched.utils.path_utils.get_annotated_filepath(image)
-        m = hashlib.sha256()
-        with open(image_path, 'rb') as f:
-            m.update(f.read())
-        return m.digest().hex()
-
-    @classmethod
-    def VALIDATE_INPUTS(s, image):
-        if not ldm_patched.utils.path_utils.exists_annotated_filepath(image):
-            return "Invalid image file: {}".format(image)
-
-        return True
-
-class ImageScale:
-    upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "lanczos"]
-    crop_methods = ["disabled", "center"]
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "image": ("IMAGE",), "upscale_method": (s.upscale_methods,),
-                              "width": ("INT", {"default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
-                              "height": ("INT", {"default": 512, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
-                              "crop": (s.crop_methods,)}}
-    RETURN_TYPES = ("IMAGE",)
-    FUNCTION = "upscale"
-
-    CATEGORY = "image/upscaling"
-
-    def upscale(self, image, upscale_method, width, height, crop):
-        if width == 0 and height == 0:
-            s = image
-        else:
-            samples = image.movedim(-1,1)
-
-            if width == 0:
-                width = max(1, round(samples.shape[3] * height / samples.shape[2]))
-            elif height == 0:
-                height = max(1, round(samples.shape[2] * width / samples.shape[3]))
-
-            s = ldm_patched.modules.utils.common_upscale(samples, width, height, upscale_method, crop)
-            s = s.movedim(1,-1)
-        return (s,)
-
-class ImageScaleBy:
-    upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "lanczos"]
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "image": ("IMAGE",), "upscale_method": (s.upscale_methods,),
-                              "scale_by": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 8.0, "step": 0.01}),}}
-    RETURN_TYPES = ("IMAGE",)
-    FUNCTION = "upscale"
-
-    CATEGORY = "image/upscaling"
-
-    def upscale(self, image, upscale_method, scale_by):
-        samples = image.movedim(-1,1)
-        width = round(samples.shape[3] * scale_by)
-        height = round(samples.shape[2] * scale_by)
-        s = ldm_patched.modules.utils.common_upscale(samples, width, height, upscale_method, "disabled")
-        s = s.movedim(1,-1)
-        return (s,)
-
-class ImageInvert:
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "image": ("IMAGE",)}}
-
-    RETURN_TYPES = ("IMAGE",)
-    FUNCTION = "invert"
-
-    CATEGORY = "image"
-
-    def invert(self, image):
-        s = 1.0 - image
-        return (s,)
-
-class ImageBatch:
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "image1": ("IMAGE",), "image2": ("IMAGE",)}}
-
-    RETURN_TYPES = ("IMAGE",)
-    FUNCTION = "batch"
-
-    CATEGORY = "image"
-
-    def batch(self, image1, image2):
-        if image1.shape[1:] != image2.shape[1:]:
-            image2 = ldm_patched.modules.utils.common_upscale(image2.movedim(-1,1), image1.shape[2], image1.shape[1], "bilinear", "center").movedim(1,-1)
-        s = torch.cat((image1, image2), dim=0)
-        return (s,)
-
-class EmptyImage:
-    def __init__(self, device="cpu"):
-        self.device = device
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "width": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}),
-                              "height": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}),
-                              "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
-                              "color": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFF, "step": 1, "display": "color"}),
-                              }}
-    RETURN_TYPES = ("IMAGE",)
-    FUNCTION = "generate"
-
-    CATEGORY = "image"
-
-    def generate(self, width, height, batch_size=1, color=0):
-        r = torch.full([batch_size, height, width, 1], ((color >> 16) & 0xFF) / 0xFF)
-        g = torch.full([batch_size, height, width, 1], ((color >> 8) & 0xFF) / 0xFF)
-        b = torch.full([batch_size, height, width, 1], ((color) & 0xFF) / 0xFF)
-        return (torch.cat((r, g, b), dim=-1), )
-
-class ImagePadForOutpaint:
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "image": ("IMAGE",),
-                "left": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
-                "top": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
-                "right": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
-                "bottom": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
-                "feathering": ("INT", {"default": 40, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
-            }
-        }
-
-    RETURN_TYPES = ("IMAGE", "MASK")
-    FUNCTION = "expand_image"
-
-    CATEGORY = "image"
-
-    def expand_image(self, image, left, top, right, bottom, feathering):
-        d1, d2, d3, d4 = image.size()
-
-        new_image = torch.ones(
-            (d1, d2 + top + bottom, d3 + left + right, d4),
-            dtype=torch.float32,
-        ) * 0.5
-
-        new_image[:, top:top + d2, left:left + d3, :] = image
-
-        mask = torch.ones(
-            (d2 + top + bottom, d3 + left + right),
-            dtype=torch.float32,
-        )
-
-        t = torch.zeros(
-            (d2, d3),
-            dtype=torch.float32
-        )
-
-        if feathering > 0 and feathering * 2 < d2 and feathering * 2 < d3:
-
-            for i in range(d2):
-                for j in range(d3):
-                    dt = i if top != 0 else d2
-                    db = d2 - i if bottom != 0 else d2
-
-                    dl = j if left != 0 else d3
-                    dr = d3 - j if right != 0 else d3
-
-                    d = min(dt, db, dl, dr)
-
-                    if d >= feathering:
-                        continue
-
-                    v = (feathering - d) / feathering
-
-                    t[i, j] = v * v
-
-        mask[top:top + d2, left:left + d3] = t
-
-        return (new_image, mask)
-
-
-NODE_CLASS_MAPPINGS = {
-    "KSampler": KSampler,
-    "CheckpointLoaderSimple": CheckpointLoaderSimple,
-    "CLIPTextEncode": CLIPTextEncode,
-    "CLIPSetLastLayer": CLIPSetLastLayer,
-    "VAEDecode": VAEDecode,
-    "VAEEncode": VAEEncode,
-    "VAEEncodeForInpaint": VAEEncodeForInpaint,
-    "VAELoader": VAELoader,
-    "EmptyLatentImage": EmptyLatentImage,
-    "LatentUpscale": LatentUpscale,
-    "LatentUpscaleBy": LatentUpscaleBy,
-    "LatentFromBatch": LatentFromBatch,
-    "RepeatLatentBatch": RepeatLatentBatch,
-    "SaveImage": SaveImage,
-    "PreviewImage": PreviewImage,
-    "LoadImage": LoadImage,
-    "LoadImageMask": LoadImageMask,
-    "ImageScale": ImageScale,
-    "ImageScaleBy": ImageScaleBy,
-    "ImageInvert": ImageInvert,
-    "ImageBatch": ImageBatch,
-    "ImagePadForOutpaint": ImagePadForOutpaint,
-    "EmptyImage": EmptyImage,
-    "ConditioningAverage": ConditioningAverage ,
-    "ConditioningCombine": ConditioningCombine,
-    "ConditioningConcat": ConditioningConcat,
-    "ConditioningSetArea": ConditioningSetArea,
-    "ConditioningSetAreaPercentage": ConditioningSetAreaPercentage,
-    "ConditioningSetMask": ConditioningSetMask,
-    "KSamplerAdvanced": KSamplerAdvanced,
-    "SetLatentNoiseMask": SetLatentNoiseMask,
-    "LatentComposite": LatentComposite,
-    "LatentBlend": LatentBlend,
-    "LatentRotate": LatentRotate,
-    "LatentFlip": LatentFlip,
-    "LatentCrop": LatentCrop,
-    "LoraLoader": LoraLoader,
-    "CLIPLoader": CLIPLoader,
-    "UNETLoader": UNETLoader,
-    "DualCLIPLoader": DualCLIPLoader,
-    "CLIPVisionEncode": CLIPVisionEncode,
-    "StyleModelApply": StyleModelApply,
-    "unCLIPConditioning": unCLIPConditioning,
-    "ControlNetApply": ControlNetApply,
-    "ControlNetApplyAdvanced": ControlNetApplyAdvanced,
-    "ControlNetLoader": ControlNetLoader,
-    "DiffControlNetLoader": DiffControlNetLoader,
-    "StyleModelLoader": StyleModelLoader,
-    "CLIPVisionLoader": CLIPVisionLoader,
-    "VAEDecodeTiled": VAEDecodeTiled,
-    "VAEEncodeTiled": VAEEncodeTiled,
-    "unCLIPCheckpointLoader": unCLIPCheckpointLoader,
-    "GLIGENLoader": GLIGENLoader,
-    "GLIGENTextBoxApply": GLIGENTextBoxApply,
-    "InpaintModelConditioning": InpaintModelConditioning,
-
-    "CheckpointLoader": CheckpointLoader,
-    "DiffusersLoader": DiffusersLoader,
-
-    "LoadLatent": LoadLatent,
-    "SaveLatent": SaveLatent,
-
-    "ConditioningZeroOut": ConditioningZeroOut,
-    "ConditioningSetTimestepRange": ConditioningSetTimestepRange,
-    "LoraLoaderModelOnly": LoraLoaderModelOnly,
-}
-
-NODE_DISPLAY_NAME_MAPPINGS = {
-    # Sampling
-    "KSampler": "KSampler",
-    "KSamplerAdvanced": "KSampler (Advanced)",
-    # Loaders
-    "CheckpointLoader": "Load Checkpoint With Config (DEPRECATED)",
-    "CheckpointLoaderSimple": "Load Checkpoint",
-    "VAELoader": "Load VAE",
-    "LoraLoader": "Load LoRA",
-    "CLIPLoader": "Load CLIP",
-    "ControlNetLoader": "Load ControlNet Model",
-    "DiffControlNetLoader": "Load ControlNet Model (diff)",
-    "StyleModelLoader": "Load Style Model",
-    "CLIPVisionLoader": "Load CLIP Vision",
-    "UpscaleModelLoader": "Load Upscale Model",
-    # Conditioning
-    "CLIPVisionEncode": "CLIP Vision Encode",
-    "StyleModelApply": "Apply Style Model",
-    "CLIPTextEncode": "CLIP Text Encode (Prompt)",
-    "CLIPSetLastLayer": "CLIP Set Last Layer",
-    "ConditioningCombine": "Conditioning (Combine)",
-    "ConditioningAverage ": "Conditioning (Average)",
-    "ConditioningConcat": "Conditioning (Concat)",
-    "ConditioningSetArea": "Conditioning (Set Area)",
-    "ConditioningSetAreaPercentage": "Conditioning (Set Area with Percentage)",
-    "ConditioningSetMask": "Conditioning (Set Mask)",
-    "ControlNetApply": "Apply ControlNet",
-    "ControlNetApplyAdvanced": "Apply ControlNet (Advanced)",
-    # Latent
-    "VAEEncodeForInpaint": "VAE Encode (for Inpainting)",
-    "SetLatentNoiseMask": "Set Latent Noise Mask",
-    "VAEDecode": "VAE Decode",
-    "VAEEncode": "VAE Encode",
-    "LatentRotate": "Rotate Latent",
-    "LatentFlip": "Flip Latent",
-    "LatentCrop": "Crop Latent",
-    "EmptyLatentImage": "Empty Latent Image",
-    "LatentUpscale": "Upscale Latent",
-    "LatentUpscaleBy": "Upscale Latent By",
-    "LatentComposite": "Latent Composite",
-    "LatentBlend": "Latent Blend",
-    "LatentFromBatch" : "Latent From Batch",
-    "RepeatLatentBatch": "Repeat Latent Batch",
-    # Image
-    "SaveImage": "Save Image",
-    "PreviewImage": "Preview Image",
-    "LoadImage": "Load Image",
-    "LoadImageMask": "Load Image (as Mask)",
-    "ImageScale": "Upscale Image",
-    "ImageScaleBy": "Upscale Image By",
-    "ImageUpscaleWithModel": "Upscale Image (using Model)",
-    "ImageInvert": "Invert Image",
-    "ImagePadForOutpaint": "Pad Image for Outpainting",
-    "ImageBatch": "Batch Images",
-    # _for_testing
-    "VAEDecodeTiled": "VAE Decode (Tiled)",
-    "VAEEncodeTiled": "VAE Encode (Tiled)",
-}
-
-EXTENSION_WEB_DIRS = {}
-
-def load_custom_node(module_path, ignore=set()):
-    module_name = os.path.basename(module_path)
-    if os.path.isfile(module_path):
-        sp = os.path.splitext(module_path)
-        module_name = sp[0]
-    try:
-        if os.path.isfile(module_path):
-            module_spec = importlib.util.spec_from_file_location(module_name, module_path)
-            module_dir = os.path.split(module_path)[0]
-        else:
-            module_spec = importlib.util.spec_from_file_location(module_name, os.path.join(module_path, "__init__.py"))
-            module_dir = module_path
-
-        module = importlib.util.module_from_spec(module_spec)
-        sys.modules[module_name] = module
-        module_spec.loader.exec_module(module)
-
-        if hasattr(module, "WEB_DIRECTORY") and getattr(module, "WEB_DIRECTORY") is not None:
-            web_dir = os.path.abspath(os.path.join(module_dir, getattr(module, "WEB_DIRECTORY")))
-            if os.path.isdir(web_dir):
-                EXTENSION_WEB_DIRS[module_name] = web_dir
-
-        if hasattr(module, "NODE_CLASS_MAPPINGS") and getattr(module, "NODE_CLASS_MAPPINGS") is not None:
-            for name in module.NODE_CLASS_MAPPINGS:
-                if name not in ignore:
-                    NODE_CLASS_MAPPINGS[name] = module.NODE_CLASS_MAPPINGS[name]
-            if hasattr(module, "NODE_DISPLAY_NAME_MAPPINGS") and getattr(module, "NODE_DISPLAY_NAME_MAPPINGS") is not None:
-                NODE_DISPLAY_NAME_MAPPINGS.update(module.NODE_DISPLAY_NAME_MAPPINGS)
-            return True
-        else:
-            print(f"Skip {module_path} module for custom nodes due to the lack of NODE_CLASS_MAPPINGS.")
-            return False
-    except Exception as e:
-        print(traceback.format_exc())
-        print(f"Cannot import {module_path} module for custom nodes:", e)
-        return False
-
-def load_custom_nodes():
-    base_node_names = set(NODE_CLASS_MAPPINGS.keys())
-    node_paths = ldm_patched.utils.path_utils.get_folder_paths("custom_nodes")
-    node_import_times = []
-    for custom_node_path in node_paths:
-        possible_modules = os.listdir(os.path.realpath(custom_node_path))
-        if "__pycache__" in possible_modules:
-            possible_modules.remove("__pycache__")
-
-        for possible_module in possible_modules:
-            module_path = os.path.join(custom_node_path, possible_module)
-            if os.path.isfile(module_path) and os.path.splitext(module_path)[1] != ".py": continue
-            if module_path.endswith(".disabled"): continue
-            time_before = time.perf_counter()
-            success = load_custom_node(module_path, base_node_names)
-            node_import_times.append((time.perf_counter() - time_before, module_path, success))
-
-    if len(node_import_times) > 0:
-        print("\nImport times for custom nodes:")
-        for n in sorted(node_import_times):
-            if n[2]:
-                import_message = ""
-            else:
-                import_message = " (IMPORT FAILED)"
-            print("{:6.1f} seconds{}:".format(n[0], import_message), n[1])
-        print()
-
-def init_custom_nodes():
-    extras_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "ldm_patched_extras")
-    extras_files = [
-        "nodes_latent.py",
-        "nodes_hypernetwork.py",
-        "nodes_upscale_model.py",
-        "nodes_post_processing.py",
-        "nodes_mask.py",
-        "nodes_compositing.py",
-        "nodes_rebatch.py",
-        "nodes_model_merging.py",
-        "nodes_tomesd.py",
-        "nodes_clip_sdxl.py",
-        "nodes_canny.py",
-        "nodes_freelunch.py",
-        "nodes_custom_sampler.py",
-        "nodes_hypertile.py",
-        "nodes_model_advanced.py",
-        "nodes_model_downscale.py",
-        "nodes_images.py",
-        "nodes_video_model.py",
-        "nodes_sag.py",
-        "nodes_perpneg.py",
-        "nodes_stable3d.py",
-        "nodes_sdupscale.py",
-        "nodes_photomaker.py",
-    ]
-
-    for node_file in extras_files:
-        load_custom_node(os.path.join(extras_dir, node_file))
-
-    load_custom_nodes()

ldm_patched/contrib/external_canny.py

@@ -1,301 +0,0 @@
-# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
-
-#From https://github.com/kornia/kornia
-import math
-
-import torch
-import torch.nn.functional as F
-import ldm_patched.modules.model_management
-
-def get_canny_nms_kernel(device=None, dtype=None):
-    """Utility function that returns 3x3 kernels for the Canny Non-maximal suppression."""
-    return torch.tensor(
-        [
-            [[[0.0, 0.0, 0.0], [0.0, 1.0, -1.0], [0.0, 0.0, 0.0]]],
-            [[[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, -1.0]]],
-            [[[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, -1.0, 0.0]]],
-            [[[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [-1.0, 0.0, 0.0]]],
-            [[[0.0, 0.0, 0.0], [-1.0, 1.0, 0.0], [0.0, 0.0, 0.0]]],
-            [[[-1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]]],
-            [[[0.0, -1.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]]],
-            [[[0.0, 0.0, -1.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]]],
-        ],
-        device=device,
-        dtype=dtype,
-    )
-
-
-def get_hysteresis_kernel(device=None, dtype=None):
-    """Utility function that returns the 3x3 kernels for the Canny hysteresis."""
-    return torch.tensor(
-        [
-            [[[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 0.0]]],
-            [[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]],
-            [[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 1.0, 0.0]]],
-            [[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]],
-            [[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 0.0]]],
-            [[[1.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]],
-            [[[0.0, 1.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]],
-            [[[0.0, 0.0, 1.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]],
-        ],
-        device=device,
-        dtype=dtype,
-    )
-
-def gaussian_blur_2d(img, kernel_size, sigma):
-    ksize_half = (kernel_size - 1) * 0.5
-
-    x = torch.linspace(-ksize_half, ksize_half, steps=kernel_size)
-
-    pdf = torch.exp(-0.5 * (x / sigma).pow(2))
-
-    x_kernel = pdf / pdf.sum()
-    x_kernel = x_kernel.to(device=img.device, dtype=img.dtype)
-
-    kernel2d = torch.mm(x_kernel[:, None], x_kernel[None, :])
-    kernel2d = kernel2d.expand(img.shape[-3], 1, kernel2d.shape[0], kernel2d.shape[1])
-
-    padding = [kernel_size // 2, kernel_size // 2, kernel_size // 2, kernel_size // 2]
-
-    img = torch.nn.functional.pad(img, padding, mode="reflect")
-    img = torch.nn.functional.conv2d(img, kernel2d, groups=img.shape[-3])
-
-    return img
-
-def get_sobel_kernel2d(device=None, dtype=None):
-    kernel_x = torch.tensor([[-1.0, 0.0, 1.0], [-2.0, 0.0, 2.0], [-1.0, 0.0, 1.0]], device=device, dtype=dtype)
-    kernel_y = kernel_x.transpose(0, 1)
-    return torch.stack([kernel_x, kernel_y])
-
-def spatial_gradient(input, normalized: bool = True):
-    r"""Compute the first order image derivative in both x and y using a Sobel operator.
-    .. image:: _static/img/spatial_gradient.png
-    Args:
-        input: input image tensor with shape :math:`(B, C, H, W)`.
-        mode: derivatives modality, can be: `sobel` or `diff`.
-        order: the order of the derivatives.
-        normalized: whether the output is normalized.
-    Return:
-        the derivatives of the input feature map. with shape :math:`(B, C, 2, H, W)`.
-    .. note::
-       See a working example `here <https://kornia.readthedocs.io/en/latest/
-       filtering_edges.html>`__.
-    Examples:
-        >>> input = torch.rand(1, 3, 4, 4)
-        >>> output = spatial_gradient(input)  # 1x3x2x4x4
-        >>> output.shape
-        torch.Size([1, 3, 2, 4, 4])
-    """
-    # KORNIA_CHECK_IS_TENSOR(input)
-    # KORNIA_CHECK_SHAPE(input, ['B', 'C', 'H', 'W'])
-
-    # allocate kernel
-    kernel = get_sobel_kernel2d(device=input.device, dtype=input.dtype)
-    if normalized:
-        kernel = normalize_kernel2d(kernel)
-
-    # prepare kernel
-    b, c, h, w = input.shape
-    tmp_kernel = kernel[:, None, ...]
-
-    # Pad with "replicate for spatial dims, but with zeros for channel
-    spatial_pad = [kernel.size(1) // 2, kernel.size(1) // 2, kernel.size(2) // 2, kernel.size(2) // 2]
-    out_channels: int = 2
-    padded_inp = torch.nn.functional.pad(input.reshape(b * c, 1, h, w), spatial_pad, 'replicate')
-    out = F.conv2d(padded_inp, tmp_kernel, groups=1, padding=0, stride=1)
-    return out.reshape(b, c, out_channels, h, w)
-
-def rgb_to_grayscale(image, rgb_weights = None):
-    r"""Convert a RGB image to grayscale version of image.
-
-    .. image:: _static/img/rgb_to_grayscale.png
-
-    The image data is assumed to be in the range of (0, 1).
-
-    Args:
-        image: RGB image to be converted to grayscale with shape :math:`(*,3,H,W)`.
-        rgb_weights: Weights that will be applied on each channel (RGB).
-            The sum of the weights should add up to one.
-    Returns:
-        grayscale version of the image with shape :math:`(*,1,H,W)`.
-
-    .. note::
-       See a working example `here <https://kornia.readthedocs.io/en/latest/
-       color_conversions.html>`__.
-
-    Example:
-        >>> input = torch.rand(2, 3, 4, 5)
-        >>> gray = rgb_to_grayscale(input) # 2x1x4x5
-    """
-
-    if len(image.shape) < 3 or image.shape[-3] != 3:
-        raise ValueError(f"Input size must have a shape of (*, 3, H, W). Got {image.shape}")
-
-    if rgb_weights is None:
-        # 8 bit images
-        if image.dtype == torch.uint8:
-            rgb_weights = torch.tensor([76, 150, 29], device=image.device, dtype=torch.uint8)
-        # floating point images
-        elif image.dtype in (torch.float16, torch.float32, torch.float64):
-            rgb_weights = torch.tensor([0.299, 0.587, 0.114], device=image.device, dtype=image.dtype)
-        else:
-            raise TypeError(f"Unknown data type: {image.dtype}")
-    else:
-        # is tensor that we make sure is in the same device/dtype
-        rgb_weights = rgb_weights.to(image)
-
-    # unpack the color image channels with RGB order
-    r: Tensor = image[..., 0:1, :, :]
-    g: Tensor = image[..., 1:2, :, :]
-    b: Tensor = image[..., 2:3, :, :]
-
-    w_r, w_g, w_b = rgb_weights.unbind()
-    return w_r * r + w_g * g + w_b * b
-
-def canny(
-    input,
-    low_threshold = 0.1,
-    high_threshold = 0.2,
-    kernel_size  = 5,
-    sigma = 1,
-    hysteresis = True,
-    eps = 1e-6,
-):
-    r"""Find edges of the input image and filters them using the Canny algorithm.
-    .. image:: _static/img/canny.png
-    Args:
-        input: input image tensor with shape :math:`(B,C,H,W)`.
-        low_threshold: lower threshold for the hysteresis procedure.
-        high_threshold: upper threshold for the hysteresis procedure.
-        kernel_size: the size of the kernel for the gaussian blur.
-        sigma: the standard deviation of the kernel for the gaussian blur.
-        hysteresis: if True, applies the hysteresis edge tracking.
-            Otherwise, the edges are divided between weak (0.5) and strong (1) edges.
-        eps: regularization number to avoid NaN during backprop.
-    Returns:
-        - the canny edge magnitudes map, shape of :math:`(B,1,H,W)`.
-        - the canny edge detection filtered by thresholds and hysteresis, shape of :math:`(B,1,H,W)`.
-    .. note::
-       See a working example `here <https://kornia.readthedocs.io/en/latest/
-       canny.html>`__.
-    Example:
-        >>> input = torch.rand(5, 3, 4, 4)
-        >>> magnitude, edges = canny(input)  # 5x3x4x4
-        >>> magnitude.shape
-        torch.Size([5, 1, 4, 4])
-        >>> edges.shape
-        torch.Size([5, 1, 4, 4])
-    """
-    # KORNIA_CHECK_IS_TENSOR(input)
-    # KORNIA_CHECK_SHAPE(input, ['B', 'C', 'H', 'W'])
-    # KORNIA_CHECK(
-    #     low_threshold <= high_threshold,
-    #     "Invalid input thresholds. low_threshold should be smaller than the high_threshold. Got: "
-    #     f"{low_threshold}>{high_threshold}",
-    # )
-    # KORNIA_CHECK(0 < low_threshold < 1, f'Invalid low threshold. Should be in range (0, 1). Got: {low_threshold}')
-    # KORNIA_CHECK(0 < high_threshold < 1, f'Invalid high threshold. Should be in range (0, 1). Got: {high_threshold}')
-
-    device = input.device
-    dtype = input.dtype
-
-    # To Grayscale
-    if input.shape[1] == 3:
-        input = rgb_to_grayscale(input)
-
-    # Gaussian filter
-    blurred: Tensor = gaussian_blur_2d(input, kernel_size, sigma)
-
-    # Compute the gradients
-    gradients: Tensor = spatial_gradient(blurred, normalized=False)
-
-    # Unpack the edges
-    gx: Tensor = gradients[:, :, 0]
-    gy: Tensor = gradients[:, :, 1]
-
-    # Compute gradient magnitude and angle
-    magnitude: Tensor = torch.sqrt(gx * gx + gy * gy + eps)
-    angle: Tensor = torch.atan2(gy, gx)
-
-    # Radians to Degrees
-    angle = 180.0 * angle / math.pi
-
-    # Round angle to the nearest 45 degree
-    angle = torch.round(angle / 45) * 45
-
-    # Non-maximal suppression
-    nms_kernels: Tensor = get_canny_nms_kernel(device, dtype)
-    nms_magnitude: Tensor = F.conv2d(magnitude, nms_kernels, padding=nms_kernels.shape[-1] // 2)
-
-    # Get the indices for both directions
-    positive_idx: Tensor = (angle / 45) % 8
-    positive_idx = positive_idx.long()
-
-    negative_idx: Tensor = ((angle / 45) + 4) % 8
-    negative_idx = negative_idx.long()
-
-    # Apply the non-maximum suppression to the different directions
-    channel_select_filtered_positive: Tensor = torch.gather(nms_magnitude, 1, positive_idx)
-    channel_select_filtered_negative: Tensor = torch.gather(nms_magnitude, 1, negative_idx)
-
-    channel_select_filtered: Tensor = torch.stack(
-        [channel_select_filtered_positive, channel_select_filtered_negative], 1
-    )
-
-    is_max: Tensor = channel_select_filtered.min(dim=1)[0] > 0.0
-
-    magnitude = magnitude * is_max
-
-    # Threshold
-    edges: Tensor = F.threshold(magnitude, low_threshold, 0.0)
-
-    low: Tensor = magnitude > low_threshold
-    high: Tensor = magnitude > high_threshold
-
-    edges = low * 0.5 + high * 0.5
-    edges = edges.to(dtype)
-
-    # Hysteresis
-    if hysteresis:
-        edges_old: Tensor = -torch.ones(edges.shape, device=edges.device, dtype=dtype)
-        hysteresis_kernels: Tensor = get_hysteresis_kernel(device, dtype)
-
-        while ((edges_old - edges).abs() != 0).any():
-            weak: Tensor = (edges == 0.5).float()
-            strong: Tensor = (edges == 1).float()
-
-            hysteresis_magnitude: Tensor = F.conv2d(
-                edges, hysteresis_kernels, padding=hysteresis_kernels.shape[-1] // 2
-            )
-            hysteresis_magnitude = (hysteresis_magnitude == 1).any(1, keepdim=True).to(dtype)
-            hysteresis_magnitude = hysteresis_magnitude * weak + strong
-
-            edges_old = edges.clone()
-            edges = hysteresis_magnitude + (hysteresis_magnitude == 0) * weak * 0.5
-
-        edges = hysteresis_magnitude
-
-    return magnitude, edges
-
-
-class Canny:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {"image": ("IMAGE",),
-                                "low_threshold": ("FLOAT", {"default": 0.4, "min": 0.01, "max": 0.99, "step": 0.01}),
-                                "high_threshold": ("FLOAT", {"default": 0.8, "min": 0.01, "max": 0.99, "step": 0.01})
-                                }}
-
-    RETURN_TYPES = ("IMAGE",)
-    FUNCTION = "detect_edge"
-
-    CATEGORY = "image/preprocessors"
-
-    def detect_edge(self, image, low_threshold, high_threshold):
-        output = canny(image.to(ldm_patched.modules.model_management.get_torch_device()).movedim(-1, 1), low_threshold, high_threshold)
-        img_out = output[1].to(ldm_patched.modules.model_management.intermediate_device()).repeat(1, 3, 1, 1).movedim(1, -1)
-        return (img_out,)
-
-NODE_CLASS_MAPPINGS = {
-    "Canny": Canny,
-}

ldm_patched/contrib/external_clip_sdxl.py

@@ -1,58 +0,0 @@
-# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
-
-import torch
-from ldm_patched.contrib.external import MAX_RESOLUTION
-
-class CLIPTextEncodeSDXLRefiner:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {
-            "ascore": ("FLOAT", {"default": 6.0, "min": 0.0, "max": 1000.0, "step": 0.01}),
-            "width": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
-            "height": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
-            "text": ("STRING", {"multiline": True}), "clip": ("CLIP", ),
-            }}
-    RETURN_TYPES = ("CONDITIONING",)
-    FUNCTION = "encode"
-
-    CATEGORY = "advanced/conditioning"
-
-    def encode(self, clip, ascore, width, height, text):
-        tokens = clip.tokenize(text)
-        cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)
-        return ([[cond, {"pooled_output": pooled, "aesthetic_score": ascore, "width": width,"height": height}]], )
-
-class CLIPTextEncodeSDXL:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {
-            "width": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
-            "height": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
-            "crop_w": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION}),
-            "crop_h": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION}),
-            "target_width": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
-            "target_height": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}),
-            "text_g": ("STRING", {"multiline": True, "default": "CLIP_G"}), "clip": ("CLIP", ),
-            "text_l": ("STRING", {"multiline": True, "default": "CLIP_L"}), "clip": ("CLIP", ),
-            }}
-    RETURN_TYPES = ("CONDITIONING",)
-    FUNCTION = "encode"
-
-    CATEGORY = "advanced/conditioning"
-
-    def encode(self, clip, width, height, crop_w, crop_h, target_width, target_height, text_g, text_l):
-        tokens = clip.tokenize(text_g)
-        tokens["l"] = clip.tokenize(text_l)["l"]
-        if len(tokens["l"]) != len(tokens["g"]):
-            empty = clip.tokenize("")
-            while len(tokens["l"]) < len(tokens["g"]):
-                tokens["l"] += empty["l"]
-            while len(tokens["l"]) > len(tokens["g"]):
-                tokens["g"] += empty["g"]
-        cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)
-        return ([[cond, {"pooled_output": pooled, "width": width, "height": height, "crop_w": crop_w, "crop_h": crop_h, "target_width": target_width, "target_height": target_height}]], )
-
-NODE_CLASS_MAPPINGS = {
-    "CLIPTextEncodeSDXLRefiner": CLIPTextEncodeSDXLRefiner,
-    "CLIPTextEncodeSDXL": CLIPTextEncodeSDXL,
-}

ldm_patched/contrib/external_compositing.py

@@ -1,204 +0,0 @@
-# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
-
-import numpy as np
-import torch
-import ldm_patched.modules.utils
-from enum import Enum
-
-def resize_mask(mask, shape):
-    return torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(shape[0], shape[1]), mode="bilinear").squeeze(1)
-
-class PorterDuffMode(Enum):
-    ADD = 0
-    CLEAR = 1
-    DARKEN = 2
-    DST = 3
-    DST_ATOP = 4
-    DST_IN = 5
-    DST_OUT = 6
-    DST_OVER = 7
-    LIGHTEN = 8
-    MULTIPLY = 9
-    OVERLAY = 10
-    SCREEN = 11
-    SRC = 12
-    SRC_ATOP = 13
-    SRC_IN = 14
-    SRC_OUT = 15
-    SRC_OVER = 16
-    XOR = 17
-
-
-def porter_duff_composite(src_image: torch.Tensor, src_alpha: torch.Tensor, dst_image: torch.Tensor, dst_alpha: torch.Tensor, mode: PorterDuffMode):
-    if mode == PorterDuffMode.ADD:
-        out_alpha = torch.clamp(src_alpha + dst_alpha, 0, 1)
-        out_image = torch.clamp(src_image + dst_image, 0, 1)
-    elif mode == PorterDuffMode.CLEAR:
-        out_alpha = torch.zeros_like(dst_alpha)
-        out_image = torch.zeros_like(dst_image)
-    elif mode == PorterDuffMode.DARKEN:
-        out_alpha = src_alpha + dst_alpha  - src_alpha * dst_alpha
-        out_image = (1 - dst_alpha) * src_image + (1 - src_alpha) * dst_image + torch.min(src_image, dst_image)
-    elif mode == PorterDuffMode.DST:
-        out_alpha = dst_alpha
-        out_image = dst_image
-    elif mode == PorterDuffMode.DST_ATOP:
-        out_alpha = src_alpha
-        out_image = src_alpha * dst_image + (1 - dst_alpha) * src_image
-    elif mode == PorterDuffMode.DST_IN:
-        out_alpha = src_alpha * dst_alpha
-        out_image = dst_image * src_alpha
-    elif mode == PorterDuffMode.DST_OUT:
-        out_alpha = (1 - src_alpha) * dst_alpha
-        out_image = (1 - src_alpha) * dst_image
-    elif mode == PorterDuffMode.DST_OVER:
-        out_alpha = dst_alpha + (1 - dst_alpha) * src_alpha
-        out_image = dst_image + (1 - dst_alpha) * src_image
-    elif mode == PorterDuffMode.LIGHTEN:
-        out_alpha = src_alpha + dst_alpha - src_alpha * dst_alpha
-        out_image = (1 - dst_alpha) * src_image + (1 - src_alpha) * dst_image + torch.max(src_image, dst_image)
-    elif mode == PorterDuffMode.MULTIPLY:
-        out_alpha = src_alpha * dst_alpha
-        out_image = src_image * dst_image
-    elif mode == PorterDuffMode.OVERLAY:
-        out_alpha = src_alpha + dst_alpha - src_alpha * dst_alpha
-        out_image = torch.where(2 * dst_image < dst_alpha, 2 * src_image * dst_image,
-            src_alpha * dst_alpha - 2 * (dst_alpha - src_image) * (src_alpha - dst_image))
-    elif mode == PorterDuffMode.SCREEN:
-        out_alpha = src_alpha + dst_alpha - src_alpha * dst_alpha
-        out_image = src_image + dst_image - src_image * dst_image
-    elif mode == PorterDuffMode.SRC:
-        out_alpha = src_alpha
-        out_image = src_image
-    elif mode == PorterDuffMode.SRC_ATOP:
-        out_alpha = dst_alpha
-        out_image = dst_alpha * src_image + (1 - src_alpha) * dst_image
-    elif mode == PorterDuffMode.SRC_IN:
-        out_alpha = src_alpha * dst_alpha
-        out_image = src_image * dst_alpha
-    elif mode == PorterDuffMode.SRC_OUT:
-        out_alpha = (1 - dst_alpha) * src_alpha
-        out_image = (1 - dst_alpha) * src_image
-    elif mode == PorterDuffMode.SRC_OVER:
-        out_alpha = src_alpha + (1 - src_alpha) * dst_alpha
-        out_image = src_image + (1 - src_alpha) * dst_image
-    elif mode == PorterDuffMode.XOR:
-        out_alpha = (1 - dst_alpha) * src_alpha + (1 - src_alpha) * dst_alpha
-        out_image = (1 - dst_alpha) * src_image + (1 - src_alpha) * dst_image
-    else:
-        out_alpha = None
-        out_image = None
-    return out_image, out_alpha
-
-
-class PorterDuffImageComposite:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "source": ("IMAGE",),
-                "source_alpha": ("MASK",),
-                "destination": ("IMAGE",),
-                "destination_alpha": ("MASK",),
-                "mode": ([mode.name for mode in PorterDuffMode], {"default": PorterDuffMode.DST.name}),
-            },
-        }
-
-    RETURN_TYPES = ("IMAGE", "MASK")
-    FUNCTION = "composite"
-    CATEGORY = "mask/compositing"
-
-    def composite(self, source: torch.Tensor, source_alpha: torch.Tensor, destination: torch.Tensor, destination_alpha: torch.Tensor, mode):
-        batch_size = min(len(source), len(source_alpha), len(destination), len(destination_alpha))
-        out_images = []
-        out_alphas = []
-
-        for i in range(batch_size):
-            src_image = source[i]
-            dst_image = destination[i]
-
-            assert src_image.shape[2] == dst_image.shape[2] # inputs need to have same number of channels
-
-            src_alpha = source_alpha[i].unsqueeze(2)
-            dst_alpha = destination_alpha[i].unsqueeze(2)
-
-            if dst_alpha.shape[:2] != dst_image.shape[:2]:
-                upscale_input = dst_alpha.unsqueeze(0).permute(0, 3, 1, 2)
-                upscale_output = ldm_patched.modules.utils.common_upscale(upscale_input, dst_image.shape[1], dst_image.shape[0], upscale_method='bicubic', crop='center')
-                dst_alpha = upscale_output.permute(0, 2, 3, 1).squeeze(0)
-            if src_image.shape != dst_image.shape:
-                upscale_input = src_image.unsqueeze(0).permute(0, 3, 1, 2)
-                upscale_output = ldm_patched.modules.utils.common_upscale(upscale_input, dst_image.shape[1], dst_image.shape[0], upscale_method='bicubic', crop='center')
-                src_image = upscale_output.permute(0, 2, 3, 1).squeeze(0)
-            if src_alpha.shape != dst_alpha.shape:
-                upscale_input = src_alpha.unsqueeze(0).permute(0, 3, 1, 2)
-                upscale_output = ldm_patched.modules.utils.common_upscale(upscale_input, dst_alpha.shape[1], dst_alpha.shape[0], upscale_method='bicubic', crop='center')
-                src_alpha = upscale_output.permute(0, 2, 3, 1).squeeze(0)
-
-            out_image, out_alpha = porter_duff_composite(src_image, src_alpha, dst_image, dst_alpha, PorterDuffMode[mode])
-
-            out_images.append(out_image)
-            out_alphas.append(out_alpha.squeeze(2))
-
-        result = (torch.stack(out_images), torch.stack(out_alphas))
-        return result
-
-
-class SplitImageWithAlpha:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-                "required": {
-                    "image": ("IMAGE",),
-                }
-        }
-
-    CATEGORY = "mask/compositing"
-    RETURN_TYPES = ("IMAGE", "MASK")
-    FUNCTION = "split_image_with_alpha"
-
-    def split_image_with_alpha(self, image: torch.Tensor):
-        out_images = [i[:,:,:3] for i in image]
-        out_alphas = [i[:,:,3] if i.shape[2] > 3 else torch.ones_like(i[:,:,0]) for i in image]
-        result = (torch.stack(out_images), 1.0 - torch.stack(out_alphas))
-        return result
-
-
-class JoinImageWithAlpha:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-                "required": {
-                    "image": ("IMAGE",),
-                    "alpha": ("MASK",),
-                }
-        }
-
-    CATEGORY = "mask/compositing"
-    RETURN_TYPES = ("IMAGE",)
-    FUNCTION = "join_image_with_alpha"
-
-    def join_image_with_alpha(self, image: torch.Tensor, alpha: torch.Tensor):
-        batch_size = min(len(image), len(alpha))
-        out_images = []
-
-        alpha = 1.0 - resize_mask(alpha, image.shape[1:])
-        for i in range(batch_size):
-           out_images.append(torch.cat((image[i][:,:,:3], alpha[i].unsqueeze(2)), dim=2))
-
-        result = (torch.stack(out_images),)
-        return result
-
-
-NODE_CLASS_MAPPINGS = {
-    "PorterDuffImageComposite": PorterDuffImageComposite,
-    "SplitImageWithAlpha": SplitImageWithAlpha,
-    "JoinImageWithAlpha": JoinImageWithAlpha,
-}
-
-
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "PorterDuffImageComposite": "Porter-Duff Image Composite",
-    "SplitImageWithAlpha": "Split Image with Alpha",
-    "JoinImageWithAlpha": "Join Image with Alpha",
-}

ldm_patched/contrib/external_custom_sampler.py

@@ -1,297 +0,0 @@
-# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
-
-import ldm_patched.modules.samplers
-import ldm_patched.modules.sample
-from ldm_patched.k_diffusion import sampling as k_diffusion_sampling
-import ldm_patched.utils.latent_visualization
-import torch
-import ldm_patched.modules.utils
-
-
-class BasicScheduler:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"model": ("MODEL",),
-                     "scheduler": (ldm_patched.modules.samplers.SCHEDULER_NAMES, ),
-                     "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
-                     "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
-                      }
-               }
-    RETURN_TYPES = ("SIGMAS",)
-    CATEGORY = "sampling/custom_sampling/schedulers"
-
-    FUNCTION = "get_sigmas"
-
-    def get_sigmas(self, model, scheduler, steps, denoise):
-        total_steps = steps
-        if denoise < 1.0:
-            total_steps = int(steps/denoise)
-
-        ldm_patched.modules.model_management.load_models_gpu([model])
-        sigmas = ldm_patched.modules.samplers.calculate_sigmas_scheduler(model.model, scheduler, total_steps).cpu()
-        sigmas = sigmas[-(steps + 1):]
-        return (sigmas, )
-
-
-class KarrasScheduler:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
-                     "sigma_max": ("FLOAT", {"default": 14.614642, "min": 0.0, "max": 1000.0, "step":0.01, "round": False}),
-                     "sigma_min": ("FLOAT", {"default": 0.0291675, "min": 0.0, "max": 1000.0, "step":0.01, "round": False}),
-                     "rho": ("FLOAT", {"default": 7.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}),
-                    }
-               }
-    RETURN_TYPES = ("SIGMAS",)
-    CATEGORY = "sampling/custom_sampling/schedulers"
-
-    FUNCTION = "get_sigmas"
-
-    def get_sigmas(self, steps, sigma_max, sigma_min, rho):
-        sigmas = k_diffusion_sampling.get_sigmas_karras(n=steps, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho)
-        return (sigmas, )
-
-class ExponentialScheduler:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
-                     "sigma_max": ("FLOAT", {"default": 14.614642, "min": 0.0, "max": 1000.0, "step":0.01, "round": False}),
-                     "sigma_min": ("FLOAT", {"default": 0.0291675, "min": 0.0, "max": 1000.0, "step":0.01, "round": False}),
-                    }
-               }
-    RETURN_TYPES = ("SIGMAS",)
-    CATEGORY = "sampling/custom_sampling/schedulers"
-
-    FUNCTION = "get_sigmas"
-
-    def get_sigmas(self, steps, sigma_max, sigma_min):
-        sigmas = k_diffusion_sampling.get_sigmas_exponential(n=steps, sigma_min=sigma_min, sigma_max=sigma_max)
-        return (sigmas, )
-
-class PolyexponentialScheduler:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
-                     "sigma_max": ("FLOAT", {"default": 14.614642, "min": 0.0, "max": 1000.0, "step":0.01, "round": False}),
-                     "sigma_min": ("FLOAT", {"default": 0.0291675, "min": 0.0, "max": 1000.0, "step":0.01, "round": False}),
-                     "rho": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}),
-                    }
-               }
-    RETURN_TYPES = ("SIGMAS",)
-    CATEGORY = "sampling/custom_sampling/schedulers"
-
-    FUNCTION = "get_sigmas"
-
-    def get_sigmas(self, steps, sigma_max, sigma_min, rho):
-        sigmas = k_diffusion_sampling.get_sigmas_polyexponential(n=steps, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho)
-        return (sigmas, )
-
-class SDTurboScheduler:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"model": ("MODEL",),
-                     "steps": ("INT", {"default": 1, "min": 1, "max": 10}),
-                     "denoise": ("FLOAT", {"default": 1.0, "min": 0, "max": 1.0, "step": 0.01}),
-                      }
-               }
-    RETURN_TYPES = ("SIGMAS",)
-    CATEGORY = "sampling/custom_sampling/schedulers"
-
-    FUNCTION = "get_sigmas"
-
-    def get_sigmas(self, model, steps, denoise):
-        start_step = 10 - int(10 * denoise)
-        timesteps = torch.flip(torch.arange(1, 11) * 100 - 1, (0,))[start_step:start_step + steps]
-        ldm_patched.modules.model_management.load_models_gpu([model])
-        sigmas = model.model.model_sampling.sigma(timesteps)
-        sigmas = torch.cat([sigmas, sigmas.new_zeros([1])])
-        return (sigmas, )
-
-class VPScheduler:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
-                     "beta_d": ("FLOAT", {"default": 19.9, "min": 0.0, "max": 1000.0, "step":0.01, "round": False}), #TODO: fix default values
-                     "beta_min": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 1000.0, "step":0.01, "round": False}),
-                     "eps_s": ("FLOAT", {"default": 0.001, "min": 0.0, "max": 1.0, "step":0.0001, "round": False}),
-                    }
-               }
-    RETURN_TYPES = ("SIGMAS",)
-    CATEGORY = "sampling/custom_sampling/schedulers"
-
-    FUNCTION = "get_sigmas"
-
-    def get_sigmas(self, steps, beta_d, beta_min, eps_s):
-        sigmas = k_diffusion_sampling.get_sigmas_vp(n=steps, beta_d=beta_d, beta_min=beta_min, eps_s=eps_s)
-        return (sigmas, )
-
-class SplitSigmas:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"sigmas": ("SIGMAS", ),
-                    "step": ("INT", {"default": 0, "min": 0, "max": 10000}),
-                     }
-                }
-    RETURN_TYPES = ("SIGMAS","SIGMAS")
-    CATEGORY = "sampling/custom_sampling/sigmas"
-
-    FUNCTION = "get_sigmas"
-
-    def get_sigmas(self, sigmas, step):
-        sigmas1 = sigmas[:step + 1]
-        sigmas2 = sigmas[step:]
-        return (sigmas1, sigmas2)
-
-class FlipSigmas:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"sigmas": ("SIGMAS", ),
-                     }
-                }
-    RETURN_TYPES = ("SIGMAS",)
-    CATEGORY = "sampling/custom_sampling/sigmas"
-
-    FUNCTION = "get_sigmas"
-
-    def get_sigmas(self, sigmas):
-        sigmas = sigmas.flip(0)
-        if sigmas[0] == 0:
-            sigmas[0] = 0.0001
-        return (sigmas,)
-
-class KSamplerSelect:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"sampler_name": (ldm_patched.modules.samplers.SAMPLER_NAMES, ),
-                      }
-               }
-    RETURN_TYPES = ("SAMPLER",)
-    CATEGORY = "sampling/custom_sampling/samplers"
-
-    FUNCTION = "get_sampler"
-
-    def get_sampler(self, sampler_name):
-        sampler = ldm_patched.modules.samplers.sampler_object(sampler_name)
-        return (sampler, )
-
-class SamplerDPMPP_2M_SDE:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"solver_type": (['midpoint', 'heun'], ),
-                     "eta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}),
-                     "s_noise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}),
-                     "noise_device": (['gpu', 'cpu'], ),
-                      }
-               }
-    RETURN_TYPES = ("SAMPLER",)
-    CATEGORY = "sampling/custom_sampling/samplers"
-
-    FUNCTION = "get_sampler"
-
-    def get_sampler(self, solver_type, eta, s_noise, noise_device):
-        if noise_device == 'cpu':
-            sampler_name = "dpmpp_2m_sde"
-        else:
-            sampler_name = "dpmpp_2m_sde_gpu"
-        sampler = ldm_patched.modules.samplers.ksampler(sampler_name, {"eta": eta, "s_noise": s_noise, "solver_type": solver_type})
-        return (sampler, )
-
-
-class SamplerDPMPP_SDE:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"eta": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}),
-                     "s_noise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.01, "round": False}),
-                     "r": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 100.0, "step":0.01, "round": False}),
-                     "noise_device": (['gpu', 'cpu'], ),
-                      }
-               }
-    RETURN_TYPES = ("SAMPLER",)
-    CATEGORY = "sampling/custom_sampling/samplers"
-
-    FUNCTION = "get_sampler"
-
-    def get_sampler(self, eta, s_noise, r, noise_device):
-        if noise_device == 'cpu':
-            sampler_name = "dpmpp_sde"
-        else:
-            sampler_name = "dpmpp_sde_gpu"
-        sampler = ldm_patched.modules.samplers.ksampler(sampler_name, {"eta": eta, "s_noise": s_noise, "r": r})
-        return (sampler, )
-
-class SamplerCustom:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"model": ("MODEL",),
-                    "add_noise": ("BOOLEAN", {"default": True}),
-                    "noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
-                    "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
-                    "positive": ("CONDITIONING", ),
-                    "negative": ("CONDITIONING", ),
-                    "sampler": ("SAMPLER", ),
-                    "sigmas": ("SIGMAS", ),
-                    "latent_image": ("LATENT", ),
-                     }
-                }
-
-    RETURN_TYPES = ("LATENT","LATENT")
-    RETURN_NAMES = ("output", "denoised_output")
-
-    FUNCTION = "sample"
-
-    CATEGORY = "sampling/custom_sampling"
-
-    def sample(self, model, add_noise, noise_seed, cfg, positive, negative, sampler, sigmas, latent_image):
-        latent = latent_image
-        latent_image = latent["samples"]
-        if not add_noise:
-            noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
-        else:
-            batch_inds = latent["batch_index"] if "batch_index" in latent else None
-            noise = ldm_patched.modules.sample.prepare_noise(latent_image, noise_seed, batch_inds)
-
-        noise_mask = None
-        if "noise_mask" in latent:
-            noise_mask = latent["noise_mask"]
-
-        x0_output = {}
-        callback = ldm_patched.utils.latent_visualization.prepare_callback(model, sigmas.shape[-1] - 1, x0_output)
-
-        disable_pbar = not ldm_patched.modules.utils.PROGRESS_BAR_ENABLED
-        samples = ldm_patched.modules.sample.sample_custom(model, noise, cfg, sampler, sigmas, positive, negative, latent_image, noise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=noise_seed)
-
-        out = latent.copy()
-        out["samples"] = samples
-        if "x0" in x0_output:
-            out_denoised = latent.copy()
-            out_denoised["samples"] = model.model.process_latent_out(x0_output["x0"].cpu())
-        else:
-            out_denoised = out
-        return (out, out_denoised)
-
-NODE_CLASS_MAPPINGS = {
-    "SamplerCustom": SamplerCustom,
-    "BasicScheduler": BasicScheduler,
-    "KarrasScheduler": KarrasScheduler,
-    "ExponentialScheduler": ExponentialScheduler,
-    "PolyexponentialScheduler": PolyexponentialScheduler,
-    "VPScheduler": VPScheduler,
-    "SDTurboScheduler": SDTurboScheduler,
-    "KSamplerSelect": KSamplerSelect,
-    "SamplerDPMPP_2M_SDE": SamplerDPMPP_2M_SDE,
-    "SamplerDPMPP_SDE": SamplerDPMPP_SDE,
-    "SplitSigmas": SplitSigmas,
-    "FlipSigmas": FlipSigmas,
-}

ldm_patched/contrib/external_freelunch.py

@@ -1,115 +0,0 @@
-# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
-
-#code originally taken from: https://github.com/ChenyangSi/FreeU (under MIT License)
-
-import torch
-
-
-def Fourier_filter(x, threshold, scale):
-    # FFT
-    x_freq = torch.fft.fftn(x.float(), dim=(-2, -1))
-    x_freq = torch.fft.fftshift(x_freq, dim=(-2, -1))
-
-    B, C, H, W = x_freq.shape
-    mask = torch.ones((B, C, H, W), device=x.device)
-
-    crow, ccol = H // 2, W //2
-    mask[..., crow - threshold:crow + threshold, ccol - threshold:ccol + threshold] = scale
-    x_freq = x_freq * mask
-
-    # IFFT
-    x_freq = torch.fft.ifftshift(x_freq, dim=(-2, -1))
-    x_filtered = torch.fft.ifftn(x_freq, dim=(-2, -1)).real
-
-    return x_filtered.to(x.dtype)
-
-
-class FreeU:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model": ("MODEL",),
-                             "b1": ("FLOAT", {"default": 1.1, "min": 0.0, "max": 10.0, "step": 0.01}),
-                             "b2": ("FLOAT", {"default": 1.2, "min": 0.0, "max": 10.0, "step": 0.01}),
-                             "s1": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 10.0, "step": 0.01}),
-                             "s2": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 10.0, "step": 0.01}),
-                              }}
-    RETURN_TYPES = ("MODEL",)
-    FUNCTION = "patch"
-
-    CATEGORY = "model_patches"
-
-    def patch(self, model, b1, b2, s1, s2):
-        model_channels = model.model.model_config.unet_config["model_channels"]
-        scale_dict = {model_channels * 4: (b1, s1), model_channels * 2: (b2, s2)}
-        on_cpu_devices = {}
-
-        def output_block_patch(h, hsp, transformer_options):
-            scale = scale_dict.get(h.shape[1], None)
-            if scale is not None:
-                h[:,:h.shape[1] // 2] = h[:,:h.shape[1] // 2] * scale[0]
-                if hsp.device not in on_cpu_devices:
-                    try:
-                        hsp = Fourier_filter(hsp, threshold=1, scale=scale[1])
-                    except:
-                        print("Device", hsp.device, "does not support the torch.fft functions used in the FreeU node, switching to CPU.")
-                        on_cpu_devices[hsp.device] = True
-                        hsp = Fourier_filter(hsp.cpu(), threshold=1, scale=scale[1]).to(hsp.device)
-                else:
-                    hsp = Fourier_filter(hsp.cpu(), threshold=1, scale=scale[1]).to(hsp.device)
-
-            return h, hsp
-
-        m = model.clone()
-        m.set_model_output_block_patch(output_block_patch)
-        return (m, )
-
-class FreeU_V2:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model": ("MODEL",),
-                             "b1": ("FLOAT", {"default": 1.3, "min": 0.0, "max": 10.0, "step": 0.01}),
-                             "b2": ("FLOAT", {"default": 1.4, "min": 0.0, "max": 10.0, "step": 0.01}),
-                             "s1": ("FLOAT", {"default": 0.9, "min": 0.0, "max": 10.0, "step": 0.01}),
-                             "s2": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 10.0, "step": 0.01}),
-                              }}
-    RETURN_TYPES = ("MODEL",)
-    FUNCTION = "patch"
-
-    CATEGORY = "model_patches"
-
-    def patch(self, model, b1, b2, s1, s2):
-        model_channels = model.model.model_config.unet_config["model_channels"]
-        scale_dict = {model_channels * 4: (b1, s1), model_channels * 2: (b2, s2)}
-        on_cpu_devices = {}
-
-        def output_block_patch(h, hsp, transformer_options):
-            scale = scale_dict.get(h.shape[1], None)
-            if scale is not None:
-                hidden_mean = h.mean(1).unsqueeze(1)
-                B = hidden_mean.shape[0]
-                hidden_max, _ = torch.max(hidden_mean.view(B, -1), dim=-1, keepdim=True)
-                hidden_min, _ = torch.min(hidden_mean.view(B, -1), dim=-1, keepdim=True)
-                hidden_mean = (hidden_mean - hidden_min.unsqueeze(2).unsqueeze(3)) / (hidden_max - hidden_min).unsqueeze(2).unsqueeze(3)
-
-                h[:,:h.shape[1] // 2] = h[:,:h.shape[1] // 2] * ((scale[0] - 1 ) * hidden_mean + 1)
-
-                if hsp.device not in on_cpu_devices:
-                    try:
-                        hsp = Fourier_filter(hsp, threshold=1, scale=scale[1])
-                    except:
-                        print("Device", hsp.device, "does not support the torch.fft functions used in the FreeU node, switching to CPU.")
-                        on_cpu_devices[hsp.device] = True
-                        hsp = Fourier_filter(hsp.cpu(), threshold=1, scale=scale[1]).to(hsp.device)
-                else:
-                    hsp = Fourier_filter(hsp.cpu(), threshold=1, scale=scale[1]).to(hsp.device)
-
-            return h, hsp
-
-        m = model.clone()
-        m.set_model_output_block_patch(output_block_patch)
-        return (m, )
-
-NODE_CLASS_MAPPINGS = {
-    "FreeU": FreeU,
-    "FreeU_V2": FreeU_V2,
-}

ldm_patched/contrib/external_hypernetwork.py

@@ -1,121 +0,0 @@
-# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
-
-import ldm_patched.modules.utils
-import ldm_patched.utils.path_utils
-import torch
-
-def load_hypernetwork_patch(path, strength):
-    sd = ldm_patched.modules.utils.load_torch_file(path, safe_load=True)
-    activation_func = sd.get('activation_func', 'linear')
-    is_layer_norm = sd.get('is_layer_norm', False)
-    use_dropout = sd.get('use_dropout', False)
-    activate_output = sd.get('activate_output', False)
-    last_layer_dropout = sd.get('last_layer_dropout', False)
-
-    valid_activation = {
-        "linear": torch.nn.Identity,
-        "relu": torch.nn.ReLU,
-        "leakyrelu": torch.nn.LeakyReLU,
-        "elu": torch.nn.ELU,
-        "swish": torch.nn.Hardswish,
-        "tanh": torch.nn.Tanh,
-        "sigmoid": torch.nn.Sigmoid,
-        "softsign": torch.nn.Softsign,
-        "mish": torch.nn.Mish,
-    }
-
-    if activation_func not in valid_activation:
-        print("Unsupported Hypernetwork format, if you report it I might implement it.", path, " ", activation_func, is_layer_norm, use_dropout, activate_output, last_layer_dropout)
-        return None
-
-    out = {}
-
-    for d in sd:
-        try:
-            dim = int(d)
-        except:
-            continue
-
-        output = []
-        for index in [0, 1]:
-            attn_weights = sd[dim][index]
-            keys = attn_weights.keys()
-
-            linears = filter(lambda a: a.endswith(".weight"), keys)
-            linears = list(map(lambda a: a[:-len(".weight")], linears))
-            layers = []
-
-            i = 0
-            while i < len(linears):
-                lin_name = linears[i]
-                last_layer = (i == (len(linears) - 1))
-                penultimate_layer = (i == (len(linears) - 2))
-
-                lin_weight = attn_weights['{}.weight'.format(lin_name)]
-                lin_bias = attn_weights['{}.bias'.format(lin_name)]
-                layer = torch.nn.Linear(lin_weight.shape[1], lin_weight.shape[0])
-                layer.load_state_dict({"weight": lin_weight, "bias": lin_bias})
-                layers.append(layer)
-                if activation_func != "linear":
-                    if (not last_layer) or (activate_output):
-                        layers.append(valid_activation[activation_func]())
-                if is_layer_norm:
-                    i += 1
-                    ln_name = linears[i]
-                    ln_weight = attn_weights['{}.weight'.format(ln_name)]
-                    ln_bias = attn_weights['{}.bias'.format(ln_name)]
-                    ln = torch.nn.LayerNorm(ln_weight.shape[0])
-                    ln.load_state_dict({"weight": ln_weight, "bias": ln_bias})
-                    layers.append(ln)
-                if use_dropout:
-                    if (not last_layer) and (not penultimate_layer or last_layer_dropout):
-                        layers.append(torch.nn.Dropout(p=0.3))
-                i += 1
-
-            output.append(torch.nn.Sequential(*layers))
-        out[dim] = torch.nn.ModuleList(output)
-
-    class hypernetwork_patch:
-        def __init__(self, hypernet, strength):
-            self.hypernet = hypernet
-            self.strength = strength
-        def __call__(self, q, k, v, extra_options):
-            dim = k.shape[-1]
-            if dim in self.hypernet:
-                hn = self.hypernet[dim]
-                k = k + hn[0](k) * self.strength
-                v = v + hn[1](v) * self.strength
-
-            return q, k, v
-
-        def to(self, device):
-            for d in self.hypernet.keys():
-                self.hypernet[d] = self.hypernet[d].to(device)
-            return self
-
-    return hypernetwork_patch(out, strength)
-
-class HypernetworkLoader:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model": ("MODEL",),
-                              "hypernetwork_name": (ldm_patched.utils.path_utils.get_filename_list("hypernetworks"), ),
-                              "strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
-                              }}
-    RETURN_TYPES = ("MODEL",)
-    FUNCTION = "load_hypernetwork"
-
-    CATEGORY = "loaders"
-
-    def load_hypernetwork(self, model, hypernetwork_name, strength):
-        hypernetwork_path = ldm_patched.utils.path_utils.get_full_path("hypernetworks", hypernetwork_name)
-        model_hypernetwork = model.clone()
-        patch = load_hypernetwork_patch(hypernetwork_path, strength)
-        if patch is not None:
-            model_hypernetwork.set_model_attn1_patch(patch)
-            model_hypernetwork.set_model_attn2_patch(patch)
-        return (model_hypernetwork,)
-
-NODE_CLASS_MAPPINGS = {
-    "HypernetworkLoader": HypernetworkLoader
-}

ldm_patched/contrib/external_hypertile.py

@@ -1,85 +0,0 @@
-# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
-
-#Taken from: https://github.com/tfernd/HyperTile/
-
-import math
-from einops import rearrange
-# Use torch rng for consistency across generations
-from torch import randint
-
-def random_divisor(value: int, min_value: int, /, max_options: int = 1) -> int:
-    min_value = min(min_value, value)
-
-    # All big divisors of value (inclusive)
-    divisors = [i for i in range(min_value, value + 1) if value % i == 0]
-
-    ns = [value // i for i in divisors[:max_options]]  # has at least 1 element
-
-    if len(ns) - 1 > 0:
-        idx = randint(low=0, high=len(ns) - 1, size=(1,)).item()
-    else:
-        idx = 0
-
-    return ns[idx]
-
-class HyperTile:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model": ("MODEL",),
-                             "tile_size": ("INT", {"default": 256, "min": 1, "max": 2048}),
-                             "swap_size": ("INT", {"default": 2, "min": 1, "max": 128}),
-                             "max_depth": ("INT", {"default": 0, "min": 0, "max": 10}),
-                             "scale_depth": ("BOOLEAN", {"default": False}),
-                              }}
-    RETURN_TYPES = ("MODEL",)
-    FUNCTION = "patch"
-
-    CATEGORY = "model_patches"
-
-    def patch(self, model, tile_size, swap_size, max_depth, scale_depth):
-        model_channels = model.model.model_config.unet_config["model_channels"]
-
-        latent_tile_size = max(32, tile_size) // 8
-        self.temp = None
-
-        def hypertile_in(q, k, v, extra_options):
-            model_chans = q.shape[-2]
-            orig_shape = extra_options['original_shape']
-            apply_to = []
-            for i in range(max_depth + 1):
-                apply_to.append((orig_shape[-2] / (2 ** i)) * (orig_shape[-1] / (2 ** i)))
-
-            if model_chans in apply_to:
-                shape = extra_options["original_shape"]
-                aspect_ratio = shape[-1] / shape[-2]
-
-                hw = q.size(1)
-                h, w = round(math.sqrt(hw * aspect_ratio)), round(math.sqrt(hw / aspect_ratio))
-
-                factor = (2 ** apply_to.index(model_chans)) if scale_depth else 1
-                nh = random_divisor(h, latent_tile_size * factor, swap_size)
-                nw = random_divisor(w, latent_tile_size * factor, swap_size)
-
-                if nh * nw > 1:
-                    q = rearrange(q, "b (nh h nw w) c -> (b nh nw) (h w) c", h=h // nh, w=w // nw, nh=nh, nw=nw)
-                    self.temp = (nh, nw, h, w)
-                return q, k, v
-
-            return q, k, v
-        def hypertile_out(out, extra_options):
-            if self.temp is not None:
-                nh, nw, h, w = self.temp
-                self.temp = None
-                out = rearrange(out, "(b nh nw) hw c -> b nh nw hw c", nh=nh, nw=nw)
-                out = rearrange(out, "b nh nw (h w) c -> b (nh h nw w) c", h=h // nh, w=w // nw)
-            return out
-
-
-        m = model.clone()
-        m.set_model_attn1_patch(hypertile_in)
-        m.set_model_attn1_output_patch(hypertile_out)
-        return (m, )
-
-NODE_CLASS_MAPPINGS = {
-    "HyperTile": HyperTile,
-}

ldm_patched/contrib/external_images.py

@@ -1,177 +0,0 @@
-# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
-
-import ldm_patched.contrib.external
-import ldm_patched.utils.path_utils
-from ldm_patched.modules.args_parser import args
-
-from PIL import Image
-from PIL.PngImagePlugin import PngInfo
-
-import numpy as np
-import json
-import os
-
-MAX_RESOLUTION = ldm_patched.contrib.external.MAX_RESOLUTION
-
-class ImageCrop:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "image": ("IMAGE",),
-                              "width": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}),
-                              "height": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}),
-                              "x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
-                              "y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
-                              }}
-    RETURN_TYPES = ("IMAGE",)
-    FUNCTION = "crop"
-
-    CATEGORY = "image/transform"
-
-    def crop(self, image, width, height, x, y):
-        x = min(x, image.shape[2] - 1)
-        y = min(y, image.shape[1] - 1)
-        to_x = width + x
-        to_y = height + y
-        img = image[:,y:to_y, x:to_x, :]
-        return (img,)
-
-class RepeatImageBatch:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "image": ("IMAGE",),
-                              "amount": ("INT", {"default": 1, "min": 1, "max": 64}),
-                              }}
-    RETURN_TYPES = ("IMAGE",)
-    FUNCTION = "repeat"
-
-    CATEGORY = "image/batch"
-
-    def repeat(self, image, amount):
-        s = image.repeat((amount, 1,1,1))
-        return (s,)
-
-class SaveAnimatedWEBP:
-    def __init__(self):
-        self.output_dir = ldm_patched.utils.path_utils.get_output_directory()
-        self.type = "output"
-        self.prefix_append = ""
-
-    methods = {"default": 4, "fastest": 0, "slowest": 6}
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"images": ("IMAGE", ),
-                     "filename_prefix": ("STRING", {"default": "ldm_patched"}),
-                     "fps": ("FLOAT", {"default": 6.0, "min": 0.01, "max": 1000.0, "step": 0.01}),
-                     "lossless": ("BOOLEAN", {"default": True}),
-                     "quality": ("INT", {"default": 80, "min": 0, "max": 100}),
-                     "method": (list(s.methods.keys()),),
-                     # "num_frames": ("INT", {"default": 0, "min": 0, "max": 8192}),
-                     },
-                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
-                }
-
-    RETURN_TYPES = ()
-    FUNCTION = "save_images"
-
-    OUTPUT_NODE = True
-
-    CATEGORY = "image/animation"
-
-    def save_images(self, images, fps, filename_prefix, lossless, quality, method, num_frames=0, prompt=None, extra_pnginfo=None):
-        method = self.methods.get(method)
-        filename_prefix += self.prefix_append
-        full_output_folder, filename, counter, subfolder, filename_prefix = ldm_patched.utils.path_utils.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0])
-        results = list()
-        pil_images = []
-        for image in images:
-            i = 255. * image.cpu().numpy()
-            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
-            pil_images.append(img)
-
-        metadata = pil_images[0].getexif()
-        if not args.disable_server_info:
-            if prompt is not None:
-                metadata[0x0110] = "prompt:{}".format(json.dumps(prompt))
-            if extra_pnginfo is not None:
-                inital_exif = 0x010f
-                for x in extra_pnginfo:
-                    metadata[inital_exif] = "{}:{}".format(x, json.dumps(extra_pnginfo[x]))
-                    inital_exif -= 1
-
-        if num_frames == 0:
-            num_frames = len(pil_images)
-
-        c = len(pil_images)
-        for i in range(0, c, num_frames):
-            file = f"{filename}_{counter:05}_.webp"
-            pil_images[i].save(os.path.join(full_output_folder, file), save_all=True, duration=int(1000.0/fps), append_images=pil_images[i + 1:i + num_frames], exif=metadata, lossless=lossless, quality=quality, method=method)
-            results.append({
-                "filename": file,
-                "subfolder": subfolder,
-                "type": self.type
-            })
-            counter += 1
-
-        animated = num_frames != 1
-        return { "ui": { "images": results, "animated": (animated,) } }
-
-class SaveAnimatedPNG:
-    def __init__(self):
-        self.output_dir = ldm_patched.utils.path_utils.get_output_directory()
-        self.type = "output"
-        self.prefix_append = ""
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required":
-                    {"images": ("IMAGE", ),
-                     "filename_prefix": ("STRING", {"default": "ldm_patched"}),
-                     "fps": ("FLOAT", {"default": 6.0, "min": 0.01, "max": 1000.0, "step": 0.01}),
-                     "compress_level": ("INT", {"default": 4, "min": 0, "max": 9})
-                     },
-                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
-                }
-
-    RETURN_TYPES = ()
-    FUNCTION = "save_images"
-
-    OUTPUT_NODE = True
-
-    CATEGORY = "image/animation"
-
-    def save_images(self, images, fps, compress_level, filename_prefix="ldm_patched", prompt=None, extra_pnginfo=None):
-        filename_prefix += self.prefix_append
-        full_output_folder, filename, counter, subfolder, filename_prefix = ldm_patched.utils.path_utils.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0])
-        results = list()
-        pil_images = []
-        for image in images:
-            i = 255. * image.cpu().numpy()
-            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
-            pil_images.append(img)
-
-        metadata = None
-        if not args.disable_server_info:
-            metadata = PngInfo()
-            if prompt is not None:
-                metadata.add(b"ldm_patched", "prompt".encode("latin-1", "strict") + b"\0" + json.dumps(prompt).encode("latin-1", "strict"), after_idat=True)
-            if extra_pnginfo is not None:
-                for x in extra_pnginfo:
-                    metadata.add(b"ldm_patched", x.encode("latin-1", "strict") + b"\0" + json.dumps(extra_pnginfo[x]).encode("latin-1", "strict"), after_idat=True)
-
-        file = f"{filename}_{counter:05}_.png"
-        pil_images[0].save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=compress_level, save_all=True, duration=int(1000.0/fps), append_images=pil_images[1:])
-        results.append({
-            "filename": file,
-            "subfolder": subfolder,
-            "type": self.type
-        })
-
-        return { "ui": { "images": results, "animated": (True,)} }
-
-NODE_CLASS_MAPPINGS = {
-    "ImageCrop": ImageCrop,
-    "RepeatImageBatch": RepeatImageBatch,
-    "SaveAnimatedWEBP": SaveAnimatedWEBP,
-    "SaveAnimatedPNG": SaveAnimatedPNG,
-}

ldm_patched/contrib/external_latent.py

@@ -1,157 +0,0 @@
-# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
-
-import ldm_patched.modules.utils
-import torch
-
-def reshape_latent_to(target_shape, latent):
-    if latent.shape[1:] != target_shape[1:]:
-        latent = ldm_patched.modules.utils.common_upscale(latent, target_shape[3], target_shape[2], "bilinear", "center")
-    return ldm_patched.modules.utils.repeat_to_batch_size(latent, target_shape[0])
-
-
-class LatentAdd:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "samples1": ("LATENT",), "samples2": ("LATENT",)}}
-
-    RETURN_TYPES = ("LATENT",)
-    FUNCTION = "op"
-
-    CATEGORY = "latent/advanced"
-
-    def op(self, samples1, samples2):
-        samples_out = samples1.copy()
-
-        s1 = samples1["samples"]
-        s2 = samples2["samples"]
-
-        s2 = reshape_latent_to(s1.shape, s2)
-        samples_out["samples"] = s1 + s2
-        return (samples_out,)
-
-class LatentSubtract:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "samples1": ("LATENT",), "samples2": ("LATENT",)}}
-
-    RETURN_TYPES = ("LATENT",)
-    FUNCTION = "op"
-
-    CATEGORY = "latent/advanced"
-
-    def op(self, samples1, samples2):
-        samples_out = samples1.copy()
-
-        s1 = samples1["samples"]
-        s2 = samples2["samples"]
-
-        s2 = reshape_latent_to(s1.shape, s2)
-        samples_out["samples"] = s1 - s2
-        return (samples_out,)
-
-class LatentMultiply:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "samples": ("LATENT",),
-                              "multiplier": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
-                             }}
-
-    RETURN_TYPES = ("LATENT",)
-    FUNCTION = "op"
-
-    CATEGORY = "latent/advanced"
-
-    def op(self, samples, multiplier):
-        samples_out = samples.copy()
-
-        s1 = samples["samples"]
-        samples_out["samples"] = s1 * multiplier
-        return (samples_out,)
-
-class LatentInterpolate:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "samples1": ("LATENT",),
-                              "samples2": ("LATENT",),
-                              "ratio": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
-                              }}
-
-    RETURN_TYPES = ("LATENT",)
-    FUNCTION = "op"
-
-    CATEGORY = "latent/advanced"
-
-    def op(self, samples1, samples2, ratio):
-        samples_out = samples1.copy()
-
-        s1 = samples1["samples"]
-        s2 = samples2["samples"]
-
-        s2 = reshape_latent_to(s1.shape, s2)
-
-        m1 = torch.linalg.vector_norm(s1, dim=(1))
-        m2 = torch.linalg.vector_norm(s2, dim=(1))
-
-        s1 = torch.nan_to_num(s1 / m1)
-        s2 = torch.nan_to_num(s2 / m2)
-
-        t = (s1 * ratio + s2 * (1.0 - ratio))
-        mt = torch.linalg.vector_norm(t, dim=(1))
-        st = torch.nan_to_num(t / mt)
-
-        samples_out["samples"] = st * (m1 * ratio + m2 * (1.0 - ratio))
-        return (samples_out,)
-
-class LatentBatch:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "samples1": ("LATENT",), "samples2": ("LATENT",)}}
-
-    RETURN_TYPES = ("LATENT",)
-    FUNCTION = "batch"
-
-    CATEGORY = "latent/batch"
-
-    def batch(self, samples1, samples2):
-        samples_out = samples1.copy()
-        s1 = samples1["samples"]
-        s2 = samples2["samples"]
-
-        if s1.shape[1:] != s2.shape[1:]:
-            s2 = ldm_patched.modules.utils.common_upscale(s2, s1.shape[3], s1.shape[2], "bilinear", "center")
-        s = torch.cat((s1, s2), dim=0)
-        samples_out["samples"] = s
-        samples_out["batch_index"] = samples1.get("batch_index", [x for x in range(0, s1.shape[0])]) + samples2.get("batch_index", [x for x in range(0, s2.shape[0])])
-        return (samples_out,)
-
-class LatentBatchSeedBehavior:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "samples": ("LATENT",),
-                              "seed_behavior": (["random", "fixed"],),}}
-
-    RETURN_TYPES = ("LATENT",)
-    FUNCTION = "op"
-
-    CATEGORY = "latent/advanced"
-
-    def op(self, samples, seed_behavior):
-        samples_out = samples.copy()
-        latent = samples["samples"]
-        if seed_behavior == "random":
-            if 'batch_index' in samples_out:
-                samples_out.pop('batch_index')
-        elif seed_behavior == "fixed":
-            batch_number = samples_out.get("batch_index", [0])[0]
-            samples_out["batch_index"] = [batch_number] * latent.shape[0]
-
-        return (samples_out,)
-
-NODE_CLASS_MAPPINGS = {
-    "LatentAdd": LatentAdd,
-    "LatentSubtract": LatentSubtract,
-    "LatentMultiply": LatentMultiply,
-    "LatentInterpolate": LatentInterpolate,
-    "LatentBatch": LatentBatch,
-    "LatentBatchSeedBehavior": LatentBatchSeedBehavior,
-}

ldm_patched/contrib/external_mask.py

@@ -1,365 +0,0 @@
-# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
-
-import numpy as np
-import scipy.ndimage
-import torch
-import ldm_patched.modules.utils
-
-from ldm_patched.contrib.external import MAX_RESOLUTION
-
-def composite(destination, source, x, y, mask = None, multiplier = 8, resize_source = False):
-    source = source.to(destination.device)
-    if resize_source:
-        source = torch.nn.functional.interpolate(source, size=(destination.shape[2], destination.shape[3]), mode="bilinear")
-
-    source = ldm_patched.modules.utils.repeat_to_batch_size(source, destination.shape[0])
-
-    x = max(-source.shape[3] * multiplier, min(x, destination.shape[3] * multiplier))
-    y = max(-source.shape[2] * multiplier, min(y, destination.shape[2] * multiplier))
-
-    left, top = (x // multiplier, y // multiplier)
-    right, bottom = (left + source.shape[3], top + source.shape[2],)
-
-    if mask is None:
-        mask = torch.ones_like(source)
-    else:
-        mask = mask.to(destination.device, copy=True)
-        mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(source.shape[2], source.shape[3]), mode="bilinear")
-        mask = ldm_patched.modules.utils.repeat_to_batch_size(mask, source.shape[0])
-
-    # calculate the bounds of the source that will be overlapping the destination
-    # this prevents the source trying to overwrite latent pixels that are out of bounds
-    # of the destination
-    visible_width, visible_height = (destination.shape[3] - left + min(0, x), destination.shape[2] - top + min(0, y),)
-
-    mask = mask[:, :, :visible_height, :visible_width]
-    inverse_mask = torch.ones_like(mask) - mask
-
-    source_portion = mask * source[:, :, :visible_height, :visible_width]
-    destination_portion = inverse_mask  * destination[:, :, top:bottom, left:right]
-
-    destination[:, :, top:bottom, left:right] = source_portion + destination_portion
-    return destination
-
-class LatentCompositeMasked:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "destination": ("LATENT",),
-                "source": ("LATENT",),
-                "x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
-                "y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 8}),
-                "resize_source": ("BOOLEAN", {"default": False}),
-            },
-            "optional": {
-                "mask": ("MASK",),
-            }
-        }
-    RETURN_TYPES = ("LATENT",)
-    FUNCTION = "composite"
-
-    CATEGORY = "latent"
-
-    def composite(self, destination, source, x, y, resize_source, mask = None):
-        output = destination.copy()
-        destination = destination["samples"].clone()
-        source = source["samples"]
-        output["samples"] = composite(destination, source, x, y, mask, 8, resize_source)
-        return (output,)
-
-class ImageCompositeMasked:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "destination": ("IMAGE",),
-                "source": ("IMAGE",),
-                "x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
-                "y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
-                "resize_source": ("BOOLEAN", {"default": False}),
-            },
-            "optional": {
-                "mask": ("MASK",),
-            }
-        }
-    RETURN_TYPES = ("IMAGE",)
-    FUNCTION = "composite"
-
-    CATEGORY = "image"
-
-    def composite(self, destination, source, x, y, resize_source, mask = None):
-        destination = destination.clone().movedim(-1, 1)
-        output = composite(destination, source.movedim(-1, 1), x, y, mask, 1, resize_source).movedim(1, -1)
-        return (output,)
-
-class MaskToImage:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-                "required": {
-                    "mask": ("MASK",),
-                }
-        }
-
-    CATEGORY = "mask"
-
-    RETURN_TYPES = ("IMAGE",)
-    FUNCTION = "mask_to_image"
-
-    def mask_to_image(self, mask):
-        result = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
-        return (result,)
-
-class ImageToMask:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-                "required": {
-                    "image": ("IMAGE",),
-                    "channel": (["red", "green", "blue", "alpha"],),
-                }
-        }
-
-    CATEGORY = "mask"
-
-    RETURN_TYPES = ("MASK",)
-    FUNCTION = "image_to_mask"
-
-    def image_to_mask(self, image, channel):
-        channels = ["red", "green", "blue", "alpha"]
-        mask = image[:, :, :, channels.index(channel)]
-        return (mask,)
-
-class ImageColorToMask:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-                "required": {
-                    "image": ("IMAGE",),
-                    "color": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFF, "step": 1, "display": "color"}),
-                }
-        }
-
-    CATEGORY = "mask"
-
-    RETURN_TYPES = ("MASK",)
-    FUNCTION = "image_to_mask"
-
-    def image_to_mask(self, image, color):
-        temp = (torch.clamp(image, 0, 1.0) * 255.0).round().to(torch.int)
-        temp = torch.bitwise_left_shift(temp[:,:,:,0], 16) + torch.bitwise_left_shift(temp[:,:,:,1], 8) + temp[:,:,:,2]
-        mask = torch.where(temp == color, 255, 0).float()
-        return (mask,)
-
-class SolidMask:
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "value": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
-                "width": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}),
-                "height": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}),
-            }
-        }
-
-    CATEGORY = "mask"
-
-    RETURN_TYPES = ("MASK",)
-
-    FUNCTION = "solid"
-
-    def solid(self, value, width, height):
-        out = torch.full((1, height, width), value, dtype=torch.float32, device="cpu")
-        return (out,)
-
-class InvertMask:
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "mask": ("MASK",),
-            }
-        }
-
-    CATEGORY = "mask"
-
-    RETURN_TYPES = ("MASK",)
-
-    FUNCTION = "invert"
-
-    def invert(self, mask):
-        out = 1.0 - mask
-        return (out,)
-
-class CropMask:
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "mask": ("MASK",),
-                "x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
-                "y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
-                "width": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}),
-                "height": ("INT", {"default": 512, "min": 1, "max": MAX_RESOLUTION, "step": 1}),
-            }
-        }
-
-    CATEGORY = "mask"
-
-    RETURN_TYPES = ("MASK",)
-
-    FUNCTION = "crop"
-
-    def crop(self, mask, x, y, width, height):
-        mask = mask.reshape((-1, mask.shape[-2], mask.shape[-1]))
-        out = mask[:, y:y + height, x:x + width]
-        return (out,)
-
-class MaskComposite:
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "destination": ("MASK",),
-                "source": ("MASK",),
-                "x": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
-                "y": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
-                "operation": (["multiply", "add", "subtract", "and", "or", "xor"],),
-            }
-        }
-
-    CATEGORY = "mask"
-
-    RETURN_TYPES = ("MASK",)
-
-    FUNCTION = "combine"
-
-    def combine(self, destination, source, x, y, operation):
-        output = destination.reshape((-1, destination.shape[-2], destination.shape[-1])).clone()
-        source = source.reshape((-1, source.shape[-2], source.shape[-1]))
-
-        left, top = (x, y,)
-        right, bottom = (min(left + source.shape[-1], destination.shape[-1]), min(top + source.shape[-2], destination.shape[-2]))
-        visible_width, visible_height = (right - left, bottom - top,)
-
-        source_portion = source[:, :visible_height, :visible_width]
-        destination_portion = destination[:, top:bottom, left:right]
-
-        if operation == "multiply":
-            output[:, top:bottom, left:right] = destination_portion * source_portion
-        elif operation == "add":
-            output[:, top:bottom, left:right] = destination_portion + source_portion
-        elif operation == "subtract":
-            output[:, top:bottom, left:right] = destination_portion - source_portion
-        elif operation == "and":
-            output[:, top:bottom, left:right] = torch.bitwise_and(destination_portion.round().bool(), source_portion.round().bool()).float()
-        elif operation == "or":
-            output[:, top:bottom, left:right] = torch.bitwise_or(destination_portion.round().bool(), source_portion.round().bool()).float()
-        elif operation == "xor":
-            output[:, top:bottom, left:right] = torch.bitwise_xor(destination_portion.round().bool(), source_portion.round().bool()).float()
-
-        output = torch.clamp(output, 0.0, 1.0)
-
-        return (output,)
-
-class FeatherMask:
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "mask": ("MASK",),
-                "left": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
-                "top": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
-                "right": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
-                "bottom": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION, "step": 1}),
-            }
-        }
-
-    CATEGORY = "mask"
-
-    RETURN_TYPES = ("MASK",)
-
-    FUNCTION = "feather"
-
-    def feather(self, mask, left, top, right, bottom):
-        output = mask.reshape((-1, mask.shape[-2], mask.shape[-1])).clone()
-
-        left = min(left, output.shape[-1])
-        right = min(right, output.shape[-1])
-        top = min(top, output.shape[-2])
-        bottom = min(bottom, output.shape[-2])
-
-        for x in range(left):
-            feather_rate = (x + 1.0) / left
-            output[:, :, x] *= feather_rate
-
-        for x in range(right):
-            feather_rate = (x + 1) / right
-            output[:, :, -x] *= feather_rate
-
-        for y in range(top):
-            feather_rate = (y + 1) / top
-            output[:, y, :] *= feather_rate
-
-        for y in range(bottom):
-            feather_rate = (y + 1) / bottom
-            output[:, -y, :] *= feather_rate
-
-        return (output,)
-    
-class GrowMask:
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {
-                "mask": ("MASK",),
-                "expand": ("INT", {"default": 0, "min": -MAX_RESOLUTION, "max": MAX_RESOLUTION, "step": 1}),
-                "tapered_corners": ("BOOLEAN", {"default": True}),
-            },
-        }
-    
-    CATEGORY = "mask"
-
-    RETURN_TYPES = ("MASK",)
-
-    FUNCTION = "expand_mask"
-
-    def expand_mask(self, mask, expand, tapered_corners):
-        c = 0 if tapered_corners else 1
-        kernel = np.array([[c, 1, c],
-                           [1, 1, 1],
-                           [c, 1, c]])
-        mask = mask.reshape((-1, mask.shape[-2], mask.shape[-1]))
-        out = []
-        for m in mask:
-            output = m.numpy()
-            for _ in range(abs(expand)):
-                if expand < 0:
-                    output = scipy.ndimage.grey_erosion(output, footprint=kernel)
-                else:
-                    output = scipy.ndimage.grey_dilation(output, footprint=kernel)
-            output = torch.from_numpy(output)
-            out.append(output)
-        return (torch.stack(out, dim=0),)
-
-
-
-NODE_CLASS_MAPPINGS = {
-    "LatentCompositeMasked": LatentCompositeMasked,
-    "ImageCompositeMasked": ImageCompositeMasked,
-    "MaskToImage": MaskToImage,
-    "ImageToMask": ImageToMask,
-    "ImageColorToMask": ImageColorToMask,
-    "SolidMask": SolidMask,
-    "InvertMask": InvertMask,
-    "CropMask": CropMask,
-    "MaskComposite": MaskComposite,
-    "FeatherMask": FeatherMask,
-    "GrowMask": GrowMask,
-}
-
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "ImageToMask": "Convert Image to Mask",
-    "MaskToImage": "Convert Mask to Image",
-}

ldm_patched/contrib/external_model_advanced.py

@@ -1,177 +0,0 @@
-# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
-
-import ldm_patched.utils.path_utils
-import ldm_patched.modules.sd
-import ldm_patched.modules.model_sampling
-import torch
-
-class LCM(ldm_patched.modules.model_sampling.EPS):
-    def calculate_denoised(self, sigma, model_output, model_input):
-        timestep = self.timestep(sigma).view(sigma.shape[:1] + (1,) * (model_output.ndim - 1))
-        sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1))
-        x0 = model_input - model_output * sigma
-
-        sigma_data = 0.5
-        scaled_timestep = timestep * 10.0 #timestep_scaling
-
-        c_skip = sigma_data**2 / (scaled_timestep**2 + sigma_data**2)
-        c_out = scaled_timestep / (scaled_timestep**2 + sigma_data**2) ** 0.5
-
-        return c_out * x0 + c_skip * model_input
-
-class ModelSamplingDiscreteDistilled(ldm_patched.modules.model_sampling.ModelSamplingDiscrete):
-    original_timesteps = 50
-
-    def __init__(self, model_config=None):
-        super().__init__(model_config)
-
-        self.skip_steps = self.num_timesteps // self.original_timesteps
-
-        sigmas_valid = torch.zeros((self.original_timesteps), dtype=torch.float32)
-        for x in range(self.original_timesteps):
-            sigmas_valid[self.original_timesteps - 1 - x] = self.sigmas[self.num_timesteps - 1 - x * self.skip_steps]
-
-        self.set_sigmas(sigmas_valid)
-
-    def timestep(self, sigma):
-        log_sigma = sigma.log()
-        dists = log_sigma.to(self.log_sigmas.device) - self.log_sigmas[:, None]
-        return (dists.abs().argmin(dim=0).view(sigma.shape) * self.skip_steps + (self.skip_steps - 1)).to(sigma.device)
-
-    def sigma(self, timestep):
-        t = torch.clamp(((timestep.float().to(self.log_sigmas.device) - (self.skip_steps - 1)) / self.skip_steps).float(), min=0, max=(len(self.sigmas) - 1))
-        low_idx = t.floor().long()
-        high_idx = t.ceil().long()
-        w = t.frac()
-        log_sigma = (1 - w) * self.log_sigmas[low_idx] + w * self.log_sigmas[high_idx]
-        return log_sigma.exp().to(timestep.device)
-
-
-def rescale_zero_terminal_snr_sigmas(sigmas):
-    alphas_cumprod = 1 / ((sigmas * sigmas) + 1)
-    alphas_bar_sqrt = alphas_cumprod.sqrt()
-
-    # Store old values.
-    alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone()
-    alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone()
-
-    # Shift so the last timestep is zero.
-    alphas_bar_sqrt -= (alphas_bar_sqrt_T)
-
-    # Scale so the first timestep is back to the old value.
-    alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T)
-
-    # Convert alphas_bar_sqrt to betas
-    alphas_bar = alphas_bar_sqrt**2  # Revert sqrt
-    alphas_bar[-1] = 4.8973451890853435e-08
-    return ((1 - alphas_bar) / alphas_bar) ** 0.5
-
-class ModelSamplingDiscrete:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model": ("MODEL",),
-                              "sampling": (["eps", "v_prediction", "lcm"],),
-                              "zsnr": ("BOOLEAN", {"default": False}),
-                              }}
-
-    RETURN_TYPES = ("MODEL",)
-    FUNCTION = "patch"
-
-    CATEGORY = "advanced/model"
-
-    def patch(self, model, sampling, zsnr):
-        m = model.clone()
-
-        sampling_base = ldm_patched.modules.model_sampling.ModelSamplingDiscrete
-        if sampling == "eps":
-            sampling_type = ldm_patched.modules.model_sampling.EPS
-        elif sampling == "v_prediction":
-            sampling_type = ldm_patched.modules.model_sampling.V_PREDICTION
-        elif sampling == "lcm":
-            sampling_type = LCM
-            sampling_base = ModelSamplingDiscreteDistilled
-
-        class ModelSamplingAdvanced(sampling_base, sampling_type):
-            pass
-
-        model_sampling = ModelSamplingAdvanced(model.model.model_config)
-        if zsnr:
-            model_sampling.set_sigmas(rescale_zero_terminal_snr_sigmas(model_sampling.sigmas))
-
-        m.add_object_patch("model_sampling", model_sampling)
-        return (m, )
-
-class ModelSamplingContinuousEDM:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model": ("MODEL",),
-                              "sampling": (["v_prediction", "eps"],),
-                              "sigma_max": ("FLOAT", {"default": 120.0, "min": 0.0, "max": 1000.0, "step":0.001, "round": False}),
-                              "sigma_min": ("FLOAT", {"default": 0.002, "min": 0.0, "max": 1000.0, "step":0.001, "round": False}),
-                              }}
-
-    RETURN_TYPES = ("MODEL",)
-    FUNCTION = "patch"
-
-    CATEGORY = "advanced/model"
-
-    def patch(self, model, sampling, sigma_max, sigma_min):
-        m = model.clone()
-
-        if sampling == "eps":
-            sampling_type = ldm_patched.modules.model_sampling.EPS
-        elif sampling == "v_prediction":
-            sampling_type = ldm_patched.modules.model_sampling.V_PREDICTION
-
-        class ModelSamplingAdvanced(ldm_patched.modules.model_sampling.ModelSamplingContinuousEDM, sampling_type):
-            pass
-
-        model_sampling = ModelSamplingAdvanced(model.model.model_config)
-        model_sampling.set_sigma_range(sigma_min, sigma_max)
-        m.add_object_patch("model_sampling", model_sampling)
-        return (m, )
-
-class RescaleCFG:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model": ("MODEL",),
-                              "multiplier": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 1.0, "step": 0.01}),
-                              }}
-    RETURN_TYPES = ("MODEL",)
-    FUNCTION = "patch"
-
-    CATEGORY = "advanced/model"
-
-    def patch(self, model, multiplier):
-        def rescale_cfg(args):
-            cond = args["cond"]
-            uncond = args["uncond"]
-            cond_scale = args["cond_scale"]
-            sigma = args["sigma"]
-            sigma = sigma.view(sigma.shape[:1] + (1,) * (cond.ndim - 1))
-            x_orig = args["input"]
-
-            #rescale cfg has to be done on v-pred model output
-            x = x_orig / (sigma * sigma + 1.0)
-            cond = ((x - (x_orig - cond)) * (sigma ** 2 + 1.0) ** 0.5) / (sigma)
-            uncond = ((x - (x_orig - uncond)) * (sigma ** 2 + 1.0) ** 0.5) / (sigma)
-
-            #rescalecfg
-            x_cfg = uncond + cond_scale * (cond - uncond)
-            ro_pos = torch.std(cond, dim=(1,2,3), keepdim=True)
-            ro_cfg = torch.std(x_cfg, dim=(1,2,3), keepdim=True)
-
-            x_rescaled = x_cfg * (ro_pos / ro_cfg)
-            x_final = multiplier * x_rescaled + (1.0 - multiplier) * x_cfg
-
-            return x_orig - (x - x_final * sigma / (sigma * sigma + 1.0) ** 0.5)
-
-        m = model.clone()
-        m.set_model_sampler_cfg_function(rescale_cfg)
-        return (m, )
-
-NODE_CLASS_MAPPINGS = {
-    "ModelSamplingDiscrete": ModelSamplingDiscrete,
-    "ModelSamplingContinuousEDM": ModelSamplingContinuousEDM,
-    "RescaleCFG": RescaleCFG,
-}

ldm_patched/contrib/external_model_downscale.py

@@ -1,55 +0,0 @@
-# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
-
-import torch
-import ldm_patched.modules.utils
-
-class PatchModelAddDownscale:
-    upscale_methods = ["bicubic", "nearest-exact", "bilinear", "area", "bislerp"]
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model": ("MODEL",),
-                              "block_number": ("INT", {"default": 3, "min": 1, "max": 32, "step": 1}),
-                              "downscale_factor": ("FLOAT", {"default": 2.0, "min": 0.1, "max": 9.0, "step": 0.001}),
-                              "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
-                              "end_percent": ("FLOAT", {"default": 0.35, "min": 0.0, "max": 1.0, "step": 0.001}),
-                              "downscale_after_skip": ("BOOLEAN", {"default": True}),
-                              "downscale_method": (s.upscale_methods,),
-                              "upscale_method": (s.upscale_methods,),
-                              }}
-    RETURN_TYPES = ("MODEL",)
-    FUNCTION = "patch"
-
-    CATEGORY = "_for_testing"
-
-    def patch(self, model, block_number, downscale_factor, start_percent, end_percent, downscale_after_skip, downscale_method, upscale_method):
-        sigma_start = model.model.model_sampling.percent_to_sigma(start_percent)
-        sigma_end = model.model.model_sampling.percent_to_sigma(end_percent)
-
-        def input_block_patch(h, transformer_options):
-            if transformer_options["block"][1] == block_number:
-                sigma = transformer_options["sigmas"][0].item()
-                if sigma <= sigma_start and sigma >= sigma_end:
-                    h = ldm_patched.modules.utils.common_upscale(h, round(h.shape[-1] * (1.0 / downscale_factor)), round(h.shape[-2] * (1.0 / downscale_factor)), downscale_method, "disabled")
-            return h
-
-        def output_block_patch(h, hsp, transformer_options):
-            if h.shape[2] != hsp.shape[2]:
-                h = ldm_patched.modules.utils.common_upscale(h, hsp.shape[-1], hsp.shape[-2], upscale_method, "disabled")
-            return h, hsp
-
-        m = model.clone()
-        if downscale_after_skip:
-            m.set_model_input_block_patch_after_skip(input_block_patch)
-        else:
-            m.set_model_input_block_patch(input_block_patch)
-        m.set_model_output_block_patch(output_block_patch)
-        return (m, )
-
-NODE_CLASS_MAPPINGS = {
-    "PatchModelAddDownscale": PatchModelAddDownscale,
-}
-
-NODE_DISPLAY_NAME_MAPPINGS = {
-    # Sampling
-    "PatchModelAddDownscale": "PatchModelAddDownscale (Kohya Deep Shrink)",
-}

ldm_patched/contrib/external_model_merging.py

@@ -1,286 +0,0 @@
-# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
-
-import ldm_patched.modules.sd
-import ldm_patched.modules.utils
-import ldm_patched.modules.model_base
-import ldm_patched.modules.model_management
-
-import ldm_patched.utils.path_utils
-import json
-import os
-
-from ldm_patched.modules.args_parser import args
-
-class ModelMergeSimple:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model1": ("MODEL",),
-                              "model2": ("MODEL",),
-                              "ratio": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
-                              }}
-    RETURN_TYPES = ("MODEL",)
-    FUNCTION = "merge"
-
-    CATEGORY = "advanced/model_merging"
-
-    def merge(self, model1, model2, ratio):
-        m = model1.clone()
-        kp = model2.get_key_patches("diffusion_model.")
-        for k in kp:
-            m.add_patches({k: kp[k]}, 1.0 - ratio, ratio)
-        return (m, )
-
-class ModelSubtract:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model1": ("MODEL",),
-                              "model2": ("MODEL",),
-                              "multiplier": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
-                              }}
-    RETURN_TYPES = ("MODEL",)
-    FUNCTION = "merge"
-
-    CATEGORY = "advanced/model_merging"
-
-    def merge(self, model1, model2, multiplier):
-        m = model1.clone()
-        kp = model2.get_key_patches("diffusion_model.")
-        for k in kp:
-            m.add_patches({k: kp[k]}, - multiplier, multiplier)
-        return (m, )
-
-class ModelAdd:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model1": ("MODEL",),
-                              "model2": ("MODEL",),
-                              }}
-    RETURN_TYPES = ("MODEL",)
-    FUNCTION = "merge"
-
-    CATEGORY = "advanced/model_merging"
-
-    def merge(self, model1, model2):
-        m = model1.clone()
-        kp = model2.get_key_patches("diffusion_model.")
-        for k in kp:
-            m.add_patches({k: kp[k]}, 1.0, 1.0)
-        return (m, )
-
-
-class CLIPMergeSimple:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "clip1": ("CLIP",),
-                              "clip2": ("CLIP",),
-                              "ratio": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
-                              }}
-    RETURN_TYPES = ("CLIP",)
-    FUNCTION = "merge"
-
-    CATEGORY = "advanced/model_merging"
-
-    def merge(self, clip1, clip2, ratio):
-        m = clip1.clone()
-        kp = clip2.get_key_patches()
-        for k in kp:
-            if k.endswith(".position_ids") or k.endswith(".logit_scale"):
-                continue
-            m.add_patches({k: kp[k]}, 1.0 - ratio, ratio)
-        return (m, )
-
-class ModelMergeBlocks:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model1": ("MODEL",),
-                              "model2": ("MODEL",),
-                              "input": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
-                              "middle": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
-                              "out": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01})
-                              }}
-    RETURN_TYPES = ("MODEL",)
-    FUNCTION = "merge"
-
-    CATEGORY = "advanced/model_merging"
-
-    def merge(self, model1, model2, **kwargs):
-        m = model1.clone()
-        kp = model2.get_key_patches("diffusion_model.")
-        default_ratio = next(iter(kwargs.values()))
-
-        for k in kp:
-            ratio = default_ratio
-            k_unet = k[len("diffusion_model."):]
-
-            last_arg_size = 0
-            for arg in kwargs:
-                if k_unet.startswith(arg) and last_arg_size < len(arg):
-                    ratio = kwargs[arg]
-                    last_arg_size = len(arg)
-
-            m.add_patches({k: kp[k]}, 1.0 - ratio, ratio)
-        return (m, )
-
-def save_checkpoint(model, clip=None, vae=None, clip_vision=None, filename_prefix=None, output_dir=None, prompt=None, extra_pnginfo=None):
-    full_output_folder, filename, counter, subfolder, filename_prefix = ldm_patched.utils.path_utils.get_save_image_path(filename_prefix, output_dir)
-    prompt_info = ""
-    if prompt is not None:
-        prompt_info = json.dumps(prompt)
-
-    metadata = {}
-
-    enable_modelspec = True
-    if isinstance(model.model, ldm_patched.modules.model_base.SDXL):
-        metadata["modelspec.architecture"] = "stable-diffusion-xl-v1-base"
-    elif isinstance(model.model, ldm_patched.modules.model_base.SDXLRefiner):
-        metadata["modelspec.architecture"] = "stable-diffusion-xl-v1-refiner"
-    else:
-        enable_modelspec = False
-
-    if enable_modelspec:
-        metadata["modelspec.sai_model_spec"] = "1.0.0"
-        metadata["modelspec.implementation"] = "sgm"
-        metadata["modelspec.title"] = "{} {}".format(filename, counter)
-
-    #TODO:
-    # "stable-diffusion-v1", "stable-diffusion-v1-inpainting", "stable-diffusion-v2-512",
-    # "stable-diffusion-v2-768-v", "stable-diffusion-v2-unclip-l", "stable-diffusion-v2-unclip-h",
-    # "v2-inpainting"
-
-    if model.model.model_type == ldm_patched.modules.model_base.ModelType.EPS:
-        metadata["modelspec.predict_key"] = "epsilon"
-    elif model.model.model_type == ldm_patched.modules.model_base.ModelType.V_PREDICTION:
-        metadata["modelspec.predict_key"] = "v"
-
-    if not args.disable_server_info:
-        metadata["prompt"] = prompt_info
-        if extra_pnginfo is not None:
-            for x in extra_pnginfo:
-                metadata[x] = json.dumps(extra_pnginfo[x])
-
-    output_checkpoint = f"{filename}_{counter:05}_.safetensors"
-    output_checkpoint = os.path.join(full_output_folder, output_checkpoint)
-
-    ldm_patched.modules.sd.save_checkpoint(output_checkpoint, model, clip, vae, clip_vision, metadata=metadata)
-
-class CheckpointSave:
-    def __init__(self):
-        self.output_dir = ldm_patched.utils.path_utils.get_output_directory()
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model": ("MODEL",),
-                              "clip": ("CLIP",),
-                              "vae": ("VAE",),
-                              "filename_prefix": ("STRING", {"default": "checkpoints/ldm_patched"}),},
-                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},}
-    RETURN_TYPES = ()
-    FUNCTION = "save"
-    OUTPUT_NODE = True
-
-    CATEGORY = "advanced/model_merging"
-
-    def save(self, model, clip, vae, filename_prefix, prompt=None, extra_pnginfo=None):
-        save_checkpoint(model, clip=clip, vae=vae, filename_prefix=filename_prefix, output_dir=self.output_dir, prompt=prompt, extra_pnginfo=extra_pnginfo)
-        return {}
-
-class CLIPSave:
-    def __init__(self):
-        self.output_dir = ldm_patched.utils.path_utils.get_output_directory()
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "clip": ("CLIP",),
-                              "filename_prefix": ("STRING", {"default": "clip/ldm_patched"}),},
-                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},}
-    RETURN_TYPES = ()
-    FUNCTION = "save"
-    OUTPUT_NODE = True
-
-    CATEGORY = "advanced/model_merging"
-
-    def save(self, clip, filename_prefix, prompt=None, extra_pnginfo=None):
-        prompt_info = ""
-        if prompt is not None:
-            prompt_info = json.dumps(prompt)
-
-        metadata = {}
-        if not args.disable_server_info:
-            metadata["prompt"] = prompt_info
-            if extra_pnginfo is not None:
-                for x in extra_pnginfo:
-                    metadata[x] = json.dumps(extra_pnginfo[x])
-
-        ldm_patched.modules.model_management.load_models_gpu([clip.load_model()])
-        clip_sd = clip.get_sd()
-
-        for prefix in ["clip_l.", "clip_g.", ""]:
-            k = list(filter(lambda a: a.startswith(prefix), clip_sd.keys()))
-            current_clip_sd = {}
-            for x in k:
-                current_clip_sd[x] = clip_sd.pop(x)
-            if len(current_clip_sd) == 0:
-                continue
-
-            p = prefix[:-1]
-            replace_prefix = {}
-            filename_prefix_ = filename_prefix
-            if len(p) > 0:
-                filename_prefix_ = "{}_{}".format(filename_prefix_, p)
-                replace_prefix[prefix] = ""
-            replace_prefix["transformer."] = ""
-
-            full_output_folder, filename, counter, subfolder, filename_prefix_ = ldm_patched.utils.path_utils.get_save_image_path(filename_prefix_, self.output_dir)
-
-            output_checkpoint = f"{filename}_{counter:05}_.safetensors"
-            output_checkpoint = os.path.join(full_output_folder, output_checkpoint)
-
-            current_clip_sd = ldm_patched.modules.utils.state_dict_prefix_replace(current_clip_sd, replace_prefix)
-
-            ldm_patched.modules.utils.save_torch_file(current_clip_sd, output_checkpoint, metadata=metadata)
-        return {}
-
-class VAESave:
-    def __init__(self):
-        self.output_dir = ldm_patched.utils.path_utils.get_output_directory()
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "vae": ("VAE",),
-                              "filename_prefix": ("STRING", {"default": "vae/ldm_patched_vae"}),},
-                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},}
-    RETURN_TYPES = ()
-    FUNCTION = "save"
-    OUTPUT_NODE = True
-
-    CATEGORY = "advanced/model_merging"
-
-    def save(self, vae, filename_prefix, prompt=None, extra_pnginfo=None):
-        full_output_folder, filename, counter, subfolder, filename_prefix = ldm_patched.utils.path_utils.get_save_image_path(filename_prefix, self.output_dir)
-        prompt_info = ""
-        if prompt is not None:
-            prompt_info = json.dumps(prompt)
-
-        metadata = {}
-        if not args.disable_server_info:
-            metadata["prompt"] = prompt_info
-            if extra_pnginfo is not None:
-                for x in extra_pnginfo:
-                    metadata[x] = json.dumps(extra_pnginfo[x])
-
-        output_checkpoint = f"{filename}_{counter:05}_.safetensors"
-        output_checkpoint = os.path.join(full_output_folder, output_checkpoint)
-
-        ldm_patched.modules.utils.save_torch_file(vae.get_sd(), output_checkpoint, metadata=metadata)
-        return {}
-
-NODE_CLASS_MAPPINGS = {
-    "ModelMergeSimple": ModelMergeSimple,
-    "ModelMergeBlocks": ModelMergeBlocks,
-    "ModelMergeSubtract": ModelSubtract,
-    "ModelMergeAdd": ModelAdd,
-    "CheckpointSave": CheckpointSave,
-    "CLIPMergeSimple": CLIPMergeSimple,
-    "CLIPSave": CLIPSave,
-    "VAESave": VAESave,
-}

ldm_patched/contrib/external_perpneg.py

@@ -1,57 +0,0 @@
-# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
-
-import torch
-import ldm_patched.modules.model_management
-import ldm_patched.modules.sample
-import ldm_patched.modules.samplers
-import ldm_patched.modules.utils
-
-
-class PerpNeg:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": {"model": ("MODEL", ),
-                             "empty_conditioning": ("CONDITIONING", ),
-                             "neg_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0}),
-                            }}
-    RETURN_TYPES = ("MODEL",)
-    FUNCTION = "patch"
-
-    CATEGORY = "_for_testing"
-
-    def patch(self, model, empty_conditioning, neg_scale):
-        m = model.clone()
-        nocond = ldm_patched.modules.sample.convert_cond(empty_conditioning)
-
-        def cfg_function(args):
-            model = args["model"]
-            noise_pred_pos = args["cond_denoised"]
-            noise_pred_neg = args["uncond_denoised"]
-            cond_scale = args["cond_scale"]
-            x = args["input"]
-            sigma = args["sigma"]
-            model_options = args["model_options"]
-            nocond_processed = ldm_patched.modules.samplers.encode_model_conds(model.extra_conds, nocond, x, x.device, "negative")
-
-            (noise_pred_nocond, _) = ldm_patched.modules.samplers.calc_cond_uncond_batch(model, nocond_processed, None, x, sigma, model_options)
-
-            pos = noise_pred_pos - noise_pred_nocond
-            neg = noise_pred_neg - noise_pred_nocond
-            perp = ((torch.mul(pos, neg).sum())/(torch.norm(neg)**2)) * neg
-            perp_neg = perp * neg_scale
-            cfg_result = noise_pred_nocond + cond_scale*(pos - perp_neg)
-            cfg_result = x - cfg_result
-            return cfg_result
-
-        m.set_model_sampler_cfg_function(cfg_function)
-
-        return (m, )
-
-
-NODE_CLASS_MAPPINGS = {
-    "PerpNeg": PerpNeg,
-}
-
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "PerpNeg": "Perp-Neg",
-}

ldm_patched/contrib/external_photomaker.py

@@ -1,189 +0,0 @@
-# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
-
-import torch
-import torch.nn as nn
-import ldm_patched.utils.path_utils
-import ldm_patched.modules.clip_model
-import ldm_patched.modules.clip_vision
-import ldm_patched.modules.ops
-
-# code for model from: https://github.com/TencentARC/PhotoMaker/blob/main/photomaker/model.py under Apache License Version 2.0
-VISION_CONFIG_DICT = {
-    "hidden_size": 1024,
-    "image_size": 224,
-    "intermediate_size": 4096,
-    "num_attention_heads": 16,
-    "num_channels": 3,
-    "num_hidden_layers": 24,
-    "patch_size": 14,
-    "projection_dim": 768,
-    "hidden_act": "quick_gelu",
-}
-
-class MLP(nn.Module):
-    def __init__(self, in_dim, out_dim, hidden_dim, use_residual=True, operations=ldm_patched.modules.ops):
-        super().__init__()
-        if use_residual:
-            assert in_dim == out_dim
-        self.layernorm = operations.LayerNorm(in_dim)
-        self.fc1 = operations.Linear(in_dim, hidden_dim)
-        self.fc2 = operations.Linear(hidden_dim, out_dim)
-        self.use_residual = use_residual
-        self.act_fn = nn.GELU()
-
-    def forward(self, x):
-        residual = x
-        x = self.layernorm(x)
-        x = self.fc1(x)
-        x = self.act_fn(x)
-        x = self.fc2(x)
-        if self.use_residual:
-            x = x + residual
-        return x
-
-
-class FuseModule(nn.Module):
-    def __init__(self, embed_dim, operations):
-        super().__init__()
-        self.mlp1 = MLP(embed_dim * 2, embed_dim, embed_dim, use_residual=False, operations=operations)
-        self.mlp2 = MLP(embed_dim, embed_dim, embed_dim, use_residual=True, operations=operations)
-        self.layer_norm = operations.LayerNorm(embed_dim)
-
-    def fuse_fn(self, prompt_embeds, id_embeds):
-        stacked_id_embeds = torch.cat([prompt_embeds, id_embeds], dim=-1)
-        stacked_id_embeds = self.mlp1(stacked_id_embeds) + prompt_embeds
-        stacked_id_embeds = self.mlp2(stacked_id_embeds)
-        stacked_id_embeds = self.layer_norm(stacked_id_embeds)
-        return stacked_id_embeds
-
-    def forward(
-        self,
-        prompt_embeds,
-        id_embeds,
-        class_tokens_mask,
-    ) -> torch.Tensor:
-        # id_embeds shape: [b, max_num_inputs, 1, 2048]
-        id_embeds = id_embeds.to(prompt_embeds.dtype)
-        num_inputs = class_tokens_mask.sum().unsqueeze(0) # TODO: check for training case
-        batch_size, max_num_inputs = id_embeds.shape[:2]
-        # seq_length: 77
-        seq_length = prompt_embeds.shape[1]
-        # flat_id_embeds shape: [b*max_num_inputs, 1, 2048]
-        flat_id_embeds = id_embeds.view(
-            -1, id_embeds.shape[-2], id_embeds.shape[-1]
-        )
-        # valid_id_mask [b*max_num_inputs]
-        valid_id_mask = (
-            torch.arange(max_num_inputs, device=flat_id_embeds.device)[None, :]
-            < num_inputs[:, None]
-        )
-        valid_id_embeds = flat_id_embeds[valid_id_mask.flatten()]
-
-        prompt_embeds = prompt_embeds.view(-1, prompt_embeds.shape[-1])
-        class_tokens_mask = class_tokens_mask.view(-1)
-        valid_id_embeds = valid_id_embeds.view(-1, valid_id_embeds.shape[-1])
-        # slice out the image token embeddings
-        image_token_embeds = prompt_embeds[class_tokens_mask]
-        stacked_id_embeds = self.fuse_fn(image_token_embeds, valid_id_embeds)
-        assert class_tokens_mask.sum() == stacked_id_embeds.shape[0], f"{class_tokens_mask.sum()} != {stacked_id_embeds.shape[0]}"
-        prompt_embeds.masked_scatter_(class_tokens_mask[:, None], stacked_id_embeds.to(prompt_embeds.dtype))
-        updated_prompt_embeds = prompt_embeds.view(batch_size, seq_length, -1)
-        return updated_prompt_embeds
-
-class PhotoMakerIDEncoder(ldm_patched.modules.clip_model.CLIPVisionModelProjection):
-    def __init__(self):
-        self.load_device = ldm_patched.modules.model_management.text_encoder_device()
-        offload_device = ldm_patched.modules.model_management.text_encoder_offload_device()
-        dtype = ldm_patched.modules.model_management.text_encoder_dtype(self.load_device)
-
-        super().__init__(VISION_CONFIG_DICT, dtype, offload_device, ldm_patched.modules.ops.manual_cast)
-        self.visual_projection_2 = ldm_patched.modules.ops.manual_cast.Linear(1024, 1280, bias=False)
-        self.fuse_module = FuseModule(2048, ldm_patched.modules.ops.manual_cast)
-
-    def forward(self, id_pixel_values, prompt_embeds, class_tokens_mask):
-        b, num_inputs, c, h, w = id_pixel_values.shape
-        id_pixel_values = id_pixel_values.view(b * num_inputs, c, h, w)
-
-        shared_id_embeds = self.vision_model(id_pixel_values)[2]
-        id_embeds = self.visual_projection(shared_id_embeds)
-        id_embeds_2 = self.visual_projection_2(shared_id_embeds)
-
-        id_embeds = id_embeds.view(b, num_inputs, 1, -1)
-        id_embeds_2 = id_embeds_2.view(b, num_inputs, 1, -1)
-
-        id_embeds = torch.cat((id_embeds, id_embeds_2), dim=-1)
-        updated_prompt_embeds = self.fuse_module(prompt_embeds, id_embeds, class_tokens_mask)
-
-        return updated_prompt_embeds
-
-
-class PhotoMakerLoader:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "photomaker_model_name": (ldm_patched.utils.path_utils.get_filename_list("photomaker"), )}}
-
-    RETURN_TYPES = ("PHOTOMAKER",)
-    FUNCTION = "load_photomaker_model"
-
-    CATEGORY = "_for_testing/photomaker"
-
-    def load_photomaker_model(self, photomaker_model_name):
-        photomaker_model_path = ldm_patched.utils.path_utils.get_full_path("photomaker", photomaker_model_name)
-        photomaker_model = PhotoMakerIDEncoder()
-        data = ldm_patched.modules.utils.load_torch_file(photomaker_model_path, safe_load=True)
-        if "id_encoder" in data:
-            data = data["id_encoder"]
-        photomaker_model.load_state_dict(data)
-        return (photomaker_model,)
-
-
-class PhotoMakerEncode:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "photomaker": ("PHOTOMAKER",),
-                              "image": ("IMAGE",),
-                              "clip": ("CLIP", ),
-                              "text": ("STRING", {"multiline": True, "default": "photograph of photomaker"}),
-                             }}
-
-    RETURN_TYPES = ("CONDITIONING",)
-    FUNCTION = "apply_photomaker"
-
-    CATEGORY = "_for_testing/photomaker"
-
-    def apply_photomaker(self, photomaker, image, clip, text):
-        special_token = "photomaker"
-        pixel_values = ldm_patched.modules.clip_vision.clip_preprocess(image.to(photomaker.load_device)).float()
-        try:
-            index = text.split(" ").index(special_token) + 1
-        except ValueError:
-            index = -1
-        tokens = clip.tokenize(text, return_word_ids=True)
-        out_tokens = {}
-        for k in tokens:
-            out_tokens[k] = []
-            for t in tokens[k]:
-                f = list(filter(lambda x: x[2] != index, t))
-                while len(f) < len(t):
-                    f.append(t[-1])
-                out_tokens[k].append(f)
-
-        cond, pooled = clip.encode_from_tokens(out_tokens, return_pooled=True)
-
-        if index > 0:
-            token_index = index - 1
-            num_id_images = 1
-            class_tokens_mask = [True if token_index <= i < token_index+num_id_images else False for i in range(77)]
-            out = photomaker(id_pixel_values=pixel_values.unsqueeze(0), prompt_embeds=cond.to(photomaker.load_device),
-                            class_tokens_mask=torch.tensor(class_tokens_mask, dtype=torch.bool, device=photomaker.load_device).unsqueeze(0))
-        else:
-            out = cond
-
-        return ([[out, {"pooled_output": pooled}]], )
-
-
-NODE_CLASS_MAPPINGS = {
-    "PhotoMakerLoader": PhotoMakerLoader,
-    "PhotoMakerEncode": PhotoMakerEncode,
-}
-

ldm_patched/contrib/external_post_processing.py

@@ -1,278 +0,0 @@
-# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-from PIL import Image
-import math
-
-import ldm_patched.modules.utils
-
-
-class Blend:
-    def __init__(self):
-        pass
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "image1": ("IMAGE",),
-                "image2": ("IMAGE",),
-                "blend_factor": ("FLOAT", {
-                    "default": 0.5,
-                    "min": 0.0,
-                    "max": 1.0,
-                    "step": 0.01
-                }),
-                "blend_mode": (["normal", "multiply", "screen", "overlay", "soft_light", "difference"],),
-            },
-        }
-
-    RETURN_TYPES = ("IMAGE",)
-    FUNCTION = "blend_images"
-
-    CATEGORY = "image/postprocessing"
-
-    def blend_images(self, image1: torch.Tensor, image2: torch.Tensor, blend_factor: float, blend_mode: str):
-        image2 = image2.to(image1.device)
-        if image1.shape != image2.shape:
-            image2 = image2.permute(0, 3, 1, 2)
-            image2 = ldm_patched.modules.utils.common_upscale(image2, image1.shape[2], image1.shape[1], upscale_method='bicubic', crop='center')
-            image2 = image2.permute(0, 2, 3, 1)
-
-        blended_image = self.blend_mode(image1, image2, blend_mode)
-        blended_image = image1 * (1 - blend_factor) + blended_image * blend_factor
-        blended_image = torch.clamp(blended_image, 0, 1)
-        return (blended_image,)
-
-    def blend_mode(self, img1, img2, mode):
-        if mode == "normal":
-            return img2
-        elif mode == "multiply":
-            return img1 * img2
-        elif mode == "screen":
-            return 1 - (1 - img1) * (1 - img2)
-        elif mode == "overlay":
-            return torch.where(img1 <= 0.5, 2 * img1 * img2, 1 - 2 * (1 - img1) * (1 - img2))
-        elif mode == "soft_light":
-            return torch.where(img2 <= 0.5, img1 - (1 - 2 * img2) * img1 * (1 - img1), img1 + (2 * img2 - 1) * (self.g(img1) - img1))
-        elif mode == "difference":
-            return img1 - img2
-        else:
-            raise ValueError(f"Unsupported blend mode: {mode}")
-
-    def g(self, x):
-        return torch.where(x <= 0.25, ((16 * x - 12) * x + 4) * x, torch.sqrt(x))
-
-def gaussian_kernel(kernel_size: int, sigma: 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 = torch.sqrt(x * x + y * y)
-    g = torch.exp(-(d * d) / (2.0 * sigma * sigma))
-    return g / g.sum()
-
-class Blur:
-    def __init__(self):
-        pass
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "image": ("IMAGE",),
-                "blur_radius": ("INT", {
-                    "default": 1,
-                    "min": 1,
-                    "max": 31,
-                    "step": 1
-                }),
-                "sigma": ("FLOAT", {
-                    "default": 1.0,
-                    "min": 0.1,
-                    "max": 10.0,
-                    "step": 0.1
-                }),
-            },
-        }
-
-    RETURN_TYPES = ("IMAGE",)
-    FUNCTION = "blur"
-
-    CATEGORY = "image/postprocessing"
-
-    def blur(self, image: torch.Tensor, blur_radius: int, sigma: float):
-        if blur_radius == 0:
-            return (image,)
-
-        batch_size, height, width, channels = image.shape
-
-        kernel_size = blur_radius * 2 + 1
-        kernel = gaussian_kernel(kernel_size, sigma, device=image.device).repeat(channels, 1, 1).unsqueeze(1)
-
-        image = image.permute(0, 3, 1, 2) # Torch wants (B, C, H, W) we use (B, H, W, C)
-        padded_image = F.pad(image, (blur_radius,blur_radius,blur_radius,blur_radius), 'reflect')
-        blurred = F.conv2d(padded_image, kernel, padding=kernel_size // 2, groups=channels)[:,:,blur_radius:-blur_radius, blur_radius:-blur_radius]
-        blurred = blurred.permute(0, 2, 3, 1)
-
-        return (blurred,)
-
-class Quantize:
-    def __init__(self):
-        pass
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "image": ("IMAGE",),
-                "colors": ("INT", {
-                    "default": 256,
-                    "min": 1,
-                    "max": 256,
-                    "step": 1
-                }),
-                "dither": (["none", "floyd-steinberg", "bayer-2", "bayer-4", "bayer-8", "bayer-16"],),
-            },
-        }
-
-    RETURN_TYPES = ("IMAGE",)
-    FUNCTION = "quantize"
-
-    CATEGORY = "image/postprocessing"
-
-    def bayer(im, pal_im, order):
-        def normalized_bayer_matrix(n):
-            if n == 0:
-                return np.zeros((1,1), "float32")
-            else:
-                q = 4 ** n
-                m = q * normalized_bayer_matrix(n - 1)
-                return np.bmat(((m-1.5, m+0.5), (m+1.5, m-0.5))) / q
-
-        num_colors = len(pal_im.getpalette()) // 3
-        spread = 2 * 256 / num_colors
-        bayer_n = int(math.log2(order))
-        bayer_matrix = torch.from_numpy(spread * normalized_bayer_matrix(bayer_n) + 0.5)
-
-        result = torch.from_numpy(np.array(im).astype(np.float32))
-        tw = math.ceil(result.shape[0] / bayer_matrix.shape[0])
-        th = math.ceil(result.shape[1] / bayer_matrix.shape[1])
-        tiled_matrix = bayer_matrix.tile(tw, th).unsqueeze(-1)
-        result.add_(tiled_matrix[:result.shape[0],:result.shape[1]]).clamp_(0, 255)
-        result = result.to(dtype=torch.uint8)
-
-        im = Image.fromarray(result.cpu().numpy())
-        im = im.quantize(palette=pal_im, dither=Image.Dither.NONE)
-        return im
-
-    def quantize(self, image: torch.Tensor, colors: int, dither: str):
-        batch_size, height, width, _ = image.shape
-        result = torch.zeros_like(image)
-
-        for b in range(batch_size):
-            im = Image.fromarray((image[b] * 255).to(torch.uint8).numpy(), mode='RGB')
-
-            pal_im = im.quantize(colors=colors) # Required as described in https://github.com/python-pillow/Pillow/issues/5836
-
-            if dither == "none":
-                quantized_image = im.quantize(palette=pal_im, dither=Image.Dither.NONE)
-            elif dither == "floyd-steinberg":
-                quantized_image = im.quantize(palette=pal_im, dither=Image.Dither.FLOYDSTEINBERG)
-            elif dither.startswith("bayer"):
-                order = int(dither.split('-')[-1])
-                quantized_image = Quantize.bayer(im, pal_im, order)
-
-            quantized_array = torch.tensor(np.array(quantized_image.convert("RGB"))).float() / 255
-            result[b] = quantized_array
-
-        return (result,)
-
-class Sharpen:
-    def __init__(self):
-        pass
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "image": ("IMAGE",),
-                "sharpen_radius": ("INT", {
-                    "default": 1,
-                    "min": 1,
-                    "max": 31,
-                    "step": 1
-                }),
-                "sigma": ("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 = "sharpen"
-
-    CATEGORY = "image/postprocessing"
-
-    def sharpen(self, image: torch.Tensor, sharpen_radius: int, sigma:float, alpha: float):
-        if sharpen_radius == 0:
-            return (image,)
-
-        batch_size, height, width, channels = image.shape
-
-        kernel_size = sharpen_radius * 2 + 1
-        kernel = gaussian_kernel(kernel_size, sigma, device=image.device) * -(alpha*10)
-        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) # Torch wants (B, C, H, W) we use (B, H, W, C)
-        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)[:,:,sharpen_radius:-sharpen_radius, sharpen_radius:-sharpen_radius]
-        sharpened = sharpened.permute(0, 2, 3, 1)
-
-        result = torch.clamp(sharpened, 0, 1)
-
-        return (result,)
-
-class ImageScaleToTotalPixels:
-    upscale_methods = ["nearest-exact", "bilinear", "area", "bicubic", "lanczos"]
-    crop_methods = ["disabled", "center"]
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "image": ("IMAGE",), "upscale_method": (s.upscale_methods,),
-                              "megapixels": ("FLOAT", {"default": 1.0, "min": 0.01, "max": 16.0, "step": 0.01}),
-                            }}
-    RETURN_TYPES = ("IMAGE",)
-    FUNCTION = "upscale"
-
-    CATEGORY = "image/upscaling"
-
-    def upscale(self, image, upscale_method, megapixels):
-        samples = image.movedim(-1,1)
-        total = int(megapixels * 1024 * 1024)
-
-        scale_by = math.sqrt(total / (samples.shape[3] * samples.shape[2]))
-        width = round(samples.shape[3] * scale_by)
-        height = round(samples.shape[2] * scale_by)
-
-        s = ldm_patched.modules.utils.common_upscale(samples, width, height, upscale_method, "disabled")
-        s = s.movedim(1,-1)
-        return (s,)
-
-NODE_CLASS_MAPPINGS = {
-    "ImageBlend": Blend,
-    "ImageBlur": Blur,
-    "ImageQuantize": Quantize,
-    "ImageSharpen": Sharpen,
-    "ImageScaleToTotalPixels": ImageScaleToTotalPixels,
-}

ldm_patched/contrib/external_rebatch.py

@@ -1,140 +0,0 @@
-# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
-
-import torch
-
-class LatentRebatch:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "latents": ("LATENT",),
-                              "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
-                              }}
-    RETURN_TYPES = ("LATENT",)
-    INPUT_IS_LIST = True
-    OUTPUT_IS_LIST = (True, )
-
-    FUNCTION = "rebatch"
-
-    CATEGORY = "latent/batch"
-
-    @staticmethod
-    def get_batch(latents, list_ind, offset):
-        '''prepare a batch out of the list of latents'''
-        samples = latents[list_ind]['samples']
-        shape = samples.shape
-        mask = latents[list_ind]['noise_mask'] if 'noise_mask' in latents[list_ind] else torch.ones((shape[0], 1, shape[2]*8, shape[3]*8), device='cpu')
-        if mask.shape[-1] != shape[-1] * 8 or mask.shape[-2] != shape[-2]:
-            torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(shape[-2]*8, shape[-1]*8), mode="bilinear")
-        if mask.shape[0] < samples.shape[0]:
-            mask = mask.repeat((shape[0] - 1) // mask.shape[0] + 1, 1, 1, 1)[:shape[0]]
-        if 'batch_index' in latents[list_ind]:
-            batch_inds = latents[list_ind]['batch_index']
-        else:
-            batch_inds = [x+offset for x in range(shape[0])]
-        return samples, mask, batch_inds
-
-    @staticmethod
-    def get_slices(indexable, num, batch_size):
-        '''divides an indexable object into num slices of length batch_size, and a remainder'''
-        slices = []
-        for i in range(num):
-            slices.append(indexable[i*batch_size:(i+1)*batch_size])
-        if num * batch_size < len(indexable):
-            return slices, indexable[num * batch_size:]
-        else:
-            return slices, None
-    
-    @staticmethod
-    def slice_batch(batch, num, batch_size):
-        result = [LatentRebatch.get_slices(x, num, batch_size) for x in batch]
-        return list(zip(*result))
-
-    @staticmethod
-    def cat_batch(batch1, batch2):
-        if batch1[0] is None:
-            return batch2
-        result = [torch.cat((b1, b2)) if torch.is_tensor(b1) else b1 + b2 for b1, b2 in zip(batch1, batch2)]
-        return result
-
-    def rebatch(self, latents, batch_size):
-        batch_size = batch_size[0]
-
-        output_list = []
-        current_batch = (None, None, None)
-        processed = 0
-
-        for i in range(len(latents)):
-            # fetch new entry of list
-            #samples, masks, indices = self.get_batch(latents, i)
-            next_batch = self.get_batch(latents, i, processed)
-            processed += len(next_batch[2])
-            # set to current if current is None
-            if current_batch[0] is None:
-                current_batch = next_batch
-            # add previous to list if dimensions do not match
-            elif next_batch[0].shape[-1] != current_batch[0].shape[-1] or next_batch[0].shape[-2] != current_batch[0].shape[-2]:
-                sliced, _ = self.slice_batch(current_batch, 1, batch_size)
-                output_list.append({'samples': sliced[0][0], 'noise_mask': sliced[1][0], 'batch_index': sliced[2][0]})
-                current_batch = next_batch
-            # cat if everything checks out
-            else:
-                current_batch = self.cat_batch(current_batch, next_batch)
-
-            # add to list if dimensions gone above target batch size
-            if current_batch[0].shape[0] > batch_size:
-                num = current_batch[0].shape[0] // batch_size
-                sliced, remainder = self.slice_batch(current_batch, num, batch_size)
-                
-                for i in range(num):
-                    output_list.append({'samples': sliced[0][i], 'noise_mask': sliced[1][i], 'batch_index': sliced[2][i]})
-
-                current_batch = remainder
-
-        #add remainder
-        if current_batch[0] is not None:
-            sliced, _ = self.slice_batch(current_batch, 1, batch_size)
-            output_list.append({'samples': sliced[0][0], 'noise_mask': sliced[1][0], 'batch_index': sliced[2][0]})
-
-        #get rid of empty masks
-        for s in output_list:
-            if s['noise_mask'].mean() == 1.0:
-                del s['noise_mask']
-
-        return (output_list,)
-
-class ImageRebatch:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "images": ("IMAGE",),
-                              "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
-                              }}
-    RETURN_TYPES = ("IMAGE",)
-    INPUT_IS_LIST = True
-    OUTPUT_IS_LIST = (True, )
-
-    FUNCTION = "rebatch"
-
-    CATEGORY = "image/batch"
-
-    def rebatch(self, images, batch_size):
-        batch_size = batch_size[0]
-
-        output_list = []
-        all_images = []
-        for img in images:
-            for i in range(img.shape[0]):
-                all_images.append(img[i:i+1])
-
-        for i in range(0, len(all_images), batch_size):
-            output_list.append(torch.cat(all_images[i:i+batch_size], dim=0))
-
-        return (output_list,)
-
-NODE_CLASS_MAPPINGS = {
-    "RebatchLatents": LatentRebatch,
-    "RebatchImages": ImageRebatch,
-}
-
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "RebatchLatents": "Rebatch Latents",
-    "RebatchImages": "Rebatch Images",
-}

ldm_patched/contrib/external_sag.py

@@ -1,172 +0,0 @@
-# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
-
-import torch
-from torch import einsum
-import torch.nn.functional as F
-import math
-
-from einops import rearrange, repeat
-import os
-from ldm_patched.ldm.modules.attention import optimized_attention, _ATTN_PRECISION
-import ldm_patched.modules.samplers
-
-# from ldm_patched.modules/ldm/modules/attention.py
-# but modified to return attention scores as well as output
-def attention_basic_with_sim(q, k, v, heads, mask=None):
-    b, _, dim_head = q.shape
-    dim_head //= heads
-    scale = dim_head ** -0.5
-
-    h = heads
-    q, k, v = map(
-        lambda t: t.unsqueeze(3)
-        .reshape(b, -1, heads, dim_head)
-        .permute(0, 2, 1, 3)
-        .reshape(b * heads, -1, dim_head)
-        .contiguous(),
-        (q, k, v),
-    )
-
-    # force cast to fp32 to avoid overflowing
-    if _ATTN_PRECISION =="fp32":
-        sim = einsum('b i d, b j d -> b i j', q.float(), k.float()) * scale
-    else:
-        sim = einsum('b i d, b j d -> b i j', q, k) * scale
-
-    del q, k
-
-    if mask is not None:
-        mask = rearrange(mask, 'b ... -> b (...)')
-        max_neg_value = -torch.finfo(sim.dtype).max
-        mask = repeat(mask, 'b j -> (b h) () j', h=h)
-        sim.masked_fill_(~mask, max_neg_value)
-
-    # attention, what we cannot get enough of
-    sim = sim.softmax(dim=-1)
-
-    out = einsum('b i j, b j d -> b i d', sim.to(v.dtype), v)
-    out = (
-        out.unsqueeze(0)
-        .reshape(b, heads, -1, dim_head)
-        .permute(0, 2, 1, 3)
-        .reshape(b, -1, heads * dim_head)
-    )
-    return (out, sim)
-
-def create_blur_map(x0, attn, sigma=3.0, threshold=1.0):
-    # reshape and GAP the attention map
-    _, hw1, hw2 = attn.shape
-    b, _, lh, lw = x0.shape
-    attn = attn.reshape(b, -1, hw1, hw2)
-    # Global Average Pool
-    mask = attn.mean(1, keepdim=False).sum(1, keepdim=False) > threshold
-    ratio = 2**(math.ceil(math.sqrt(lh * lw / hw1)) - 1).bit_length()
-    mid_shape = [math.ceil(lh / ratio), math.ceil(lw / ratio)]
-
-    # Reshape
-    mask = (
-        mask.reshape(b, *mid_shape)
-        .unsqueeze(1)
-        .type(attn.dtype)
-    )
-    # Upsample
-    mask = F.interpolate(mask, (lh, lw))
-
-    blurred = gaussian_blur_2d(x0, kernel_size=9, sigma=sigma)
-    blurred = blurred * mask + x0 * (1 - mask)
-    return blurred
-
-def gaussian_blur_2d(img, kernel_size, sigma):
-    ksize_half = (kernel_size - 1) * 0.5
-
-    x = torch.linspace(-ksize_half, ksize_half, steps=kernel_size)
-
-    pdf = torch.exp(-0.5 * (x / sigma).pow(2))
-
-    x_kernel = pdf / pdf.sum()
-    x_kernel = x_kernel.to(device=img.device, dtype=img.dtype)
-
-    kernel2d = torch.mm(x_kernel[:, None], x_kernel[None, :])
-    kernel2d = kernel2d.expand(img.shape[-3], 1, kernel2d.shape[0], kernel2d.shape[1])
-
-    padding = [kernel_size // 2, kernel_size // 2, kernel_size // 2, kernel_size // 2]
-
-    img = F.pad(img, padding, mode="reflect")
-    img = F.conv2d(img, kernel2d, groups=img.shape[-3])
-    return img
-
-class SelfAttentionGuidance:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model": ("MODEL",),
-                             "scale": ("FLOAT", {"default": 0.5, "min": -2.0, "max": 5.0, "step": 0.1}),
-                             "blur_sigma": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 10.0, "step": 0.1}),
-                              }}
-    RETURN_TYPES = ("MODEL",)
-    FUNCTION = "patch"
-
-    CATEGORY = "_for_testing"
-
-    def patch(self, model, scale, blur_sigma):
-        m = model.clone()
-
-        attn_scores = None
-
-        # TODO: make this work properly with chunked batches
-        #       currently, we can only save the attn from one UNet call
-        def attn_and_record(q, k, v, extra_options):
-            nonlocal attn_scores
-            # if uncond, save the attention scores
-            heads = extra_options["n_heads"]
-            cond_or_uncond = extra_options["cond_or_uncond"]
-            b = q.shape[0] // len(cond_or_uncond)
-            if 1 in cond_or_uncond:
-                uncond_index = cond_or_uncond.index(1)
-                # do the entire attention operation, but save the attention scores to attn_scores
-                (out, sim) = attention_basic_with_sim(q, k, v, heads=heads)
-                # when using a higher batch size, I BELIEVE the result batch dimension is [uc1, ... ucn, c1, ... cn]
-                n_slices = heads * b
-                attn_scores = sim[n_slices * uncond_index:n_slices * (uncond_index+1)]
-                return out
-            else:
-                return optimized_attention(q, k, v, heads=heads)
-
-        def post_cfg_function(args):
-            nonlocal attn_scores
-            uncond_attn = attn_scores
-
-            sag_scale = scale
-            sag_sigma = blur_sigma
-            sag_threshold = 1.0
-            model = args["model"]
-            uncond_pred = args["uncond_denoised"]
-            uncond = args["uncond"]
-            cfg_result = args["denoised"]
-            sigma = args["sigma"]
-            model_options = args["model_options"]
-            x = args["input"]
-            if min(cfg_result.shape[2:]) <= 4: #skip when too small to add padding
-                return cfg_result
-
-            # create the adversarially blurred image
-            degraded = create_blur_map(uncond_pred, uncond_attn, sag_sigma, sag_threshold)
-            degraded_noised = degraded + x - uncond_pred
-            # call into the UNet
-            (sag, _) = ldm_patched.modules.samplers.calc_cond_uncond_batch(model, uncond, None, degraded_noised, sigma, model_options)
-            return cfg_result + (degraded - sag) * sag_scale
-
-        m.set_model_sampler_post_cfg_function(post_cfg_function, disable_cfg1_optimization=True)
-
-        # from diffusers:
-        # unet.mid_block.attentions[0].transformer_blocks[0].attn1.patch
-        m.set_model_attn1_replace(attn_and_record, "middle", 0, 0)
-
-        return (m, )
-
-NODE_CLASS_MAPPINGS = {
-    "SelfAttentionGuidance": SelfAttentionGuidance,
-}
-
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "SelfAttentionGuidance": "Self-Attention Guidance",
-}

ldm_patched/contrib/external_sdupscale.py

@@ -1,49 +0,0 @@
-# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
-
-import torch
-import ldm_patched.contrib.external
-import ldm_patched.modules.utils
-
-class SD_4XUpscale_Conditioning:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "images": ("IMAGE",),
-                              "positive": ("CONDITIONING",),
-                              "negative": ("CONDITIONING",),
-                              "scale_ratio": ("FLOAT", {"default": 4.0, "min": 0.0, "max": 10.0, "step": 0.01}),
-                              "noise_augmentation": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
-                             }}
-    RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
-    RETURN_NAMES = ("positive", "negative", "latent")
-
-    FUNCTION = "encode"
-
-    CATEGORY = "conditioning/upscale_diffusion"
-
-    def encode(self, images, positive, negative, scale_ratio, noise_augmentation):
-        width = max(1, round(images.shape[-2] * scale_ratio))
-        height = max(1, round(images.shape[-3] * scale_ratio))
-
-        pixels = ldm_patched.modules.utils.common_upscale((images.movedim(-1,1) * 2.0) - 1.0, width // 4, height // 4, "bilinear", "center")
-
-        out_cp = []
-        out_cn = []
-
-        for t in positive:
-            n = [t[0], t[1].copy()]
-            n[1]['concat_image'] = pixels
-            n[1]['noise_augmentation'] = noise_augmentation
-            out_cp.append(n)
-
-        for t in negative:
-            n = [t[0], t[1].copy()]
-            n[1]['concat_image'] = pixels
-            n[1]['noise_augmentation'] = noise_augmentation
-            out_cn.append(n)
-
-        latent = torch.zeros([images.shape[0], 4, height // 4, width // 4])
-        return (out_cp, out_cn, {"samples":latent})
-
-NODE_CLASS_MAPPINGS = {
-    "SD_4XUpscale_Conditioning": SD_4XUpscale_Conditioning,
-}

ldm_patched/contrib/external_stable3d.py

@@ -1,104 +0,0 @@
-# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
-
-import torch
-import ldm_patched.contrib.external
-import ldm_patched.modules.utils
-
-def camera_embeddings(elevation, azimuth):
-    elevation = torch.as_tensor([elevation])
-    azimuth = torch.as_tensor([azimuth])
-    embeddings = torch.stack(
-        [
-                torch.deg2rad(
-                    (90 - elevation) - (90)
-                ),  # Zero123 polar is 90-elevation
-                torch.sin(torch.deg2rad(azimuth)),
-                torch.cos(torch.deg2rad(azimuth)),
-                torch.deg2rad(
-                    90 - torch.full_like(elevation, 0)
-                ),
-        ], dim=-1).unsqueeze(1)
-
-    return embeddings
-
-
-class StableZero123_Conditioning:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "clip_vision": ("CLIP_VISION",),
-                              "init_image": ("IMAGE",),
-                              "vae": ("VAE",),
-                              "width": ("INT", {"default": 256, "min": 16, "max": ldm_patched.contrib.external.MAX_RESOLUTION, "step": 8}),
-                              "height": ("INT", {"default": 256, "min": 16, "max": ldm_patched.contrib.external.MAX_RESOLUTION, "step": 8}),
-                              "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
-                              "elevation": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0}),
-                              "azimuth": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0}),
-                             }}
-    RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
-    RETURN_NAMES = ("positive", "negative", "latent")
-
-    FUNCTION = "encode"
-
-    CATEGORY = "conditioning/3d_models"
-
-    def encode(self, clip_vision, init_image, vae, width, height, batch_size, elevation, azimuth):
-        output = clip_vision.encode_image(init_image)
-        pooled = output.image_embeds.unsqueeze(0)
-        pixels = ldm_patched.modules.utils.common_upscale(init_image.movedim(-1,1), width, height, "bilinear", "center").movedim(1,-1)
-        encode_pixels = pixels[:,:,:,:3]
-        t = vae.encode(encode_pixels)
-        cam_embeds = camera_embeddings(elevation, azimuth)
-        cond = torch.cat([pooled, cam_embeds.to(pooled.device).repeat((pooled.shape[0], 1, 1))], dim=-1)
-
-        positive = [[cond, {"concat_latent_image": t}]]
-        negative = [[torch.zeros_like(pooled), {"concat_latent_image": torch.zeros_like(t)}]]
-        latent = torch.zeros([batch_size, 4, height // 8, width // 8])
-        return (positive, negative, {"samples":latent})
-
-class StableZero123_Conditioning_Batched:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "clip_vision": ("CLIP_VISION",),
-                              "init_image": ("IMAGE",),
-                              "vae": ("VAE",),
-                              "width": ("INT", {"default": 256, "min": 16, "max": ldm_patched.contrib.external.MAX_RESOLUTION, "step": 8}),
-                              "height": ("INT", {"default": 256, "min": 16, "max": ldm_patched.contrib.external.MAX_RESOLUTION, "step": 8}),
-                              "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}),
-                              "elevation": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0}),
-                              "azimuth": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0}),
-                              "elevation_batch_increment": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0}),
-                              "azimuth_batch_increment": ("FLOAT", {"default": 0.0, "min": -180.0, "max": 180.0}),
-                             }}
-    RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
-    RETURN_NAMES = ("positive", "negative", "latent")
-
-    FUNCTION = "encode"
-
-    CATEGORY = "conditioning/3d_models"
-
-    def encode(self, clip_vision, init_image, vae, width, height, batch_size, elevation, azimuth, elevation_batch_increment, azimuth_batch_increment):
-        output = clip_vision.encode_image(init_image)
-        pooled = output.image_embeds.unsqueeze(0)
-        pixels = ldm_patched.modules.utils.common_upscale(init_image.movedim(-1,1), width, height, "bilinear", "center").movedim(1,-1)
-        encode_pixels = pixels[:,:,:,:3]
-        t = vae.encode(encode_pixels)
-
-        cam_embeds = []
-        for i in range(batch_size):
-            cam_embeds.append(camera_embeddings(elevation, azimuth))
-            elevation += elevation_batch_increment
-            azimuth += azimuth_batch_increment
-
-        cam_embeds = torch.cat(cam_embeds, dim=0)
-        cond = torch.cat([ldm_patched.modules.utils.repeat_to_batch_size(pooled, batch_size), cam_embeds], dim=-1)
-
-        positive = [[cond, {"concat_latent_image": t}]]
-        negative = [[torch.zeros_like(pooled), {"concat_latent_image": torch.zeros_like(t)}]]
-        latent = torch.zeros([batch_size, 4, height // 8, width // 8])
-        return (positive, negative, {"samples":latent, "batch_index": [0] * batch_size})
-
-
-NODE_CLASS_MAPPINGS = {
-    "StableZero123_Conditioning": StableZero123_Conditioning,
-    "StableZero123_Conditioning_Batched": StableZero123_Conditioning_Batched,
-}

ldm_patched/contrib/external_tomesd.py

@@ -1,179 +0,0 @@
-# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
-
-#Taken from: https://github.com/dbolya/tomesd
-
-import torch
-from typing import Tuple, Callable
-import math
-
-def do_nothing(x: torch.Tensor, mode:str=None):
-    return x
-
-
-def mps_gather_workaround(input, dim, index):
-    if input.shape[-1] == 1:
-        return torch.gather(
-            input.unsqueeze(-1),
-            dim - 1 if dim < 0 else dim,
-            index.unsqueeze(-1)
-        ).squeeze(-1)
-    else:
-        return torch.gather(input, dim, index)
-
-
-def bipartite_soft_matching_random2d(metric: torch.Tensor,
-                                     w: int, h: int, sx: int, sy: int, r: int,
-                                     no_rand: bool = False) -> Tuple[Callable, Callable]:
-    """
-    Partitions the tokens into src and dst and merges r tokens from src to dst.
-    Dst tokens are partitioned by choosing one randomy in each (sx, sy) region.
-    Args:
-     - metric [B, N, C]: metric to use for similarity
-     - w: image width in tokens
-     - h: image height in tokens
-     - sx: stride in the x dimension for dst, must divide w
-     - sy: stride in the y dimension for dst, must divide h
-     - r: number of tokens to remove (by merging)
-     - no_rand: if true, disable randomness (use top left corner only)
-    """
-    B, N, _ = metric.shape
-
-    if r <= 0 or w == 1 or h == 1:
-        return do_nothing, do_nothing
-
-    gather = mps_gather_workaround if metric.device.type == "mps" else torch.gather
-    
-    with torch.no_grad():
-        
-        hsy, wsx = h // sy, w // sx
-
-        # For each sy by sx kernel, randomly assign one token to be dst and the rest src
-        if no_rand:
-            rand_idx = torch.zeros(hsy, wsx, 1, device=metric.device, dtype=torch.int64)
-        else:
-            rand_idx = torch.randint(sy*sx, size=(hsy, wsx, 1), device=metric.device)
-        
-        # The image might not divide sx and sy, so we need to work on a view of the top left if the idx buffer instead
-        idx_buffer_view = torch.zeros(hsy, wsx, sy*sx, device=metric.device, dtype=torch.int64)
-        idx_buffer_view.scatter_(dim=2, index=rand_idx, src=-torch.ones_like(rand_idx, dtype=rand_idx.dtype))
-        idx_buffer_view = idx_buffer_view.view(hsy, wsx, sy, sx).transpose(1, 2).reshape(hsy * sy, wsx * sx)
-
-        # Image is not divisible by sx or sy so we need to move it into a new buffer
-        if (hsy * sy) < h or (wsx * sx) < w:
-            idx_buffer = torch.zeros(h, w, device=metric.device, dtype=torch.int64)
-            idx_buffer[:(hsy * sy), :(wsx * sx)] = idx_buffer_view
-        else:
-            idx_buffer = idx_buffer_view
-
-        # We set dst tokens to be -1 and src to be 0, so an argsort gives us dst|src indices
-        rand_idx = idx_buffer.reshape(1, -1, 1).argsort(dim=1)
-
-        # We're finished with these
-        del idx_buffer, idx_buffer_view
-
-        # rand_idx is currently dst|src, so split them
-        num_dst = hsy * wsx
-        a_idx = rand_idx[:, num_dst:, :] # src
-        b_idx = rand_idx[:, :num_dst, :] # dst
-
-        def split(x):
-            C = x.shape[-1]
-            src = gather(x, dim=1, index=a_idx.expand(B, N - num_dst, C))
-            dst = gather(x, dim=1, index=b_idx.expand(B, num_dst, C))
-            return src, dst
-
-        # Cosine similarity between A and B
-        metric = metric / metric.norm(dim=-1, keepdim=True)
-        a, b = split(metric)
-        scores = a @ b.transpose(-1, -2)
-
-        # Can't reduce more than the # tokens in src
-        r = min(a.shape[1], r)
-
-        # Find the most similar greedily
-        node_max, node_idx = scores.max(dim=-1)
-        edge_idx = node_max.argsort(dim=-1, descending=True)[..., None]
-
-        unm_idx = edge_idx[..., r:, :]  # Unmerged Tokens
-        src_idx = edge_idx[..., :r, :]  # Merged Tokens
-        dst_idx = gather(node_idx[..., None], dim=-2, index=src_idx)
-
-    def merge(x: torch.Tensor, mode="mean") -> torch.Tensor:
-        src, dst = split(x)
-        n, t1, c = src.shape
-        
-        unm = gather(src, dim=-2, index=unm_idx.expand(n, t1 - r, c))
-        src = gather(src, dim=-2, index=src_idx.expand(n, r, c))
-        dst = dst.scatter_reduce(-2, dst_idx.expand(n, r, c), src, reduce=mode)
-
-        return torch.cat([unm, dst], dim=1)
-
-    def unmerge(x: torch.Tensor) -> torch.Tensor:
-        unm_len = unm_idx.shape[1]
-        unm, dst = x[..., :unm_len, :], x[..., unm_len:, :]
-        _, _, c = unm.shape
-
-        src = gather(dst, dim=-2, index=dst_idx.expand(B, r, c))
-
-        # Combine back to the original shape
-        out = torch.zeros(B, N, c, device=x.device, dtype=x.dtype)
-        out.scatter_(dim=-2, index=b_idx.expand(B, num_dst, c), src=dst)
-        out.scatter_(dim=-2, index=gather(a_idx.expand(B, a_idx.shape[1], 1), dim=1, index=unm_idx).expand(B, unm_len, c), src=unm)
-        out.scatter_(dim=-2, index=gather(a_idx.expand(B, a_idx.shape[1], 1), dim=1, index=src_idx).expand(B, r, c), src=src)
-
-        return out
-
-    return merge, unmerge
-
-
-def get_functions(x, ratio, original_shape):
-    b, c, original_h, original_w = original_shape
-    original_tokens = original_h * original_w
-    downsample = int(math.ceil(math.sqrt(original_tokens // x.shape[1])))
-    stride_x = 2
-    stride_y = 2
-    max_downsample = 1
-
-    if downsample <= max_downsample:
-        w = int(math.ceil(original_w / downsample))
-        h = int(math.ceil(original_h / downsample))
-        r = int(x.shape[1] * ratio)
-        no_rand = False
-        m, u = bipartite_soft_matching_random2d(x, w, h, stride_x, stride_y, r, no_rand)
-        return m, u
-
-    nothing = lambda y: y
-    return nothing, nothing
-
-
-
-class TomePatchModel:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model": ("MODEL",),
-                              "ratio": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01}),
-                              }}
-    RETURN_TYPES = ("MODEL",)
-    FUNCTION = "patch"
-
-    CATEGORY = "_for_testing"
-
-    def patch(self, model, ratio):
-        self.u = None
-        def tomesd_m(q, k, v, extra_options):
-            #NOTE: In the reference code get_functions takes x (input of the transformer block) as the argument instead of q
-            #however from my basic testing it seems that using q instead gives better results
-            m, self.u = get_functions(q, ratio, extra_options["original_shape"])
-            return m(q), k, v
-        def tomesd_u(n, extra_options):
-            return self.u(n)
-
-        m = model.clone()
-        m.set_model_attn1_patch(tomesd_m)
-        m.set_model_attn1_output_patch(tomesd_u)
-        return (m, )
-
-
-NODE_CLASS_MAPPINGS = {
-    "TomePatchModel": TomePatchModel,
-}

ldm_patched/contrib/external_upscale_model.py

@@ -1,68 +0,0 @@
-# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
-
-import os
-from ldm_patched.pfn import model_loading
-from ldm_patched.modules import model_management
-import torch
-import ldm_patched.modules.utils
-import ldm_patched.utils.path_utils
-
-class UpscaleModelLoader:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model_name": (ldm_patched.utils.path_utils.get_filename_list("upscale_models"), ),
-                             }}
-    RETURN_TYPES = ("UPSCALE_MODEL",)
-    FUNCTION = "load_model"
-
-    CATEGORY = "loaders"
-
-    def load_model(self, model_name):
-        model_path = ldm_patched.utils.path_utils.get_full_path("upscale_models", model_name)
-        sd = ldm_patched.modules.utils.load_torch_file(model_path, safe_load=True)
-        if "module.layers.0.residual_group.blocks.0.norm1.weight" in sd:
-            sd = ldm_patched.modules.utils.state_dict_prefix_replace(sd, {"module.":""})
-        out = model_loading.load_state_dict(sd).eval()
-        return (out, )
-
-
-class ImageUpscaleWithModel:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "upscale_model": ("UPSCALE_MODEL",),
-                              "image": ("IMAGE",),
-                              }}
-    RETURN_TYPES = ("IMAGE",)
-    FUNCTION = "upscale"
-
-    CATEGORY = "image/upscaling"
-
-    def upscale(self, upscale_model, image):
-        device = model_management.get_torch_device()
-        upscale_model.to(device)
-        in_img = image.movedim(-1,-3).to(device)
-        free_memory = model_management.get_free_memory(device)
-
-        tile = 512
-        overlap = 32
-
-        oom = True
-        while oom:
-            try:
-                steps = in_img.shape[0] * ldm_patched.modules.utils.get_tiled_scale_steps(in_img.shape[3], in_img.shape[2], tile_x=tile, tile_y=tile, overlap=overlap)
-                pbar = ldm_patched.modules.utils.ProgressBar(steps)
-                s = ldm_patched.modules.utils.tiled_scale(in_img, lambda a: upscale_model(a), tile_x=tile, tile_y=tile, overlap=overlap, upscale_amount=upscale_model.scale, pbar=pbar)
-                oom = False
-            except model_management.OOM_EXCEPTION as e:
-                tile //= 2
-                if tile < 128:
-                    raise e
-
-        upscale_model.cpu()
-        s = torch.clamp(s.movedim(-3,-1), min=0, max=1.0)
-        return (s,)
-
-NODE_CLASS_MAPPINGS = {
-    "UpscaleModelLoader": UpscaleModelLoader,
-    "ImageUpscaleWithModel": ImageUpscaleWithModel
-}

ldm_patched/contrib/external_video_model.py

@@ -1,108 +0,0 @@
-# https://github.com/comfyanonymous/ComfyUI/blob/master/nodes.py 
-
-import ldm_patched.contrib.external
-import torch
-import ldm_patched.modules.utils
-import ldm_patched.modules.sd
-import ldm_patched.utils.path_utils
-import ldm_patched.contrib.external_model_merging
-
-
-class ImageOnlyCheckpointLoader:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "ckpt_name": (ldm_patched.utils.path_utils.get_filename_list("checkpoints"), ),
-                             }}
-    RETURN_TYPES = ("MODEL", "CLIP_VISION", "VAE")
-    FUNCTION = "load_checkpoint"
-
-    CATEGORY = "loaders/video_models"
-
-    def load_checkpoint(self, ckpt_name, output_vae=True, output_clip=True):
-        ckpt_path = ldm_patched.utils.path_utils.get_full_path("checkpoints", ckpt_name)
-        out = ldm_patched.modules.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=False, output_clipvision=True, embedding_directory=ldm_patched.utils.path_utils.get_folder_paths("embeddings"))
-        return (out[0], out[3], out[2])
-
-
-class SVD_img2vid_Conditioning:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "clip_vision": ("CLIP_VISION",),
-                              "init_image": ("IMAGE",),
-                              "vae": ("VAE",),
-                              "width": ("INT", {"default": 1024, "min": 16, "max": ldm_patched.contrib.external.MAX_RESOLUTION, "step": 8}),
-                              "height": ("INT", {"default": 576, "min": 16, "max": ldm_patched.contrib.external.MAX_RESOLUTION, "step": 8}),
-                              "video_frames": ("INT", {"default": 14, "min": 1, "max": 4096}),
-                              "motion_bucket_id": ("INT", {"default": 127, "min": 1, "max": 1023}),
-                              "fps": ("INT", {"default": 6, "min": 1, "max": 1024}),
-                              "augmentation_level": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 10.0, "step": 0.01})
-                             }}
-    RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "LATENT")
-    RETURN_NAMES = ("positive", "negative", "latent")
-
-    FUNCTION = "encode"
-
-    CATEGORY = "conditioning/video_models"
-
-    def encode(self, clip_vision, init_image, vae, width, height, video_frames, motion_bucket_id, fps, augmentation_level):
-        output = clip_vision.encode_image(init_image)
-        pooled = output.image_embeds.unsqueeze(0)
-        pixels = ldm_patched.modules.utils.common_upscale(init_image.movedim(-1,1), width, height, "bilinear", "center").movedim(1,-1)
-        encode_pixels = pixels[:,:,:,:3]
-        if augmentation_level > 0:
-            encode_pixels += torch.randn_like(pixels) * augmentation_level
-        t = vae.encode(encode_pixels)
-        positive = [[pooled, {"motion_bucket_id": motion_bucket_id, "fps": fps, "augmentation_level": augmentation_level, "concat_latent_image": t}]]
-        negative = [[torch.zeros_like(pooled), {"motion_bucket_id": motion_bucket_id, "fps": fps, "augmentation_level": augmentation_level, "concat_latent_image": torch.zeros_like(t)}]]
-        latent = torch.zeros([video_frames, 4, height // 8, width // 8])
-        return (positive, negative, {"samples":latent})
-
-class VideoLinearCFGGuidance:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model": ("MODEL",),
-                              "min_cfg": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step":0.5, "round": 0.01}),
-                              }}
-    RETURN_TYPES = ("MODEL",)
-    FUNCTION = "patch"
-
-    CATEGORY = "sampling/video_models"
-
-    def patch(self, model, min_cfg):
-        def linear_cfg(args):
-            cond = args["cond"]
-            uncond = args["uncond"]
-            cond_scale = args["cond_scale"]
-
-            scale = torch.linspace(min_cfg, cond_scale, cond.shape[0], device=cond.device).reshape((cond.shape[0], 1, 1, 1))
-            return uncond + scale * (cond - uncond)
-
-        m = model.clone()
-        m.set_model_sampler_cfg_function(linear_cfg)
-        return (m, )
-
-class ImageOnlyCheckpointSave(ldm_patched.contrib.external_model_merging.CheckpointSave):
-    CATEGORY = "_for_testing"
-
-    @classmethod
-    def INPUT_TYPES(s):
-        return {"required": { "model": ("MODEL",),
-                              "clip_vision": ("CLIP_VISION",),
-                              "vae": ("VAE",),
-                              "filename_prefix": ("STRING", {"default": "checkpoints/ldm_patched"}),},
-                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},}
-
-    def save(self, model, clip_vision, vae, filename_prefix, prompt=None, extra_pnginfo=None):
-        ldm_patched.contrib.external_model_merging.save_checkpoint(model, clip_vision=clip_vision, vae=vae, filename_prefix=filename_prefix, output_dir=self.output_dir, prompt=prompt, extra_pnginfo=extra_pnginfo)
-        return {}
-
-NODE_CLASS_MAPPINGS = {
-    "ImageOnlyCheckpointLoader": ImageOnlyCheckpointLoader,
-    "SVD_img2vid_Conditioning": SVD_img2vid_Conditioning,
-    "VideoLinearCFGGuidance": VideoLinearCFGGuidance,
-    "ImageOnlyCheckpointSave": ImageOnlyCheckpointSave,
-}
-
-NODE_DISPLAY_NAME_MAPPINGS = {
-    "ImageOnlyCheckpointLoader": "Image Only Checkpoint Loader (img2vid model)",
-}

ldm_patched/controlnet/cldm.py

@@ -1,312 +0,0 @@
-#taken from: https://github.com/lllyasviel/ControlNet
-#and modified
-
-import torch
-import torch as th
-import torch.nn as nn
-
-from ldm_patched.ldm.modules.diffusionmodules.util import (
-    zero_module,
-    timestep_embedding,
-)
-
-from ldm_patched.ldm.modules.attention import SpatialTransformer
-from ldm_patched.ldm.modules.diffusionmodules.openaimodel import UNetModel, TimestepEmbedSequential, ResBlock, Downsample
-from ldm_patched.ldm.util import exists
-import ldm_patched.modules.ops
-
-class ControlledUnetModel(UNetModel):
-    #implemented in the ldm unet
-    pass
-
-class ControlNet(nn.Module):
-    def __init__(
-        self,
-        image_size,
-        in_channels,
-        model_channels,
-        hint_channels,
-        num_res_blocks,
-        dropout=0,
-        channel_mult=(1, 2, 4, 8),
-        conv_resample=True,
-        dims=2,
-        num_classes=None,
-        use_checkpoint=False,
-        dtype=torch.float32,
-        num_heads=-1,
-        num_head_channels=-1,
-        num_heads_upsample=-1,
-        use_scale_shift_norm=False,
-        resblock_updown=False,
-        use_new_attention_order=False,
-        use_spatial_transformer=False,    # custom transformer support
-        transformer_depth=1,              # custom transformer support
-        context_dim=None,                 # custom transformer support
-        n_embed=None,                     # custom support for prediction of discrete ids into codebook of first stage vq model
-        legacy=True,
-        disable_self_attentions=None,
-        num_attention_blocks=None,
-        disable_middle_self_attn=False,
-        use_linear_in_transformer=False,
-        adm_in_channels=None,
-        transformer_depth_middle=None,
-        transformer_depth_output=None,
-        device=None,
-        operations=ldm_patched.modules.ops.disable_weight_init,
-        **kwargs,
-    ):
-        super().__init__()
-        assert use_spatial_transformer == True, "use_spatial_transformer has to be true"
-        if use_spatial_transformer:
-            assert context_dim is not None, 'Fool!! You forgot to include the dimension of your cross-attention conditioning...'
-
-        if context_dim is not None:
-            assert use_spatial_transformer, 'Fool!! You forgot to use the spatial transformer for your cross-attention conditioning...'
-            # from omegaconf.listconfig import ListConfig
-            # if type(context_dim) == ListConfig:
-            #     context_dim = list(context_dim)
-
-        if num_heads_upsample == -1:
-            num_heads_upsample = num_heads
-
-        if num_heads == -1:
-            assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set'
-
-        if num_head_channels == -1:
-            assert num_heads != -1, 'Either num_heads or num_head_channels has to be set'
-
-        self.dims = dims
-        self.image_size = image_size
-        self.in_channels = in_channels
-        self.model_channels = model_channels
-
-        if isinstance(num_res_blocks, int):
-            self.num_res_blocks = len(channel_mult) * [num_res_blocks]
-        else:
-            if len(num_res_blocks) != len(channel_mult):
-                raise ValueError("provide num_res_blocks either as an int (globally constant) or "
-                                 "as a list/tuple (per-level) with the same length as channel_mult")
-            self.num_res_blocks = num_res_blocks
-
-        if disable_self_attentions is not None:
-            # should be a list of booleans, indicating whether to disable self-attention in TransformerBlocks or not
-            assert len(disable_self_attentions) == len(channel_mult)
-        if num_attention_blocks is not None:
-            assert len(num_attention_blocks) == len(self.num_res_blocks)
-            assert all(map(lambda i: self.num_res_blocks[i] >= num_attention_blocks[i], range(len(num_attention_blocks))))
-
-        transformer_depth = transformer_depth[:]
-
-        self.dropout = dropout
-        self.channel_mult = channel_mult
-        self.conv_resample = conv_resample
-        self.num_classes = num_classes
-        self.use_checkpoint = use_checkpoint
-        self.dtype = dtype
-        self.num_heads = num_heads
-        self.num_head_channels = num_head_channels
-        self.num_heads_upsample = num_heads_upsample
-        self.predict_codebook_ids = n_embed is not None
-
-        time_embed_dim = model_channels * 4
-        self.time_embed = nn.Sequential(
-            operations.Linear(model_channels, time_embed_dim, dtype=self.dtype, device=device),
-            nn.SiLU(),
-            operations.Linear(time_embed_dim, time_embed_dim, dtype=self.dtype, device=device),
-        )
-
-        if self.num_classes is not None:
-            if isinstance(self.num_classes, int):
-                self.label_emb = nn.Embedding(num_classes, time_embed_dim)
-            elif self.num_classes == "continuous":
-                print("setting up linear c_adm embedding layer")
-                self.label_emb = nn.Linear(1, time_embed_dim)
-            elif self.num_classes == "sequential":
-                assert adm_in_channels is not None
-                self.label_emb = nn.Sequential(
-                    nn.Sequential(
-                        operations.Linear(adm_in_channels, time_embed_dim, dtype=self.dtype, device=device),
-                        nn.SiLU(),
-                        operations.Linear(time_embed_dim, time_embed_dim, dtype=self.dtype, device=device),
-                    )
-                )
-            else:
-                raise ValueError()
-
-        self.input_blocks = nn.ModuleList(
-            [
-                TimestepEmbedSequential(
-                    operations.conv_nd(dims, in_channels, model_channels, 3, padding=1, dtype=self.dtype, device=device)
-                )
-            ]
-        )
-        self.zero_convs = nn.ModuleList([self.make_zero_conv(model_channels, operations=operations, dtype=self.dtype, device=device)])
-
-        self.input_hint_block = TimestepEmbedSequential(
-                    operations.conv_nd(dims, hint_channels, 16, 3, padding=1, dtype=self.dtype, device=device),
-                    nn.SiLU(),
-                    operations.conv_nd(dims, 16, 16, 3, padding=1, dtype=self.dtype, device=device),
-                    nn.SiLU(),
-                    operations.conv_nd(dims, 16, 32, 3, padding=1, stride=2, dtype=self.dtype, device=device),
-                    nn.SiLU(),
-                    operations.conv_nd(dims, 32, 32, 3, padding=1, dtype=self.dtype, device=device),
-                    nn.SiLU(),
-                    operations.conv_nd(dims, 32, 96, 3, padding=1, stride=2, dtype=self.dtype, device=device),
-                    nn.SiLU(),
-                    operations.conv_nd(dims, 96, 96, 3, padding=1, dtype=self.dtype, device=device),
-                    nn.SiLU(),
-                    operations.conv_nd(dims, 96, 256, 3, padding=1, stride=2, dtype=self.dtype, device=device),
-                    nn.SiLU(),
-                    operations.conv_nd(dims, 256, model_channels, 3, padding=1, dtype=self.dtype, device=device)
-        )
-
-        self._feature_size = model_channels
-        input_block_chans = [model_channels]
-        ch = model_channels
-        ds = 1
-        for level, mult in enumerate(channel_mult):
-            for nr in range(self.num_res_blocks[level]):
-                layers = [
-                    ResBlock(
-                        ch,
-                        time_embed_dim,
-                        dropout,
-                        out_channels=mult * model_channels,
-                        dims=dims,
-                        use_checkpoint=use_checkpoint,
-                        use_scale_shift_norm=use_scale_shift_norm,
-                        dtype=self.dtype,
-                        device=device,
-                        operations=operations,
-                    )
-                ]
-                ch = mult * model_channels
-                num_transformers = transformer_depth.pop(0)
-                if num_transformers > 0:
-                    if num_head_channels == -1:
-                        dim_head = ch // num_heads
-                    else:
-                        num_heads = ch // num_head_channels
-                        dim_head = num_head_channels
-                    if legacy:
-                        #num_heads = 1
-                        dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
-                    if exists(disable_self_attentions):
-                        disabled_sa = disable_self_attentions[level]
-                    else:
-                        disabled_sa = False
-
-                    if not exists(num_attention_blocks) or nr < num_attention_blocks[level]:
-                        layers.append(
-                            SpatialTransformer(
-                                ch, num_heads, dim_head, depth=num_transformers, context_dim=context_dim,
-                                disable_self_attn=disabled_sa, use_linear=use_linear_in_transformer,
-                                use_checkpoint=use_checkpoint, dtype=self.dtype, device=device, operations=operations
-                            )
-                        )
-                self.input_blocks.append(TimestepEmbedSequential(*layers))
-                self.zero_convs.append(self.make_zero_conv(ch, operations=operations, dtype=self.dtype, device=device))
-                self._feature_size += ch
-                input_block_chans.append(ch)
-            if level != len(channel_mult) - 1:
-                out_ch = ch
-                self.input_blocks.append(
-                    TimestepEmbedSequential(
-                        ResBlock(
-                            ch,
-                            time_embed_dim,
-                            dropout,
-                            out_channels=out_ch,
-                            dims=dims,
-                            use_checkpoint=use_checkpoint,
-                            use_scale_shift_norm=use_scale_shift_norm,
-                            down=True,
-                            dtype=self.dtype,
-                            device=device,
-                            operations=operations
-                        )
-                        if resblock_updown
-                        else Downsample(
-                            ch, conv_resample, dims=dims, out_channels=out_ch, dtype=self.dtype, device=device, operations=operations
-                        )
-                    )
-                )
-                ch = out_ch
-                input_block_chans.append(ch)
-                self.zero_convs.append(self.make_zero_conv(ch, operations=operations, dtype=self.dtype, device=device))
-                ds *= 2
-                self._feature_size += ch
-
-        if num_head_channels == -1:
-            dim_head = ch // num_heads
-        else:
-            num_heads = ch // num_head_channels
-            dim_head = num_head_channels
-        if legacy:
-            #num_heads = 1
-            dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
-        mid_block = [
-            ResBlock(
-                ch,
-                time_embed_dim,
-                dropout,
-                dims=dims,
-                use_checkpoint=use_checkpoint,
-                use_scale_shift_norm=use_scale_shift_norm,
-                dtype=self.dtype,
-                device=device,
-                operations=operations
-            )]
-        if transformer_depth_middle >= 0:
-            mid_block += [SpatialTransformer(  # always uses a self-attn
-                            ch, num_heads, dim_head, depth=transformer_depth_middle, context_dim=context_dim,
-                            disable_self_attn=disable_middle_self_attn, use_linear=use_linear_in_transformer,
-                            use_checkpoint=use_checkpoint, dtype=self.dtype, device=device, operations=operations
-                        ),
-            ResBlock(
-                ch,
-                time_embed_dim,
-                dropout,
-                dims=dims,
-                use_checkpoint=use_checkpoint,
-                use_scale_shift_norm=use_scale_shift_norm,
-                dtype=self.dtype,
-                device=device,
-                operations=operations
-            )]
-        self.middle_block = TimestepEmbedSequential(*mid_block)
-        self.middle_block_out = self.make_zero_conv(ch, operations=operations, dtype=self.dtype, device=device)
-        self._feature_size += ch
-
-    def make_zero_conv(self, channels, operations=None, dtype=None, device=None):
-        return TimestepEmbedSequential(operations.conv_nd(self.dims, channels, channels, 1, padding=0, dtype=dtype, device=device))
-
-    def forward(self, x, hint, timesteps, context, y=None, **kwargs):
-        t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype)
-        emb = self.time_embed(t_emb)
-
-        guided_hint = self.input_hint_block(hint, emb, context)
-
-        outs = []
-
-        hs = []
-        if self.num_classes is not None:
-            assert y.shape[0] == x.shape[0]
-            emb = emb + self.label_emb(y)
-
-        h = x
-        for module, zero_conv in zip(self.input_blocks, self.zero_convs):
-            if guided_hint is not None:
-                h = module(h, emb, context)
-                h += guided_hint
-                guided_hint = None
-            else:
-                h = module(h, emb, context)
-            outs.append(zero_conv(h, emb, context))
-
-        h = self.middle_block(h, emb, context)
-        outs.append(self.middle_block_out(h, emb, context))
-
-        return outs
-

ldm_patched/k_diffusion/sampling.py

@@ -1,810 +0,0 @@
-import math
-
-from scipy import integrate
-import torch
-from torch import nn
-import torchsde
-from tqdm.auto import trange, tqdm
-
-from . import utils
-
-
-def append_zero(x):
-    return torch.cat([x, x.new_zeros([1])])
-
-
-def get_sigmas_karras(n, sigma_min, sigma_max, rho=7., device='cpu'):
-    """Constructs the noise schedule of Karras et al. (2022)."""
-    ramp = torch.linspace(0, 1, n, device=device)
-    min_inv_rho = sigma_min ** (1 / rho)
-    max_inv_rho = sigma_max ** (1 / rho)
-    sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
-    return append_zero(sigmas).to(device)
-
-
-def get_sigmas_exponential(n, sigma_min, sigma_max, device='cpu'):
-    """Constructs an exponential noise schedule."""
-    sigmas = torch.linspace(math.log(sigma_max), math.log(sigma_min), n, device=device).exp()
-    return append_zero(sigmas)
-
-
-def get_sigmas_polyexponential(n, sigma_min, sigma_max, rho=1., device='cpu'):
-    """Constructs an polynomial in log sigma noise schedule."""
-    ramp = torch.linspace(1, 0, n, device=device) ** rho
-    sigmas = torch.exp(ramp * (math.log(sigma_max) - math.log(sigma_min)) + math.log(sigma_min))
-    return append_zero(sigmas)
-
-
-def get_sigmas_vp(n, beta_d=19.9, beta_min=0.1, eps_s=1e-3, device='cpu'):
-    """Constructs a continuous VP noise schedule."""
-    t = torch.linspace(1, eps_s, n, device=device)
-    sigmas = torch.sqrt(torch.exp(beta_d * t ** 2 / 2 + beta_min * t) - 1)
-    return append_zero(sigmas)
-
-
-def to_d(x, sigma, denoised):
-    """Converts a denoiser output to a Karras ODE derivative."""
-    return (x - denoised) / utils.append_dims(sigma, x.ndim)
-
-
-def get_ancestral_step(sigma_from, sigma_to, eta=1.):
-    """Calculates the noise level (sigma_down) to step down to and the amount
-    of noise to add (sigma_up) when doing an ancestral sampling step."""
-    if not eta:
-        return sigma_to, 0.
-    sigma_up = min(sigma_to, eta * (sigma_to ** 2 * (sigma_from ** 2 - sigma_to ** 2) / sigma_from ** 2) ** 0.5)
-    sigma_down = (sigma_to ** 2 - sigma_up ** 2) ** 0.5
-    return sigma_down, sigma_up
-
-
-def default_noise_sampler(x):
-    return lambda sigma, sigma_next: torch.randn_like(x)
-
-
-class BatchedBrownianTree:
-    """A wrapper around torchsde.BrownianTree that enables batches of entropy."""
-
-    def __init__(self, x, t0, t1, seed=None, **kwargs):
-        self.cpu_tree = True
-        if "cpu" in kwargs:
-            self.cpu_tree = kwargs.pop("cpu")
-        t0, t1, self.sign = self.sort(t0, t1)
-        w0 = kwargs.get('w0', torch.zeros_like(x))
-        if seed is None:
-            seed = torch.randint(0, 2 ** 63 - 1, []).item()
-        self.batched = True
-        try:
-            assert len(seed) == x.shape[0]
-            w0 = w0[0]
-        except TypeError:
-            seed = [seed]
-            self.batched = False
-        if self.cpu_tree:
-            self.trees = [torchsde.BrownianTree(t0.cpu(), w0.cpu(), t1.cpu(), entropy=s, **kwargs) for s in seed]
-        else:
-            self.trees = [torchsde.BrownianTree(t0, w0, t1, entropy=s, **kwargs) for s in seed]
-
-    @staticmethod
-    def sort(a, b):
-        return (a, b, 1) if a < b else (b, a, -1)
-
-    def __call__(self, t0, t1):
-        t0, t1, sign = self.sort(t0, t1)
-        if self.cpu_tree:
-            w = torch.stack([tree(t0.cpu().float(), t1.cpu().float()).to(t0.dtype).to(t0.device) for tree in self.trees]) * (self.sign * sign)
-        else:
-            w = torch.stack([tree(t0, t1) for tree in self.trees]) * (self.sign * sign)
-
-        return w if self.batched else w[0]
-
-
-class BrownianTreeNoiseSampler:
-    """A noise sampler backed by a torchsde.BrownianTree.
-
-    Args:
-        x (Tensor): The tensor whose shape, device and dtype to use to generate
-            random samples.
-        sigma_min (float): The low end of the valid interval.
-        sigma_max (float): The high end of the valid interval.
-        seed (int or List[int]): The random seed. If a list of seeds is
-            supplied instead of a single integer, then the noise sampler will
-            use one BrownianTree per batch item, each with its own seed.
-        transform (callable): A function that maps sigma to the sampler's
-            internal timestep.
-    """
-
-    def __init__(self, x, sigma_min, sigma_max, seed=None, transform=lambda x: x, cpu=False):
-        self.transform = transform
-        t0, t1 = self.transform(torch.as_tensor(sigma_min)), self.transform(torch.as_tensor(sigma_max))
-        self.tree = BatchedBrownianTree(x, t0, t1, seed, cpu=cpu)
-
-    def __call__(self, sigma, sigma_next):
-        t0, t1 = self.transform(torch.as_tensor(sigma)), self.transform(torch.as_tensor(sigma_next))
-        return self.tree(t0, t1) / (t1 - t0).abs().sqrt()
-
-
-@torch.no_grad()
-def sample_euler(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
-    """Implements Algorithm 2 (Euler steps) from Karras et al. (2022)."""
-    extra_args = {} if extra_args is None else extra_args
-    s_in = x.new_ones([x.shape[0]])
-    for i in trange(len(sigmas) - 1, disable=disable):
-        gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
-        sigma_hat = sigmas[i] * (gamma + 1)
-        if gamma > 0:
-            eps = torch.randn_like(x) * s_noise
-            x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
-        denoised = model(x, sigma_hat * s_in, **extra_args)
-        d = to_d(x, sigma_hat, denoised)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
-        dt = sigmas[i + 1] - sigma_hat
-        # Euler method
-        x = x + d * dt
-    return x
-
-
-@torch.no_grad()
-def sample_euler_ancestral(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
-    """Ancestral sampling with Euler method steps."""
-    extra_args = {} if extra_args is None else extra_args
-    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
-    s_in = x.new_ones([x.shape[0]])
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-        d = to_d(x, sigmas[i], denoised)
-        # Euler method
-        dt = sigma_down - sigmas[i]
-        x = x + d * dt
-        if sigmas[i + 1] > 0:
-            x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
-    return x
-
-
-@torch.no_grad()
-def sample_heun(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
-    """Implements Algorithm 2 (Heun steps) from Karras et al. (2022)."""
-    extra_args = {} if extra_args is None else extra_args
-    s_in = x.new_ones([x.shape[0]])
-    for i in trange(len(sigmas) - 1, disable=disable):
-        gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
-        sigma_hat = sigmas[i] * (gamma + 1)
-        if gamma > 0:
-            eps = torch.randn_like(x) * s_noise
-            x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
-        denoised = model(x, sigma_hat * s_in, **extra_args)
-        d = to_d(x, sigma_hat, denoised)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
-        dt = sigmas[i + 1] - sigma_hat
-        if sigmas[i + 1] == 0:
-            # Euler method
-            x = x + d * dt
-        else:
-            # Heun's method
-            x_2 = x + d * dt
-            denoised_2 = model(x_2, sigmas[i + 1] * s_in, **extra_args)
-            d_2 = to_d(x_2, sigmas[i + 1], denoised_2)
-            d_prime = (d + d_2) / 2
-            x = x + d_prime * dt
-    return x
-
-
-@torch.no_grad()
-def sample_dpm_2(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
-    """A sampler inspired by DPM-Solver-2 and Algorithm 2 from Karras et al. (2022)."""
-    extra_args = {} if extra_args is None else extra_args
-    s_in = x.new_ones([x.shape[0]])
-    for i in trange(len(sigmas) - 1, disable=disable):
-        gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
-        sigma_hat = sigmas[i] * (gamma + 1)
-        if gamma > 0:
-            eps = torch.randn_like(x) * s_noise
-            x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
-        denoised = model(x, sigma_hat * s_in, **extra_args)
-        d = to_d(x, sigma_hat, denoised)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
-        if sigmas[i + 1] == 0:
-            # Euler method
-            dt = sigmas[i + 1] - sigma_hat
-            x = x + d * dt
-        else:
-            # DPM-Solver-2
-            sigma_mid = sigma_hat.log().lerp(sigmas[i + 1].log(), 0.5).exp()
-            dt_1 = sigma_mid - sigma_hat
-            dt_2 = sigmas[i + 1] - sigma_hat
-            x_2 = x + d * dt_1
-            denoised_2 = model(x_2, sigma_mid * s_in, **extra_args)
-            d_2 = to_d(x_2, sigma_mid, denoised_2)
-            x = x + d_2 * dt_2
-    return x
-
-
-@torch.no_grad()
-def sample_dpm_2_ancestral(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
-    """Ancestral sampling with DPM-Solver second-order steps."""
-    extra_args = {} if extra_args is None else extra_args
-    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
-    s_in = x.new_ones([x.shape[0]])
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-        d = to_d(x, sigmas[i], denoised)
-        if sigma_down == 0:
-            # Euler method
-            dt = sigma_down - sigmas[i]
-            x = x + d * dt
-        else:
-            # DPM-Solver-2
-            sigma_mid = sigmas[i].log().lerp(sigma_down.log(), 0.5).exp()
-            dt_1 = sigma_mid - sigmas[i]
-            dt_2 = sigma_down - sigmas[i]
-            x_2 = x + d * dt_1
-            denoised_2 = model(x_2, sigma_mid * s_in, **extra_args)
-            d_2 = to_d(x_2, sigma_mid, denoised_2)
-            x = x + d_2 * dt_2
-            x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
-    return x
-
-
-def linear_multistep_coeff(order, t, i, j):
-    if order - 1 > i:
-        raise ValueError(f'Order {order} too high for step {i}')
-    def fn(tau):
-        prod = 1.
-        for k in range(order):
-            if j == k:
-                continue
-            prod *= (tau - t[i - k]) / (t[i - j] - t[i - k])
-        return prod
-    return integrate.quad(fn, t[i], t[i + 1], epsrel=1e-4)[0]
-
-
-@torch.no_grad()
-def sample_lms(model, x, sigmas, extra_args=None, callback=None, disable=None, order=4):
-    extra_args = {} if extra_args is None else extra_args
-    s_in = x.new_ones([x.shape[0]])
-    sigmas_cpu = sigmas.detach().cpu().numpy()
-    ds = []
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        d = to_d(x, sigmas[i], denoised)
-        ds.append(d)
-        if len(ds) > order:
-            ds.pop(0)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-        cur_order = min(i + 1, order)
-        coeffs = [linear_multistep_coeff(cur_order, sigmas_cpu, i, j) for j in range(cur_order)]
-        x = x + sum(coeff * d for coeff, d in zip(coeffs, reversed(ds)))
-    return x
-
-
-class PIDStepSizeController:
-    """A PID controller for ODE adaptive step size control."""
-    def __init__(self, h, pcoeff, icoeff, dcoeff, order=1, accept_safety=0.81, eps=1e-8):
-        self.h = h
-        self.b1 = (pcoeff + icoeff + dcoeff) / order
-        self.b2 = -(pcoeff + 2 * dcoeff) / order
-        self.b3 = dcoeff / order
-        self.accept_safety = accept_safety
-        self.eps = eps
-        self.errs = []
-
-    def limiter(self, x):
-        return 1 + math.atan(x - 1)
-
-    def propose_step(self, error):
-        inv_error = 1 / (float(error) + self.eps)
-        if not self.errs:
-            self.errs = [inv_error, inv_error, inv_error]
-        self.errs[0] = inv_error
-        factor = self.errs[0] ** self.b1 * self.errs[1] ** self.b2 * self.errs[2] ** self.b3
-        factor = self.limiter(factor)
-        accept = factor >= self.accept_safety
-        if accept:
-            self.errs[2] = self.errs[1]
-            self.errs[1] = self.errs[0]
-        self.h *= factor
-        return accept
-
-
-class DPMSolver(nn.Module):
-    """DPM-Solver. See https://arxiv.org/abs/2206.00927."""
-
-    def __init__(self, model, extra_args=None, eps_callback=None, info_callback=None):
-        super().__init__()
-        self.model = model
-        self.extra_args = {} if extra_args is None else extra_args
-        self.eps_callback = eps_callback
-        self.info_callback = info_callback
-
-    def t(self, sigma):
-        return -sigma.log()
-
-    def sigma(self, t):
-        return t.neg().exp()
-
-    def eps(self, eps_cache, key, x, t, *args, **kwargs):
-        if key in eps_cache:
-            return eps_cache[key], eps_cache
-        sigma = self.sigma(t) * x.new_ones([x.shape[0]])
-        eps = (x - self.model(x, sigma, *args, **self.extra_args, **kwargs)) / self.sigma(t)
-        if self.eps_callback is not None:
-            self.eps_callback()
-        return eps, {key: eps, **eps_cache}
-
-    def dpm_solver_1_step(self, x, t, t_next, eps_cache=None):
-        eps_cache = {} if eps_cache is None else eps_cache
-        h = t_next - t
-        eps, eps_cache = self.eps(eps_cache, 'eps', x, t)
-        x_1 = x - self.sigma(t_next) * h.expm1() * eps
-        return x_1, eps_cache
-
-    def dpm_solver_2_step(self, x, t, t_next, r1=1 / 2, eps_cache=None):
-        eps_cache = {} if eps_cache is None else eps_cache
-        h = t_next - t
-        eps, eps_cache = self.eps(eps_cache, 'eps', x, t)
-        s1 = t + r1 * h
-        u1 = x - self.sigma(s1) * (r1 * h).expm1() * eps
-        eps_r1, eps_cache = self.eps(eps_cache, 'eps_r1', u1, s1)
-        x_2 = x - self.sigma(t_next) * h.expm1() * eps - self.sigma(t_next) / (2 * r1) * h.expm1() * (eps_r1 - eps)
-        return x_2, eps_cache
-
-    def dpm_solver_3_step(self, x, t, t_next, r1=1 / 3, r2=2 / 3, eps_cache=None):
-        eps_cache = {} if eps_cache is None else eps_cache
-        h = t_next - t
-        eps, eps_cache = self.eps(eps_cache, 'eps', x, t)
-        s1 = t + r1 * h
-        s2 = t + r2 * h
-        u1 = x - self.sigma(s1) * (r1 * h).expm1() * eps
-        eps_r1, eps_cache = self.eps(eps_cache, 'eps_r1', u1, s1)
-        u2 = x - self.sigma(s2) * (r2 * h).expm1() * eps - self.sigma(s2) * (r2 / r1) * ((r2 * h).expm1() / (r2 * h) - 1) * (eps_r1 - eps)
-        eps_r2, eps_cache = self.eps(eps_cache, 'eps_r2', u2, s2)
-        x_3 = x - self.sigma(t_next) * h.expm1() * eps - self.sigma(t_next) / r2 * (h.expm1() / h - 1) * (eps_r2 - eps)
-        return x_3, eps_cache
-
-    def dpm_solver_fast(self, x, t_start, t_end, nfe, eta=0., s_noise=1., noise_sampler=None):
-        noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
-        if not t_end > t_start and eta:
-            raise ValueError('eta must be 0 for reverse sampling')
-
-        m = math.floor(nfe / 3) + 1
-        ts = torch.linspace(t_start, t_end, m + 1, device=x.device)
-
-        if nfe % 3 == 0:
-            orders = [3] * (m - 2) + [2, 1]
-        else:
-            orders = [3] * (m - 1) + [nfe % 3]
-
-        for i in range(len(orders)):
-            eps_cache = {}
-            t, t_next = ts[i], ts[i + 1]
-            if eta:
-                sd, su = get_ancestral_step(self.sigma(t), self.sigma(t_next), eta)
-                t_next_ = torch.minimum(t_end, self.t(sd))
-                su = (self.sigma(t_next) ** 2 - self.sigma(t_next_) ** 2) ** 0.5
-            else:
-                t_next_, su = t_next, 0.
-
-            eps, eps_cache = self.eps(eps_cache, 'eps', x, t)
-            denoised = x - self.sigma(t) * eps
-            if self.info_callback is not None:
-                self.info_callback({'x': x, 'i': i, 't': ts[i], 't_up': t, 'denoised': denoised})
-
-            if orders[i] == 1:
-                x, eps_cache = self.dpm_solver_1_step(x, t, t_next_, eps_cache=eps_cache)
-            elif orders[i] == 2:
-                x, eps_cache = self.dpm_solver_2_step(x, t, t_next_, eps_cache=eps_cache)
-            else:
-                x, eps_cache = self.dpm_solver_3_step(x, t, t_next_, eps_cache=eps_cache)
-
-            x = x + su * s_noise * noise_sampler(self.sigma(t), self.sigma(t_next))
-
-        return x
-
-    def dpm_solver_adaptive(self, x, t_start, t_end, order=3, rtol=0.05, atol=0.0078, h_init=0.05, pcoeff=0., icoeff=1., dcoeff=0., accept_safety=0.81, eta=0., s_noise=1., noise_sampler=None):
-        noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
-        if order not in {2, 3}:
-            raise ValueError('order should be 2 or 3')
-        forward = t_end > t_start
-        if not forward and eta:
-            raise ValueError('eta must be 0 for reverse sampling')
-        h_init = abs(h_init) * (1 if forward else -1)
-        atol = torch.tensor(atol)
-        rtol = torch.tensor(rtol)
-        s = t_start
-        x_prev = x
-        accept = True
-        pid = PIDStepSizeController(h_init, pcoeff, icoeff, dcoeff, 1.5 if eta else order, accept_safety)
-        info = {'steps': 0, 'nfe': 0, 'n_accept': 0, 'n_reject': 0}
-
-        while s < t_end - 1e-5 if forward else s > t_end + 1e-5:
-            eps_cache = {}
-            t = torch.minimum(t_end, s + pid.h) if forward else torch.maximum(t_end, s + pid.h)
-            if eta:
-                sd, su = get_ancestral_step(self.sigma(s), self.sigma(t), eta)
-                t_ = torch.minimum(t_end, self.t(sd))
-                su = (self.sigma(t) ** 2 - self.sigma(t_) ** 2) ** 0.5
-            else:
-                t_, su = t, 0.
-
-            eps, eps_cache = self.eps(eps_cache, 'eps', x, s)
-            denoised = x - self.sigma(s) * eps
-
-            if order == 2:
-                x_low, eps_cache = self.dpm_solver_1_step(x, s, t_, eps_cache=eps_cache)
-                x_high, eps_cache = self.dpm_solver_2_step(x, s, t_, eps_cache=eps_cache)
-            else:
-                x_low, eps_cache = self.dpm_solver_2_step(x, s, t_, r1=1 / 3, eps_cache=eps_cache)
-                x_high, eps_cache = self.dpm_solver_3_step(x, s, t_, eps_cache=eps_cache)
-            delta = torch.maximum(atol, rtol * torch.maximum(x_low.abs(), x_prev.abs()))
-            error = torch.linalg.norm((x_low - x_high) / delta) / x.numel() ** 0.5
-            accept = pid.propose_step(error)
-            if accept:
-                x_prev = x_low
-                x = x_high + su * s_noise * noise_sampler(self.sigma(s), self.sigma(t))
-                s = t
-                info['n_accept'] += 1
-            else:
-                info['n_reject'] += 1
-            info['nfe'] += order
-            info['steps'] += 1
-
-            if self.info_callback is not None:
-                self.info_callback({'x': x, 'i': info['steps'] - 1, 't': s, 't_up': s, 'denoised': denoised, 'error': error, 'h': pid.h, **info})
-
-        return x, info
-
-
-@torch.no_grad()
-def sample_dpm_fast(model, x, sigma_min, sigma_max, n, extra_args=None, callback=None, disable=None, eta=0., s_noise=1., noise_sampler=None):
-    """DPM-Solver-Fast (fixed step size). See https://arxiv.org/abs/2206.00927."""
-    if sigma_min <= 0 or sigma_max <= 0:
-        raise ValueError('sigma_min and sigma_max must not be 0')
-    with tqdm(total=n, disable=disable) as pbar:
-        dpm_solver = DPMSolver(model, extra_args, eps_callback=pbar.update)
-        if callback is not None:
-            dpm_solver.info_callback = lambda info: callback({'sigma': dpm_solver.sigma(info['t']), 'sigma_hat': dpm_solver.sigma(info['t_up']), **info})
-        return dpm_solver.dpm_solver_fast(x, dpm_solver.t(torch.tensor(sigma_max)), dpm_solver.t(torch.tensor(sigma_min)), n, eta, s_noise, noise_sampler)
-
-
-@torch.no_grad()
-def sample_dpm_adaptive(model, x, sigma_min, sigma_max, extra_args=None, callback=None, disable=None, order=3, rtol=0.05, atol=0.0078, h_init=0.05, pcoeff=0., icoeff=1., dcoeff=0., accept_safety=0.81, eta=0., s_noise=1., noise_sampler=None, return_info=False):
-    """DPM-Solver-12 and 23 (adaptive step size). See https://arxiv.org/abs/2206.00927."""
-    if sigma_min <= 0 or sigma_max <= 0:
-        raise ValueError('sigma_min and sigma_max must not be 0')
-    with tqdm(disable=disable) as pbar:
-        dpm_solver = DPMSolver(model, extra_args, eps_callback=pbar.update)
-        if callback is not None:
-            dpm_solver.info_callback = lambda info: callback({'sigma': dpm_solver.sigma(info['t']), 'sigma_hat': dpm_solver.sigma(info['t_up']), **info})
-        x, info = dpm_solver.dpm_solver_adaptive(x, dpm_solver.t(torch.tensor(sigma_max)), dpm_solver.t(torch.tensor(sigma_min)), order, rtol, atol, h_init, pcoeff, icoeff, dcoeff, accept_safety, eta, s_noise, noise_sampler)
-    if return_info:
-        return x, info
-    return x
-
-
-@torch.no_grad()
-def sample_dpmpp_2s_ancestral(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
-    """Ancestral sampling with DPM-Solver++(2S) second-order steps."""
-    extra_args = {} if extra_args is None else extra_args
-    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
-    s_in = x.new_ones([x.shape[0]])
-    sigma_fn = lambda t: t.neg().exp()
-    t_fn = lambda sigma: sigma.log().neg()
-
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        sigma_down, sigma_up = get_ancestral_step(sigmas[i], sigmas[i + 1], eta=eta)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-        if sigma_down == 0:
-            # Euler method
-            d = to_d(x, sigmas[i], denoised)
-            dt = sigma_down - sigmas[i]
-            x = x + d * dt
-        else:
-            # DPM-Solver++(2S)
-            t, t_next = t_fn(sigmas[i]), t_fn(sigma_down)
-            r = 1 / 2
-            h = t_next - t
-            s = t + r * h
-            x_2 = (sigma_fn(s) / sigma_fn(t)) * x - (-h * r).expm1() * denoised
-            denoised_2 = model(x_2, sigma_fn(s) * s_in, **extra_args)
-            x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised_2
-        # Noise addition
-        if sigmas[i + 1] > 0:
-            x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
-    return x
-
-
-@torch.no_grad()
-def sample_dpmpp_sde(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, r=1 / 2):
-    """DPM-Solver++ (stochastic)."""
-    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
-    seed = extra_args.get("seed", None)
-    noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=True) if noise_sampler is None else noise_sampler
-    extra_args = {} if extra_args is None else extra_args
-    s_in = x.new_ones([x.shape[0]])
-    sigma_fn = lambda t: t.neg().exp()
-    t_fn = lambda sigma: sigma.log().neg()
-
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-        if sigmas[i + 1] == 0:
-            # Euler method
-            d = to_d(x, sigmas[i], denoised)
-            dt = sigmas[i + 1] - sigmas[i]
-            x = x + d * dt
-        else:
-            # DPM-Solver++
-            t, t_next = t_fn(sigmas[i]), t_fn(sigmas[i + 1])
-            h = t_next - t
-            s = t + h * r
-            fac = 1 / (2 * r)
-
-            # Step 1
-            sd, su = get_ancestral_step(sigma_fn(t), sigma_fn(s), eta)
-            s_ = t_fn(sd)
-            x_2 = (sigma_fn(s_) / sigma_fn(t)) * x - (t - s_).expm1() * denoised
-            x_2 = x_2 + noise_sampler(sigma_fn(t), sigma_fn(s)) * s_noise * su
-            denoised_2 = model(x_2, sigma_fn(s) * s_in, **extra_args)
-
-            # Step 2
-            sd, su = get_ancestral_step(sigma_fn(t), sigma_fn(t_next), eta)
-            t_next_ = t_fn(sd)
-            denoised_d = (1 - fac) * denoised + fac * denoised_2
-            x = (sigma_fn(t_next_) / sigma_fn(t)) * x - (t - t_next_).expm1() * denoised_d
-            x = x + noise_sampler(sigma_fn(t), sigma_fn(t_next)) * s_noise * su
-    return x
-
-
-@torch.no_grad()
-def sample_dpmpp_2m(model, x, sigmas, extra_args=None, callback=None, disable=None):
-    """DPM-Solver++(2M)."""
-    extra_args = {} if extra_args is None else extra_args
-    s_in = x.new_ones([x.shape[0]])
-    sigma_fn = lambda t: t.neg().exp()
-    t_fn = lambda sigma: sigma.log().neg()
-    old_denoised = None
-
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-        t, t_next = t_fn(sigmas[i]), t_fn(sigmas[i + 1])
-        h = t_next - t
-        if old_denoised is None or sigmas[i + 1] == 0:
-            x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised
-        else:
-            h_last = t - t_fn(sigmas[i - 1])
-            r = h_last / h
-            denoised_d = (1 + 1 / (2 * r)) * denoised - (1 / (2 * r)) * old_denoised
-            x = (sigma_fn(t_next) / sigma_fn(t)) * x - (-h).expm1() * denoised_d
-        old_denoised = denoised
-    return x
-
-@torch.no_grad()
-def sample_dpmpp_2m_sde(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, solver_type='midpoint'):
-    """DPM-Solver++(2M) SDE."""
-
-    if solver_type not in {'heun', 'midpoint'}:
-        raise ValueError('solver_type must be \'heun\' or \'midpoint\'')
-
-    seed = extra_args.get("seed", None)
-    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
-    noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=True) if noise_sampler is None else noise_sampler
-    extra_args = {} if extra_args is None else extra_args
-    s_in = x.new_ones([x.shape[0]])
-
-    old_denoised = None
-    h_last = None
-    h = None
-
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-        if sigmas[i + 1] == 0:
-            # Denoising step
-            x = denoised
-        else:
-            # DPM-Solver++(2M) SDE
-            t, s = -sigmas[i].log(), -sigmas[i + 1].log()
-            h = s - t
-            eta_h = eta * h
-
-            x = sigmas[i + 1] / sigmas[i] * (-eta_h).exp() * x + (-h - eta_h).expm1().neg() * denoised
-
-            if old_denoised is not None:
-                r = h_last / h
-                if solver_type == 'heun':
-                    x = x + ((-h - eta_h).expm1().neg() / (-h - eta_h) + 1) * (1 / r) * (denoised - old_denoised)
-                elif solver_type == 'midpoint':
-                    x = x + 0.5 * (-h - eta_h).expm1().neg() * (1 / r) * (denoised - old_denoised)
-
-            if eta:
-                x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * eta_h).expm1().neg().sqrt() * s_noise
-
-        old_denoised = denoised
-        h_last = h
-    return x
-
-@torch.no_grad()
-def sample_dpmpp_3m_sde(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
-    """DPM-Solver++(3M) SDE."""
-
-    seed = extra_args.get("seed", None)
-    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
-    noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=seed, cpu=True) if noise_sampler is None else noise_sampler
-    extra_args = {} if extra_args is None else extra_args
-    s_in = x.new_ones([x.shape[0]])
-
-    denoised_1, denoised_2 = None, None
-    h, h_1, h_2 = None, None, None
-
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-        if sigmas[i + 1] == 0:
-            # Denoising step
-            x = denoised
-        else:
-            t, s = -sigmas[i].log(), -sigmas[i + 1].log()
-            h = s - t
-            h_eta = h * (eta + 1)
-
-            x = torch.exp(-h_eta) * x + (-h_eta).expm1().neg() * denoised
-
-            if h_2 is not None:
-                r0 = h_1 / h
-                r1 = h_2 / h
-                d1_0 = (denoised - denoised_1) / r0
-                d1_1 = (denoised_1 - denoised_2) / r1
-                d1 = d1_0 + (d1_0 - d1_1) * r0 / (r0 + r1)
-                d2 = (d1_0 - d1_1) / (r0 + r1)
-                phi_2 = h_eta.neg().expm1() / h_eta + 1
-                phi_3 = phi_2 / h_eta - 0.5
-                x = x + phi_2 * d1 - phi_3 * d2
-            elif h_1 is not None:
-                r = h_1 / h
-                d = (denoised - denoised_1) / r
-                phi_2 = h_eta.neg().expm1() / h_eta + 1
-                x = x + phi_2 * d
-
-            if eta:
-                x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * sigmas[i + 1] * (-2 * h * eta).expm1().neg().sqrt() * s_noise
-
-        denoised_1, denoised_2 = denoised, denoised_1
-        h_1, h_2 = h, h_1
-    return x
-
-@torch.no_grad()
-def sample_dpmpp_3m_sde_gpu(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None):
-    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
-    noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=extra_args.get("seed", None), cpu=False) if noise_sampler is None else noise_sampler
-    return sample_dpmpp_3m_sde(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler)
-
-@torch.no_grad()
-def sample_dpmpp_2m_sde_gpu(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, solver_type='midpoint'):
-    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
-    noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=extra_args.get("seed", None), cpu=False) if noise_sampler is None else noise_sampler
-    return sample_dpmpp_2m_sde(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler, solver_type=solver_type)
-
-@torch.no_grad()
-def sample_dpmpp_sde_gpu(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, r=1 / 2):
-    sigma_min, sigma_max = sigmas[sigmas > 0].min(), sigmas.max()
-    noise_sampler = BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=extra_args.get("seed", None), cpu=False) if noise_sampler is None else noise_sampler
-    return sample_dpmpp_sde(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, eta=eta, s_noise=s_noise, noise_sampler=noise_sampler, r=r)
-
-
-def DDPMSampler_step(x, sigma, sigma_prev, noise, noise_sampler):
-    alpha_cumprod = 1 / ((sigma * sigma) + 1)
-    alpha_cumprod_prev = 1 / ((sigma_prev * sigma_prev) + 1)
-    alpha = (alpha_cumprod / alpha_cumprod_prev)
-
-    mu = (1.0 / alpha).sqrt() * (x - (1 - alpha) * noise / (1 - alpha_cumprod).sqrt())
-    if sigma_prev > 0:
-        mu += ((1 - alpha) * (1. - alpha_cumprod_prev) / (1. - alpha_cumprod)).sqrt() * noise_sampler(sigma, sigma_prev)
-    return mu
-
-def generic_step_sampler(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler=None, step_function=None):
-    extra_args = {} if extra_args is None else extra_args
-    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
-    s_in = x.new_ones([x.shape[0]])
-
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-        x = step_function(x / torch.sqrt(1.0 + sigmas[i] ** 2.0), sigmas[i], sigmas[i + 1], (x - denoised) / sigmas[i], noise_sampler)
-        if sigmas[i + 1] != 0:
-            x *= torch.sqrt(1.0 + sigmas[i + 1] ** 2.0)
-    return x
-
-
-@torch.no_grad()
-def sample_ddpm(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler=None):
-    return generic_step_sampler(model, x, sigmas, extra_args, callback, disable, noise_sampler, DDPMSampler_step)
-
-@torch.no_grad()
-def sample_lcm(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler=None):
-    extra_args = {} if extra_args is None else extra_args
-    noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
-    s_in = x.new_ones([x.shape[0]])
-    for i in trange(len(sigmas) - 1, disable=disable):
-        denoised = model(x, sigmas[i] * s_in, **extra_args)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
-
-        x = denoised
-        if sigmas[i + 1] > 0:
-            x += sigmas[i + 1] * noise_sampler(sigmas[i], sigmas[i + 1])
-    return x
-
-
-
-@torch.no_grad()
-def sample_heunpp2(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
-    # From MIT licensed: https://github.com/Carzit/sd-webui-samplers-scheduler/
-    extra_args = {} if extra_args is None else extra_args
-    s_in = x.new_ones([x.shape[0]])
-    s_end = sigmas[-1]
-    for i in trange(len(sigmas) - 1, disable=disable):
-        gamma = min(s_churn / (len(sigmas) - 1), 2 ** 0.5 - 1) if s_tmin <= sigmas[i] <= s_tmax else 0.
-        eps = torch.randn_like(x) * s_noise
-        sigma_hat = sigmas[i] * (gamma + 1)
-        if gamma > 0:
-            x = x + eps * (sigma_hat ** 2 - sigmas[i] ** 2) ** 0.5
-        denoised = model(x, sigma_hat * s_in, **extra_args)
-        d = to_d(x, sigma_hat, denoised)
-        if callback is not None:
-            callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigma_hat, 'denoised': denoised})
-        dt = sigmas[i + 1] - sigma_hat
-        if sigmas[i + 1] == s_end:
-            # Euler method
-            x = x + d * dt
-        elif sigmas[i + 2] == s_end:
-
-            # Heun's method
-            x_2 = x + d * dt
-            denoised_2 = model(x_2, sigmas[i + 1] * s_in, **extra_args)
-            d_2 = to_d(x_2, sigmas[i + 1], denoised_2)
-
-            w = 2 * sigmas[0]
-            w2 = sigmas[i+1]/w
-            w1 = 1 - w2
-
-            d_prime = d * w1 + d_2 * w2
-
-
-            x = x + d_prime * dt
-
-        else:
-            # Heun++
-            x_2 = x + d * dt
-            denoised_2 = model(x_2, sigmas[i + 1] * s_in, **extra_args)
-            d_2 = to_d(x_2, sigmas[i + 1], denoised_2)
-            dt_2 = sigmas[i + 2] - sigmas[i + 1]
-
-            x_3 = x_2 + d_2 * dt_2
-            denoised_3 = model(x_3, sigmas[i + 2] * s_in, **extra_args)
-            d_3 = to_d(x_3, sigmas[i + 2], denoised_3)
-
-            w = 3 * sigmas[0]
-            w2 = sigmas[i + 1] / w
-            w3 = sigmas[i + 2] / w
-            w1 = 1 - w2 - w3
-
-            d_prime = w1 * d + w2 * d_2 + w3 * d_3
-            x = x + d_prime * dt
-    return x

ldm_patched/k_diffusion/utils.py

@@ -1,313 +0,0 @@
-from contextlib import contextmanager
-import hashlib
-import math
-from pathlib import Path
-import shutil
-import urllib
-import warnings
-
-from PIL import Image
-import torch
-from torch import nn, optim
-from torch.utils import data
-
-
-def hf_datasets_augs_helper(examples, transform, image_key, mode='RGB'):
-    """Apply passed in transforms for HuggingFace Datasets."""
-    images = [transform(image.convert(mode)) for image in examples[image_key]]
-    return {image_key: images}
-
-
-def append_dims(x, target_dims):
-    """Appends dimensions to the end of a tensor until it has target_dims dimensions."""
-    dims_to_append = target_dims - x.ndim
-    if dims_to_append < 0:
-        raise ValueError(f'input has {x.ndim} dims but target_dims is {target_dims}, which is less')
-    expanded = x[(...,) + (None,) * dims_to_append]
-    # MPS will get inf values if it tries to index into the new axes, but detaching fixes this.
-    # https://github.com/pytorch/pytorch/issues/84364
-    return expanded.detach().clone() if expanded.device.type == 'mps' else expanded
-
-
-def n_params(module):
-    """Returns the number of trainable parameters in a module."""
-    return sum(p.numel() for p in module.parameters())
-
-
-def download_file(path, url, digest=None):
-    """Downloads a file if it does not exist, optionally checking its SHA-256 hash."""
-    path = Path(path)
-    path.parent.mkdir(parents=True, exist_ok=True)
-    if not path.exists():
-        with urllib.request.urlopen(url) as response, open(path, 'wb') as f:
-            shutil.copyfileobj(response, f)
-    if digest is not None:
-        file_digest = hashlib.sha256(open(path, 'rb').read()).hexdigest()
-        if digest != file_digest:
-            raise OSError(f'hash of {path} (url: {url}) failed to validate')
-    return path
-
-
-@contextmanager
-def train_mode(model, mode=True):
-    """A context manager that places a model into training mode and restores
-    the previous mode on exit."""
-    modes = [module.training for module in model.modules()]
-    try:
-        yield model.train(mode)
-    finally:
-        for i, module in enumerate(model.modules()):
-            module.training = modes[i]
-
-
-def eval_mode(model):
-    """A context manager that places a model into evaluation mode and restores
-    the previous mode on exit."""
-    return train_mode(model, False)
-
-
-@torch.no_grad()
-def ema_update(model, averaged_model, decay):
-    """Incorporates updated model parameters into an exponential moving averaged
-    version of a model. It should be called after each optimizer step."""
-    model_params = dict(model.named_parameters())
-    averaged_params = dict(averaged_model.named_parameters())
-    assert model_params.keys() == averaged_params.keys()
-
-    for name, param in model_params.items():
-        averaged_params[name].mul_(decay).add_(param, alpha=1 - decay)
-
-    model_buffers = dict(model.named_buffers())
-    averaged_buffers = dict(averaged_model.named_buffers())
-    assert model_buffers.keys() == averaged_buffers.keys()
-
-    for name, buf in model_buffers.items():
-        averaged_buffers[name].copy_(buf)
-
-
-class EMAWarmup:
-    """Implements an EMA warmup using an inverse decay schedule.
-    If inv_gamma=1 and power=1, implements a simple average. inv_gamma=1, power=2/3 are
-    good values for models you plan to train for a million or more steps (reaches decay
-    factor 0.999 at 31.6K steps, 0.9999 at 1M steps), inv_gamma=1, power=3/4 for models
-    you plan to train for less (reaches decay factor 0.999 at 10K steps, 0.9999 at
-    215.4k steps).
-    Args:
-        inv_gamma (float): Inverse multiplicative factor of EMA warmup. Default: 1.
-        power (float): Exponential factor of EMA warmup. Default: 1.
-        min_value (float): The minimum EMA decay rate. Default: 0.
-        max_value (float): The maximum EMA decay rate. Default: 1.
-        start_at (int): The epoch to start averaging at. Default: 0.
-        last_epoch (int): The index of last epoch. Default: 0.
-    """
-
-    def __init__(self, inv_gamma=1., power=1., min_value=0., max_value=1., start_at=0,
-                 last_epoch=0):
-        self.inv_gamma = inv_gamma
-        self.power = power
-        self.min_value = min_value
-        self.max_value = max_value
-        self.start_at = start_at
-        self.last_epoch = last_epoch
-
-    def state_dict(self):
-        """Returns the state of the class as a :class:`dict`."""
-        return dict(self.__dict__.items())
-
-    def load_state_dict(self, state_dict):
-        """Loads the class's state.
-        Args:
-            state_dict (dict): scaler state. Should be an object returned
-                from a call to :meth:`state_dict`.
-        """
-        self.__dict__.update(state_dict)
-
-    def get_value(self):
-        """Gets the current EMA decay rate."""
-        epoch = max(0, self.last_epoch - self.start_at)
-        value = 1 - (1 + epoch / self.inv_gamma) ** -self.power
-        return 0. if epoch < 0 else min(self.max_value, max(self.min_value, value))
-
-    def step(self):
-        """Updates the step count."""
-        self.last_epoch += 1
-
-
-class InverseLR(optim.lr_scheduler._LRScheduler):
-    """Implements an inverse decay learning rate schedule with an optional exponential
-    warmup. When last_epoch=-1, sets initial lr as lr.
-    inv_gamma is the number of steps/epochs required for the learning rate to decay to
-    (1 / 2)**power of its original value.
-    Args:
-        optimizer (Optimizer): Wrapped optimizer.
-        inv_gamma (float): Inverse multiplicative factor of learning rate decay. Default: 1.
-        power (float): Exponential factor of learning rate decay. Default: 1.
-        warmup (float): Exponential warmup factor (0 <= warmup < 1, 0 to disable)
-            Default: 0.
-        min_lr (float): The minimum learning rate. Default: 0.
-        last_epoch (int): The index of last epoch. Default: -1.
-        verbose (bool): If ``True``, prints a message to stdout for
-            each update. Default: ``False``.
-    """
-
-    def __init__(self, optimizer, inv_gamma=1., power=1., warmup=0., min_lr=0.,
-                 last_epoch=-1, verbose=False):
-        self.inv_gamma = inv_gamma
-        self.power = power
-        if not 0. <= warmup < 1:
-            raise ValueError('Invalid value for warmup')
-        self.warmup = warmup
-        self.min_lr = min_lr
-        super().__init__(optimizer, last_epoch, verbose)
-
-    def get_lr(self):
-        if not self._get_lr_called_within_step:
-            warnings.warn("To get the last learning rate computed by the scheduler, "
-                          "please use `get_last_lr()`.")
-
-        return self._get_closed_form_lr()
-
-    def _get_closed_form_lr(self):
-        warmup = 1 - self.warmup ** (self.last_epoch + 1)
-        lr_mult = (1 + self.last_epoch / self.inv_gamma) ** -self.power
-        return [warmup * max(self.min_lr, base_lr * lr_mult)
-                for base_lr in self.base_lrs]
-
-
-class ExponentialLR(optim.lr_scheduler._LRScheduler):
-    """Implements an exponential learning rate schedule with an optional exponential
-    warmup. When last_epoch=-1, sets initial lr as lr. Decays the learning rate
-    continuously by decay (default 0.5) every num_steps steps.
-    Args:
-        optimizer (Optimizer): Wrapped optimizer.
-        num_steps (float): The number of steps to decay the learning rate by decay in.
-        decay (float): The factor by which to decay the learning rate every num_steps
-            steps. Default: 0.5.
-        warmup (float): Exponential warmup factor (0 <= warmup < 1, 0 to disable)
-            Default: 0.
-        min_lr (float): The minimum learning rate. Default: 0.
-        last_epoch (int): The index of last epoch. Default: -1.
-        verbose (bool): If ``True``, prints a message to stdout for
-            each update. Default: ``False``.
-    """
-
-    def __init__(self, optimizer, num_steps, decay=0.5, warmup=0., min_lr=0.,
-                 last_epoch=-1, verbose=False):
-        self.num_steps = num_steps
-        self.decay = decay
-        if not 0. <= warmup < 1:
-            raise ValueError('Invalid value for warmup')
-        self.warmup = warmup
-        self.min_lr = min_lr
-        super().__init__(optimizer, last_epoch, verbose)
-
-    def get_lr(self):
-        if not self._get_lr_called_within_step:
-            warnings.warn("To get the last learning rate computed by the scheduler, "
-                          "please use `get_last_lr()`.")
-
-        return self._get_closed_form_lr()
-
-    def _get_closed_form_lr(self):
-        warmup = 1 - self.warmup ** (self.last_epoch + 1)
-        lr_mult = (self.decay ** (1 / self.num_steps)) ** self.last_epoch
-        return [warmup * max(self.min_lr, base_lr * lr_mult)
-                for base_lr in self.base_lrs]
-
-
-def rand_log_normal(shape, loc=0., scale=1., device='cpu', dtype=torch.float32):
-    """Draws samples from an lognormal distribution."""
-    return (torch.randn(shape, device=device, dtype=dtype) * scale + loc).exp()
-
-
-def rand_log_logistic(shape, loc=0., scale=1., min_value=0., max_value=float('inf'), device='cpu', dtype=torch.float32):
-    """Draws samples from an optionally truncated log-logistic distribution."""
-    min_value = torch.as_tensor(min_value, device=device, dtype=torch.float64)
-    max_value = torch.as_tensor(max_value, device=device, dtype=torch.float64)
-    min_cdf = min_value.log().sub(loc).div(scale).sigmoid()
-    max_cdf = max_value.log().sub(loc).div(scale).sigmoid()
-    u = torch.rand(shape, device=device, dtype=torch.float64) * (max_cdf - min_cdf) + min_cdf
-    return u.logit().mul(scale).add(loc).exp().to(dtype)
-
-
-def rand_log_uniform(shape, min_value, max_value, device='cpu', dtype=torch.float32):
-    """Draws samples from an log-uniform distribution."""
-    min_value = math.log(min_value)
-    max_value = math.log(max_value)
-    return (torch.rand(shape, device=device, dtype=dtype) * (max_value - min_value) + min_value).exp()
-
-
-def rand_v_diffusion(shape, sigma_data=1., min_value=0., max_value=float('inf'), device='cpu', dtype=torch.float32):
-    """Draws samples from a truncated v-diffusion training timestep distribution."""
-    min_cdf = math.atan(min_value / sigma_data) * 2 / math.pi
-    max_cdf = math.atan(max_value / sigma_data) * 2 / math.pi
-    u = torch.rand(shape, device=device, dtype=dtype) * (max_cdf - min_cdf) + min_cdf
-    return torch.tan(u * math.pi / 2) * sigma_data
-
-
-def rand_split_log_normal(shape, loc, scale_1, scale_2, device='cpu', dtype=torch.float32):
-    """Draws samples from a split lognormal distribution."""
-    n = torch.randn(shape, device=device, dtype=dtype).abs()
-    u = torch.rand(shape, device=device, dtype=dtype)
-    n_left = n * -scale_1 + loc
-    n_right = n * scale_2 + loc
-    ratio = scale_1 / (scale_1 + scale_2)
-    return torch.where(u < ratio, n_left, n_right).exp()
-
-
-class FolderOfImages(data.Dataset):
-    """Recursively finds all images in a directory. It does not support
-    classes/targets."""
-
-    IMG_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp'}
-
-    def __init__(self, root, transform=None):
-        super().__init__()
-        self.root = Path(root)
-        self.transform = nn.Identity() if transform is None else transform
-        self.paths = sorted(path for path in self.root.rglob('*') if path.suffix.lower() in self.IMG_EXTENSIONS)
-
-    def __repr__(self):
-        return f'FolderOfImages(root="{self.root}", len: {len(self)})'
-
-    def __len__(self):
-        return len(self.paths)
-
-    def __getitem__(self, key):
-        path = self.paths[key]
-        with open(path, 'rb') as f:
-            image = Image.open(f).convert('RGB')
-        image = self.transform(image)
-        return image,
-
-
-class CSVLogger:
-    def __init__(self, filename, columns):
-        self.filename = Path(filename)
-        self.columns = columns
-        if self.filename.exists():
-            self.file = open(self.filename, 'a')
-        else:
-            self.file = open(self.filename, 'w')
-            self.write(*self.columns)
-
-    def write(self, *args):
-        print(*args, sep=',', file=self.file, flush=True)
-
-
-@contextmanager
-def tf32_mode(cudnn=None, matmul=None):
-    """A context manager that sets whether TF32 is allowed on cuDNN or matmul."""
-    cudnn_old = torch.backends.cudnn.allow_tf32
-    matmul_old = torch.backends.cuda.matmul.allow_tf32
-    try:
-        if cudnn is not None:
-            torch.backends.cudnn.allow_tf32 = cudnn
-        if matmul is not None:
-            torch.backends.cuda.matmul.allow_tf32 = matmul
-        yield
-    finally:
-        if cudnn is not None:
-            torch.backends.cudnn.allow_tf32 = cudnn_old
-        if matmul is not None:
-            torch.backends.cuda.matmul.allow_tf32 = matmul_old

ldm_patched/ldm/models/autoencoder.py

@@ -1,228 +0,0 @@
-import torch
-# import pytorch_lightning as pl
-import torch.nn.functional as F
-from contextlib import contextmanager
-from typing import Any, Dict, List, Optional, Tuple, Union
-
-from ldm_patched.ldm.modules.distributions.distributions import DiagonalGaussianDistribution
-
-from ldm_patched.ldm.util import instantiate_from_config
-from ldm_patched.ldm.modules.ema import LitEma
-import ldm_patched.modules.ops
-
-class DiagonalGaussianRegularizer(torch.nn.Module):
-    def __init__(self, sample: bool = True):
-        super().__init__()
-        self.sample = sample
-
-    def get_trainable_parameters(self) -> Any:
-        yield from ()
-
-    def forward(self, z: torch.Tensor) -> Tuple[torch.Tensor, dict]:
-        log = dict()
-        posterior = DiagonalGaussianDistribution(z)
-        if self.sample:
-            z = posterior.sample()
-        else:
-            z = posterior.mode()
-        kl_loss = posterior.kl()
-        kl_loss = torch.sum(kl_loss) / kl_loss.shape[0]
-        log["kl_loss"] = kl_loss
-        return z, log
-
-
-class AbstractAutoencoder(torch.nn.Module):
-    """
-    This is the base class for all autoencoders, including image autoencoders, image autoencoders with discriminators,
-    unCLIP models, etc. Hence, it is fairly general, and specific features
-    (e.g. discriminator training, encoding, decoding) must be implemented in subclasses.
-    """
-
-    def __init__(
-        self,
-        ema_decay: Union[None, float] = None,
-        monitor: Union[None, str] = None,
-        input_key: str = "jpg",
-        **kwargs,
-    ):
-        super().__init__()
-
-        self.input_key = input_key
-        self.use_ema = ema_decay is not None
-        if monitor is not None:
-            self.monitor = monitor
-
-        if self.use_ema:
-            self.model_ema = LitEma(self, decay=ema_decay)
-            logpy.info(f"Keeping EMAs of {len(list(self.model_ema.buffers()))}.")
-
-    def get_input(self, batch) -> Any:
-        raise NotImplementedError()
-
-    def on_train_batch_end(self, *args, **kwargs):
-        # for EMA computation
-        if self.use_ema:
-            self.model_ema(self)
-
-    @contextmanager
-    def ema_scope(self, context=None):
-        if self.use_ema:
-            self.model_ema.store(self.parameters())
-            self.model_ema.copy_to(self)
-            if context is not None:
-                logpy.info(f"{context}: Switched to EMA weights")
-        try:
-            yield None
-        finally:
-            if self.use_ema:
-                self.model_ema.restore(self.parameters())
-                if context is not None:
-                    logpy.info(f"{context}: Restored training weights")
-
-    def encode(self, *args, **kwargs) -> torch.Tensor:
-        raise NotImplementedError("encode()-method of abstract base class called")
-
-    def decode(self, *args, **kwargs) -> torch.Tensor:
-        raise NotImplementedError("decode()-method of abstract base class called")
-
-    def instantiate_optimizer_from_config(self, params, lr, cfg):
-        logpy.info(f"loading >>> {cfg['target']} <<< optimizer from config")
-        return get_obj_from_str(cfg["target"])(
-            params, lr=lr, **cfg.get("params", dict())
-        )
-
-    def configure_optimizers(self) -> Any:
-        raise NotImplementedError()
-
-
-class AutoencodingEngine(AbstractAutoencoder):
-    """
-    Base class for all image autoencoders that we train, like VQGAN or AutoencoderKL
-    (we also restore them explicitly as special cases for legacy reasons).
-    Regularizations such as KL or VQ are moved to the regularizer class.
-    """
-
-    def __init__(
-        self,
-        *args,
-        encoder_config: Dict,
-        decoder_config: Dict,
-        regularizer_config: Dict,
-        **kwargs,
-    ):
-        super().__init__(*args, **kwargs)
-
-        self.encoder: torch.nn.Module = instantiate_from_config(encoder_config)
-        self.decoder: torch.nn.Module = instantiate_from_config(decoder_config)
-        self.regularization: AbstractRegularizer = instantiate_from_config(
-            regularizer_config
-        )
-
-    def get_last_layer(self):
-        return self.decoder.get_last_layer()
-
-    def encode(
-        self,
-        x: torch.Tensor,
-        return_reg_log: bool = False,
-        unregularized: bool = False,
-    ) -> Union[torch.Tensor, Tuple[torch.Tensor, dict]]:
-        z = self.encoder(x)
-        if unregularized:
-            return z, dict()
-        z, reg_log = self.regularization(z)
-        if return_reg_log:
-            return z, reg_log
-        return z
-
-    def decode(self, z: torch.Tensor, **kwargs) -> torch.Tensor:
-        x = self.decoder(z, **kwargs)
-        return x
-
-    def forward(
-        self, x: torch.Tensor, **additional_decode_kwargs
-    ) -> Tuple[torch.Tensor, torch.Tensor, dict]:
-        z, reg_log = self.encode(x, return_reg_log=True)
-        dec = self.decode(z, **additional_decode_kwargs)
-        return z, dec, reg_log
-
-
-class AutoencodingEngineLegacy(AutoencodingEngine):
-    def __init__(self, embed_dim: int, **kwargs):
-        self.max_batch_size = kwargs.pop("max_batch_size", None)
-        ddconfig = kwargs.pop("ddconfig")
-        super().__init__(
-            encoder_config={
-                "target": "ldm_patched.ldm.modules.diffusionmodules.model.Encoder",
-                "params": ddconfig,
-            },
-            decoder_config={
-                "target": "ldm_patched.ldm.modules.diffusionmodules.model.Decoder",
-                "params": ddconfig,
-            },
-            **kwargs,
-        )
-        self.quant_conv = ldm_patched.modules.ops.disable_weight_init.Conv2d(
-            (1 + ddconfig["double_z"]) * ddconfig["z_channels"],
-            (1 + ddconfig["double_z"]) * embed_dim,
-            1,
-        )
-        self.post_quant_conv = ldm_patched.modules.ops.disable_weight_init.Conv2d(embed_dim, ddconfig["z_channels"], 1)
-        self.embed_dim = embed_dim
-
-    def get_autoencoder_params(self) -> list:
-        params = super().get_autoencoder_params()
-        return params
-
-    def encode(
-        self, x: torch.Tensor, return_reg_log: bool = False
-    ) -> Union[torch.Tensor, Tuple[torch.Tensor, dict]]:
-        if self.max_batch_size is None:
-            z = self.encoder(x)
-            z = self.quant_conv(z)
-        else:
-            N = x.shape[0]
-            bs = self.max_batch_size
-            n_batches = int(math.ceil(N / bs))
-            z = list()
-            for i_batch in range(n_batches):
-                z_batch = self.encoder(x[i_batch * bs : (i_batch + 1) * bs])
-                z_batch = self.quant_conv(z_batch)
-                z.append(z_batch)
-            z = torch.cat(z, 0)
-
-        z, reg_log = self.regularization(z)
-        if return_reg_log:
-            return z, reg_log
-        return z
-
-    def decode(self, z: torch.Tensor, **decoder_kwargs) -> torch.Tensor:
-        if self.max_batch_size is None:
-            dec = self.post_quant_conv(z)
-            dec = self.decoder(dec, **decoder_kwargs)
-        else:
-            N = z.shape[0]
-            bs = self.max_batch_size
-            n_batches = int(math.ceil(N / bs))
-            dec = list()
-            for i_batch in range(n_batches):
-                dec_batch = self.post_quant_conv(z[i_batch * bs : (i_batch + 1) * bs])
-                dec_batch = self.decoder(dec_batch, **decoder_kwargs)
-                dec.append(dec_batch)
-            dec = torch.cat(dec, 0)
-
-        return dec
-
-
-class AutoencoderKL(AutoencodingEngineLegacy):
-    def __init__(self, **kwargs):
-        if "lossconfig" in kwargs:
-            kwargs["loss_config"] = kwargs.pop("lossconfig")
-        super().__init__(
-            regularizer_config={
-                "target": (
-                    "ldm_patched.ldm.models.autoencoder.DiagonalGaussianRegularizer"
-                )
-            },
-            **kwargs,
-        )

ldm_patched/ldm/modules/attention.py

@@ -1,781 +0,0 @@
-import math
-import torch
-import torch.nn.functional as F
-from torch import nn, einsum
-from einops import rearrange, repeat
-from typing import Optional, Any
-
-from .diffusionmodules.util import checkpoint, AlphaBlender, timestep_embedding
-from .sub_quadratic_attention import efficient_dot_product_attention
-
-from ldm_patched.modules import model_management
-
-if model_management.xformers_enabled():
-    import xformers
-    import xformers.ops
-
-from ldm_patched.modules.args_parser import args
-import ldm_patched.modules.ops
-ops = ldm_patched.modules.ops.disable_weight_init
-
-# CrossAttn precision handling
-if args.disable_attention_upcast:
-    print("disabling upcasting of attention")
-    _ATTN_PRECISION = "fp16"
-else:
-    _ATTN_PRECISION = "fp32"
-
-
-def exists(val):
-    return val is not None
-
-
-def uniq(arr):
-    return{el: True for el in arr}.keys()
-
-
-def default(val, d):
-    if exists(val):
-        return val
-    return d
-
-
-def max_neg_value(t):
-    return -torch.finfo(t.dtype).max
-
-
-def init_(tensor):
-    dim = tensor.shape[-1]
-    std = 1 / math.sqrt(dim)
-    tensor.uniform_(-std, std)
-    return tensor
-
-
-# feedforward
-class GEGLU(nn.Module):
-    def __init__(self, dim_in, dim_out, dtype=None, device=None, operations=ops):
-        super().__init__()
-        self.proj = operations.Linear(dim_in, dim_out * 2, dtype=dtype, device=device)
-
-    def forward(self, x):
-        x, gate = self.proj(x).chunk(2, dim=-1)
-        return x * F.gelu(gate)
-
-
-class FeedForward(nn.Module):
-    def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0., dtype=None, device=None, operations=ops):
-        super().__init__()
-        inner_dim = int(dim * mult)
-        dim_out = default(dim_out, dim)
-        project_in = nn.Sequential(
-            operations.Linear(dim, inner_dim, dtype=dtype, device=device),
-            nn.GELU()
-        ) if not glu else GEGLU(dim, inner_dim, dtype=dtype, device=device, operations=operations)
-
-        self.net = nn.Sequential(
-            project_in,
-            nn.Dropout(dropout),
-            operations.Linear(inner_dim, dim_out, dtype=dtype, device=device)
-        )
-
-    def forward(self, x):
-        return self.net(x)
-
-def Normalize(in_channels, dtype=None, device=None):
-    return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True, dtype=dtype, device=device)
-
-def attention_basic(q, k, v, heads, mask=None):
-    b, _, dim_head = q.shape
-    dim_head //= heads
-    scale = dim_head ** -0.5
-
-    h = heads
-    q, k, v = map(
-        lambda t: t.unsqueeze(3)
-        .reshape(b, -1, heads, dim_head)
-        .permute(0, 2, 1, 3)
-        .reshape(b * heads, -1, dim_head)
-        .contiguous(),
-        (q, k, v),
-    )
-
-    # force cast to fp32 to avoid overflowing
-    if _ATTN_PRECISION =="fp32":
-        sim = einsum('b i d, b j d -> b i j', q.float(), k.float()) * scale
-    else:
-        sim = einsum('b i d, b j d -> b i j', q, k) * scale
-
-    del q, k
-
-    if exists(mask):
-        if mask.dtype == torch.bool:
-            mask = rearrange(mask, 'b ... -> b (...)') #TODO: check if this bool part matches pytorch attention
-            max_neg_value = -torch.finfo(sim.dtype).max
-            mask = repeat(mask, 'b j -> (b h) () j', h=h)
-            sim.masked_fill_(~mask, max_neg_value)
-        else:
-            sim += mask
-
-    # attention, what we cannot get enough of
-    sim = sim.softmax(dim=-1)
-
-    out = einsum('b i j, b j d -> b i d', sim.to(v.dtype), v)
-    out = (
-        out.unsqueeze(0)
-        .reshape(b, heads, -1, dim_head)
-        .permute(0, 2, 1, 3)
-        .reshape(b, -1, heads * dim_head)
-    )
-    return out
-
-
-def attention_sub_quad(query, key, value, heads, mask=None):
-    b, _, dim_head = query.shape
-    dim_head //= heads
-
-    scale = dim_head ** -0.5
-    query = query.unsqueeze(3).reshape(b, -1, heads, dim_head).permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head)
-    value = value.unsqueeze(3).reshape(b, -1, heads, dim_head).permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head)
-
-    key = key.unsqueeze(3).reshape(b, -1, heads, dim_head).permute(0, 2, 3, 1).reshape(b * heads, dim_head, -1)
-
-    dtype = query.dtype
-    upcast_attention = _ATTN_PRECISION =="fp32" and query.dtype != torch.float32
-    if upcast_attention:
-        bytes_per_token = torch.finfo(torch.float32).bits//8
-    else:
-        bytes_per_token = torch.finfo(query.dtype).bits//8
-    batch_x_heads, q_tokens, _ = query.shape
-    _, _, k_tokens = key.shape
-    qk_matmul_size_bytes = batch_x_heads * bytes_per_token * q_tokens * k_tokens
-
-    mem_free_total, mem_free_torch = model_management.get_free_memory(query.device, True)
-
-    kv_chunk_size_min = None
-    kv_chunk_size = None
-    query_chunk_size = None
-
-    for x in [4096, 2048, 1024, 512, 256]:
-        count = mem_free_total / (batch_x_heads * bytes_per_token * x * 4.0)
-        if count >= k_tokens:
-            kv_chunk_size = k_tokens
-            query_chunk_size = x
-            break
-
-    if query_chunk_size is None:
-        query_chunk_size = 512
-
-    hidden_states = efficient_dot_product_attention(
-        query,
-        key,
-        value,
-        query_chunk_size=query_chunk_size,
-        kv_chunk_size=kv_chunk_size,
-        kv_chunk_size_min=kv_chunk_size_min,
-        use_checkpoint=False,
-        upcast_attention=upcast_attention,
-        mask=mask,
-    )
-
-    hidden_states = hidden_states.to(dtype)
-
-    hidden_states = hidden_states.unflatten(0, (-1, heads)).transpose(1,2).flatten(start_dim=2)
-    return hidden_states
-
-def attention_split(q, k, v, heads, mask=None):
-    b, _, dim_head = q.shape
-    dim_head //= heads
-    scale = dim_head ** -0.5
-
-    h = heads
-    q, k, v = map(
-        lambda t: t.unsqueeze(3)
-        .reshape(b, -1, heads, dim_head)
-        .permute(0, 2, 1, 3)
-        .reshape(b * heads, -1, dim_head)
-        .contiguous(),
-        (q, k, v),
-    )
-
-    r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype)
-
-    mem_free_total = model_management.get_free_memory(q.device)
-
-    if _ATTN_PRECISION =="fp32":
-        element_size = 4
-    else:
-        element_size = q.element_size()
-
-    gb = 1024 ** 3
-    tensor_size = q.shape[0] * q.shape[1] * k.shape[1] * element_size
-    modifier = 3
-    mem_required = tensor_size * modifier
-    steps = 1
-
-
-    if mem_required > mem_free_total:
-        steps = 2**(math.ceil(math.log(mem_required / mem_free_total, 2)))
-        # print(f"Expected tensor size:{tensor_size/gb:0.1f}GB, cuda free:{mem_free_cuda/gb:0.1f}GB "
-        #      f"torch free:{mem_free_torch/gb:0.1f} total:{mem_free_total/gb:0.1f} steps:{steps}")
-
-    if steps > 64:
-        max_res = math.floor(math.sqrt(math.sqrt(mem_free_total / 2.5)) / 8) * 64
-        raise RuntimeError(f'Not enough memory, use lower resolution (max approx. {max_res}x{max_res}). '
-                            f'Need: {mem_required/64/gb:0.1f}GB free, Have:{mem_free_total/gb:0.1f}GB free')
-
-    # print("steps", steps, mem_required, mem_free_total, modifier, q.element_size(), tensor_size)
-    first_op_done = False
-    cleared_cache = False
-    while True:
-        try:
-            slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1]
-            for i in range(0, q.shape[1], slice_size):
-                end = i + slice_size
-                if _ATTN_PRECISION =="fp32":
-                    with torch.autocast(enabled=False, device_type = 'cuda'):
-                        s1 = einsum('b i d, b j d -> b i j', q[:, i:end].float(), k.float()) * scale
-                else:
-                    s1 = einsum('b i d, b j d -> b i j', q[:, i:end], k) * scale
-
-                if mask is not None:
-                    if len(mask.shape) == 2:
-                        s1 += mask[i:end]
-                    else:
-                        s1 += mask[:, i:end]
-
-                s2 = s1.softmax(dim=-1).to(v.dtype)
-                del s1
-                first_op_done = True
-
-                r1[:, i:end] = einsum('b i j, b j d -> b i d', s2, v)
-                del s2
-            break
-        except model_management.OOM_EXCEPTION as e:
-            if first_op_done == False:
-                model_management.soft_empty_cache(True)
-                if cleared_cache == False:
-                    cleared_cache = True
-                    print("out of memory error, emptying cache and trying again")
-                    continue
-                steps *= 2
-                if steps > 64:
-                    raise e
-                print("out of memory error, increasing steps and trying again", steps)
-            else:
-                raise e
-
-    del q, k, v
-
-    r1 = (
-        r1.unsqueeze(0)
-        .reshape(b, heads, -1, dim_head)
-        .permute(0, 2, 1, 3)
-        .reshape(b, -1, heads * dim_head)
-    )
-    return r1
-
-BROKEN_XFORMERS = False
-try:
-    x_vers = xformers.__version__
-    #I think 0.0.23 is also broken (q with bs bigger than 65535 gives CUDA error)
-    BROKEN_XFORMERS = x_vers.startswith("0.0.21") or x_vers.startswith("0.0.22") or x_vers.startswith("0.0.23")
-except:
-    pass
-
-def attention_xformers(q, k, v, heads, mask=None):
-    b, _, dim_head = q.shape
-    dim_head //= heads
-    if BROKEN_XFORMERS:
-        if b * heads > 65535:
-            return attention_pytorch(q, k, v, heads, mask)
-
-    q, k, v = map(
-        lambda t: t.unsqueeze(3)
-        .reshape(b, -1, heads, dim_head)
-        .permute(0, 2, 1, 3)
-        .reshape(b * heads, -1, dim_head)
-        .contiguous(),
-        (q, k, v),
-    )
-
-    if mask is not None:
-        pad = 8 - q.shape[1] % 8
-        mask_out = torch.empty([q.shape[0], q.shape[1], q.shape[1] + pad], dtype=q.dtype, device=q.device)
-        mask_out[:, :, :mask.shape[-1]] = mask
-        mask = mask_out[:, :, :mask.shape[-1]]
-
-    out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=mask)
-
-    out = (
-        out.unsqueeze(0)
-        .reshape(b, heads, -1, dim_head)
-        .permute(0, 2, 1, 3)
-        .reshape(b, -1, heads * dim_head)
-    )
-    return out
-
-def attention_pytorch(q, k, v, heads, mask=None):
-    b, _, dim_head = q.shape
-    dim_head //= heads
-    q, k, v = map(
-        lambda t: t.view(b, -1, heads, dim_head).transpose(1, 2),
-        (q, k, v),
-    )
-
-    out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False)
-    out = (
-        out.transpose(1, 2).reshape(b, -1, heads * dim_head)
-    )
-    return out
-
-
-optimized_attention = attention_basic
-
-if model_management.xformers_enabled():
-    print("Using xformers cross attention")
-    optimized_attention = attention_xformers
-elif model_management.pytorch_attention_enabled():
-    print("Using pytorch cross attention")
-    optimized_attention = attention_pytorch
-else:
-    if args.attention_split:
-        print("Using split optimization for cross attention")
-        optimized_attention = attention_split
-    else:
-        print("Using sub quadratic optimization for cross attention, if you have memory or speed issues try using: --attention-split")
-        optimized_attention = attention_sub_quad
-
-optimized_attention_masked = optimized_attention
-
-def optimized_attention_for_device(device, mask=False, small_input=False):
-    if small_input:
-        if model_management.pytorch_attention_enabled():
-            return attention_pytorch #TODO: need to confirm but this is probably slightly faster for small inputs in all cases
-        else:
-            return attention_basic
-
-    if device == torch.device("cpu"):
-        return attention_sub_quad
-
-    if mask:
-        return optimized_attention_masked
-
-    return optimized_attention
-
-
-class CrossAttention(nn.Module):
-    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0., dtype=None, device=None, operations=ops):
-        super().__init__()
-        inner_dim = dim_head * heads
-        context_dim = default(context_dim, query_dim)
-
-        self.heads = heads
-        self.dim_head = dim_head
-
-        self.to_q = operations.Linear(query_dim, inner_dim, bias=False, dtype=dtype, device=device)
-        self.to_k = operations.Linear(context_dim, inner_dim, bias=False, dtype=dtype, device=device)
-        self.to_v = operations.Linear(context_dim, inner_dim, bias=False, dtype=dtype, device=device)
-
-        self.to_out = nn.Sequential(operations.Linear(inner_dim, query_dim, dtype=dtype, device=device), nn.Dropout(dropout))
-
-    def forward(self, x, context=None, value=None, mask=None):
-        q = self.to_q(x)
-        context = default(context, x)
-        k = self.to_k(context)
-        if value is not None:
-            v = self.to_v(value)
-            del value
-        else:
-            v = self.to_v(context)
-
-        if mask is None:
-            out = optimized_attention(q, k, v, self.heads)
-        else:
-            out = optimized_attention_masked(q, k, v, self.heads, mask)
-        return self.to_out(out)
-
-
-class BasicTransformerBlock(nn.Module):
-    def __init__(self, dim, n_heads, d_head, dropout=0., context_dim=None, gated_ff=True, checkpoint=True, ff_in=False, inner_dim=None,
-                 disable_self_attn=False, disable_temporal_crossattention=False, switch_temporal_ca_to_sa=False, dtype=None, device=None, operations=ops):
-        super().__init__()
-
-        self.ff_in = ff_in or inner_dim is not None
-        if inner_dim is None:
-            inner_dim = dim
-
-        self.is_res = inner_dim == dim
-
-        if self.ff_in:
-            self.norm_in = operations.LayerNorm(dim, dtype=dtype, device=device)
-            self.ff_in = FeedForward(dim, dim_out=inner_dim, dropout=dropout, glu=gated_ff, dtype=dtype, device=device, operations=operations)
-
-        self.disable_self_attn = disable_self_attn
-        self.attn1 = CrossAttention(query_dim=inner_dim, heads=n_heads, dim_head=d_head, dropout=dropout,
-                              context_dim=context_dim if self.disable_self_attn else None, dtype=dtype, device=device, operations=operations)  # is a self-attention if not self.disable_self_attn
-        self.ff = FeedForward(inner_dim, dim_out=dim, dropout=dropout, glu=gated_ff, dtype=dtype, device=device, operations=operations)
-
-        if disable_temporal_crossattention:
-            if switch_temporal_ca_to_sa:
-                raise ValueError
-            else:
-                self.attn2 = None
-        else:
-            context_dim_attn2 = None
-            if not switch_temporal_ca_to_sa:
-                context_dim_attn2 = context_dim
-
-            self.attn2 = CrossAttention(query_dim=inner_dim, context_dim=context_dim_attn2,
-                                heads=n_heads, dim_head=d_head, dropout=dropout, dtype=dtype, device=device, operations=operations)  # is self-attn if context is none
-            self.norm2 = operations.LayerNorm(inner_dim, dtype=dtype, device=device)
-
-        self.norm1 = operations.LayerNorm(inner_dim, dtype=dtype, device=device)
-        self.norm3 = operations.LayerNorm(inner_dim, dtype=dtype, device=device)
-        self.checkpoint = checkpoint
-        self.n_heads = n_heads
-        self.d_head = d_head
-        self.switch_temporal_ca_to_sa = switch_temporal_ca_to_sa
-
-    def forward(self, x, context=None, transformer_options={}):
-        return checkpoint(self._forward, (x, context, transformer_options), self.parameters(), self.checkpoint)
-
-    def _forward(self, x, context=None, transformer_options={}):
-        extra_options = {}
-        block = transformer_options.get("block", None)
-        block_index = transformer_options.get("block_index", 0)
-        transformer_patches = {}
-        transformer_patches_replace = {}
-
-        for k in transformer_options:
-            if k == "patches":
-                transformer_patches = transformer_options[k]
-            elif k == "patches_replace":
-                transformer_patches_replace = transformer_options[k]
-            else:
-                extra_options[k] = transformer_options[k]
-
-        extra_options["n_heads"] = self.n_heads
-        extra_options["dim_head"] = self.d_head
-
-        if self.ff_in:
-            x_skip = x
-            x = self.ff_in(self.norm_in(x))
-            if self.is_res:
-                x += x_skip
-
-        n = self.norm1(x)
-        if self.disable_self_attn:
-            context_attn1 = context
-        else:
-            context_attn1 = None
-        value_attn1 = None
-
-        if "attn1_patch" in transformer_patches:
-            patch = transformer_patches["attn1_patch"]
-            if context_attn1 is None:
-                context_attn1 = n
-            value_attn1 = context_attn1
-            for p in patch:
-                n, context_attn1, value_attn1 = p(n, context_attn1, value_attn1, extra_options)
-
-        if block is not None:
-            transformer_block = (block[0], block[1], block_index)
-        else:
-            transformer_block = None
-        attn1_replace_patch = transformer_patches_replace.get("attn1", {})
-        block_attn1 = transformer_block
-        if block_attn1 not in attn1_replace_patch:
-            block_attn1 = block
-
-        if block_attn1 in attn1_replace_patch:
-            if context_attn1 is None:
-                context_attn1 = n
-                value_attn1 = n
-            n = self.attn1.to_q(n)
-            context_attn1 = self.attn1.to_k(context_attn1)
-            value_attn1 = self.attn1.to_v(value_attn1)
-            n = attn1_replace_patch[block_attn1](n, context_attn1, value_attn1, extra_options)
-            n = self.attn1.to_out(n)
-        else:
-            n = self.attn1(n, context=context_attn1, value=value_attn1)
-
-        if "attn1_output_patch" in transformer_patches:
-            patch = transformer_patches["attn1_output_patch"]
-            for p in patch:
-                n = p(n, extra_options)
-
-        x += n
-        if "middle_patch" in transformer_patches:
-            patch = transformer_patches["middle_patch"]
-            for p in patch:
-                x = p(x, extra_options)
-
-        if self.attn2 is not None:
-            n = self.norm2(x)
-            if self.switch_temporal_ca_to_sa:
-                context_attn2 = n
-            else:
-                context_attn2 = context
-            value_attn2 = None
-            if "attn2_patch" in transformer_patches:
-                patch = transformer_patches["attn2_patch"]
-                value_attn2 = context_attn2
-                for p in patch:
-                    n, context_attn2, value_attn2 = p(n, context_attn2, value_attn2, extra_options)
-
-            attn2_replace_patch = transformer_patches_replace.get("attn2", {})
-            block_attn2 = transformer_block
-            if block_attn2 not in attn2_replace_patch:
-                block_attn2 = block
-
-            if block_attn2 in attn2_replace_patch:
-                if value_attn2 is None:
-                    value_attn2 = context_attn2
-                n = self.attn2.to_q(n)
-                context_attn2 = self.attn2.to_k(context_attn2)
-                value_attn2 = self.attn2.to_v(value_attn2)
-                n = attn2_replace_patch[block_attn2](n, context_attn2, value_attn2, extra_options)
-                n = self.attn2.to_out(n)
-            else:
-                n = self.attn2(n, context=context_attn2, value=value_attn2)
-
-        if "attn2_output_patch" in transformer_patches:
-            patch = transformer_patches["attn2_output_patch"]
-            for p in patch:
-                n = p(n, extra_options)
-
-        x += n
-        if self.is_res:
-            x_skip = x
-        x = self.ff(self.norm3(x))
-        if self.is_res:
-            x += x_skip
-
-        return x
-
-
-class SpatialTransformer(nn.Module):
-    """
-    Transformer block for image-like data.
-    First, project the input (aka embedding)
-    and reshape to b, t, d.
-    Then apply standard transformer action.
-    Finally, reshape to image
-    NEW: use_linear for more efficiency instead of the 1x1 convs
-    """
-    def __init__(self, in_channels, n_heads, d_head,
-                 depth=1, dropout=0., context_dim=None,
-                 disable_self_attn=False, use_linear=False,
-                 use_checkpoint=True, dtype=None, device=None, operations=ops):
-        super().__init__()
-        if exists(context_dim) and not isinstance(context_dim, list):
-            context_dim = [context_dim] * depth
-        self.in_channels = in_channels
-        inner_dim = n_heads * d_head
-        self.norm = operations.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True, dtype=dtype, device=device)
-        if not use_linear:
-            self.proj_in = operations.Conv2d(in_channels,
-                                     inner_dim,
-                                     kernel_size=1,
-                                     stride=1,
-                                     padding=0, dtype=dtype, device=device)
-        else:
-            self.proj_in = operations.Linear(in_channels, inner_dim, dtype=dtype, device=device)
-
-        self.transformer_blocks = nn.ModuleList(
-            [BasicTransformerBlock(inner_dim, n_heads, d_head, dropout=dropout, context_dim=context_dim[d],
-                                   disable_self_attn=disable_self_attn, checkpoint=use_checkpoint, dtype=dtype, device=device, operations=operations)
-                for d in range(depth)]
-        )
-        if not use_linear:
-            self.proj_out = operations.Conv2d(inner_dim,in_channels,
-                                                  kernel_size=1,
-                                                  stride=1,
-                                                  padding=0, dtype=dtype, device=device)
-        else:
-            self.proj_out = operations.Linear(in_channels, inner_dim, dtype=dtype, device=device)
-        self.use_linear = use_linear
-
-    def forward(self, x, context=None, transformer_options={}):
-        # note: if no context is given, cross-attention defaults to self-attention
-        if not isinstance(context, list):
-            context = [context] * len(self.transformer_blocks)
-        b, c, h, w = x.shape
-        x_in = x
-        x = self.norm(x)
-        if not self.use_linear:
-            x = self.proj_in(x)
-        x = rearrange(x, 'b c h w -> b (h w) c').contiguous()
-        if self.use_linear:
-            x = self.proj_in(x)
-        for i, block in enumerate(self.transformer_blocks):
-            transformer_options["block_index"] = i
-            x = block(x, context=context[i], transformer_options=transformer_options)
-        if self.use_linear:
-            x = self.proj_out(x)
-        x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w).contiguous()
-        if not self.use_linear:
-            x = self.proj_out(x)
-        return x + x_in
-
-
-class SpatialVideoTransformer(SpatialTransformer):
-    def __init__(
-        self,
-        in_channels,
-        n_heads,
-        d_head,
-        depth=1,
-        dropout=0.0,
-        use_linear=False,
-        context_dim=None,
-        use_spatial_context=False,
-        timesteps=None,
-        merge_strategy: str = "fixed",
-        merge_factor: float = 0.5,
-        time_context_dim=None,
-        ff_in=False,
-        checkpoint=False,
-        time_depth=1,
-        disable_self_attn=False,
-        disable_temporal_crossattention=False,
-        max_time_embed_period: int = 10000,
-        dtype=None, device=None, operations=ops
-    ):
-        super().__init__(
-            in_channels,
-            n_heads,
-            d_head,
-            depth=depth,
-            dropout=dropout,
-            use_checkpoint=checkpoint,
-            context_dim=context_dim,
-            use_linear=use_linear,
-            disable_self_attn=disable_self_attn,
-            dtype=dtype, device=device, operations=operations
-        )
-        self.time_depth = time_depth
-        self.depth = depth
-        self.max_time_embed_period = max_time_embed_period
-
-        time_mix_d_head = d_head
-        n_time_mix_heads = n_heads
-
-        time_mix_inner_dim = int(time_mix_d_head * n_time_mix_heads)
-
-        inner_dim = n_heads * d_head
-        if use_spatial_context:
-            time_context_dim = context_dim
-
-        self.time_stack = nn.ModuleList(
-            [
-                BasicTransformerBlock(
-                    inner_dim,
-                    n_time_mix_heads,
-                    time_mix_d_head,
-                    dropout=dropout,
-                    context_dim=time_context_dim,
-                    # timesteps=timesteps,
-                    checkpoint=checkpoint,
-                    ff_in=ff_in,
-                    inner_dim=time_mix_inner_dim,
-                    disable_self_attn=disable_self_attn,
-                    disable_temporal_crossattention=disable_temporal_crossattention,
-                    dtype=dtype, device=device, operations=operations
-                )
-                for _ in range(self.depth)
-            ]
-        )
-
-        assert len(self.time_stack) == len(self.transformer_blocks)
-
-        self.use_spatial_context = use_spatial_context
-        self.in_channels = in_channels
-
-        time_embed_dim = self.in_channels * 4
-        self.time_pos_embed = nn.Sequential(
-            operations.Linear(self.in_channels, time_embed_dim, dtype=dtype, device=device),
-            nn.SiLU(),
-            operations.Linear(time_embed_dim, self.in_channels, dtype=dtype, device=device),
-        )
-
-        self.time_mixer = AlphaBlender(
-            alpha=merge_factor, merge_strategy=merge_strategy
-        )
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        context: Optional[torch.Tensor] = None,
-        time_context: Optional[torch.Tensor] = None,
-        timesteps: Optional[int] = None,
-        image_only_indicator: Optional[torch.Tensor] = None,
-        transformer_options={}
-    ) -> torch.Tensor:
-        _, _, h, w = x.shape
-        x_in = x
-        spatial_context = None
-        if exists(context):
-            spatial_context = context
-
-        if self.use_spatial_context:
-            assert (
-                context.ndim == 3
-            ), f"n dims of spatial context should be 3 but are {context.ndim}"
-
-            if time_context is None:
-                time_context = context
-            time_context_first_timestep = time_context[::timesteps]
-            time_context = repeat(
-                time_context_first_timestep, "b ... -> (b n) ...", n=h * w
-            )
-        elif time_context is not None and not self.use_spatial_context:
-            time_context = repeat(time_context, "b ... -> (b n) ...", n=h * w)
-            if time_context.ndim == 2:
-                time_context = rearrange(time_context, "b c -> b 1 c")
-
-        x = self.norm(x)
-        if not self.use_linear:
-            x = self.proj_in(x)
-        x = rearrange(x, "b c h w -> b (h w) c")
-        if self.use_linear:
-            x = self.proj_in(x)
-
-        num_frames = torch.arange(timesteps, device=x.device)
-        num_frames = repeat(num_frames, "t -> b t", b=x.shape[0] // timesteps)
-        num_frames = rearrange(num_frames, "b t -> (b t)")
-        t_emb = timestep_embedding(num_frames, self.in_channels, repeat_only=False, max_period=self.max_time_embed_period).to(x.dtype)
-        emb = self.time_pos_embed(t_emb)
-        emb = emb[:, None, :]
-
-        for it_, (block, mix_block) in enumerate(
-            zip(self.transformer_blocks, self.time_stack)
-        ):
-            transformer_options["block_index"] = it_
-            x = block(
-                x,
-                context=spatial_context,
-                transformer_options=transformer_options,
-            )
-
-            x_mix = x
-            x_mix = x_mix + emb
-
-            B, S, C = x_mix.shape
-            x_mix = rearrange(x_mix, "(b t) s c -> (b s) t c", t=timesteps)
-            x_mix = mix_block(x_mix, context=time_context) #TODO: transformer_options
-            x_mix = rearrange(
-                x_mix, "(b s) t c -> (b t) s c", s=S, b=B // timesteps, c=C, t=timesteps
-            )
-
-            x = self.time_mixer(x_spatial=x, x_temporal=x_mix, image_only_indicator=image_only_indicator)
-
-        if self.use_linear:
-            x = self.proj_out(x)
-        x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w)
-        if not self.use_linear:
-            x = self.proj_out(x)
-        out = x + x_in
-        return out
-
-

ldm_patched/ldm/modules/diffusionmodules/model.py

@@ -1,650 +0,0 @@
-# pytorch_diffusion + derived encoder decoder
-import math
-import torch
-import torch.nn as nn
-import numpy as np
-from einops import rearrange
-from typing import Optional, Any
-
-from ldm_patched.modules import model_management
-import ldm_patched.modules.ops
-ops = ldm_patched.modules.ops.disable_weight_init
-
-if model_management.xformers_enabled_vae():
-    import xformers
-    import xformers.ops
-
-def get_timestep_embedding(timesteps, embedding_dim):
-    """
-    This matches the implementation in Denoising Diffusion Probabilistic Models:
-    From Fairseq.
-    Build sinusoidal embeddings.
-    This matches the implementation in tensor2tensor, but differs slightly
-    from the description in Section 3.5 of "Attention Is All You Need".
-    """
-    assert len(timesteps.shape) == 1
-
-    half_dim = embedding_dim // 2
-    emb = math.log(10000) / (half_dim - 1)
-    emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb)
-    emb = emb.to(device=timesteps.device)
-    emb = timesteps.float()[:, None] * emb[None, :]
-    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
-    if embedding_dim % 2 == 1:  # zero pad
-        emb = torch.nn.functional.pad(emb, (0,1,0,0))
-    return emb
-
-
-def nonlinearity(x):
-    # swish
-    return x*torch.sigmoid(x)
-
-
-def Normalize(in_channels, num_groups=32):
-    return ops.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
-
-
-class Upsample(nn.Module):
-    def __init__(self, in_channels, with_conv):
-        super().__init__()
-        self.with_conv = with_conv
-        if self.with_conv:
-            self.conv = ops.Conv2d(in_channels,
-                                        in_channels,
-                                        kernel_size=3,
-                                        stride=1,
-                                        padding=1)
-
-    def forward(self, x):
-        try:
-            x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
-        except: #operation not implemented for bf16
-            b, c, h, w = x.shape
-            out = torch.empty((b, c, h*2, w*2), dtype=x.dtype, layout=x.layout, device=x.device)
-            split = 8
-            l = out.shape[1] // split
-            for i in range(0, out.shape[1], l):
-                out[:,i:i+l] = torch.nn.functional.interpolate(x[:,i:i+l].to(torch.float32), scale_factor=2.0, mode="nearest").to(x.dtype)
-            del x
-            x = out
-
-        if self.with_conv:
-            x = self.conv(x)
-        return x
-
-
-class Downsample(nn.Module):
-    def __init__(self, in_channels, with_conv):
-        super().__init__()
-        self.with_conv = with_conv
-        if self.with_conv:
-            # no asymmetric padding in torch conv, must do it ourselves
-            self.conv = ops.Conv2d(in_channels,
-                                        in_channels,
-                                        kernel_size=3,
-                                        stride=2,
-                                        padding=0)
-
-    def forward(self, x):
-        if self.with_conv:
-            pad = (0,1,0,1)
-            x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
-            x = self.conv(x)
-        else:
-            x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
-        return x
-
-
-class ResnetBlock(nn.Module):
-    def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False,
-                 dropout, temb_channels=512):
-        super().__init__()
-        self.in_channels = in_channels
-        out_channels = in_channels if out_channels is None else out_channels
-        self.out_channels = out_channels
-        self.use_conv_shortcut = conv_shortcut
-
-        self.swish = torch.nn.SiLU(inplace=True)
-        self.norm1 = Normalize(in_channels)
-        self.conv1 = ops.Conv2d(in_channels,
-                                     out_channels,
-                                     kernel_size=3,
-                                     stride=1,
-                                     padding=1)
-        if temb_channels > 0:
-            self.temb_proj = ops.Linear(temb_channels,
-                                             out_channels)
-        self.norm2 = Normalize(out_channels)
-        self.dropout = torch.nn.Dropout(dropout, inplace=True)
-        self.conv2 = ops.Conv2d(out_channels,
-                                     out_channels,
-                                     kernel_size=3,
-                                     stride=1,
-                                     padding=1)
-        if self.in_channels != self.out_channels:
-            if self.use_conv_shortcut:
-                self.conv_shortcut = ops.Conv2d(in_channels,
-                                                     out_channels,
-                                                     kernel_size=3,
-                                                     stride=1,
-                                                     padding=1)
-            else:
-                self.nin_shortcut = ops.Conv2d(in_channels,
-                                                    out_channels,
-                                                    kernel_size=1,
-                                                    stride=1,
-                                                    padding=0)
-
-    def forward(self, x, temb):
-        h = x
-        h = self.norm1(h)
-        h = self.swish(h)
-        h = self.conv1(h)
-
-        if temb is not None:
-            h = h + self.temb_proj(self.swish(temb))[:,:,None,None]
-
-        h = self.norm2(h)
-        h = self.swish(h)
-        h = self.dropout(h)
-        h = self.conv2(h)
-
-        if self.in_channels != self.out_channels:
-            if self.use_conv_shortcut:
-                x = self.conv_shortcut(x)
-            else:
-                x = self.nin_shortcut(x)
-
-        return x+h
-
-def slice_attention(q, k, v):
-    r1 = torch.zeros_like(k, device=q.device)
-    scale = (int(q.shape[-1])**(-0.5))
-
-    mem_free_total = model_management.get_free_memory(q.device)
-
-    gb = 1024 ** 3
-    tensor_size = q.shape[0] * q.shape[1] * k.shape[2] * q.element_size()
-    modifier = 3 if q.element_size() == 2 else 2.5
-    mem_required = tensor_size * modifier
-    steps = 1
-
-    if mem_required > mem_free_total:
-        steps = 2**(math.ceil(math.log(mem_required / mem_free_total, 2)))
-
-    while True:
-        try:
-            slice_size = q.shape[1] // steps if (q.shape[1] % steps) == 0 else q.shape[1]
-            for i in range(0, q.shape[1], slice_size):
-                end = i + slice_size
-                s1 = torch.bmm(q[:, i:end], k) * scale
-
-                s2 = torch.nn.functional.softmax(s1, dim=2).permute(0,2,1)
-                del s1
-
-                r1[:, :, i:end] = torch.bmm(v, s2)
-                del s2
-            break
-        except model_management.OOM_EXCEPTION as e:
-            model_management.soft_empty_cache(True)
-            steps *= 2
-            if steps > 128:
-                raise e
-            print("out of memory error, increasing steps and trying again", steps)
-
-    return r1
-
-def normal_attention(q, k, v):
-    # compute attention
-    b,c,h,w = q.shape
-
-    q = q.reshape(b,c,h*w)
-    q = q.permute(0,2,1)   # b,hw,c
-    k = k.reshape(b,c,h*w) # b,c,hw
-    v = v.reshape(b,c,h*w)
-
-    r1 = slice_attention(q, k, v)
-    h_ = r1.reshape(b,c,h,w)
-    del r1
-    return h_
-
-def xformers_attention(q, k, v):
-    # compute attention
-    B, C, H, W = q.shape
-    q, k, v = map(
-        lambda t: t.view(B, C, -1).transpose(1, 2).contiguous(),
-        (q, k, v),
-    )
-
-    try:
-        out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None)
-        out = out.transpose(1, 2).reshape(B, C, H, W)
-    except NotImplementedError as e:
-        out = slice_attention(q.view(B, -1, C), k.view(B, -1, C).transpose(1, 2), v.view(B, -1, C).transpose(1, 2)).reshape(B, C, H, W)
-    return out
-
-def pytorch_attention(q, k, v):
-    # compute attention
-    B, C, H, W = q.shape
-    q, k, v = map(
-        lambda t: t.view(B, 1, C, -1).transpose(2, 3).contiguous(),
-        (q, k, v),
-    )
-
-    try:
-        out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False)
-        out = out.transpose(2, 3).reshape(B, C, H, W)
-    except model_management.OOM_EXCEPTION as e:
-        print("scaled_dot_product_attention OOMed: switched to slice attention")
-        out = slice_attention(q.view(B, -1, C), k.view(B, -1, C).transpose(1, 2), v.view(B, -1, C).transpose(1, 2)).reshape(B, C, H, W)
-    return out
-
-
-class AttnBlock(nn.Module):
-    def __init__(self, in_channels):
-        super().__init__()
-        self.in_channels = in_channels
-
-        self.norm = Normalize(in_channels)
-        self.q = ops.Conv2d(in_channels,
-                                 in_channels,
-                                 kernel_size=1,
-                                 stride=1,
-                                 padding=0)
-        self.k = ops.Conv2d(in_channels,
-                                 in_channels,
-                                 kernel_size=1,
-                                 stride=1,
-                                 padding=0)
-        self.v = ops.Conv2d(in_channels,
-                                 in_channels,
-                                 kernel_size=1,
-                                 stride=1,
-                                 padding=0)
-        self.proj_out = ops.Conv2d(in_channels,
-                                        in_channels,
-                                        kernel_size=1,
-                                        stride=1,
-                                        padding=0)
-
-        if model_management.xformers_enabled_vae():
-            print("Using xformers attention in VAE")
-            self.optimized_attention = xformers_attention
-        elif model_management.pytorch_attention_enabled():
-            print("Using pytorch attention in VAE")
-            self.optimized_attention = pytorch_attention
-        else:
-            print("Using split attention in VAE")
-            self.optimized_attention = normal_attention
-
-    def forward(self, x):
-        h_ = x
-        h_ = self.norm(h_)
-        q = self.q(h_)
-        k = self.k(h_)
-        v = self.v(h_)
-
-        h_ = self.optimized_attention(q, k, v)
-
-        h_ = self.proj_out(h_)
-
-        return x+h_
-
-
-def make_attn(in_channels, attn_type="vanilla", attn_kwargs=None):
-    return AttnBlock(in_channels)
-
-
-class Model(nn.Module):
-    def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
-                 attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
-                 resolution, use_timestep=True, use_linear_attn=False, attn_type="vanilla"):
-        super().__init__()
-        if use_linear_attn: attn_type = "linear"
-        self.ch = ch
-        self.temb_ch = self.ch*4
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        self.resolution = resolution
-        self.in_channels = in_channels
-
-        self.use_timestep = use_timestep
-        if self.use_timestep:
-            # timestep embedding
-            self.temb = nn.Module()
-            self.temb.dense = nn.ModuleList([
-                ops.Linear(self.ch,
-                                self.temb_ch),
-                ops.Linear(self.temb_ch,
-                                self.temb_ch),
-            ])
-
-        # downsampling
-        self.conv_in = ops.Conv2d(in_channels,
-                                       self.ch,
-                                       kernel_size=3,
-                                       stride=1,
-                                       padding=1)
-
-        curr_res = resolution
-        in_ch_mult = (1,)+tuple(ch_mult)
-        self.down = nn.ModuleList()
-        for i_level in range(self.num_resolutions):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_in = ch*in_ch_mult[i_level]
-            block_out = ch*ch_mult[i_level]
-            for i_block in range(self.num_res_blocks):
-                block.append(ResnetBlock(in_channels=block_in,
-                                         out_channels=block_out,
-                                         temb_channels=self.temb_ch,
-                                         dropout=dropout))
-                block_in = block_out
-                if curr_res in attn_resolutions:
-                    attn.append(make_attn(block_in, attn_type=attn_type))
-            down = nn.Module()
-            down.block = block
-            down.attn = attn
-            if i_level != self.num_resolutions-1:
-                down.downsample = Downsample(block_in, resamp_with_conv)
-                curr_res = curr_res // 2
-            self.down.append(down)
-
-        # middle
-        self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(in_channels=block_in,
-                                       out_channels=block_in,
-                                       temb_channels=self.temb_ch,
-                                       dropout=dropout)
-        self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
-        self.mid.block_2 = ResnetBlock(in_channels=block_in,
-                                       out_channels=block_in,
-                                       temb_channels=self.temb_ch,
-                                       dropout=dropout)
-
-        # upsampling
-        self.up = nn.ModuleList()
-        for i_level in reversed(range(self.num_resolutions)):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_out = ch*ch_mult[i_level]
-            skip_in = ch*ch_mult[i_level]
-            for i_block in range(self.num_res_blocks+1):
-                if i_block == self.num_res_blocks:
-                    skip_in = ch*in_ch_mult[i_level]
-                block.append(ResnetBlock(in_channels=block_in+skip_in,
-                                         out_channels=block_out,
-                                         temb_channels=self.temb_ch,
-                                         dropout=dropout))
-                block_in = block_out
-                if curr_res in attn_resolutions:
-                    attn.append(make_attn(block_in, attn_type=attn_type))
-            up = nn.Module()
-            up.block = block
-            up.attn = attn
-            if i_level != 0:
-                up.upsample = Upsample(block_in, resamp_with_conv)
-                curr_res = curr_res * 2
-            self.up.insert(0, up) # prepend to get consistent order
-
-        # end
-        self.norm_out = Normalize(block_in)
-        self.conv_out = ops.Conv2d(block_in,
-                                        out_ch,
-                                        kernel_size=3,
-                                        stride=1,
-                                        padding=1)
-
-    def forward(self, x, t=None, context=None):
-        #assert x.shape[2] == x.shape[3] == self.resolution
-        if context is not None:
-            # assume aligned context, cat along channel axis
-            x = torch.cat((x, context), dim=1)
-        if self.use_timestep:
-            # timestep embedding
-            assert t is not None
-            temb = get_timestep_embedding(t, self.ch)
-            temb = self.temb.dense[0](temb)
-            temb = nonlinearity(temb)
-            temb = self.temb.dense[1](temb)
-        else:
-            temb = None
-
-        # downsampling
-        hs = [self.conv_in(x)]
-        for i_level in range(self.num_resolutions):
-            for i_block in range(self.num_res_blocks):
-                h = self.down[i_level].block[i_block](hs[-1], temb)
-                if len(self.down[i_level].attn) > 0:
-                    h = self.down[i_level].attn[i_block](h)
-                hs.append(h)
-            if i_level != self.num_resolutions-1:
-                hs.append(self.down[i_level].downsample(hs[-1]))
-
-        # middle
-        h = hs[-1]
-        h = self.mid.block_1(h, temb)
-        h = self.mid.attn_1(h)
-        h = self.mid.block_2(h, temb)
-
-        # upsampling
-        for i_level in reversed(range(self.num_resolutions)):
-            for i_block in range(self.num_res_blocks+1):
-                h = self.up[i_level].block[i_block](
-                    torch.cat([h, hs.pop()], dim=1), temb)
-                if len(self.up[i_level].attn) > 0:
-                    h = self.up[i_level].attn[i_block](h)
-            if i_level != 0:
-                h = self.up[i_level].upsample(h)
-
-        # end
-        h = self.norm_out(h)
-        h = nonlinearity(h)
-        h = self.conv_out(h)
-        return h
-
-    def get_last_layer(self):
-        return self.conv_out.weight
-
-
-class Encoder(nn.Module):
-    def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
-                 attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
-                 resolution, z_channels, double_z=True, use_linear_attn=False, attn_type="vanilla",
-                 **ignore_kwargs):
-        super().__init__()
-        if use_linear_attn: attn_type = "linear"
-        self.ch = ch
-        self.temb_ch = 0
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        self.resolution = resolution
-        self.in_channels = in_channels
-
-        # downsampling
-        self.conv_in = ops.Conv2d(in_channels,
-                                       self.ch,
-                                       kernel_size=3,
-                                       stride=1,
-                                       padding=1)
-
-        curr_res = resolution
-        in_ch_mult = (1,)+tuple(ch_mult)
-        self.in_ch_mult = in_ch_mult
-        self.down = nn.ModuleList()
-        for i_level in range(self.num_resolutions):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_in = ch*in_ch_mult[i_level]
-            block_out = ch*ch_mult[i_level]
-            for i_block in range(self.num_res_blocks):
-                block.append(ResnetBlock(in_channels=block_in,
-                                         out_channels=block_out,
-                                         temb_channels=self.temb_ch,
-                                         dropout=dropout))
-                block_in = block_out
-                if curr_res in attn_resolutions:
-                    attn.append(make_attn(block_in, attn_type=attn_type))
-            down = nn.Module()
-            down.block = block
-            down.attn = attn
-            if i_level != self.num_resolutions-1:
-                down.downsample = Downsample(block_in, resamp_with_conv)
-                curr_res = curr_res // 2
-            self.down.append(down)
-
-        # middle
-        self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(in_channels=block_in,
-                                       out_channels=block_in,
-                                       temb_channels=self.temb_ch,
-                                       dropout=dropout)
-        self.mid.attn_1 = make_attn(block_in, attn_type=attn_type)
-        self.mid.block_2 = ResnetBlock(in_channels=block_in,
-                                       out_channels=block_in,
-                                       temb_channels=self.temb_ch,
-                                       dropout=dropout)
-
-        # end
-        self.norm_out = Normalize(block_in)
-        self.conv_out = ops.Conv2d(block_in,
-                                        2*z_channels if double_z else z_channels,
-                                        kernel_size=3,
-                                        stride=1,
-                                        padding=1)
-
-    def forward(self, x):
-        # timestep embedding
-        temb = None
-        # downsampling
-        h = self.conv_in(x)
-        for i_level in range(self.num_resolutions):
-            for i_block in range(self.num_res_blocks):
-                h = self.down[i_level].block[i_block](h, temb)
-                if len(self.down[i_level].attn) > 0:
-                    h = self.down[i_level].attn[i_block](h)
-            if i_level != self.num_resolutions-1:
-                h = self.down[i_level].downsample(h)
-
-        # middle
-        h = self.mid.block_1(h, temb)
-        h = self.mid.attn_1(h)
-        h = self.mid.block_2(h, temb)
-
-        # end
-        h = self.norm_out(h)
-        h = nonlinearity(h)
-        h = self.conv_out(h)
-        return h
-
-
-class Decoder(nn.Module):
-    def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
-                 attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
-                 resolution, z_channels, give_pre_end=False, tanh_out=False, use_linear_attn=False,
-                 conv_out_op=ops.Conv2d,
-                 resnet_op=ResnetBlock,
-                 attn_op=AttnBlock,
-                **ignorekwargs):
-        super().__init__()
-        if use_linear_attn: attn_type = "linear"
-        self.ch = ch
-        self.temb_ch = 0
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        self.resolution = resolution
-        self.in_channels = in_channels
-        self.give_pre_end = give_pre_end
-        self.tanh_out = tanh_out
-
-        # compute in_ch_mult, block_in and curr_res at lowest res
-        in_ch_mult = (1,)+tuple(ch_mult)
-        block_in = ch*ch_mult[self.num_resolutions-1]
-        curr_res = resolution // 2**(self.num_resolutions-1)
-        self.z_shape = (1,z_channels,curr_res,curr_res)
-        print("Working with z of shape {} = {} dimensions.".format(
-            self.z_shape, np.prod(self.z_shape)))
-
-        # z to block_in
-        self.conv_in = ops.Conv2d(z_channels,
-                                       block_in,
-                                       kernel_size=3,
-                                       stride=1,
-                                       padding=1)
-
-        # middle
-        self.mid = nn.Module()
-        self.mid.block_1 = resnet_op(in_channels=block_in,
-                                       out_channels=block_in,
-                                       temb_channels=self.temb_ch,
-                                       dropout=dropout)
-        self.mid.attn_1 = attn_op(block_in)
-        self.mid.block_2 = resnet_op(in_channels=block_in,
-                                       out_channels=block_in,
-                                       temb_channels=self.temb_ch,
-                                       dropout=dropout)
-
-        # upsampling
-        self.up = nn.ModuleList()
-        for i_level in reversed(range(self.num_resolutions)):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_out = ch*ch_mult[i_level]
-            for i_block in range(self.num_res_blocks+1):
-                block.append(resnet_op(in_channels=block_in,
-                                         out_channels=block_out,
-                                         temb_channels=self.temb_ch,
-                                         dropout=dropout))
-                block_in = block_out
-                if curr_res in attn_resolutions:
-                    attn.append(attn_op(block_in))
-            up = nn.Module()
-            up.block = block
-            up.attn = attn
-            if i_level != 0:
-                up.upsample = Upsample(block_in, resamp_with_conv)
-                curr_res = curr_res * 2
-            self.up.insert(0, up) # prepend to get consistent order
-
-        # end
-        self.norm_out = Normalize(block_in)
-        self.conv_out = conv_out_op(block_in,
-                                        out_ch,
-                                        kernel_size=3,
-                                        stride=1,
-                                        padding=1)
-
-    def forward(self, z, **kwargs):
-        #assert z.shape[1:] == self.z_shape[1:]
-        self.last_z_shape = z.shape
-
-        # timestep embedding
-        temb = None
-
-        # z to block_in
-        h = self.conv_in(z)
-
-        # middle
-        h = self.mid.block_1(h, temb, **kwargs)
-        h = self.mid.attn_1(h, **kwargs)
-        h = self.mid.block_2(h, temb, **kwargs)
-
-        # upsampling
-        for i_level in reversed(range(self.num_resolutions)):
-            for i_block in range(self.num_res_blocks+1):
-                h = self.up[i_level].block[i_block](h, temb, **kwargs)
-                if len(self.up[i_level].attn) > 0:
-                    h = self.up[i_level].attn[i_block](h, **kwargs)
-            if i_level != 0:
-                h = self.up[i_level].upsample(h)
-
-        # end
-        if self.give_pre_end:
-            return h
-
-        h = self.norm_out(h)
-        h = nonlinearity(h)
-        h = self.conv_out(h, **kwargs)
-        if self.tanh_out:
-            h = torch.tanh(h)
-        return h

ldm_patched/ldm/modules/diffusionmodules/openaimodel.py

@@ -1,886 +0,0 @@
-from abc import abstractmethod
-
-import torch as th
-import torch.nn as nn
-import torch.nn.functional as F
-from einops import rearrange
-
-from .util import (
-    checkpoint,
-    avg_pool_nd,
-    zero_module,
-    timestep_embedding,
-    AlphaBlender,
-)
-from ..attention import SpatialTransformer, SpatialVideoTransformer, default
-from ldm_patched.ldm.util import exists
-import ldm_patched.modules.ops
-ops = ldm_patched.modules.ops.disable_weight_init
-
-class TimestepBlock(nn.Module):
-    """
-    Any module where forward() takes timestep embeddings as a second argument.
-    """
-
-    @abstractmethod
-    def forward(self, x, emb):
-        """
-        Apply the module to `x` given `emb` timestep embeddings.
-        """
-
-#This is needed because accelerate makes a copy of transformer_options which breaks "transformer_index"
-def forward_timestep_embed(ts, x, emb, context=None, transformer_options={}, output_shape=None, time_context=None, num_video_frames=None, image_only_indicator=None):
-    for layer in ts:
-        if isinstance(layer, VideoResBlock):
-            x = layer(x, emb, num_video_frames, image_only_indicator)
-        elif isinstance(layer, TimestepBlock):
-            x = layer(x, emb)
-        elif isinstance(layer, SpatialVideoTransformer):
-            x = layer(x, context, time_context, num_video_frames, image_only_indicator, transformer_options)
-            if "transformer_index" in transformer_options:
-                transformer_options["transformer_index"] += 1
-        elif isinstance(layer, SpatialTransformer):
-            x = layer(x, context, transformer_options)
-            if "transformer_index" in transformer_options:
-                transformer_options["transformer_index"] += 1
-        elif isinstance(layer, Upsample):
-            x = layer(x, output_shape=output_shape)
-        else:
-            x = layer(x)
-    return x
-
-class TimestepEmbedSequential(nn.Sequential, TimestepBlock):
-    """
-    A sequential module that passes timestep embeddings to the children that
-    support it as an extra input.
-    """
-
-    def forward(self, *args, **kwargs):
-        return forward_timestep_embed(self, *args, **kwargs)
-
-class Upsample(nn.Module):
-    """
-    An upsampling layer with an optional convolution.
-    :param channels: channels in the inputs and outputs.
-    :param use_conv: a bool determining if a convolution is applied.
-    :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
-                 upsampling occurs in the inner-two dimensions.
-    """
-
-    def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1, dtype=None, device=None, operations=ops):
-        super().__init__()
-        self.channels = channels
-        self.out_channels = out_channels or channels
-        self.use_conv = use_conv
-        self.dims = dims
-        if use_conv:
-            self.conv = operations.conv_nd(dims, self.channels, self.out_channels, 3, padding=padding, dtype=dtype, device=device)
-
-    def forward(self, x, output_shape=None):
-        assert x.shape[1] == self.channels
-        if self.dims == 3:
-            shape = [x.shape[2], x.shape[3] * 2, x.shape[4] * 2]
-            if output_shape is not None:
-                shape[1] = output_shape[3]
-                shape[2] = output_shape[4]
-        else:
-            shape = [x.shape[2] * 2, x.shape[3] * 2]
-            if output_shape is not None:
-                shape[0] = output_shape[2]
-                shape[1] = output_shape[3]
-
-        x = F.interpolate(x, size=shape, mode="nearest")
-        if self.use_conv:
-            x = self.conv(x)
-        return x
-
-class Downsample(nn.Module):
-    """
-    A downsampling layer with an optional convolution.
-    :param channels: channels in the inputs and outputs.
-    :param use_conv: a bool determining if a convolution is applied.
-    :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
-                 downsampling occurs in the inner-two dimensions.
-    """
-
-    def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1, dtype=None, device=None, operations=ops):
-        super().__init__()
-        self.channels = channels
-        self.out_channels = out_channels or channels
-        self.use_conv = use_conv
-        self.dims = dims
-        stride = 2 if dims != 3 else (1, 2, 2)
-        if use_conv:
-            self.op = operations.conv_nd(
-                dims, self.channels, self.out_channels, 3, stride=stride, padding=padding, dtype=dtype, device=device
-            )
-        else:
-            assert self.channels == self.out_channels
-            self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)
-
-    def forward(self, x):
-        assert x.shape[1] == self.channels
-        return self.op(x)
-
-
-class ResBlock(TimestepBlock):
-    """
-    A residual block that can optionally change the number of channels.
-    :param channels: the number of input channels.
-    :param emb_channels: the number of timestep embedding channels.
-    :param dropout: the rate of dropout.
-    :param out_channels: if specified, the number of out channels.
-    :param use_conv: if True and out_channels is specified, use a spatial
-        convolution instead of a smaller 1x1 convolution to change the
-        channels in the skip connection.
-    :param dims: determines if the signal is 1D, 2D, or 3D.
-    :param use_checkpoint: if True, use gradient checkpointing on this module.
-    :param up: if True, use this block for upsampling.
-    :param down: if True, use this block for downsampling.
-    """
-
-    def __init__(
-        self,
-        channels,
-        emb_channels,
-        dropout,
-        out_channels=None,
-        use_conv=False,
-        use_scale_shift_norm=False,
-        dims=2,
-        use_checkpoint=False,
-        up=False,
-        down=False,
-        kernel_size=3,
-        exchange_temb_dims=False,
-        skip_t_emb=False,
-        dtype=None,
-        device=None,
-        operations=ops
-    ):
-        super().__init__()
-        self.channels = channels
-        self.emb_channels = emb_channels
-        self.dropout = dropout
-        self.out_channels = out_channels or channels
-        self.use_conv = use_conv
-        self.use_checkpoint = use_checkpoint
-        self.use_scale_shift_norm = use_scale_shift_norm
-        self.exchange_temb_dims = exchange_temb_dims
-
-        if isinstance(kernel_size, list):
-            padding = [k // 2 for k in kernel_size]
-        else:
-            padding = kernel_size // 2
-
-        self.in_layers = nn.Sequential(
-            operations.GroupNorm(32, channels, dtype=dtype, device=device),
-            nn.SiLU(),
-            operations.conv_nd(dims, channels, self.out_channels, kernel_size, padding=padding, dtype=dtype, device=device),
-        )
-
-        self.updown = up or down
-
-        if up:
-            self.h_upd = Upsample(channels, False, dims, dtype=dtype, device=device)
-            self.x_upd = Upsample(channels, False, dims, dtype=dtype, device=device)
-        elif down:
-            self.h_upd = Downsample(channels, False, dims, dtype=dtype, device=device)
-            self.x_upd = Downsample(channels, False, dims, dtype=dtype, device=device)
-        else:
-            self.h_upd = self.x_upd = nn.Identity()
-
-        self.skip_t_emb = skip_t_emb
-        if self.skip_t_emb:
-            self.emb_layers = None
-            self.exchange_temb_dims = False
-        else:
-            self.emb_layers = nn.Sequential(
-                nn.SiLU(),
-                operations.Linear(
-                    emb_channels,
-                    2 * self.out_channels if use_scale_shift_norm else self.out_channels, dtype=dtype, device=device
-                ),
-            )
-        self.out_layers = nn.Sequential(
-            operations.GroupNorm(32, self.out_channels, dtype=dtype, device=device),
-            nn.SiLU(),
-            nn.Dropout(p=dropout),
-            operations.conv_nd(dims, self.out_channels, self.out_channels, kernel_size, padding=padding, dtype=dtype, device=device)
-            ,
-        )
-
-        if self.out_channels == channels:
-            self.skip_connection = nn.Identity()
-        elif use_conv:
-            self.skip_connection = operations.conv_nd(
-                dims, channels, self.out_channels, kernel_size, padding=padding, dtype=dtype, device=device
-            )
-        else:
-            self.skip_connection = operations.conv_nd(dims, channels, self.out_channels, 1, dtype=dtype, device=device)
-
-    def forward(self, x, emb):
-        """
-        Apply the block to a Tensor, conditioned on a timestep embedding.
-        :param x: an [N x C x ...] Tensor of features.
-        :param emb: an [N x emb_channels] Tensor of timestep embeddings.
-        :return: an [N x C x ...] Tensor of outputs.
-        """
-        return checkpoint(
-            self._forward, (x, emb), self.parameters(), self.use_checkpoint
-        )
-
-
-    def _forward(self, x, emb):
-        if self.updown:
-            in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
-            h = in_rest(x)
-            h = self.h_upd(h)
-            x = self.x_upd(x)
-            h = in_conv(h)
-        else:
-            h = self.in_layers(x)
-
-        emb_out = None
-        if not self.skip_t_emb:
-            emb_out = self.emb_layers(emb).type(h.dtype)
-            while len(emb_out.shape) < len(h.shape):
-                emb_out = emb_out[..., None]
-        if self.use_scale_shift_norm:
-            out_norm, out_rest = self.out_layers[0], self.out_layers[1:]
-            h = out_norm(h)
-            if emb_out is not None:
-                scale, shift = th.chunk(emb_out, 2, dim=1)
-                h *= (1 + scale)
-                h += shift
-            h = out_rest(h)
-        else:
-            if emb_out is not None:
-                if self.exchange_temb_dims:
-                    emb_out = rearrange(emb_out, "b t c ... -> b c t ...")
-                h = h + emb_out
-            h = self.out_layers(h)
-        return self.skip_connection(x) + h
-
-
-class VideoResBlock(ResBlock):
-    def __init__(
-        self,
-        channels: int,
-        emb_channels: int,
-        dropout: float,
-        video_kernel_size=3,
-        merge_strategy: str = "fixed",
-        merge_factor: float = 0.5,
-        out_channels=None,
-        use_conv: bool = False,
-        use_scale_shift_norm: bool = False,
-        dims: int = 2,
-        use_checkpoint: bool = False,
-        up: bool = False,
-        down: bool = False,
-        dtype=None,
-        device=None,
-        operations=ops
-    ):
-        super().__init__(
-            channels,
-            emb_channels,
-            dropout,
-            out_channels=out_channels,
-            use_conv=use_conv,
-            use_scale_shift_norm=use_scale_shift_norm,
-            dims=dims,
-            use_checkpoint=use_checkpoint,
-            up=up,
-            down=down,
-            dtype=dtype,
-            device=device,
-            operations=operations
-        )
-
-        self.time_stack = ResBlock(
-            default(out_channels, channels),
-            emb_channels,
-            dropout=dropout,
-            dims=3,
-            out_channels=default(out_channels, channels),
-            use_scale_shift_norm=False,
-            use_conv=False,
-            up=False,
-            down=False,
-            kernel_size=video_kernel_size,
-            use_checkpoint=use_checkpoint,
-            exchange_temb_dims=True,
-            dtype=dtype,
-            device=device,
-            operations=operations
-        )
-        self.time_mixer = AlphaBlender(
-            alpha=merge_factor,
-            merge_strategy=merge_strategy,
-            rearrange_pattern="b t -> b 1 t 1 1",
-        )
-
-    def forward(
-        self,
-        x: th.Tensor,
-        emb: th.Tensor,
-        num_video_frames: int,
-        image_only_indicator = None,
-    ) -> th.Tensor:
-        x = super().forward(x, emb)
-
-        x_mix = rearrange(x, "(b t) c h w -> b c t h w", t=num_video_frames)
-        x = rearrange(x, "(b t) c h w -> b c t h w", t=num_video_frames)
-
-        x = self.time_stack(
-            x, rearrange(emb, "(b t) ... -> b t ...", t=num_video_frames)
-        )
-        x = self.time_mixer(
-            x_spatial=x_mix, x_temporal=x, image_only_indicator=image_only_indicator
-        )
-        x = rearrange(x, "b c t h w -> (b t) c h w")
-        return x
-
-
-class Timestep(nn.Module):
-    def __init__(self, dim):
-        super().__init__()
-        self.dim = dim
-
-    def forward(self, t):
-        return timestep_embedding(t, self.dim)
-
-def apply_control(h, control, name):
-    if control is not None and name in control and len(control[name]) > 0:
-        ctrl = control[name].pop()
-        if ctrl is not None:
-            try:
-                h += ctrl
-            except:
-                print("warning control could not be applied", h.shape, ctrl.shape)
-    return h
-
-class UNetModel(nn.Module):
-    """
-    The full UNet model with attention and timestep embedding.
-    :param in_channels: channels in the input Tensor.
-    :param model_channels: base channel count for the model.
-    :param out_channels: channels in the output Tensor.
-    :param num_res_blocks: number of residual blocks per downsample.
-    :param dropout: the dropout probability.
-    :param channel_mult: channel multiplier for each level of the UNet.
-    :param conv_resample: if True, use learned convolutions for upsampling and
-        downsampling.
-    :param dims: determines if the signal is 1D, 2D, or 3D.
-    :param num_classes: if specified (as an int), then this model will be
-        class-conditional with `num_classes` classes.
-    :param use_checkpoint: use gradient checkpointing to reduce memory usage.
-    :param num_heads: the number of attention heads in each attention layer.
-    :param num_heads_channels: if specified, ignore num_heads and instead use
-                               a fixed channel width per attention head.
-    :param num_heads_upsample: works with num_heads to set a different number
-                               of heads for upsampling. Deprecated.
-    :param use_scale_shift_norm: use a FiLM-like conditioning mechanism.
-    :param resblock_updown: use residual blocks for up/downsampling.
-    :param use_new_attention_order: use a different attention pattern for potentially
-                                    increased efficiency.
-    """
-
-    def __init__(
-        self,
-        image_size,
-        in_channels,
-        model_channels,
-        out_channels,
-        num_res_blocks,
-        dropout=0,
-        channel_mult=(1, 2, 4, 8),
-        conv_resample=True,
-        dims=2,
-        num_classes=None,
-        use_checkpoint=False,
-        dtype=th.float32,
-        num_heads=-1,
-        num_head_channels=-1,
-        num_heads_upsample=-1,
-        use_scale_shift_norm=False,
-        resblock_updown=False,
-        use_new_attention_order=False,
-        use_spatial_transformer=False,    # custom transformer support
-        transformer_depth=1,              # custom transformer support
-        context_dim=None,                 # custom transformer support
-        n_embed=None,                     # custom support for prediction of discrete ids into codebook of first stage vq model
-        legacy=True,
-        disable_self_attentions=None,
-        num_attention_blocks=None,
-        disable_middle_self_attn=False,
-        use_linear_in_transformer=False,
-        adm_in_channels=None,
-        transformer_depth_middle=None,
-        transformer_depth_output=None,
-        use_temporal_resblock=False,
-        use_temporal_attention=False,
-        time_context_dim=None,
-        extra_ff_mix_layer=False,
-        use_spatial_context=False,
-        merge_strategy=None,
-        merge_factor=0.0,
-        video_kernel_size=None,
-        disable_temporal_crossattention=False,
-        max_ddpm_temb_period=10000,
-        device=None,
-        operations=ops,
-    ):
-        super().__init__()
-
-        if context_dim is not None:
-            assert use_spatial_transformer, 'Fool!! You forgot to use the spatial transformer for your cross-attention conditioning...'
-            # from omegaconf.listconfig import ListConfig
-            # if type(context_dim) == ListConfig:
-            #     context_dim = list(context_dim)
-
-        if num_heads_upsample == -1:
-            num_heads_upsample = num_heads
-
-        if num_heads == -1:
-            assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set'
-
-        if num_head_channels == -1:
-            assert num_heads != -1, 'Either num_heads or num_head_channels has to be set'
-
-        self.in_channels = in_channels
-        self.model_channels = model_channels
-        self.out_channels = out_channels
-
-        if isinstance(num_res_blocks, int):
-            self.num_res_blocks = len(channel_mult) * [num_res_blocks]
-        else:
-            if len(num_res_blocks) != len(channel_mult):
-                raise ValueError("provide num_res_blocks either as an int (globally constant) or "
-                                 "as a list/tuple (per-level) with the same length as channel_mult")
-            self.num_res_blocks = num_res_blocks
-
-        if disable_self_attentions is not None:
-            # should be a list of booleans, indicating whether to disable self-attention in TransformerBlocks or not
-            assert len(disable_self_attentions) == len(channel_mult)
-        if num_attention_blocks is not None:
-            assert len(num_attention_blocks) == len(self.num_res_blocks)
-
-        transformer_depth = transformer_depth[:]
-        transformer_depth_output = transformer_depth_output[:]
-
-        self.dropout = dropout
-        self.channel_mult = channel_mult
-        self.conv_resample = conv_resample
-        self.num_classes = num_classes
-        self.use_checkpoint = use_checkpoint
-        self.dtype = dtype
-        self.num_heads = num_heads
-        self.num_head_channels = num_head_channels
-        self.num_heads_upsample = num_heads_upsample
-        self.use_temporal_resblocks = use_temporal_resblock
-        self.predict_codebook_ids = n_embed is not None
-
-        self.default_num_video_frames = None
-        self.default_image_only_indicator = None
-
-        time_embed_dim = model_channels * 4
-        self.time_embed = nn.Sequential(
-            operations.Linear(model_channels, time_embed_dim, dtype=self.dtype, device=device),
-            nn.SiLU(),
-            operations.Linear(time_embed_dim, time_embed_dim, dtype=self.dtype, device=device),
-        )
-
-        if self.num_classes is not None:
-            if isinstance(self.num_classes, int):
-                self.label_emb = nn.Embedding(num_classes, time_embed_dim, dtype=self.dtype, device=device)
-            elif self.num_classes == "continuous":
-                print("setting up linear c_adm embedding layer")
-                self.label_emb = nn.Linear(1, time_embed_dim)
-            elif self.num_classes == "sequential":
-                assert adm_in_channels is not None
-                self.label_emb = nn.Sequential(
-                    nn.Sequential(
-                        operations.Linear(adm_in_channels, time_embed_dim, dtype=self.dtype, device=device),
-                        nn.SiLU(),
-                        operations.Linear(time_embed_dim, time_embed_dim, dtype=self.dtype, device=device),
-                    )
-                )
-            else:
-                raise ValueError()
-
-        self.input_blocks = nn.ModuleList(
-            [
-                TimestepEmbedSequential(
-                    operations.conv_nd(dims, in_channels, model_channels, 3, padding=1, dtype=self.dtype, device=device)
-                )
-            ]
-        )
-        self._feature_size = model_channels
-        input_block_chans = [model_channels]
-        ch = model_channels
-        ds = 1
-
-        def get_attention_layer(
-            ch,
-            num_heads,
-            dim_head,
-            depth=1,
-            context_dim=None,
-            use_checkpoint=False,
-            disable_self_attn=False,
-        ):
-            if use_temporal_attention:
-                return SpatialVideoTransformer(
-                    ch,
-                    num_heads,
-                    dim_head,
-                    depth=depth,
-                    context_dim=context_dim,
-                    time_context_dim=time_context_dim,
-                    dropout=dropout,
-                    ff_in=extra_ff_mix_layer,
-                    use_spatial_context=use_spatial_context,
-                    merge_strategy=merge_strategy,
-                    merge_factor=merge_factor,
-                    checkpoint=use_checkpoint,
-                    use_linear=use_linear_in_transformer,
-                    disable_self_attn=disable_self_attn,
-                    disable_temporal_crossattention=disable_temporal_crossattention,
-                    max_time_embed_period=max_ddpm_temb_period,
-                    dtype=self.dtype, device=device, operations=operations
-                )
-            else:
-                return SpatialTransformer(
-                                ch, num_heads, dim_head, depth=depth, context_dim=context_dim,
-                                disable_self_attn=disable_self_attn, use_linear=use_linear_in_transformer,
-                                use_checkpoint=use_checkpoint, dtype=self.dtype, device=device, operations=operations
-                            )
-
-        def get_resblock(
-            merge_factor,
-            merge_strategy,
-            video_kernel_size,
-            ch,
-            time_embed_dim,
-            dropout,
-            out_channels,
-            dims,
-            use_checkpoint,
-            use_scale_shift_norm,
-            down=False,
-            up=False,
-            dtype=None,
-            device=None,
-            operations=ops
-        ):
-            if self.use_temporal_resblocks:
-                return VideoResBlock(
-                    merge_factor=merge_factor,
-                    merge_strategy=merge_strategy,
-                    video_kernel_size=video_kernel_size,
-                    channels=ch,
-                    emb_channels=time_embed_dim,
-                    dropout=dropout,
-                    out_channels=out_channels,
-                    dims=dims,
-                    use_checkpoint=use_checkpoint,
-                    use_scale_shift_norm=use_scale_shift_norm,
-                    down=down,
-                    up=up,
-                    dtype=dtype,
-                    device=device,
-                    operations=operations
-                )
-            else:
-                return ResBlock(
-                    channels=ch,
-                    emb_channels=time_embed_dim,
-                    dropout=dropout,
-                    out_channels=out_channels,
-                    use_checkpoint=use_checkpoint,
-                    dims=dims,
-                    use_scale_shift_norm=use_scale_shift_norm,
-                    down=down,
-                    up=up,
-                    dtype=dtype,
-                    device=device,
-                    operations=operations
-                )
-
-        for level, mult in enumerate(channel_mult):
-            for nr in range(self.num_res_blocks[level]):
-                layers = [
-                    get_resblock(
-                        merge_factor=merge_factor,
-                        merge_strategy=merge_strategy,
-                        video_kernel_size=video_kernel_size,
-                        ch=ch,
-                        time_embed_dim=time_embed_dim,
-                        dropout=dropout,
-                        out_channels=mult * model_channels,
-                        dims=dims,
-                        use_checkpoint=use_checkpoint,
-                        use_scale_shift_norm=use_scale_shift_norm,
-                        dtype=self.dtype,
-                        device=device,
-                        operations=operations,
-                    )
-                ]
-                ch = mult * model_channels
-                num_transformers = transformer_depth.pop(0)
-                if num_transformers > 0:
-                    if num_head_channels == -1:
-                        dim_head = ch // num_heads
-                    else:
-                        num_heads = ch // num_head_channels
-                        dim_head = num_head_channels
-                    if legacy:
-                        #num_heads = 1
-                        dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
-                    if exists(disable_self_attentions):
-                        disabled_sa = disable_self_attentions[level]
-                    else:
-                        disabled_sa = False
-
-                    if not exists(num_attention_blocks) or nr < num_attention_blocks[level]:
-                        layers.append(get_attention_layer(
-                                ch, num_heads, dim_head, depth=num_transformers, context_dim=context_dim,
-                                disable_self_attn=disabled_sa, use_checkpoint=use_checkpoint)
-                        )
-                self.input_blocks.append(TimestepEmbedSequential(*layers))
-                self._feature_size += ch
-                input_block_chans.append(ch)
-            if level != len(channel_mult) - 1:
-                out_ch = ch
-                self.input_blocks.append(
-                    TimestepEmbedSequential(
-                        get_resblock(
-                            merge_factor=merge_factor,
-                            merge_strategy=merge_strategy,
-                            video_kernel_size=video_kernel_size,
-                            ch=ch,
-                            time_embed_dim=time_embed_dim,
-                            dropout=dropout,
-                            out_channels=out_ch,
-                            dims=dims,
-                            use_checkpoint=use_checkpoint,
-                            use_scale_shift_norm=use_scale_shift_norm,
-                            down=True,
-                            dtype=self.dtype,
-                            device=device,
-                            operations=operations
-                        )
-                        if resblock_updown
-                        else Downsample(
-                            ch, conv_resample, dims=dims, out_channels=out_ch, dtype=self.dtype, device=device, operations=operations
-                        )
-                    )
-                )
-                ch = out_ch
-                input_block_chans.append(ch)
-                ds *= 2
-                self._feature_size += ch
-
-        if num_head_channels == -1:
-            dim_head = ch // num_heads
-        else:
-            num_heads = ch // num_head_channels
-            dim_head = num_head_channels
-        if legacy:
-            #num_heads = 1
-            dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
-        mid_block = [
-            get_resblock(
-                merge_factor=merge_factor,
-                merge_strategy=merge_strategy,
-                video_kernel_size=video_kernel_size,
-                ch=ch,
-                time_embed_dim=time_embed_dim,
-                dropout=dropout,
-                out_channels=None,
-                dims=dims,
-                use_checkpoint=use_checkpoint,
-                use_scale_shift_norm=use_scale_shift_norm,
-                dtype=self.dtype,
-                device=device,
-                operations=operations
-            )]
-        if transformer_depth_middle >= 0:
-            mid_block += [get_attention_layer(  # always uses a self-attn
-                            ch, num_heads, dim_head, depth=transformer_depth_middle, context_dim=context_dim,
-                            disable_self_attn=disable_middle_self_attn, use_checkpoint=use_checkpoint
-                        ),
-            get_resblock(
-                merge_factor=merge_factor,
-                merge_strategy=merge_strategy,
-                video_kernel_size=video_kernel_size,
-                ch=ch,
-                time_embed_dim=time_embed_dim,
-                dropout=dropout,
-                out_channels=None,
-                dims=dims,
-                use_checkpoint=use_checkpoint,
-                use_scale_shift_norm=use_scale_shift_norm,
-                dtype=self.dtype,
-                device=device,
-                operations=operations
-            )]
-        self.middle_block = TimestepEmbedSequential(*mid_block)
-        self._feature_size += ch
-
-        self.output_blocks = nn.ModuleList([])
-        for level, mult in list(enumerate(channel_mult))[::-1]:
-            for i in range(self.num_res_blocks[level] + 1):
-                ich = input_block_chans.pop()
-                layers = [
-                    get_resblock(
-                        merge_factor=merge_factor,
-                        merge_strategy=merge_strategy,
-                        video_kernel_size=video_kernel_size,
-                        ch=ch + ich,
-                        time_embed_dim=time_embed_dim,
-                        dropout=dropout,
-                        out_channels=model_channels * mult,
-                        dims=dims,
-                        use_checkpoint=use_checkpoint,
-                        use_scale_shift_norm=use_scale_shift_norm,
-                        dtype=self.dtype,
-                        device=device,
-                        operations=operations
-                    )
-                ]
-                ch = model_channels * mult
-                num_transformers = transformer_depth_output.pop()
-                if num_transformers > 0:
-                    if num_head_channels == -1:
-                        dim_head = ch // num_heads
-                    else:
-                        num_heads = ch // num_head_channels
-                        dim_head = num_head_channels
-                    if legacy:
-                        #num_heads = 1
-                        dim_head = ch // num_heads if use_spatial_transformer else num_head_channels
-                    if exists(disable_self_attentions):
-                        disabled_sa = disable_self_attentions[level]
-                    else:
-                        disabled_sa = False
-
-                    if not exists(num_attention_blocks) or i < num_attention_blocks[level]:
-                        layers.append(
-                            get_attention_layer(
-                                ch, num_heads, dim_head, depth=num_transformers, context_dim=context_dim,
-                                disable_self_attn=disabled_sa, use_checkpoint=use_checkpoint
-                            )
-                        )
-                if level and i == self.num_res_blocks[level]:
-                    out_ch = ch
-                    layers.append(
-                        get_resblock(
-                            merge_factor=merge_factor,
-                            merge_strategy=merge_strategy,
-                            video_kernel_size=video_kernel_size,
-                            ch=ch,
-                            time_embed_dim=time_embed_dim,
-                            dropout=dropout,
-                            out_channels=out_ch,
-                            dims=dims,
-                            use_checkpoint=use_checkpoint,
-                            use_scale_shift_norm=use_scale_shift_norm,
-                            up=True,
-                            dtype=self.dtype,
-                            device=device,
-                            operations=operations
-                        )
-                        if resblock_updown
-                        else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch, dtype=self.dtype, device=device, operations=operations)
-                    )
-                    ds //= 2
-                self.output_blocks.append(TimestepEmbedSequential(*layers))
-                self._feature_size += ch
-
-        self.out = nn.Sequential(
-            operations.GroupNorm(32, ch, dtype=self.dtype, device=device),
-            nn.SiLU(),
-            zero_module(operations.conv_nd(dims, model_channels, out_channels, 3, padding=1, dtype=self.dtype, device=device)),
-        )
-        if self.predict_codebook_ids:
-            self.id_predictor = nn.Sequential(
-            operations.GroupNorm(32, ch, dtype=self.dtype, device=device),
-            operations.conv_nd(dims, model_channels, n_embed, 1, dtype=self.dtype, device=device),
-            #nn.LogSoftmax(dim=1)  # change to cross_entropy and produce non-normalized logits
-        )
-
-    def forward(self, x, timesteps=None, context=None, y=None, control=None, transformer_options={}, **kwargs):
-        """
-        Apply the model to an input batch.
-        :param x: an [N x C x ...] Tensor of inputs.
-        :param timesteps: a 1-D batch of timesteps.
-        :param context: conditioning plugged in via crossattn
-        :param y: an [N] Tensor of labels, if class-conditional.
-        :return: an [N x C x ...] Tensor of outputs.
-        """
-        transformer_options["original_shape"] = list(x.shape)
-        transformer_options["transformer_index"] = 0
-        transformer_patches = transformer_options.get("patches", {})
-
-        num_video_frames = kwargs.get("num_video_frames", self.default_num_video_frames)
-        image_only_indicator = kwargs.get("image_only_indicator", self.default_image_only_indicator)
-        time_context = kwargs.get("time_context", None)
-
-        assert (y is not None) == (
-            self.num_classes is not None
-        ), "must specify y if and only if the model is class-conditional"
-        hs = []
-        t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype)
-        emb = self.time_embed(t_emb)
-
-        if self.num_classes is not None:
-            assert y.shape[0] == x.shape[0]
-            emb = emb + self.label_emb(y)
-
-        h = x
-        for id, module in enumerate(self.input_blocks):
-            transformer_options["block"] = ("input", id)
-            h = forward_timestep_embed(module, h, emb, context, transformer_options, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator)
-            h = apply_control(h, control, 'input')
-            if "input_block_patch" in transformer_patches:
-                patch = transformer_patches["input_block_patch"]
-                for p in patch:
-                    h = p(h, transformer_options)
-
-            hs.append(h)
-            if "input_block_patch_after_skip" in transformer_patches:
-                patch = transformer_patches["input_block_patch_after_skip"]
-                for p in patch:
-                    h = p(h, transformer_options)
-
-        transformer_options["block"] = ("middle", 0)
-        h = forward_timestep_embed(self.middle_block, h, emb, context, transformer_options, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator)
-        h = apply_control(h, control, 'middle')
-
-
-        for id, module in enumerate(self.output_blocks):
-            transformer_options["block"] = ("output", id)
-            hsp = hs.pop()
-            hsp = apply_control(hsp, control, 'output')
-
-            if "output_block_patch" in transformer_patches:
-                patch = transformer_patches["output_block_patch"]
-                for p in patch:
-                    h, hsp = p(h, hsp, transformer_options)
-
-            h = th.cat([h, hsp], dim=1)
-            del hsp
-            if len(hs) > 0:
-                output_shape = hs[-1].shape
-            else:
-                output_shape = None
-            h = forward_timestep_embed(module, h, emb, context, transformer_options, output_shape, time_context=time_context, num_video_frames=num_video_frames, image_only_indicator=image_only_indicator)
-        h = h.type(x.dtype)
-        if self.predict_codebook_ids:
-            return self.id_predictor(h)
-        else:
-            return self.out(h)