ldm_patched
This commit is contained in:
parent
7ad2a5c195
commit
7cb6178d47
1953
ldm_patched/contrib/external.py
Normal file
1953
ldm_patched/contrib/external.py
Normal file
File diff suppressed because it is too large
Load Diff
301
ldm_patched/contrib/external_canny.py
Normal file
301
ldm_patched/contrib/external_canny.py
Normal file
@ -0,0 +1,301 @@
|
|||||||
|
# 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-tutorials.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-tutorials.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-tutorials.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,
|
||||||
|
}
|
58
ldm_patched/contrib/external_clip_sdxl.py
Normal file
58
ldm_patched/contrib/external_clip_sdxl.py
Normal file
@ -0,0 +1,58 @@
|
|||||||
|
# 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,
|
||||||
|
}
|
204
ldm_patched/contrib/external_compositing.py
Normal file
204
ldm_patched/contrib/external_compositing.py
Normal file
@ -0,0 +1,204 @@
|
|||||||
|
# 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",
|
||||||
|
}
|
298
ldm_patched/contrib/external_custom_sampler.py
Normal file
298
ldm_patched/contrib/external_custom_sampler.py
Normal file
@ -0,0 +1,298 @@
|
|||||||
|
# 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]
|
||||||
|
inner_model = model.patch_model(patch_weights=False)
|
||||||
|
sigmas = inner_model.model_sampling.sigma(timesteps)
|
||||||
|
model.unpatch_model()
|
||||||
|
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,
|
||||||
|
}
|
115
ldm_patched/contrib/external_freelunch.py
Normal file
115
ldm_patched/contrib/external_freelunch.py
Normal file
@ -0,0 +1,115 @@
|
|||||||
|
# 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 = "_for_testing"
|
||||||
|
|
||||||
|
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 = "_for_testing"
|
||||||
|
|
||||||
|
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,
|
||||||
|
}
|
121
ldm_patched/contrib/external_hypernetwork.py
Normal file
121
ldm_patched/contrib/external_hypernetwork.py
Normal file
@ -0,0 +1,121 @@
|
|||||||
|
# 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
|
||||||
|
}
|
85
ldm_patched/contrib/external_hypertile.py
Normal file
85
ldm_patched/contrib/external_hypertile.py
Normal file
@ -0,0 +1,85 @@
|
|||||||
|
# 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 = "_for_testing"
|
||||||
|
|
||||||
|
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,
|
||||||
|
}
|
177
ldm_patched/contrib/external_images.py
Normal file
177
ldm_patched/contrib/external_images.py
Normal file
@ -0,0 +1,177 @@
|
|||||||
|
# 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,
|
||||||
|
}
|
133
ldm_patched/contrib/external_latent.py
Normal file
133
ldm_patched/contrib/external_latent.py
Normal file
@ -0,0 +1,133 @@
|
|||||||
|
# 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,)
|
||||||
|
|
||||||
|
NODE_CLASS_MAPPINGS = {
|
||||||
|
"LatentAdd": LatentAdd,
|
||||||
|
"LatentSubtract": LatentSubtract,
|
||||||
|
"LatentMultiply": LatentMultiply,
|
||||||
|
"LatentInterpolate": LatentInterpolate,
|
||||||
|
"LatentBatch": LatentBatch,
|
||||||
|
}
|
365
ldm_patched/contrib/external_mask.py
Normal file
365
ldm_patched/contrib/external_mask.py
Normal file
@ -0,0 +1,365 @@
|
|||||||
|
# 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",
|
||||||
|
}
|
177
ldm_patched/contrib/external_model_advanced.py
Normal file
177
ldm_patched/contrib/external_model_advanced.py
Normal file
@ -0,0 +1,177 @@
|
|||||||
|
# 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,
|
||||||
|
}
|
55
ldm_patched/contrib/external_model_downscale.py
Normal file
55
ldm_patched/contrib/external_model_downscale.py
Normal file
@ -0,0 +1,55 @@
|
|||||||
|
# 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)",
|
||||||
|
}
|
283
ldm_patched/contrib/external_model_merging.py
Normal file
283
ldm_patched/contrib/external_model_merging.py
Normal file
@ -0,0 +1,283 @@
|
|||||||
|
# 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, )
|
||||||
|
|
||||||
|
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):
|
||||||
|
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 = {}
|
||||||
|
|
||||||
|
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, metadata=metadata)
|
||||||
|
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,
|
||||||
|
}
|
57
ldm_patched/contrib/external_perpneg.py
Normal file
57
ldm_patched/contrib/external_perpneg.py
Normal file
@ -0,0 +1,57 @@
|
|||||||
|
# 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",
|
||||||
|
}
|
277
ldm_patched/contrib/external_post_processing.py
Normal file
277
ldm_patched/contrib/external_post_processing.py
Normal file
@ -0,0 +1,277 @@
|
|||||||
|
# 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):
|
||||||
|
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,
|
||||||
|
}
|
140
ldm_patched/contrib/external_rebatch.py
Normal file
140
ldm_patched/contrib/external_rebatch.py
Normal file
@ -0,0 +1,140 @@
|
|||||||
|
# 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",
|
||||||
|
}
|
172
ldm_patched/contrib/external_sag.py
Normal file
172
ldm_patched/contrib/external_sag.py
Normal file
@ -0,0 +1,172 @@
|
|||||||
|
# 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",
|
||||||
|
}
|
49
ldm_patched/contrib/external_sdupscale.py
Normal file
49
ldm_patched/contrib/external_sdupscale.py
Normal file
@ -0,0 +1,49 @@
|
|||||||
|
# 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,
|
||||||
|
}
|
104
ldm_patched/contrib/external_stable3d.py
Normal file
104
ldm_patched/contrib/external_stable3d.py
Normal file
@ -0,0 +1,104 @@
|
|||||||
|
# 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.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,
|
||||||
|
}
|
179
ldm_patched/contrib/external_tomesd.py
Normal file
179
ldm_patched/contrib/external_tomesd.py
Normal file
@ -0,0 +1,179 @@
|
|||||||
|
# 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,
|
||||||
|
}
|
68
ldm_patched/contrib/external_upscale_model.py
Normal file
68
ldm_patched/contrib/external_upscale_model.py
Normal file
@ -0,0 +1,68 @@
|
|||||||
|
# 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
|
||||||
|
}
|
91
ldm_patched/contrib/external_video_model.py
Normal file
91
ldm_patched/contrib/external_video_model.py
Normal file
@ -0,0 +1,91 @@
|
|||||||
|
# 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
|
||||||
|
|
||||||
|
|
||||||
|
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, )
|
||||||
|
|
||||||
|
NODE_CLASS_MAPPINGS = {
|
||||||
|
"ImageOnlyCheckpointLoader": ImageOnlyCheckpointLoader,
|
||||||
|
"SVD_img2vid_Conditioning": SVD_img2vid_Conditioning,
|
||||||
|
"VideoLinearCFGGuidance": VideoLinearCFGGuidance,
|
||||||
|
}
|
||||||
|
|
||||||
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||||
|
"ImageOnlyCheckpointLoader": "Image Only Checkpoint Loader (img2vid model)",
|
||||||
|
}
|
312
ldm_patched/controlnet/cldm.py
Normal file
312
ldm_patched/controlnet/cldm.py
Normal file
@ -0,0 +1,312 @@
|
|||||||
|
#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
|
||||||
|
|
810
ldm_patched/k_diffusion/sampling.py
Normal file
810
ldm_patched/k_diffusion/sampling.py
Normal file
@ -0,0 +1,810 @@
|
|||||||
|
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
|
313
ldm_patched/k_diffusion/utils.py
Normal file
313
ldm_patched/k_diffusion/utils.py
Normal file
@ -0,0 +1,313 @@
|
|||||||
|
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
|
228
ldm_patched/ldm/models/autoencoder.py
Normal file
228
ldm_patched/ldm/models/autoencoder.py
Normal file
@ -0,0 +1,228 @@
|
|||||||
|
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,
|
||||||
|
)
|
784
ldm_patched/ldm/modules/attention.py
Normal file
784
ldm_patched/ldm/modules/attention.py
Normal file
@ -0,0 +1,784 @@
|
|||||||
|
from inspect import isfunction
|
||||||
|
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 functools import partial
|
||||||
|
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
|
|
650
ldm_patched/ldm/modules/diffusionmodules/model.py
Normal file
650
ldm_patched/ldm/modules/diffusionmodules/model.py
Normal file
@ -0,0 +1,650 @@
|
|||||||
|
# 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
|
889
ldm_patched/ldm/modules/diffusionmodules/openaimodel.py
Normal file
889
ldm_patched/ldm/modules/diffusionmodules/openaimodel.py
Normal file
@ -0,0 +1,889 @@
|
|||||||
|
from abc import abstractmethod
|
||||||
|
import math
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import torch as th
|
||||||
|
import torch.nn as nn
|
||||||
|
import torch.nn.functional as F
|
||||||
|
from einops import rearrange
|
||||||
|
from functools import partial
|
||||||
|
|
||||||
|
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)
|
85
ldm_patched/ldm/modules/diffusionmodules/upscaling.py
Normal file
85
ldm_patched/ldm/modules/diffusionmodules/upscaling.py
Normal file
@ -0,0 +1,85 @@
|
|||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
import numpy as np
|
||||||
|
from functools import partial
|
||||||
|
|
||||||
|
from .util import extract_into_tensor, make_beta_schedule
|
||||||
|
from ldm_patched.ldm.util import default
|
||||||
|
|
||||||
|
|
||||||
|
class AbstractLowScaleModel(nn.Module):
|
||||||
|
# for concatenating a downsampled image to the latent representation
|
||||||
|
def __init__(self, noise_schedule_config=None):
|
||||||
|
super(AbstractLowScaleModel, self).__init__()
|
||||||
|
if noise_schedule_config is not None:
|
||||||
|
self.register_schedule(**noise_schedule_config)
|
||||||
|
|
||||||
|
def register_schedule(self, beta_schedule="linear", timesteps=1000,
|
||||||
|
linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
|
||||||
|
betas = make_beta_schedule(beta_schedule, timesteps, linear_start=linear_start, linear_end=linear_end,
|
||||||
|
cosine_s=cosine_s)
|
||||||
|
alphas = 1. - betas
|
||||||
|
alphas_cumprod = np.cumprod(alphas, axis=0)
|
||||||
|
alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
|
||||||
|
|
||||||
|
timesteps, = betas.shape
|
||||||
|
self.num_timesteps = int(timesteps)
|
||||||
|
self.linear_start = linear_start
|
||||||
|
self.linear_end = linear_end
|
||||||
|
assert alphas_cumprod.shape[0] == self.num_timesteps, 'alphas have to be defined for each timestep'
|
||||||
|
|
||||||
|
to_torch = partial(torch.tensor, dtype=torch.float32)
|
||||||
|
|
||||||
|
self.register_buffer('betas', to_torch(betas))
|
||||||
|
self.register_buffer('alphas_cumprod', to_torch(alphas_cumprod))
|
||||||
|
self.register_buffer('alphas_cumprod_prev', to_torch(alphas_cumprod_prev))
|
||||||
|
|
||||||
|
# calculations for diffusion q(x_t | x_{t-1}) and others
|
||||||
|
self.register_buffer('sqrt_alphas_cumprod', to_torch(np.sqrt(alphas_cumprod)))
|
||||||
|
self.register_buffer('sqrt_one_minus_alphas_cumprod', to_torch(np.sqrt(1. - alphas_cumprod)))
|
||||||
|
self.register_buffer('log_one_minus_alphas_cumprod', to_torch(np.log(1. - alphas_cumprod)))
|
||||||
|
self.register_buffer('sqrt_recip_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod)))
|
||||||
|
self.register_buffer('sqrt_recipm1_alphas_cumprod', to_torch(np.sqrt(1. / alphas_cumprod - 1)))
|
||||||
|
|
||||||
|
def q_sample(self, x_start, t, noise=None, seed=None):
|
||||||
|
if noise is None:
|
||||||
|
if seed is None:
|
||||||
|
noise = torch.randn_like(x_start)
|
||||||
|
else:
|
||||||
|
noise = torch.randn(x_start.size(), dtype=x_start.dtype, layout=x_start.layout, generator=torch.manual_seed(seed)).to(x_start.device)
|
||||||
|
return (extract_into_tensor(self.sqrt_alphas_cumprod.to(x_start.device), t, x_start.shape) * x_start +
|
||||||
|
extract_into_tensor(self.sqrt_one_minus_alphas_cumprod.to(x_start.device), t, x_start.shape) * noise)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
return x, None
|
||||||
|
|
||||||
|
def decode(self, x):
|
||||||
|
return x
|
||||||
|
|
||||||
|
|
||||||
|
class SimpleImageConcat(AbstractLowScaleModel):
|
||||||
|
# no noise level conditioning
|
||||||
|
def __init__(self):
|
||||||
|
super(SimpleImageConcat, self).__init__(noise_schedule_config=None)
|
||||||
|
self.max_noise_level = 0
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
# fix to constant noise level
|
||||||
|
return x, torch.zeros(x.shape[0], device=x.device).long()
|
||||||
|
|
||||||
|
|
||||||
|
class ImageConcatWithNoiseAugmentation(AbstractLowScaleModel):
|
||||||
|
def __init__(self, noise_schedule_config, max_noise_level=1000, to_cuda=False):
|
||||||
|
super().__init__(noise_schedule_config=noise_schedule_config)
|
||||||
|
self.max_noise_level = max_noise_level
|
||||||
|
|
||||||
|
def forward(self, x, noise_level=None, seed=None):
|
||||||
|
if noise_level is None:
|
||||||
|
noise_level = torch.randint(0, self.max_noise_level, (x.shape[0],), device=x.device).long()
|
||||||
|
else:
|
||||||
|
assert isinstance(noise_level, torch.Tensor)
|
||||||
|
z = self.q_sample(x, noise_level, seed=seed)
|
||||||
|
return z, noise_level
|
||||||
|
|
||||||
|
|
||||||
|
|
304
ldm_patched/ldm/modules/diffusionmodules/util.py
Normal file
304
ldm_patched/ldm/modules/diffusionmodules/util.py
Normal file
@ -0,0 +1,304 @@
|
|||||||
|
# adopted from
|
||||||
|
# https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py
|
||||||
|
# and
|
||||||
|
# https://github.com/lucidrains/denoising-diffusion-pytorch/blob/7706bdfc6f527f58d33f84b7b522e61e6e3164b3/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py
|
||||||
|
# and
|
||||||
|
# https://github.com/openai/guided-diffusion/blob/0ba878e517b276c45d1195eb29f6f5f72659a05b/guided_diffusion/nn.py
|
||||||
|
#
|
||||||
|
# thanks!
|
||||||
|
|
||||||
|
|
||||||
|
import os
|
||||||
|
import math
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
import numpy as np
|
||||||
|
from einops import repeat, rearrange
|
||||||
|
|
||||||
|
from ldm_patched.ldm.util import instantiate_from_config
|
||||||
|
|
||||||
|
class AlphaBlender(nn.Module):
|
||||||
|
strategies = ["learned", "fixed", "learned_with_images"]
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
alpha: float,
|
||||||
|
merge_strategy: str = "learned_with_images",
|
||||||
|
rearrange_pattern: str = "b t -> (b t) 1 1",
|
||||||
|
):
|
||||||
|
super().__init__()
|
||||||
|
self.merge_strategy = merge_strategy
|
||||||
|
self.rearrange_pattern = rearrange_pattern
|
||||||
|
|
||||||
|
assert (
|
||||||
|
merge_strategy in self.strategies
|
||||||
|
), f"merge_strategy needs to be in {self.strategies}"
|
||||||
|
|
||||||
|
if self.merge_strategy == "fixed":
|
||||||
|
self.register_buffer("mix_factor", torch.Tensor([alpha]))
|
||||||
|
elif (
|
||||||
|
self.merge_strategy == "learned"
|
||||||
|
or self.merge_strategy == "learned_with_images"
|
||||||
|
):
|
||||||
|
self.register_parameter(
|
||||||
|
"mix_factor", torch.nn.Parameter(torch.Tensor([alpha]))
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
raise ValueError(f"unknown merge strategy {self.merge_strategy}")
|
||||||
|
|
||||||
|
def get_alpha(self, image_only_indicator: torch.Tensor) -> torch.Tensor:
|
||||||
|
# skip_time_mix = rearrange(repeat(skip_time_mix, 'b -> (b t) () () ()', t=t), '(b t) 1 ... -> b 1 t ...', t=t)
|
||||||
|
if self.merge_strategy == "fixed":
|
||||||
|
# make shape compatible
|
||||||
|
# alpha = repeat(self.mix_factor, '1 -> b () t () ()', t=t, b=bs)
|
||||||
|
alpha = self.mix_factor.to(image_only_indicator.device)
|
||||||
|
elif self.merge_strategy == "learned":
|
||||||
|
alpha = torch.sigmoid(self.mix_factor.to(image_only_indicator.device))
|
||||||
|
# make shape compatible
|
||||||
|
# alpha = repeat(alpha, '1 -> s () ()', s = t * bs)
|
||||||
|
elif self.merge_strategy == "learned_with_images":
|
||||||
|
assert image_only_indicator is not None, "need image_only_indicator ..."
|
||||||
|
alpha = torch.where(
|
||||||
|
image_only_indicator.bool(),
|
||||||
|
torch.ones(1, 1, device=image_only_indicator.device),
|
||||||
|
rearrange(torch.sigmoid(self.mix_factor.to(image_only_indicator.device)), "... -> ... 1"),
|
||||||
|
)
|
||||||
|
alpha = rearrange(alpha, self.rearrange_pattern)
|
||||||
|
# make shape compatible
|
||||||
|
# alpha = repeat(alpha, '1 -> s () ()', s = t * bs)
|
||||||
|
else:
|
||||||
|
raise NotImplementedError()
|
||||||
|
return alpha
|
||||||
|
|
||||||
|
def forward(
|
||||||
|
self,
|
||||||
|
x_spatial,
|
||||||
|
x_temporal,
|
||||||
|
image_only_indicator=None,
|
||||||
|
) -> torch.Tensor:
|
||||||
|
alpha = self.get_alpha(image_only_indicator)
|
||||||
|
x = (
|
||||||
|
alpha.to(x_spatial.dtype) * x_spatial
|
||||||
|
+ (1.0 - alpha).to(x_spatial.dtype) * x_temporal
|
||||||
|
)
|
||||||
|
return x
|
||||||
|
|
||||||
|
|
||||||
|
def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
|
||||||
|
if schedule == "linear":
|
||||||
|
betas = (
|
||||||
|
torch.linspace(linear_start ** 0.5, linear_end ** 0.5, n_timestep, dtype=torch.float64) ** 2
|
||||||
|
)
|
||||||
|
|
||||||
|
elif schedule == "cosine":
|
||||||
|
timesteps = (
|
||||||
|
torch.arange(n_timestep + 1, dtype=torch.float64) / n_timestep + cosine_s
|
||||||
|
)
|
||||||
|
alphas = timesteps / (1 + cosine_s) * np.pi / 2
|
||||||
|
alphas = torch.cos(alphas).pow(2)
|
||||||
|
alphas = alphas / alphas[0]
|
||||||
|
betas = 1 - alphas[1:] / alphas[:-1]
|
||||||
|
betas = np.clip(betas, a_min=0, a_max=0.999)
|
||||||
|
|
||||||
|
elif schedule == "squaredcos_cap_v2": # used for karlo prior
|
||||||
|
# return early
|
||||||
|
return betas_for_alpha_bar(
|
||||||
|
n_timestep,
|
||||||
|
lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2,
|
||||||
|
)
|
||||||
|
|
||||||
|
elif schedule == "sqrt_linear":
|
||||||
|
betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64)
|
||||||
|
elif schedule == "sqrt":
|
||||||
|
betas = torch.linspace(linear_start, linear_end, n_timestep, dtype=torch.float64) ** 0.5
|
||||||
|
else:
|
||||||
|
raise ValueError(f"schedule '{schedule}' unknown.")
|
||||||
|
return betas.numpy()
|
||||||
|
|
||||||
|
|
||||||
|
def make_ddim_timesteps(ddim_discr_method, num_ddim_timesteps, num_ddpm_timesteps, verbose=True):
|
||||||
|
if ddim_discr_method == 'uniform':
|
||||||
|
c = num_ddpm_timesteps // num_ddim_timesteps
|
||||||
|
ddim_timesteps = np.asarray(list(range(0, num_ddpm_timesteps, c)))
|
||||||
|
elif ddim_discr_method == 'quad':
|
||||||
|
ddim_timesteps = ((np.linspace(0, np.sqrt(num_ddpm_timesteps * .8), num_ddim_timesteps)) ** 2).astype(int)
|
||||||
|
else:
|
||||||
|
raise NotImplementedError(f'There is no ddim discretization method called "{ddim_discr_method}"')
|
||||||
|
|
||||||
|
# assert ddim_timesteps.shape[0] == num_ddim_timesteps
|
||||||
|
# add one to get the final alpha values right (the ones from first scale to data during sampling)
|
||||||
|
steps_out = ddim_timesteps + 1
|
||||||
|
if verbose:
|
||||||
|
print(f'Selected timesteps for ddim sampler: {steps_out}')
|
||||||
|
return steps_out
|
||||||
|
|
||||||
|
|
||||||
|
def make_ddim_sampling_parameters(alphacums, ddim_timesteps, eta, verbose=True):
|
||||||
|
# select alphas for computing the variance schedule
|
||||||
|
alphas = alphacums[ddim_timesteps]
|
||||||
|
alphas_prev = np.asarray([alphacums[0]] + alphacums[ddim_timesteps[:-1]].tolist())
|
||||||
|
|
||||||
|
# according the the formula provided in https://arxiv.org/abs/2010.02502
|
||||||
|
sigmas = eta * np.sqrt((1 - alphas_prev) / (1 - alphas) * (1 - alphas / alphas_prev))
|
||||||
|
if verbose:
|
||||||
|
print(f'Selected alphas for ddim sampler: a_t: {alphas}; a_(t-1): {alphas_prev}')
|
||||||
|
print(f'For the chosen value of eta, which is {eta}, '
|
||||||
|
f'this results in the following sigma_t schedule for ddim sampler {sigmas}')
|
||||||
|
return sigmas, alphas, alphas_prev
|
||||||
|
|
||||||
|
|
||||||
|
def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
|
||||||
|
"""
|
||||||
|
Create a beta schedule that discretizes the given alpha_t_bar function,
|
||||||
|
which defines the cumulative product of (1-beta) over time from t = [0,1].
|
||||||
|
:param num_diffusion_timesteps: the number of betas to produce.
|
||||||
|
:param alpha_bar: a lambda that takes an argument t from 0 to 1 and
|
||||||
|
produces the cumulative product of (1-beta) up to that
|
||||||
|
part of the diffusion process.
|
||||||
|
:param max_beta: the maximum beta to use; use values lower than 1 to
|
||||||
|
prevent singularities.
|
||||||
|
"""
|
||||||
|
betas = []
|
||||||
|
for i in range(num_diffusion_timesteps):
|
||||||
|
t1 = i / num_diffusion_timesteps
|
||||||
|
t2 = (i + 1) / num_diffusion_timesteps
|
||||||
|
betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
|
||||||
|
return np.array(betas)
|
||||||
|
|
||||||
|
|
||||||
|
def extract_into_tensor(a, t, x_shape):
|
||||||
|
b, *_ = t.shape
|
||||||
|
out = a.gather(-1, t)
|
||||||
|
return out.reshape(b, *((1,) * (len(x_shape) - 1)))
|
||||||
|
|
||||||
|
|
||||||
|
def checkpoint(func, inputs, params, flag):
|
||||||
|
"""
|
||||||
|
Evaluate a function without caching intermediate activations, allowing for
|
||||||
|
reduced memory at the expense of extra compute in the backward pass.
|
||||||
|
:param func: the function to evaluate.
|
||||||
|
:param inputs: the argument sequence to pass to `func`.
|
||||||
|
:param params: a sequence of parameters `func` depends on but does not
|
||||||
|
explicitly take as arguments.
|
||||||
|
:param flag: if False, disable gradient checkpointing.
|
||||||
|
"""
|
||||||
|
if flag:
|
||||||
|
args = tuple(inputs) + tuple(params)
|
||||||
|
return CheckpointFunction.apply(func, len(inputs), *args)
|
||||||
|
else:
|
||||||
|
return func(*inputs)
|
||||||
|
|
||||||
|
|
||||||
|
class CheckpointFunction(torch.autograd.Function):
|
||||||
|
@staticmethod
|
||||||
|
def forward(ctx, run_function, length, *args):
|
||||||
|
ctx.run_function = run_function
|
||||||
|
ctx.input_tensors = list(args[:length])
|
||||||
|
ctx.input_params = list(args[length:])
|
||||||
|
ctx.gpu_autocast_kwargs = {"enabled": torch.is_autocast_enabled(),
|
||||||
|
"dtype": torch.get_autocast_gpu_dtype(),
|
||||||
|
"cache_enabled": torch.is_autocast_cache_enabled()}
|
||||||
|
with torch.no_grad():
|
||||||
|
output_tensors = ctx.run_function(*ctx.input_tensors)
|
||||||
|
return output_tensors
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def backward(ctx, *output_grads):
|
||||||
|
ctx.input_tensors = [x.detach().requires_grad_(True) for x in ctx.input_tensors]
|
||||||
|
with torch.enable_grad(), \
|
||||||
|
torch.cuda.amp.autocast(**ctx.gpu_autocast_kwargs):
|
||||||
|
# Fixes a bug where the first op in run_function modifies the
|
||||||
|
# Tensor storage in place, which is not allowed for detach()'d
|
||||||
|
# Tensors.
|
||||||
|
shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
|
||||||
|
output_tensors = ctx.run_function(*shallow_copies)
|
||||||
|
input_grads = torch.autograd.grad(
|
||||||
|
output_tensors,
|
||||||
|
ctx.input_tensors + ctx.input_params,
|
||||||
|
output_grads,
|
||||||
|
allow_unused=True,
|
||||||
|
)
|
||||||
|
del ctx.input_tensors
|
||||||
|
del ctx.input_params
|
||||||
|
del output_tensors
|
||||||
|
return (None, None) + input_grads
|
||||||
|
|
||||||
|
|
||||||
|
def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):
|
||||||
|
"""
|
||||||
|
Create sinusoidal timestep embeddings.
|
||||||
|
:param timesteps: a 1-D Tensor of N indices, one per batch element.
|
||||||
|
These may be fractional.
|
||||||
|
:param dim: the dimension of the output.
|
||||||
|
:param max_period: controls the minimum frequency of the embeddings.
|
||||||
|
:return: an [N x dim] Tensor of positional embeddings.
|
||||||
|
"""
|
||||||
|
if not repeat_only:
|
||||||
|
half = dim // 2
|
||||||
|
freqs = torch.exp(
|
||||||
|
-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32, device=timesteps.device) / half
|
||||||
|
)
|
||||||
|
args = timesteps[:, None].float() * freqs[None]
|
||||||
|
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
|
||||||
|
if dim % 2:
|
||||||
|
embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
|
||||||
|
else:
|
||||||
|
embedding = repeat(timesteps, 'b -> b d', d=dim)
|
||||||
|
return embedding
|
||||||
|
|
||||||
|
|
||||||
|
def zero_module(module):
|
||||||
|
"""
|
||||||
|
Zero out the parameters of a module and return it.
|
||||||
|
"""
|
||||||
|
for p in module.parameters():
|
||||||
|
p.detach().zero_()
|
||||||
|
return module
|
||||||
|
|
||||||
|
|
||||||
|
def scale_module(module, scale):
|
||||||
|
"""
|
||||||
|
Scale the parameters of a module and return it.
|
||||||
|
"""
|
||||||
|
for p in module.parameters():
|
||||||
|
p.detach().mul_(scale)
|
||||||
|
return module
|
||||||
|
|
||||||
|
|
||||||
|
def mean_flat(tensor):
|
||||||
|
"""
|
||||||
|
Take the mean over all non-batch dimensions.
|
||||||
|
"""
|
||||||
|
return tensor.mean(dim=list(range(1, len(tensor.shape))))
|
||||||
|
|
||||||
|
|
||||||
|
def avg_pool_nd(dims, *args, **kwargs):
|
||||||
|
"""
|
||||||
|
Create a 1D, 2D, or 3D average pooling module.
|
||||||
|
"""
|
||||||
|
if dims == 1:
|
||||||
|
return nn.AvgPool1d(*args, **kwargs)
|
||||||
|
elif dims == 2:
|
||||||
|
return nn.AvgPool2d(*args, **kwargs)
|
||||||
|
elif dims == 3:
|
||||||
|
return nn.AvgPool3d(*args, **kwargs)
|
||||||
|
raise ValueError(f"unsupported dimensions: {dims}")
|
||||||
|
|
||||||
|
|
||||||
|
class HybridConditioner(nn.Module):
|
||||||
|
|
||||||
|
def __init__(self, c_concat_config, c_crossattn_config):
|
||||||
|
super().__init__()
|
||||||
|
self.concat_conditioner = instantiate_from_config(c_concat_config)
|
||||||
|
self.crossattn_conditioner = instantiate_from_config(c_crossattn_config)
|
||||||
|
|
||||||
|
def forward(self, c_concat, c_crossattn):
|
||||||
|
c_concat = self.concat_conditioner(c_concat)
|
||||||
|
c_crossattn = self.crossattn_conditioner(c_crossattn)
|
||||||
|
return {'c_concat': [c_concat], 'c_crossattn': [c_crossattn]}
|
||||||
|
|
||||||
|
|
||||||
|
def noise_like(shape, device, repeat=False):
|
||||||
|
repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
|
||||||
|
noise = lambda: torch.randn(shape, device=device)
|
||||||
|
return repeat_noise() if repeat else noise()
|
0
ldm_patched/ldm/modules/distributions/__init__.py
Normal file
0
ldm_patched/ldm/modules/distributions/__init__.py
Normal file
92
ldm_patched/ldm/modules/distributions/distributions.py
Normal file
92
ldm_patched/ldm/modules/distributions/distributions.py
Normal file
@ -0,0 +1,92 @@
|
|||||||
|
import torch
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
|
||||||
|
class AbstractDistribution:
|
||||||
|
def sample(self):
|
||||||
|
raise NotImplementedError()
|
||||||
|
|
||||||
|
def mode(self):
|
||||||
|
raise NotImplementedError()
|
||||||
|
|
||||||
|
|
||||||
|
class DiracDistribution(AbstractDistribution):
|
||||||
|
def __init__(self, value):
|
||||||
|
self.value = value
|
||||||
|
|
||||||
|
def sample(self):
|
||||||
|
return self.value
|
||||||
|
|
||||||
|
def mode(self):
|
||||||
|
return self.value
|
||||||
|
|
||||||
|
|
||||||
|
class DiagonalGaussianDistribution(object):
|
||||||
|
def __init__(self, parameters, deterministic=False):
|
||||||
|
self.parameters = parameters
|
||||||
|
self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
|
||||||
|
self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
|
||||||
|
self.deterministic = deterministic
|
||||||
|
self.std = torch.exp(0.5 * self.logvar)
|
||||||
|
self.var = torch.exp(self.logvar)
|
||||||
|
if self.deterministic:
|
||||||
|
self.var = self.std = torch.zeros_like(self.mean).to(device=self.parameters.device)
|
||||||
|
|
||||||
|
def sample(self):
|
||||||
|
x = self.mean + self.std * torch.randn(self.mean.shape).to(device=self.parameters.device)
|
||||||
|
return x
|
||||||
|
|
||||||
|
def kl(self, other=None):
|
||||||
|
if self.deterministic:
|
||||||
|
return torch.Tensor([0.])
|
||||||
|
else:
|
||||||
|
if other is None:
|
||||||
|
return 0.5 * torch.sum(torch.pow(self.mean, 2)
|
||||||
|
+ self.var - 1.0 - self.logvar,
|
||||||
|
dim=[1, 2, 3])
|
||||||
|
else:
|
||||||
|
return 0.5 * torch.sum(
|
||||||
|
torch.pow(self.mean - other.mean, 2) / other.var
|
||||||
|
+ self.var / other.var - 1.0 - self.logvar + other.logvar,
|
||||||
|
dim=[1, 2, 3])
|
||||||
|
|
||||||
|
def nll(self, sample, dims=[1,2,3]):
|
||||||
|
if self.deterministic:
|
||||||
|
return torch.Tensor([0.])
|
||||||
|
logtwopi = np.log(2.0 * np.pi)
|
||||||
|
return 0.5 * torch.sum(
|
||||||
|
logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var,
|
||||||
|
dim=dims)
|
||||||
|
|
||||||
|
def mode(self):
|
||||||
|
return self.mean
|
||||||
|
|
||||||
|
|
||||||
|
def normal_kl(mean1, logvar1, mean2, logvar2):
|
||||||
|
"""
|
||||||
|
source: https://github.com/openai/guided-diffusion/blob/27c20a8fab9cb472df5d6bdd6c8d11c8f430b924/guided_diffusion/losses.py#L12
|
||||||
|
Compute the KL divergence between two gaussians.
|
||||||
|
Shapes are automatically broadcasted, so batches can be compared to
|
||||||
|
scalars, among other use cases.
|
||||||
|
"""
|
||||||
|
tensor = None
|
||||||
|
for obj in (mean1, logvar1, mean2, logvar2):
|
||||||
|
if isinstance(obj, torch.Tensor):
|
||||||
|
tensor = obj
|
||||||
|
break
|
||||||
|
assert tensor is not None, "at least one argument must be a Tensor"
|
||||||
|
|
||||||
|
# Force variances to be Tensors. Broadcasting helps convert scalars to
|
||||||
|
# Tensors, but it does not work for torch.exp().
|
||||||
|
logvar1, logvar2 = [
|
||||||
|
x if isinstance(x, torch.Tensor) else torch.tensor(x).to(tensor)
|
||||||
|
for x in (logvar1, logvar2)
|
||||||
|
]
|
||||||
|
|
||||||
|
return 0.5 * (
|
||||||
|
-1.0
|
||||||
|
+ logvar2
|
||||||
|
- logvar1
|
||||||
|
+ torch.exp(logvar1 - logvar2)
|
||||||
|
+ ((mean1 - mean2) ** 2) * torch.exp(-logvar2)
|
||||||
|
)
|
80
ldm_patched/ldm/modules/ema.py
Normal file
80
ldm_patched/ldm/modules/ema.py
Normal file
@ -0,0 +1,80 @@
|
|||||||
|
import torch
|
||||||
|
from torch import nn
|
||||||
|
|
||||||
|
|
||||||
|
class LitEma(nn.Module):
|
||||||
|
def __init__(self, model, decay=0.9999, use_num_upates=True):
|
||||||
|
super().__init__()
|
||||||
|
if decay < 0.0 or decay > 1.0:
|
||||||
|
raise ValueError('Decay must be between 0 and 1')
|
||||||
|
|
||||||
|
self.m_name2s_name = {}
|
||||||
|
self.register_buffer('decay', torch.tensor(decay, dtype=torch.float32))
|
||||||
|
self.register_buffer('num_updates', torch.tensor(0, dtype=torch.int) if use_num_upates
|
||||||
|
else torch.tensor(-1, dtype=torch.int))
|
||||||
|
|
||||||
|
for name, p in model.named_parameters():
|
||||||
|
if p.requires_grad:
|
||||||
|
# remove as '.'-character is not allowed in buffers
|
||||||
|
s_name = name.replace('.', '')
|
||||||
|
self.m_name2s_name.update({name: s_name})
|
||||||
|
self.register_buffer(s_name, p.clone().detach().data)
|
||||||
|
|
||||||
|
self.collected_params = []
|
||||||
|
|
||||||
|
def reset_num_updates(self):
|
||||||
|
del self.num_updates
|
||||||
|
self.register_buffer('num_updates', torch.tensor(0, dtype=torch.int))
|
||||||
|
|
||||||
|
def forward(self, model):
|
||||||
|
decay = self.decay
|
||||||
|
|
||||||
|
if self.num_updates >= 0:
|
||||||
|
self.num_updates += 1
|
||||||
|
decay = min(self.decay, (1 + self.num_updates) / (10 + self.num_updates))
|
||||||
|
|
||||||
|
one_minus_decay = 1.0 - decay
|
||||||
|
|
||||||
|
with torch.no_grad():
|
||||||
|
m_param = dict(model.named_parameters())
|
||||||
|
shadow_params = dict(self.named_buffers())
|
||||||
|
|
||||||
|
for key in m_param:
|
||||||
|
if m_param[key].requires_grad:
|
||||||
|
sname = self.m_name2s_name[key]
|
||||||
|
shadow_params[sname] = shadow_params[sname].type_as(m_param[key])
|
||||||
|
shadow_params[sname].sub_(one_minus_decay * (shadow_params[sname] - m_param[key]))
|
||||||
|
else:
|
||||||
|
assert not key in self.m_name2s_name
|
||||||
|
|
||||||
|
def copy_to(self, model):
|
||||||
|
m_param = dict(model.named_parameters())
|
||||||
|
shadow_params = dict(self.named_buffers())
|
||||||
|
for key in m_param:
|
||||||
|
if m_param[key].requires_grad:
|
||||||
|
m_param[key].data.copy_(shadow_params[self.m_name2s_name[key]].data)
|
||||||
|
else:
|
||||||
|
assert not key in self.m_name2s_name
|
||||||
|
|
||||||
|
def store(self, parameters):
|
||||||
|
"""
|
||||||
|
Save the current parameters for restoring later.
|
||||||
|
Args:
|
||||||
|
parameters: Iterable of `torch.nn.Parameter`; the parameters to be
|
||||||
|
temporarily stored.
|
||||||
|
"""
|
||||||
|
self.collected_params = [param.clone() for param in parameters]
|
||||||
|
|
||||||
|
def restore(self, parameters):
|
||||||
|
"""
|
||||||
|
Restore the parameters stored with the `store` method.
|
||||||
|
Useful to validate the model with EMA parameters without affecting the
|
||||||
|
original optimization process. Store the parameters before the
|
||||||
|
`copy_to` method. After validation (or model saving), use this to
|
||||||
|
restore the former parameters.
|
||||||
|
Args:
|
||||||
|
parameters: Iterable of `torch.nn.Parameter`; the parameters to be
|
||||||
|
updated with the stored parameters.
|
||||||
|
"""
|
||||||
|
for c_param, param in zip(self.collected_params, parameters):
|
||||||
|
param.data.copy_(c_param.data)
|
0
ldm_patched/ldm/modules/encoders/__init__.py
Normal file
0
ldm_patched/ldm/modules/encoders/__init__.py
Normal file
35
ldm_patched/ldm/modules/encoders/noise_aug_modules.py
Normal file
35
ldm_patched/ldm/modules/encoders/noise_aug_modules.py
Normal file
@ -0,0 +1,35 @@
|
|||||||
|
from ..diffusionmodules.upscaling import ImageConcatWithNoiseAugmentation
|
||||||
|
from ..diffusionmodules.openaimodel import Timestep
|
||||||
|
import torch
|
||||||
|
|
||||||
|
class CLIPEmbeddingNoiseAugmentation(ImageConcatWithNoiseAugmentation):
|
||||||
|
def __init__(self, *args, clip_stats_path=None, timestep_dim=256, **kwargs):
|
||||||
|
super().__init__(*args, **kwargs)
|
||||||
|
if clip_stats_path is None:
|
||||||
|
clip_mean, clip_std = torch.zeros(timestep_dim), torch.ones(timestep_dim)
|
||||||
|
else:
|
||||||
|
clip_mean, clip_std = torch.load(clip_stats_path, map_location="cpu")
|
||||||
|
self.register_buffer("data_mean", clip_mean[None, :], persistent=False)
|
||||||
|
self.register_buffer("data_std", clip_std[None, :], persistent=False)
|
||||||
|
self.time_embed = Timestep(timestep_dim)
|
||||||
|
|
||||||
|
def scale(self, x):
|
||||||
|
# re-normalize to centered mean and unit variance
|
||||||
|
x = (x - self.data_mean.to(x.device)) * 1. / self.data_std.to(x.device)
|
||||||
|
return x
|
||||||
|
|
||||||
|
def unscale(self, x):
|
||||||
|
# back to original data stats
|
||||||
|
x = (x * self.data_std.to(x.device)) + self.data_mean.to(x.device)
|
||||||
|
return x
|
||||||
|
|
||||||
|
def forward(self, x, noise_level=None):
|
||||||
|
if noise_level is None:
|
||||||
|
noise_level = torch.randint(0, self.max_noise_level, (x.shape[0],), device=x.device).long()
|
||||||
|
else:
|
||||||
|
assert isinstance(noise_level, torch.Tensor)
|
||||||
|
x = self.scale(x)
|
||||||
|
z = self.q_sample(x, noise_level)
|
||||||
|
z = self.unscale(z)
|
||||||
|
noise_level = self.time_embed(noise_level)
|
||||||
|
return z, noise_level
|
273
ldm_patched/ldm/modules/sub_quadratic_attention.py
Normal file
273
ldm_patched/ldm/modules/sub_quadratic_attention.py
Normal file
@ -0,0 +1,273 @@
|
|||||||
|
# original source:
|
||||||
|
# https://github.com/AminRezaei0x443/memory-efficient-attention/blob/1bc0d9e6ac5f82ea43a375135c4e1d3896ee1694/memory_efficient_attention/attention_torch.py
|
||||||
|
# license:
|
||||||
|
# MIT
|
||||||
|
# credit:
|
||||||
|
# Amin Rezaei (original author)
|
||||||
|
# Alex Birch (optimized algorithm for 3D tensors, at the expense of removing bias, masking and callbacks)
|
||||||
|
# implementation of:
|
||||||
|
# Self-attention Does Not Need O(n2) Memory":
|
||||||
|
# https://arxiv.org/abs/2112.05682v2
|
||||||
|
|
||||||
|
from functools import partial
|
||||||
|
import torch
|
||||||
|
from torch import Tensor
|
||||||
|
from torch.utils.checkpoint import checkpoint
|
||||||
|
import math
|
||||||
|
|
||||||
|
try:
|
||||||
|
from typing import Optional, NamedTuple, List, Protocol
|
||||||
|
except ImportError:
|
||||||
|
from typing import Optional, NamedTuple, List
|
||||||
|
from typing_extensions import Protocol
|
||||||
|
|
||||||
|
from torch import Tensor
|
||||||
|
from typing import List
|
||||||
|
|
||||||
|
from ldm_patched.modules import model_management
|
||||||
|
|
||||||
|
def dynamic_slice(
|
||||||
|
x: Tensor,
|
||||||
|
starts: List[int],
|
||||||
|
sizes: List[int],
|
||||||
|
) -> Tensor:
|
||||||
|
slicing = [slice(start, start + size) for start, size in zip(starts, sizes)]
|
||||||
|
return x[slicing]
|
||||||
|
|
||||||
|
class AttnChunk(NamedTuple):
|
||||||
|
exp_values: Tensor
|
||||||
|
exp_weights_sum: Tensor
|
||||||
|
max_score: Tensor
|
||||||
|
|
||||||
|
class SummarizeChunk(Protocol):
|
||||||
|
@staticmethod
|
||||||
|
def __call__(
|
||||||
|
query: Tensor,
|
||||||
|
key_t: Tensor,
|
||||||
|
value: Tensor,
|
||||||
|
) -> AttnChunk: ...
|
||||||
|
|
||||||
|
class ComputeQueryChunkAttn(Protocol):
|
||||||
|
@staticmethod
|
||||||
|
def __call__(
|
||||||
|
query: Tensor,
|
||||||
|
key_t: Tensor,
|
||||||
|
value: Tensor,
|
||||||
|
) -> Tensor: ...
|
||||||
|
|
||||||
|
def _summarize_chunk(
|
||||||
|
query: Tensor,
|
||||||
|
key_t: Tensor,
|
||||||
|
value: Tensor,
|
||||||
|
scale: float,
|
||||||
|
upcast_attention: bool,
|
||||||
|
mask,
|
||||||
|
) -> AttnChunk:
|
||||||
|
if upcast_attention:
|
||||||
|
with torch.autocast(enabled=False, device_type = 'cuda'):
|
||||||
|
query = query.float()
|
||||||
|
key_t = key_t.float()
|
||||||
|
attn_weights = torch.baddbmm(
|
||||||
|
torch.empty(1, 1, 1, device=query.device, dtype=query.dtype),
|
||||||
|
query,
|
||||||
|
key_t,
|
||||||
|
alpha=scale,
|
||||||
|
beta=0,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
attn_weights = torch.baddbmm(
|
||||||
|
torch.empty(1, 1, 1, device=query.device, dtype=query.dtype),
|
||||||
|
query,
|
||||||
|
key_t,
|
||||||
|
alpha=scale,
|
||||||
|
beta=0,
|
||||||
|
)
|
||||||
|
max_score, _ = torch.max(attn_weights, -1, keepdim=True)
|
||||||
|
max_score = max_score.detach()
|
||||||
|
attn_weights -= max_score
|
||||||
|
if mask is not None:
|
||||||
|
attn_weights += mask
|
||||||
|
torch.exp(attn_weights, out=attn_weights)
|
||||||
|
exp_weights = attn_weights.to(value.dtype)
|
||||||
|
exp_values = torch.bmm(exp_weights, value)
|
||||||
|
max_score = max_score.squeeze(-1)
|
||||||
|
return AttnChunk(exp_values, exp_weights.sum(dim=-1), max_score)
|
||||||
|
|
||||||
|
def _query_chunk_attention(
|
||||||
|
query: Tensor,
|
||||||
|
key_t: Tensor,
|
||||||
|
value: Tensor,
|
||||||
|
summarize_chunk: SummarizeChunk,
|
||||||
|
kv_chunk_size: int,
|
||||||
|
mask,
|
||||||
|
) -> Tensor:
|
||||||
|
batch_x_heads, k_channels_per_head, k_tokens = key_t.shape
|
||||||
|
_, _, v_channels_per_head = value.shape
|
||||||
|
|
||||||
|
def chunk_scanner(chunk_idx: int, mask) -> AttnChunk:
|
||||||
|
key_chunk = dynamic_slice(
|
||||||
|
key_t,
|
||||||
|
(0, 0, chunk_idx),
|
||||||
|
(batch_x_heads, k_channels_per_head, kv_chunk_size)
|
||||||
|
)
|
||||||
|
value_chunk = dynamic_slice(
|
||||||
|
value,
|
||||||
|
(0, chunk_idx, 0),
|
||||||
|
(batch_x_heads, kv_chunk_size, v_channels_per_head)
|
||||||
|
)
|
||||||
|
if mask is not None:
|
||||||
|
mask = mask[:,:,chunk_idx:chunk_idx + kv_chunk_size]
|
||||||
|
|
||||||
|
return summarize_chunk(query, key_chunk, value_chunk, mask=mask)
|
||||||
|
|
||||||
|
chunks: List[AttnChunk] = [
|
||||||
|
chunk_scanner(chunk, mask) for chunk in torch.arange(0, k_tokens, kv_chunk_size)
|
||||||
|
]
|
||||||
|
acc_chunk = AttnChunk(*map(torch.stack, zip(*chunks)))
|
||||||
|
chunk_values, chunk_weights, chunk_max = acc_chunk
|
||||||
|
|
||||||
|
global_max, _ = torch.max(chunk_max, 0, keepdim=True)
|
||||||
|
max_diffs = torch.exp(chunk_max - global_max)
|
||||||
|
chunk_values *= torch.unsqueeze(max_diffs, -1)
|
||||||
|
chunk_weights *= max_diffs
|
||||||
|
|
||||||
|
all_values = chunk_values.sum(dim=0)
|
||||||
|
all_weights = torch.unsqueeze(chunk_weights, -1).sum(dim=0)
|
||||||
|
return all_values / all_weights
|
||||||
|
|
||||||
|
# TODO: refactor CrossAttention#get_attention_scores to share code with this
|
||||||
|
def _get_attention_scores_no_kv_chunking(
|
||||||
|
query: Tensor,
|
||||||
|
key_t: Tensor,
|
||||||
|
value: Tensor,
|
||||||
|
scale: float,
|
||||||
|
upcast_attention: bool,
|
||||||
|
mask,
|
||||||
|
) -> Tensor:
|
||||||
|
if upcast_attention:
|
||||||
|
with torch.autocast(enabled=False, device_type = 'cuda'):
|
||||||
|
query = query.float()
|
||||||
|
key_t = key_t.float()
|
||||||
|
attn_scores = torch.baddbmm(
|
||||||
|
torch.empty(1, 1, 1, device=query.device, dtype=query.dtype),
|
||||||
|
query,
|
||||||
|
key_t,
|
||||||
|
alpha=scale,
|
||||||
|
beta=0,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
attn_scores = torch.baddbmm(
|
||||||
|
torch.empty(1, 1, 1, device=query.device, dtype=query.dtype),
|
||||||
|
query,
|
||||||
|
key_t,
|
||||||
|
alpha=scale,
|
||||||
|
beta=0,
|
||||||
|
)
|
||||||
|
|
||||||
|
if mask is not None:
|
||||||
|
attn_scores += mask
|
||||||
|
try:
|
||||||
|
attn_probs = attn_scores.softmax(dim=-1)
|
||||||
|
del attn_scores
|
||||||
|
except model_management.OOM_EXCEPTION:
|
||||||
|
print("ran out of memory while running softmax in _get_attention_scores_no_kv_chunking, trying slower in place softmax instead")
|
||||||
|
attn_scores -= attn_scores.max(dim=-1, keepdim=True).values
|
||||||
|
torch.exp(attn_scores, out=attn_scores)
|
||||||
|
summed = torch.sum(attn_scores, dim=-1, keepdim=True)
|
||||||
|
attn_scores /= summed
|
||||||
|
attn_probs = attn_scores
|
||||||
|
|
||||||
|
hidden_states_slice = torch.bmm(attn_probs.to(value.dtype), value)
|
||||||
|
return hidden_states_slice
|
||||||
|
|
||||||
|
class ScannedChunk(NamedTuple):
|
||||||
|
chunk_idx: int
|
||||||
|
attn_chunk: AttnChunk
|
||||||
|
|
||||||
|
def efficient_dot_product_attention(
|
||||||
|
query: Tensor,
|
||||||
|
key_t: Tensor,
|
||||||
|
value: Tensor,
|
||||||
|
query_chunk_size=1024,
|
||||||
|
kv_chunk_size: Optional[int] = None,
|
||||||
|
kv_chunk_size_min: Optional[int] = None,
|
||||||
|
use_checkpoint=True,
|
||||||
|
upcast_attention=False,
|
||||||
|
mask = None,
|
||||||
|
):
|
||||||
|
"""Computes efficient dot-product attention given query, transposed key, and value.
|
||||||
|
This is efficient version of attention presented in
|
||||||
|
https://arxiv.org/abs/2112.05682v2 which comes with O(sqrt(n)) memory requirements.
|
||||||
|
Args:
|
||||||
|
query: queries for calculating attention with shape of
|
||||||
|
`[batch * num_heads, tokens, channels_per_head]`.
|
||||||
|
key_t: keys for calculating attention with shape of
|
||||||
|
`[batch * num_heads, channels_per_head, tokens]`.
|
||||||
|
value: values to be used in attention with shape of
|
||||||
|
`[batch * num_heads, tokens, channels_per_head]`.
|
||||||
|
query_chunk_size: int: query chunks size
|
||||||
|
kv_chunk_size: Optional[int]: key/value chunks size. if None: defaults to sqrt(key_tokens)
|
||||||
|
kv_chunk_size_min: Optional[int]: key/value minimum chunk size. only considered when kv_chunk_size is None. changes `sqrt(key_tokens)` into `max(sqrt(key_tokens), kv_chunk_size_min)`, to ensure our chunk sizes don't get too small (smaller chunks = more chunks = less concurrent work done).
|
||||||
|
use_checkpoint: bool: whether to use checkpointing (recommended True for training, False for inference)
|
||||||
|
Returns:
|
||||||
|
Output of shape `[batch * num_heads, query_tokens, channels_per_head]`.
|
||||||
|
"""
|
||||||
|
batch_x_heads, q_tokens, q_channels_per_head = query.shape
|
||||||
|
_, _, k_tokens = key_t.shape
|
||||||
|
scale = q_channels_per_head ** -0.5
|
||||||
|
|
||||||
|
kv_chunk_size = min(kv_chunk_size or int(math.sqrt(k_tokens)), k_tokens)
|
||||||
|
if kv_chunk_size_min is not None:
|
||||||
|
kv_chunk_size = max(kv_chunk_size, kv_chunk_size_min)
|
||||||
|
|
||||||
|
if mask is not None and len(mask.shape) == 2:
|
||||||
|
mask = mask.unsqueeze(0)
|
||||||
|
|
||||||
|
def get_query_chunk(chunk_idx: int) -> Tensor:
|
||||||
|
return dynamic_slice(
|
||||||
|
query,
|
||||||
|
(0, chunk_idx, 0),
|
||||||
|
(batch_x_heads, min(query_chunk_size, q_tokens), q_channels_per_head)
|
||||||
|
)
|
||||||
|
|
||||||
|
def get_mask_chunk(chunk_idx: int) -> Tensor:
|
||||||
|
if mask is None:
|
||||||
|
return None
|
||||||
|
chunk = min(query_chunk_size, q_tokens)
|
||||||
|
return mask[:,chunk_idx:chunk_idx + chunk]
|
||||||
|
|
||||||
|
summarize_chunk: SummarizeChunk = partial(_summarize_chunk, scale=scale, upcast_attention=upcast_attention)
|
||||||
|
summarize_chunk: SummarizeChunk = partial(checkpoint, summarize_chunk) if use_checkpoint else summarize_chunk
|
||||||
|
compute_query_chunk_attn: ComputeQueryChunkAttn = partial(
|
||||||
|
_get_attention_scores_no_kv_chunking,
|
||||||
|
scale=scale,
|
||||||
|
upcast_attention=upcast_attention
|
||||||
|
) if k_tokens <= kv_chunk_size else (
|
||||||
|
# fast-path for when there's just 1 key-value chunk per query chunk (this is just sliced attention btw)
|
||||||
|
partial(
|
||||||
|
_query_chunk_attention,
|
||||||
|
kv_chunk_size=kv_chunk_size,
|
||||||
|
summarize_chunk=summarize_chunk,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
if q_tokens <= query_chunk_size:
|
||||||
|
# fast-path for when there's just 1 query chunk
|
||||||
|
return compute_query_chunk_attn(
|
||||||
|
query=query,
|
||||||
|
key_t=key_t,
|
||||||
|
value=value,
|
||||||
|
mask=mask,
|
||||||
|
)
|
||||||
|
|
||||||
|
# TODO: maybe we should use torch.empty_like(query) to allocate storage in-advance,
|
||||||
|
# and pass slices to be mutated, instead of torch.cat()ing the returned slices
|
||||||
|
res = torch.cat([
|
||||||
|
compute_query_chunk_attn(
|
||||||
|
query=get_query_chunk(i * query_chunk_size),
|
||||||
|
key_t=key_t,
|
||||||
|
value=value,
|
||||||
|
mask=get_mask_chunk(i * query_chunk_size)
|
||||||
|
) for i in range(math.ceil(q_tokens / query_chunk_size))
|
||||||
|
], dim=1)
|
||||||
|
return res
|
245
ldm_patched/ldm/modules/temporal_ae.py
Normal file
245
ldm_patched/ldm/modules/temporal_ae.py
Normal file
@ -0,0 +1,245 @@
|
|||||||
|
import functools
|
||||||
|
from typing import Callable, Iterable, Union
|
||||||
|
|
||||||
|
import torch
|
||||||
|
from einops import rearrange, repeat
|
||||||
|
|
||||||
|
import ldm_patched.modules.ops
|
||||||
|
ops = ldm_patched.modules.ops.disable_weight_init
|
||||||
|
|
||||||
|
from .diffusionmodules.model import (
|
||||||
|
AttnBlock,
|
||||||
|
Decoder,
|
||||||
|
ResnetBlock,
|
||||||
|
)
|
||||||
|
from .diffusionmodules.openaimodel import ResBlock, timestep_embedding
|
||||||
|
from .attention import BasicTransformerBlock
|
||||||
|
|
||||||
|
def partialclass(cls, *args, **kwargs):
|
||||||
|
class NewCls(cls):
|
||||||
|
__init__ = functools.partialmethod(cls.__init__, *args, **kwargs)
|
||||||
|
|
||||||
|
return NewCls
|
||||||
|
|
||||||
|
|
||||||
|
class VideoResBlock(ResnetBlock):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
out_channels,
|
||||||
|
*args,
|
||||||
|
dropout=0.0,
|
||||||
|
video_kernel_size=3,
|
||||||
|
alpha=0.0,
|
||||||
|
merge_strategy="learned",
|
||||||
|
**kwargs,
|
||||||
|
):
|
||||||
|
super().__init__(out_channels=out_channels, dropout=dropout, *args, **kwargs)
|
||||||
|
if video_kernel_size is None:
|
||||||
|
video_kernel_size = [3, 1, 1]
|
||||||
|
self.time_stack = ResBlock(
|
||||||
|
channels=out_channels,
|
||||||
|
emb_channels=0,
|
||||||
|
dropout=dropout,
|
||||||
|
dims=3,
|
||||||
|
use_scale_shift_norm=False,
|
||||||
|
use_conv=False,
|
||||||
|
up=False,
|
||||||
|
down=False,
|
||||||
|
kernel_size=video_kernel_size,
|
||||||
|
use_checkpoint=False,
|
||||||
|
skip_t_emb=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.merge_strategy = merge_strategy
|
||||||
|
if self.merge_strategy == "fixed":
|
||||||
|
self.register_buffer("mix_factor", torch.Tensor([alpha]))
|
||||||
|
elif self.merge_strategy == "learned":
|
||||||
|
self.register_parameter(
|
||||||
|
"mix_factor", torch.nn.Parameter(torch.Tensor([alpha]))
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
raise ValueError(f"unknown merge strategy {self.merge_strategy}")
|
||||||
|
|
||||||
|
def get_alpha(self, bs):
|
||||||
|
if self.merge_strategy == "fixed":
|
||||||
|
return self.mix_factor
|
||||||
|
elif self.merge_strategy == "learned":
|
||||||
|
return torch.sigmoid(self.mix_factor)
|
||||||
|
else:
|
||||||
|
raise NotImplementedError()
|
||||||
|
|
||||||
|
def forward(self, x, temb, skip_video=False, timesteps=None):
|
||||||
|
b, c, h, w = x.shape
|
||||||
|
if timesteps is None:
|
||||||
|
timesteps = b
|
||||||
|
|
||||||
|
x = super().forward(x, temb)
|
||||||
|
|
||||||
|
if not skip_video:
|
||||||
|
x_mix = rearrange(x, "(b t) c h w -> b c t h w", t=timesteps)
|
||||||
|
|
||||||
|
x = rearrange(x, "(b t) c h w -> b c t h w", t=timesteps)
|
||||||
|
|
||||||
|
x = self.time_stack(x, temb)
|
||||||
|
|
||||||
|
alpha = self.get_alpha(bs=b // timesteps).to(x.device)
|
||||||
|
x = alpha * x + (1.0 - alpha) * x_mix
|
||||||
|
|
||||||
|
x = rearrange(x, "b c t h w -> (b t) c h w")
|
||||||
|
return x
|
||||||
|
|
||||||
|
|
||||||
|
class AE3DConv(ops.Conv2d):
|
||||||
|
def __init__(self, in_channels, out_channels, video_kernel_size=3, *args, **kwargs):
|
||||||
|
super().__init__(in_channels, out_channels, *args, **kwargs)
|
||||||
|
if isinstance(video_kernel_size, Iterable):
|
||||||
|
padding = [int(k // 2) for k in video_kernel_size]
|
||||||
|
else:
|
||||||
|
padding = int(video_kernel_size // 2)
|
||||||
|
|
||||||
|
self.time_mix_conv = ops.Conv3d(
|
||||||
|
in_channels=out_channels,
|
||||||
|
out_channels=out_channels,
|
||||||
|
kernel_size=video_kernel_size,
|
||||||
|
padding=padding,
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(self, input, timesteps=None, skip_video=False):
|
||||||
|
if timesteps is None:
|
||||||
|
timesteps = input.shape[0]
|
||||||
|
x = super().forward(input)
|
||||||
|
if skip_video:
|
||||||
|
return x
|
||||||
|
x = rearrange(x, "(b t) c h w -> b c t h w", t=timesteps)
|
||||||
|
x = self.time_mix_conv(x)
|
||||||
|
return rearrange(x, "b c t h w -> (b t) c h w")
|
||||||
|
|
||||||
|
|
||||||
|
class AttnVideoBlock(AttnBlock):
|
||||||
|
def __init__(
|
||||||
|
self, in_channels: int, alpha: float = 0, merge_strategy: str = "learned"
|
||||||
|
):
|
||||||
|
super().__init__(in_channels)
|
||||||
|
# no context, single headed, as in base class
|
||||||
|
self.time_mix_block = BasicTransformerBlock(
|
||||||
|
dim=in_channels,
|
||||||
|
n_heads=1,
|
||||||
|
d_head=in_channels,
|
||||||
|
checkpoint=False,
|
||||||
|
ff_in=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
time_embed_dim = self.in_channels * 4
|
||||||
|
self.video_time_embed = torch.nn.Sequential(
|
||||||
|
ops.Linear(self.in_channels, time_embed_dim),
|
||||||
|
torch.nn.SiLU(),
|
||||||
|
ops.Linear(time_embed_dim, self.in_channels),
|
||||||
|
)
|
||||||
|
|
||||||
|
self.merge_strategy = merge_strategy
|
||||||
|
if self.merge_strategy == "fixed":
|
||||||
|
self.register_buffer("mix_factor", torch.Tensor([alpha]))
|
||||||
|
elif self.merge_strategy == "learned":
|
||||||
|
self.register_parameter(
|
||||||
|
"mix_factor", torch.nn.Parameter(torch.Tensor([alpha]))
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
raise ValueError(f"unknown merge strategy {self.merge_strategy}")
|
||||||
|
|
||||||
|
def forward(self, x, timesteps=None, skip_time_block=False):
|
||||||
|
if skip_time_block:
|
||||||
|
return super().forward(x)
|
||||||
|
|
||||||
|
if timesteps is None:
|
||||||
|
timesteps = x.shape[0]
|
||||||
|
|
||||||
|
x_in = x
|
||||||
|
x = self.attention(x)
|
||||||
|
h, w = x.shape[2:]
|
||||||
|
x = rearrange(x, "b c h w -> b (h w) c")
|
||||||
|
|
||||||
|
x_mix = 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)
|
||||||
|
emb = self.video_time_embed(t_emb) # b, n_channels
|
||||||
|
emb = emb[:, None, :]
|
||||||
|
x_mix = x_mix + emb
|
||||||
|
|
||||||
|
alpha = self.get_alpha().to(x.device)
|
||||||
|
x_mix = self.time_mix_block(x_mix, timesteps=timesteps)
|
||||||
|
x = alpha * x + (1.0 - alpha) * x_mix # alpha merge
|
||||||
|
|
||||||
|
x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w)
|
||||||
|
x = self.proj_out(x)
|
||||||
|
|
||||||
|
return x_in + x
|
||||||
|
|
||||||
|
def get_alpha(
|
||||||
|
self,
|
||||||
|
):
|
||||||
|
if self.merge_strategy == "fixed":
|
||||||
|
return self.mix_factor
|
||||||
|
elif self.merge_strategy == "learned":
|
||||||
|
return torch.sigmoid(self.mix_factor)
|
||||||
|
else:
|
||||||
|
raise NotImplementedError(f"unknown merge strategy {self.merge_strategy}")
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
def make_time_attn(
|
||||||
|
in_channels,
|
||||||
|
attn_type="vanilla",
|
||||||
|
attn_kwargs=None,
|
||||||
|
alpha: float = 0,
|
||||||
|
merge_strategy: str = "learned",
|
||||||
|
):
|
||||||
|
return partialclass(
|
||||||
|
AttnVideoBlock, in_channels, alpha=alpha, merge_strategy=merge_strategy
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class Conv2DWrapper(torch.nn.Conv2d):
|
||||||
|
def forward(self, input: torch.Tensor, **kwargs) -> torch.Tensor:
|
||||||
|
return super().forward(input)
|
||||||
|
|
||||||
|
|
||||||
|
class VideoDecoder(Decoder):
|
||||||
|
available_time_modes = ["all", "conv-only", "attn-only"]
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
*args,
|
||||||
|
video_kernel_size: Union[int, list] = 3,
|
||||||
|
alpha: float = 0.0,
|
||||||
|
merge_strategy: str = "learned",
|
||||||
|
time_mode: str = "conv-only",
|
||||||
|
**kwargs,
|
||||||
|
):
|
||||||
|
self.video_kernel_size = video_kernel_size
|
||||||
|
self.alpha = alpha
|
||||||
|
self.merge_strategy = merge_strategy
|
||||||
|
self.time_mode = time_mode
|
||||||
|
assert (
|
||||||
|
self.time_mode in self.available_time_modes
|
||||||
|
), f"time_mode parameter has to be in {self.available_time_modes}"
|
||||||
|
|
||||||
|
if self.time_mode != "attn-only":
|
||||||
|
kwargs["conv_out_op"] = partialclass(AE3DConv, video_kernel_size=self.video_kernel_size)
|
||||||
|
if self.time_mode not in ["conv-only", "only-last-conv"]:
|
||||||
|
kwargs["attn_op"] = partialclass(make_time_attn, alpha=self.alpha, merge_strategy=self.merge_strategy)
|
||||||
|
if self.time_mode not in ["attn-only", "only-last-conv"]:
|
||||||
|
kwargs["resnet_op"] = partialclass(VideoResBlock, video_kernel_size=self.video_kernel_size, alpha=self.alpha, merge_strategy=self.merge_strategy)
|
||||||
|
|
||||||
|
super().__init__(*args, **kwargs)
|
||||||
|
|
||||||
|
def get_last_layer(self, skip_time_mix=False, **kwargs):
|
||||||
|
if self.time_mode == "attn-only":
|
||||||
|
raise NotImplementedError("TODO")
|
||||||
|
else:
|
||||||
|
return (
|
||||||
|
self.conv_out.time_mix_conv.weight
|
||||||
|
if not skip_time_mix
|
||||||
|
else self.conv_out.weight
|
||||||
|
)
|
197
ldm_patched/ldm/util.py
Normal file
197
ldm_patched/ldm/util.py
Normal file
@ -0,0 +1,197 @@
|
|||||||
|
import importlib
|
||||||
|
|
||||||
|
import torch
|
||||||
|
from torch import optim
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
from inspect import isfunction
|
||||||
|
from PIL import Image, ImageDraw, ImageFont
|
||||||
|
|
||||||
|
|
||||||
|
def log_txt_as_img(wh, xc, size=10):
|
||||||
|
# wh a tuple of (width, height)
|
||||||
|
# xc a list of captions to plot
|
||||||
|
b = len(xc)
|
||||||
|
txts = list()
|
||||||
|
for bi in range(b):
|
||||||
|
txt = Image.new("RGB", wh, color="white")
|
||||||
|
draw = ImageDraw.Draw(txt)
|
||||||
|
font = ImageFont.truetype('data/DejaVuSans.ttf', size=size)
|
||||||
|
nc = int(40 * (wh[0] / 256))
|
||||||
|
lines = "\n".join(xc[bi][start:start + nc] for start in range(0, len(xc[bi]), nc))
|
||||||
|
|
||||||
|
try:
|
||||||
|
draw.text((0, 0), lines, fill="black", font=font)
|
||||||
|
except UnicodeEncodeError:
|
||||||
|
print("Cant encode string for logging. Skipping.")
|
||||||
|
|
||||||
|
txt = np.array(txt).transpose(2, 0, 1) / 127.5 - 1.0
|
||||||
|
txts.append(txt)
|
||||||
|
txts = np.stack(txts)
|
||||||
|
txts = torch.tensor(txts)
|
||||||
|
return txts
|
||||||
|
|
||||||
|
|
||||||
|
def ismap(x):
|
||||||
|
if not isinstance(x, torch.Tensor):
|
||||||
|
return False
|
||||||
|
return (len(x.shape) == 4) and (x.shape[1] > 3)
|
||||||
|
|
||||||
|
|
||||||
|
def isimage(x):
|
||||||
|
if not isinstance(x,torch.Tensor):
|
||||||
|
return False
|
||||||
|
return (len(x.shape) == 4) and (x.shape[1] == 3 or x.shape[1] == 1)
|
||||||
|
|
||||||
|
|
||||||
|
def exists(x):
|
||||||
|
return x is not None
|
||||||
|
|
||||||
|
|
||||||
|
def default(val, d):
|
||||||
|
if exists(val):
|
||||||
|
return val
|
||||||
|
return d() if isfunction(d) else d
|
||||||
|
|
||||||
|
|
||||||
|
def mean_flat(tensor):
|
||||||
|
"""
|
||||||
|
https://github.com/openai/guided-diffusion/blob/27c20a8fab9cb472df5d6bdd6c8d11c8f430b924/guided_diffusion/nn.py#L86
|
||||||
|
Take the mean over all non-batch dimensions.
|
||||||
|
"""
|
||||||
|
return tensor.mean(dim=list(range(1, len(tensor.shape))))
|
||||||
|
|
||||||
|
|
||||||
|
def count_params(model, verbose=False):
|
||||||
|
total_params = sum(p.numel() for p in model.parameters())
|
||||||
|
if verbose:
|
||||||
|
print(f"{model.__class__.__name__} has {total_params*1.e-6:.2f} M params.")
|
||||||
|
return total_params
|
||||||
|
|
||||||
|
|
||||||
|
def instantiate_from_config(config):
|
||||||
|
if not "target" in config:
|
||||||
|
if config == '__is_first_stage__':
|
||||||
|
return None
|
||||||
|
elif config == "__is_unconditional__":
|
||||||
|
return None
|
||||||
|
raise KeyError("Expected key `target` to instantiate.")
|
||||||
|
return get_obj_from_str(config["target"])(**config.get("params", dict()))
|
||||||
|
|
||||||
|
|
||||||
|
def get_obj_from_str(string, reload=False):
|
||||||
|
module, cls = string.rsplit(".", 1)
|
||||||
|
if reload:
|
||||||
|
module_imp = importlib.import_module(module)
|
||||||
|
importlib.reload(module_imp)
|
||||||
|
return getattr(importlib.import_module(module, package=None), cls)
|
||||||
|
|
||||||
|
|
||||||
|
class AdamWwithEMAandWings(optim.Optimizer):
|
||||||
|
# credit to https://gist.github.com/crowsonkb/65f7265353f403714fce3b2595e0b298
|
||||||
|
def __init__(self, params, lr=1.e-3, betas=(0.9, 0.999), eps=1.e-8, # TODO: check hyperparameters before using
|
||||||
|
weight_decay=1.e-2, amsgrad=False, ema_decay=0.9999, # ema decay to match previous code
|
||||||
|
ema_power=1., param_names=()):
|
||||||
|
"""AdamW that saves EMA versions of the parameters."""
|
||||||
|
if not 0.0 <= lr:
|
||||||
|
raise ValueError("Invalid learning rate: {}".format(lr))
|
||||||
|
if not 0.0 <= eps:
|
||||||
|
raise ValueError("Invalid epsilon value: {}".format(eps))
|
||||||
|
if not 0.0 <= betas[0] < 1.0:
|
||||||
|
raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
|
||||||
|
if not 0.0 <= betas[1] < 1.0:
|
||||||
|
raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
|
||||||
|
if not 0.0 <= weight_decay:
|
||||||
|
raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
|
||||||
|
if not 0.0 <= ema_decay <= 1.0:
|
||||||
|
raise ValueError("Invalid ema_decay value: {}".format(ema_decay))
|
||||||
|
defaults = dict(lr=lr, betas=betas, eps=eps,
|
||||||
|
weight_decay=weight_decay, amsgrad=amsgrad, ema_decay=ema_decay,
|
||||||
|
ema_power=ema_power, param_names=param_names)
|
||||||
|
super().__init__(params, defaults)
|
||||||
|
|
||||||
|
def __setstate__(self, state):
|
||||||
|
super().__setstate__(state)
|
||||||
|
for group in self.param_groups:
|
||||||
|
group.setdefault('amsgrad', False)
|
||||||
|
|
||||||
|
@torch.no_grad()
|
||||||
|
def step(self, closure=None):
|
||||||
|
"""Performs a single optimization step.
|
||||||
|
Args:
|
||||||
|
closure (callable, optional): A closure that reevaluates the model
|
||||||
|
and returns the loss.
|
||||||
|
"""
|
||||||
|
loss = None
|
||||||
|
if closure is not None:
|
||||||
|
with torch.enable_grad():
|
||||||
|
loss = closure()
|
||||||
|
|
||||||
|
for group in self.param_groups:
|
||||||
|
params_with_grad = []
|
||||||
|
grads = []
|
||||||
|
exp_avgs = []
|
||||||
|
exp_avg_sqs = []
|
||||||
|
ema_params_with_grad = []
|
||||||
|
state_sums = []
|
||||||
|
max_exp_avg_sqs = []
|
||||||
|
state_steps = []
|
||||||
|
amsgrad = group['amsgrad']
|
||||||
|
beta1, beta2 = group['betas']
|
||||||
|
ema_decay = group['ema_decay']
|
||||||
|
ema_power = group['ema_power']
|
||||||
|
|
||||||
|
for p in group['params']:
|
||||||
|
if p.grad is None:
|
||||||
|
continue
|
||||||
|
params_with_grad.append(p)
|
||||||
|
if p.grad.is_sparse:
|
||||||
|
raise RuntimeError('AdamW does not support sparse gradients')
|
||||||
|
grads.append(p.grad)
|
||||||
|
|
||||||
|
state = self.state[p]
|
||||||
|
|
||||||
|
# State initialization
|
||||||
|
if len(state) == 0:
|
||||||
|
state['step'] = 0
|
||||||
|
# Exponential moving average of gradient values
|
||||||
|
state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format)
|
||||||
|
# Exponential moving average of squared gradient values
|
||||||
|
state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
|
||||||
|
if amsgrad:
|
||||||
|
# Maintains max of all exp. moving avg. of sq. grad. values
|
||||||
|
state['max_exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)
|
||||||
|
# Exponential moving average of parameter values
|
||||||
|
state['param_exp_avg'] = p.detach().float().clone()
|
||||||
|
|
||||||
|
exp_avgs.append(state['exp_avg'])
|
||||||
|
exp_avg_sqs.append(state['exp_avg_sq'])
|
||||||
|
ema_params_with_grad.append(state['param_exp_avg'])
|
||||||
|
|
||||||
|
if amsgrad:
|
||||||
|
max_exp_avg_sqs.append(state['max_exp_avg_sq'])
|
||||||
|
|
||||||
|
# update the steps for each param group update
|
||||||
|
state['step'] += 1
|
||||||
|
# record the step after step update
|
||||||
|
state_steps.append(state['step'])
|
||||||
|
|
||||||
|
optim._functional.adamw(params_with_grad,
|
||||||
|
grads,
|
||||||
|
exp_avgs,
|
||||||
|
exp_avg_sqs,
|
||||||
|
max_exp_avg_sqs,
|
||||||
|
state_steps,
|
||||||
|
amsgrad=amsgrad,
|
||||||
|
beta1=beta1,
|
||||||
|
beta2=beta2,
|
||||||
|
lr=group['lr'],
|
||||||
|
weight_decay=group['weight_decay'],
|
||||||
|
eps=group['eps'],
|
||||||
|
maximize=False)
|
||||||
|
|
||||||
|
cur_ema_decay = min(ema_decay, 1 - state['step'] ** -ema_power)
|
||||||
|
for param, ema_param in zip(params_with_grad, ema_params_with_grad):
|
||||||
|
ema_param.mul_(cur_ema_decay).add_(param.float(), alpha=1 - cur_ema_decay)
|
||||||
|
|
||||||
|
return loss
|
126
ldm_patched/modules/args_parser.py
Normal file
126
ldm_patched/modules/args_parser.py
Normal file
@ -0,0 +1,126 @@
|
|||||||
|
import argparse
|
||||||
|
import enum
|
||||||
|
import ldm_patched.modules.options
|
||||||
|
|
||||||
|
class EnumAction(argparse.Action):
|
||||||
|
"""
|
||||||
|
Argparse action for handling Enums
|
||||||
|
"""
|
||||||
|
def __init__(self, **kwargs):
|
||||||
|
# Pop off the type value
|
||||||
|
enum_type = kwargs.pop("type", None)
|
||||||
|
|
||||||
|
# Ensure an Enum subclass is provided
|
||||||
|
if enum_type is None:
|
||||||
|
raise ValueError("type must be assigned an Enum when using EnumAction")
|
||||||
|
if not issubclass(enum_type, enum.Enum):
|
||||||
|
raise TypeError("type must be an Enum when using EnumAction")
|
||||||
|
|
||||||
|
# Generate choices from the Enum
|
||||||
|
choices = tuple(e.value for e in enum_type)
|
||||||
|
kwargs.setdefault("choices", choices)
|
||||||
|
kwargs.setdefault("metavar", f"[{','.join(list(choices))}]")
|
||||||
|
|
||||||
|
super(EnumAction, self).__init__(**kwargs)
|
||||||
|
|
||||||
|
self._enum = enum_type
|
||||||
|
|
||||||
|
def __call__(self, parser, namespace, values, option_string=None):
|
||||||
|
# Convert value back into an Enum
|
||||||
|
value = self._enum(values)
|
||||||
|
setattr(namespace, self.dest, value)
|
||||||
|
|
||||||
|
|
||||||
|
parser = argparse.ArgumentParser()
|
||||||
|
|
||||||
|
parser.add_argument("--listen", type=str, default="127.0.0.1", metavar="IP", nargs="?", const="0.0.0.0")
|
||||||
|
parser.add_argument("--port", type=int, default=8188)
|
||||||
|
parser.add_argument("--disable-header-check", type=str, default=None, metavar="ORIGIN", nargs="?", const="*")
|
||||||
|
parser.add_argument("--web-upload-size", type=float, default=100)
|
||||||
|
|
||||||
|
parser.add_argument("--external-working-path", type=str, default=None, metavar="PATH", nargs='+', action='append')
|
||||||
|
parser.add_argument("--output-path", type=str, default=None)
|
||||||
|
parser.add_argument("--temp-path", type=str, default=None)
|
||||||
|
parser.add_argument("--cache-path", type=str, default=None)
|
||||||
|
parser.add_argument("--in-browser", action="store_true")
|
||||||
|
parser.add_argument("--disable-in-browser", action="store_true")
|
||||||
|
parser.add_argument("--gpu-device-id", type=int, default=None, metavar="DEVICE_ID")
|
||||||
|
cm_group = parser.add_mutually_exclusive_group()
|
||||||
|
cm_group.add_argument("--async-cuda-allocation", action="store_true")
|
||||||
|
cm_group.add_argument("--disable-async-cuda-allocation", action="store_true")
|
||||||
|
|
||||||
|
parser.add_argument("--disable-attention-upcast", action="store_true")
|
||||||
|
|
||||||
|
fp_group = parser.add_mutually_exclusive_group()
|
||||||
|
fp_group.add_argument("--all-in-fp32", action="store_true")
|
||||||
|
fp_group.add_argument("--all-in-fp16", action="store_true")
|
||||||
|
|
||||||
|
fpunet_group = parser.add_mutually_exclusive_group()
|
||||||
|
fpunet_group.add_argument("--unet-in-bf16", action="store_true")
|
||||||
|
fpunet_group.add_argument("--unet-in-fp16", action="store_true")
|
||||||
|
fpunet_group.add_argument("--unet-in-fp8-e4m3fn", action="store_true")
|
||||||
|
fpunet_group.add_argument("--unet-in-fp8-e5m2", action="store_true")
|
||||||
|
|
||||||
|
fpvae_group = parser.add_mutually_exclusive_group()
|
||||||
|
fpvae_group.add_argument("--vae-in-fp16", action="store_true")
|
||||||
|
fpvae_group.add_argument("--vae-in-fp32", action="store_true")
|
||||||
|
fpvae_group.add_argument("--vae-in-bf16", action="store_true")
|
||||||
|
|
||||||
|
parser.add_argument("--vae-in-cpu", action="store_true")
|
||||||
|
|
||||||
|
fpte_group = parser.add_mutually_exclusive_group()
|
||||||
|
fpte_group.add_argument("--clip-in-fp8-e4m3fn", action="store_true")
|
||||||
|
fpte_group.add_argument("--clip-in-fp8-e5m2", action="store_true")
|
||||||
|
fpte_group.add_argument("--clip-in-fp16", action="store_true")
|
||||||
|
fpte_group.add_argument("--clip-in-fp32", action="store_true")
|
||||||
|
|
||||||
|
|
||||||
|
parser.add_argument("--directml", type=int, nargs="?", metavar="DIRECTML_DEVICE", const=-1)
|
||||||
|
|
||||||
|
parser.add_argument("--disable-ipex-hijack", action="store_true")
|
||||||
|
|
||||||
|
class LatentPreviewMethod(enum.Enum):
|
||||||
|
NoPreviews = "none"
|
||||||
|
Auto = "auto"
|
||||||
|
Latent2RGB = "fast"
|
||||||
|
TAESD = "taesd"
|
||||||
|
|
||||||
|
parser.add_argument("--preview-option", type=LatentPreviewMethod, default=LatentPreviewMethod.NoPreviews, action=EnumAction)
|
||||||
|
|
||||||
|
attn_group = parser.add_mutually_exclusive_group()
|
||||||
|
attn_group.add_argument("--attention-split", action="store_true")
|
||||||
|
attn_group.add_argument("--attention-quad", action="store_true")
|
||||||
|
attn_group.add_argument("--attention-pytorch", action="store_true")
|
||||||
|
|
||||||
|
parser.add_argument("--disable-xformers", action="store_true")
|
||||||
|
|
||||||
|
vram_group = parser.add_mutually_exclusive_group()
|
||||||
|
vram_group.add_argument("--always-gpu", action="store_true")
|
||||||
|
vram_group.add_argument("--always-high-vram", action="store_true")
|
||||||
|
vram_group.add_argument("--always-normal-vram", action="store_true")
|
||||||
|
vram_group.add_argument("--always-low-vram", action="store_true")
|
||||||
|
vram_group.add_argument("--always-no-vram", action="store_true")
|
||||||
|
vram_group.add_argument("--always-cpu", action="store_true")
|
||||||
|
|
||||||
|
|
||||||
|
parser.add_argument("--always-offload-from-vram", action="store_true")
|
||||||
|
parser.add_argument("--pytorch-deterministic", action="store_true")
|
||||||
|
|
||||||
|
parser.add_argument("--disable-server-log", action="store_true")
|
||||||
|
parser.add_argument("--debug-mode", action="store_true")
|
||||||
|
parser.add_argument("--is-windows-embedded-python", action="store_true")
|
||||||
|
|
||||||
|
parser.add_argument("--disable-server-info", action="store_true")
|
||||||
|
|
||||||
|
parser.add_argument("--multi-user", action="store_true")
|
||||||
|
|
||||||
|
if ldm_patched.modules.options.args_parsing:
|
||||||
|
args = parser.parse_args([])
|
||||||
|
else:
|
||||||
|
args = parser.parse_args([])
|
||||||
|
|
||||||
|
if args.is_windows_embedded_python:
|
||||||
|
args.in_browser = True
|
||||||
|
|
||||||
|
if args.disable_in_browser:
|
||||||
|
args.in_browser = False
|
13
ldm_patched/modules/checkpoint_pickle.py
Normal file
13
ldm_patched/modules/checkpoint_pickle.py
Normal file
@ -0,0 +1,13 @@
|
|||||||
|
import pickle
|
||||||
|
|
||||||
|
load = pickle.load
|
||||||
|
|
||||||
|
class Empty:
|
||||||
|
pass
|
||||||
|
|
||||||
|
class Unpickler(pickle.Unpickler):
|
||||||
|
def find_class(self, module, name):
|
||||||
|
#TODO: safe unpickle
|
||||||
|
if module.startswith("pytorch_lightning"):
|
||||||
|
return Empty
|
||||||
|
return super().find_class(module, name)
|
23
ldm_patched/modules/clip_config_bigg.json
Normal file
23
ldm_patched/modules/clip_config_bigg.json
Normal file
@ -0,0 +1,23 @@
|
|||||||
|
{
|
||||||
|
"architectures": [
|
||||||
|
"CLIPTextModel"
|
||||||
|
],
|
||||||
|
"attention_dropout": 0.0,
|
||||||
|
"bos_token_id": 0,
|
||||||
|
"dropout": 0.0,
|
||||||
|
"eos_token_id": 2,
|
||||||
|
"hidden_act": "gelu",
|
||||||
|
"hidden_size": 1280,
|
||||||
|
"initializer_factor": 1.0,
|
||||||
|
"initializer_range": 0.02,
|
||||||
|
"intermediate_size": 5120,
|
||||||
|
"layer_norm_eps": 1e-05,
|
||||||
|
"max_position_embeddings": 77,
|
||||||
|
"model_type": "clip_text_model",
|
||||||
|
"num_attention_heads": 20,
|
||||||
|
"num_hidden_layers": 32,
|
||||||
|
"pad_token_id": 1,
|
||||||
|
"projection_dim": 1280,
|
||||||
|
"torch_dtype": "float32",
|
||||||
|
"vocab_size": 49408
|
||||||
|
}
|
188
ldm_patched/modules/clip_model.py
Normal file
188
ldm_patched/modules/clip_model.py
Normal file
@ -0,0 +1,188 @@
|
|||||||
|
import torch
|
||||||
|
from ldm_patched.ldm.modules.attention import optimized_attention_for_device
|
||||||
|
|
||||||
|
class CLIPAttention(torch.nn.Module):
|
||||||
|
def __init__(self, embed_dim, heads, dtype, device, operations):
|
||||||
|
super().__init__()
|
||||||
|
|
||||||
|
self.heads = heads
|
||||||
|
self.q_proj = operations.Linear(embed_dim, embed_dim, bias=True, dtype=dtype, device=device)
|
||||||
|
self.k_proj = operations.Linear(embed_dim, embed_dim, bias=True, dtype=dtype, device=device)
|
||||||
|
self.v_proj = operations.Linear(embed_dim, embed_dim, bias=True, dtype=dtype, device=device)
|
||||||
|
|
||||||
|
self.out_proj = operations.Linear(embed_dim, embed_dim, bias=True, dtype=dtype, device=device)
|
||||||
|
|
||||||
|
def forward(self, x, mask=None, optimized_attention=None):
|
||||||
|
q = self.q_proj(x)
|
||||||
|
k = self.k_proj(x)
|
||||||
|
v = self.v_proj(x)
|
||||||
|
|
||||||
|
out = optimized_attention(q, k, v, self.heads, mask)
|
||||||
|
return self.out_proj(out)
|
||||||
|
|
||||||
|
ACTIVATIONS = {"quick_gelu": lambda a: a * torch.sigmoid(1.702 * a),
|
||||||
|
"gelu": torch.nn.functional.gelu,
|
||||||
|
}
|
||||||
|
|
||||||
|
class CLIPMLP(torch.nn.Module):
|
||||||
|
def __init__(self, embed_dim, intermediate_size, activation, dtype, device, operations):
|
||||||
|
super().__init__()
|
||||||
|
self.fc1 = operations.Linear(embed_dim, intermediate_size, bias=True, dtype=dtype, device=device)
|
||||||
|
self.activation = ACTIVATIONS[activation]
|
||||||
|
self.fc2 = operations.Linear(intermediate_size, embed_dim, bias=True, dtype=dtype, device=device)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
x = self.fc1(x)
|
||||||
|
x = self.activation(x)
|
||||||
|
x = self.fc2(x)
|
||||||
|
return x
|
||||||
|
|
||||||
|
class CLIPLayer(torch.nn.Module):
|
||||||
|
def __init__(self, embed_dim, heads, intermediate_size, intermediate_activation, dtype, device, operations):
|
||||||
|
super().__init__()
|
||||||
|
self.layer_norm1 = operations.LayerNorm(embed_dim, dtype=dtype, device=device)
|
||||||
|
self.self_attn = CLIPAttention(embed_dim, heads, dtype, device, operations)
|
||||||
|
self.layer_norm2 = operations.LayerNorm(embed_dim, dtype=dtype, device=device)
|
||||||
|
self.mlp = CLIPMLP(embed_dim, intermediate_size, intermediate_activation, dtype, device, operations)
|
||||||
|
|
||||||
|
def forward(self, x, mask=None, optimized_attention=None):
|
||||||
|
x += self.self_attn(self.layer_norm1(x), mask, optimized_attention)
|
||||||
|
x += self.mlp(self.layer_norm2(x))
|
||||||
|
return x
|
||||||
|
|
||||||
|
|
||||||
|
class CLIPEncoder(torch.nn.Module):
|
||||||
|
def __init__(self, num_layers, embed_dim, heads, intermediate_size, intermediate_activation, dtype, device, operations):
|
||||||
|
super().__init__()
|
||||||
|
self.layers = torch.nn.ModuleList([CLIPLayer(embed_dim, heads, intermediate_size, intermediate_activation, dtype, device, operations) for i in range(num_layers)])
|
||||||
|
|
||||||
|
def forward(self, x, mask=None, intermediate_output=None):
|
||||||
|
optimized_attention = optimized_attention_for_device(x.device, mask=mask is not None, small_input=True)
|
||||||
|
|
||||||
|
if intermediate_output is not None:
|
||||||
|
if intermediate_output < 0:
|
||||||
|
intermediate_output = len(self.layers) + intermediate_output
|
||||||
|
|
||||||
|
intermediate = None
|
||||||
|
for i, l in enumerate(self.layers):
|
||||||
|
x = l(x, mask, optimized_attention)
|
||||||
|
if i == intermediate_output:
|
||||||
|
intermediate = x.clone()
|
||||||
|
return x, intermediate
|
||||||
|
|
||||||
|
class CLIPEmbeddings(torch.nn.Module):
|
||||||
|
def __init__(self, embed_dim, vocab_size=49408, num_positions=77, dtype=None, device=None):
|
||||||
|
super().__init__()
|
||||||
|
self.token_embedding = torch.nn.Embedding(vocab_size, embed_dim, dtype=dtype, device=device)
|
||||||
|
self.position_embedding = torch.nn.Embedding(num_positions, embed_dim, dtype=dtype, device=device)
|
||||||
|
|
||||||
|
def forward(self, input_tokens):
|
||||||
|
return self.token_embedding(input_tokens) + self.position_embedding.weight
|
||||||
|
|
||||||
|
|
||||||
|
class CLIPTextModel_(torch.nn.Module):
|
||||||
|
def __init__(self, config_dict, dtype, device, operations):
|
||||||
|
num_layers = config_dict["num_hidden_layers"]
|
||||||
|
embed_dim = config_dict["hidden_size"]
|
||||||
|
heads = config_dict["num_attention_heads"]
|
||||||
|
intermediate_size = config_dict["intermediate_size"]
|
||||||
|
intermediate_activation = config_dict["hidden_act"]
|
||||||
|
|
||||||
|
super().__init__()
|
||||||
|
self.embeddings = CLIPEmbeddings(embed_dim, dtype=torch.float32, device=device)
|
||||||
|
self.encoder = CLIPEncoder(num_layers, embed_dim, heads, intermediate_size, intermediate_activation, dtype, device, operations)
|
||||||
|
self.final_layer_norm = operations.LayerNorm(embed_dim, dtype=dtype, device=device)
|
||||||
|
|
||||||
|
def forward(self, input_tokens, attention_mask=None, intermediate_output=None, final_layer_norm_intermediate=True):
|
||||||
|
x = self.embeddings(input_tokens)
|
||||||
|
mask = None
|
||||||
|
if attention_mask is not None:
|
||||||
|
mask = 1.0 - attention_mask.to(x.dtype).unsqueeze(1).unsqueeze(1).expand(attention_mask.shape[0], 1, attention_mask.shape[-1], attention_mask.shape[-1])
|
||||||
|
mask = mask.masked_fill(mask.to(torch.bool), float("-inf"))
|
||||||
|
|
||||||
|
causal_mask = torch.empty(x.shape[1], x.shape[1], dtype=x.dtype, device=x.device).fill_(float("-inf")).triu_(1)
|
||||||
|
if mask is not None:
|
||||||
|
mask += causal_mask
|
||||||
|
else:
|
||||||
|
mask = causal_mask
|
||||||
|
|
||||||
|
x, i = self.encoder(x, mask=mask, intermediate_output=intermediate_output)
|
||||||
|
x = self.final_layer_norm(x)
|
||||||
|
if i is not None and final_layer_norm_intermediate:
|
||||||
|
i = self.final_layer_norm(i)
|
||||||
|
|
||||||
|
pooled_output = x[torch.arange(x.shape[0], device=x.device), input_tokens.to(dtype=torch.int, device=x.device).argmax(dim=-1),]
|
||||||
|
return x, i, pooled_output
|
||||||
|
|
||||||
|
class CLIPTextModel(torch.nn.Module):
|
||||||
|
def __init__(self, config_dict, dtype, device, operations):
|
||||||
|
super().__init__()
|
||||||
|
self.num_layers = config_dict["num_hidden_layers"]
|
||||||
|
self.text_model = CLIPTextModel_(config_dict, dtype, device, operations)
|
||||||
|
self.dtype = dtype
|
||||||
|
|
||||||
|
def get_input_embeddings(self):
|
||||||
|
return self.text_model.embeddings.token_embedding
|
||||||
|
|
||||||
|
def set_input_embeddings(self, embeddings):
|
||||||
|
self.text_model.embeddings.token_embedding = embeddings
|
||||||
|
|
||||||
|
def forward(self, *args, **kwargs):
|
||||||
|
return self.text_model(*args, **kwargs)
|
||||||
|
|
||||||
|
class CLIPVisionEmbeddings(torch.nn.Module):
|
||||||
|
def __init__(self, embed_dim, num_channels=3, patch_size=14, image_size=224, dtype=None, device=None, operations=None):
|
||||||
|
super().__init__()
|
||||||
|
self.class_embedding = torch.nn.Parameter(torch.empty(embed_dim, dtype=dtype, device=device))
|
||||||
|
|
||||||
|
self.patch_embedding = operations.Conv2d(
|
||||||
|
in_channels=num_channels,
|
||||||
|
out_channels=embed_dim,
|
||||||
|
kernel_size=patch_size,
|
||||||
|
stride=patch_size,
|
||||||
|
bias=False,
|
||||||
|
dtype=dtype,
|
||||||
|
device=device
|
||||||
|
)
|
||||||
|
|
||||||
|
num_patches = (image_size // patch_size) ** 2
|
||||||
|
num_positions = num_patches + 1
|
||||||
|
self.position_embedding = torch.nn.Embedding(num_positions, embed_dim, dtype=dtype, device=device)
|
||||||
|
|
||||||
|
def forward(self, pixel_values):
|
||||||
|
embeds = self.patch_embedding(pixel_values).flatten(2).transpose(1, 2)
|
||||||
|
return torch.cat([self.class_embedding.to(embeds.device).expand(pixel_values.shape[0], 1, -1), embeds], dim=1) + self.position_embedding.weight.to(embeds.device)
|
||||||
|
|
||||||
|
|
||||||
|
class CLIPVision(torch.nn.Module):
|
||||||
|
def __init__(self, config_dict, dtype, device, operations):
|
||||||
|
super().__init__()
|
||||||
|
num_layers = config_dict["num_hidden_layers"]
|
||||||
|
embed_dim = config_dict["hidden_size"]
|
||||||
|
heads = config_dict["num_attention_heads"]
|
||||||
|
intermediate_size = config_dict["intermediate_size"]
|
||||||
|
intermediate_activation = config_dict["hidden_act"]
|
||||||
|
|
||||||
|
self.embeddings = CLIPVisionEmbeddings(embed_dim, config_dict["num_channels"], config_dict["patch_size"], config_dict["image_size"], dtype=torch.float32, device=device, operations=operations)
|
||||||
|
self.pre_layrnorm = operations.LayerNorm(embed_dim)
|
||||||
|
self.encoder = CLIPEncoder(num_layers, embed_dim, heads, intermediate_size, intermediate_activation, dtype, device, operations)
|
||||||
|
self.post_layernorm = operations.LayerNorm(embed_dim)
|
||||||
|
|
||||||
|
def forward(self, pixel_values, attention_mask=None, intermediate_output=None):
|
||||||
|
x = self.embeddings(pixel_values)
|
||||||
|
x = self.pre_layrnorm(x)
|
||||||
|
#TODO: attention_mask?
|
||||||
|
x, i = self.encoder(x, mask=None, intermediate_output=intermediate_output)
|
||||||
|
pooled_output = self.post_layernorm(x[:, 0, :])
|
||||||
|
return x, i, pooled_output
|
||||||
|
|
||||||
|
class CLIPVisionModelProjection(torch.nn.Module):
|
||||||
|
def __init__(self, config_dict, dtype, device, operations):
|
||||||
|
super().__init__()
|
||||||
|
self.vision_model = CLIPVision(config_dict, dtype, device, operations)
|
||||||
|
self.visual_projection = operations.Linear(config_dict["hidden_size"], config_dict["projection_dim"], bias=False)
|
||||||
|
|
||||||
|
def forward(self, *args, **kwargs):
|
||||||
|
x = self.vision_model(*args, **kwargs)
|
||||||
|
out = self.visual_projection(x[2])
|
||||||
|
return (x[0], x[1], out)
|
110
ldm_patched/modules/clip_vision.py
Normal file
110
ldm_patched/modules/clip_vision.py
Normal file
@ -0,0 +1,110 @@
|
|||||||
|
from .utils import load_torch_file, transformers_convert, common_upscale
|
||||||
|
import os
|
||||||
|
import torch
|
||||||
|
import contextlib
|
||||||
|
import json
|
||||||
|
|
||||||
|
import ldm_patched.modules.ops
|
||||||
|
import ldm_patched.modules.model_patcher
|
||||||
|
import ldm_patched.modules.model_management
|
||||||
|
import ldm_patched.modules.utils
|
||||||
|
import ldm_patched.modules.clip_model
|
||||||
|
|
||||||
|
class Output:
|
||||||
|
def __getitem__(self, key):
|
||||||
|
return getattr(self, key)
|
||||||
|
def __setitem__(self, key, item):
|
||||||
|
setattr(self, key, item)
|
||||||
|
|
||||||
|
def clip_preprocess(image, size=224):
|
||||||
|
mean = torch.tensor([ 0.48145466,0.4578275,0.40821073], device=image.device, dtype=image.dtype)
|
||||||
|
std = torch.tensor([0.26862954,0.26130258,0.27577711], device=image.device, dtype=image.dtype)
|
||||||
|
image = image.movedim(-1, 1)
|
||||||
|
if not (image.shape[2] == size and image.shape[3] == size):
|
||||||
|
scale = (size / min(image.shape[2], image.shape[3]))
|
||||||
|
image = torch.nn.functional.interpolate(image, size=(round(scale * image.shape[2]), round(scale * image.shape[3])), mode="bicubic", antialias=True)
|
||||||
|
h = (image.shape[2] - size)//2
|
||||||
|
w = (image.shape[3] - size)//2
|
||||||
|
image = image[:,:,h:h+size,w:w+size]
|
||||||
|
image = torch.clip((255. * image), 0, 255).round() / 255.0
|
||||||
|
return (image - mean.view([3,1,1])) / std.view([3,1,1])
|
||||||
|
|
||||||
|
class ClipVisionModel():
|
||||||
|
def __init__(self, json_config):
|
||||||
|
with open(json_config) as f:
|
||||||
|
config = json.load(f)
|
||||||
|
|
||||||
|
self.load_device = ldm_patched.modules.model_management.text_encoder_device()
|
||||||
|
offload_device = ldm_patched.modules.model_management.text_encoder_offload_device()
|
||||||
|
self.dtype = ldm_patched.modules.model_management.text_encoder_dtype(self.load_device)
|
||||||
|
self.model = ldm_patched.modules.clip_model.CLIPVisionModelProjection(config, self.dtype, offload_device, ldm_patched.modules.ops.manual_cast)
|
||||||
|
self.model.eval()
|
||||||
|
|
||||||
|
self.patcher = ldm_patched.modules.model_patcher.ModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device)
|
||||||
|
def load_sd(self, sd):
|
||||||
|
return self.model.load_state_dict(sd, strict=False)
|
||||||
|
|
||||||
|
def encode_image(self, image):
|
||||||
|
ldm_patched.modules.model_management.load_model_gpu(self.patcher)
|
||||||
|
pixel_values = clip_preprocess(image.to(self.load_device)).float()
|
||||||
|
out = self.model(pixel_values=pixel_values, intermediate_output=-2)
|
||||||
|
|
||||||
|
outputs = Output()
|
||||||
|
outputs["last_hidden_state"] = out[0].to(ldm_patched.modules.model_management.intermediate_device())
|
||||||
|
outputs["image_embeds"] = out[2].to(ldm_patched.modules.model_management.intermediate_device())
|
||||||
|
outputs["penultimate_hidden_states"] = out[1].to(ldm_patched.modules.model_management.intermediate_device())
|
||||||
|
return outputs
|
||||||
|
|
||||||
|
def convert_to_transformers(sd, prefix):
|
||||||
|
sd_k = sd.keys()
|
||||||
|
if "{}transformer.resblocks.0.attn.in_proj_weight".format(prefix) in sd_k:
|
||||||
|
keys_to_replace = {
|
||||||
|
"{}class_embedding".format(prefix): "vision_model.embeddings.class_embedding",
|
||||||
|
"{}conv1.weight".format(prefix): "vision_model.embeddings.patch_embedding.weight",
|
||||||
|
"{}positional_embedding".format(prefix): "vision_model.embeddings.position_embedding.weight",
|
||||||
|
"{}ln_post.bias".format(prefix): "vision_model.post_layernorm.bias",
|
||||||
|
"{}ln_post.weight".format(prefix): "vision_model.post_layernorm.weight",
|
||||||
|
"{}ln_pre.bias".format(prefix): "vision_model.pre_layrnorm.bias",
|
||||||
|
"{}ln_pre.weight".format(prefix): "vision_model.pre_layrnorm.weight",
|
||||||
|
}
|
||||||
|
|
||||||
|
for x in keys_to_replace:
|
||||||
|
if x in sd_k:
|
||||||
|
sd[keys_to_replace[x]] = sd.pop(x)
|
||||||
|
|
||||||
|
if "{}proj".format(prefix) in sd_k:
|
||||||
|
sd['visual_projection.weight'] = sd.pop("{}proj".format(prefix)).transpose(0, 1)
|
||||||
|
|
||||||
|
sd = transformers_convert(sd, prefix, "vision_model.", 48)
|
||||||
|
return sd
|
||||||
|
|
||||||
|
def load_clipvision_from_sd(sd, prefix="", convert_keys=False):
|
||||||
|
if convert_keys:
|
||||||
|
sd = convert_to_transformers(sd, prefix)
|
||||||
|
if "vision_model.encoder.layers.47.layer_norm1.weight" in sd:
|
||||||
|
json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "clip_vision_config_g.json")
|
||||||
|
elif "vision_model.encoder.layers.30.layer_norm1.weight" in sd:
|
||||||
|
json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "clip_vision_config_h.json")
|
||||||
|
elif "vision_model.encoder.layers.22.layer_norm1.weight" in sd:
|
||||||
|
json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "clip_vision_config_vitl.json")
|
||||||
|
else:
|
||||||
|
return None
|
||||||
|
|
||||||
|
clip = ClipVisionModel(json_config)
|
||||||
|
m, u = clip.load_sd(sd)
|
||||||
|
if len(m) > 0:
|
||||||
|
print("extra clip vision:", m)
|
||||||
|
u = set(u)
|
||||||
|
keys = list(sd.keys())
|
||||||
|
for k in keys:
|
||||||
|
if k not in u:
|
||||||
|
t = sd.pop(k)
|
||||||
|
del t
|
||||||
|
return clip
|
||||||
|
|
||||||
|
def load(ckpt_path):
|
||||||
|
sd = load_torch_file(ckpt_path)
|
||||||
|
if "visual.transformer.resblocks.0.attn.in_proj_weight" in sd:
|
||||||
|
return load_clipvision_from_sd(sd, prefix="visual.", convert_keys=True)
|
||||||
|
else:
|
||||||
|
return load_clipvision_from_sd(sd)
|
18
ldm_patched/modules/clip_vision_config_g.json
Normal file
18
ldm_patched/modules/clip_vision_config_g.json
Normal file
@ -0,0 +1,18 @@
|
|||||||
|
{
|
||||||
|
"attention_dropout": 0.0,
|
||||||
|
"dropout": 0.0,
|
||||||
|
"hidden_act": "gelu",
|
||||||
|
"hidden_size": 1664,
|
||||||
|
"image_size": 224,
|
||||||
|
"initializer_factor": 1.0,
|
||||||
|
"initializer_range": 0.02,
|
||||||
|
"intermediate_size": 8192,
|
||||||
|
"layer_norm_eps": 1e-05,
|
||||||
|
"model_type": "clip_vision_model",
|
||||||
|
"num_attention_heads": 16,
|
||||||
|
"num_channels": 3,
|
||||||
|
"num_hidden_layers": 48,
|
||||||
|
"patch_size": 14,
|
||||||
|
"projection_dim": 1280,
|
||||||
|
"torch_dtype": "float32"
|
||||||
|
}
|
18
ldm_patched/modules/clip_vision_config_h.json
Normal file
18
ldm_patched/modules/clip_vision_config_h.json
Normal file
@ -0,0 +1,18 @@
|
|||||||
|
{
|
||||||
|
"attention_dropout": 0.0,
|
||||||
|
"dropout": 0.0,
|
||||||
|
"hidden_act": "gelu",
|
||||||
|
"hidden_size": 1280,
|
||||||
|
"image_size": 224,
|
||||||
|
"initializer_factor": 1.0,
|
||||||
|
"initializer_range": 0.02,
|
||||||
|
"intermediate_size": 5120,
|
||||||
|
"layer_norm_eps": 1e-05,
|
||||||
|
"model_type": "clip_vision_model",
|
||||||
|
"num_attention_heads": 16,
|
||||||
|
"num_channels": 3,
|
||||||
|
"num_hidden_layers": 32,
|
||||||
|
"patch_size": 14,
|
||||||
|
"projection_dim": 1024,
|
||||||
|
"torch_dtype": "float32"
|
||||||
|
}
|
18
ldm_patched/modules/clip_vision_config_vitl.json
Normal file
18
ldm_patched/modules/clip_vision_config_vitl.json
Normal file
@ -0,0 +1,18 @@
|
|||||||
|
{
|
||||||
|
"attention_dropout": 0.0,
|
||||||
|
"dropout": 0.0,
|
||||||
|
"hidden_act": "quick_gelu",
|
||||||
|
"hidden_size": 1024,
|
||||||
|
"image_size": 224,
|
||||||
|
"initializer_factor": 1.0,
|
||||||
|
"initializer_range": 0.02,
|
||||||
|
"intermediate_size": 4096,
|
||||||
|
"layer_norm_eps": 1e-05,
|
||||||
|
"model_type": "clip_vision_model",
|
||||||
|
"num_attention_heads": 16,
|
||||||
|
"num_channels": 3,
|
||||||
|
"num_hidden_layers": 24,
|
||||||
|
"patch_size": 14,
|
||||||
|
"projection_dim": 768,
|
||||||
|
"torch_dtype": "float32"
|
||||||
|
}
|
79
ldm_patched/modules/conds.py
Normal file
79
ldm_patched/modules/conds.py
Normal file
@ -0,0 +1,79 @@
|
|||||||
|
import enum
|
||||||
|
import torch
|
||||||
|
import math
|
||||||
|
import ldm_patched.modules.utils
|
||||||
|
|
||||||
|
|
||||||
|
def lcm(a, b): #TODO: eventually replace by math.lcm (added in python3.9)
|
||||||
|
return abs(a*b) // math.gcd(a, b)
|
||||||
|
|
||||||
|
class CONDRegular:
|
||||||
|
def __init__(self, cond):
|
||||||
|
self.cond = cond
|
||||||
|
|
||||||
|
def _copy_with(self, cond):
|
||||||
|
return self.__class__(cond)
|
||||||
|
|
||||||
|
def process_cond(self, batch_size, device, **kwargs):
|
||||||
|
return self._copy_with(ldm_patched.modules.utils.repeat_to_batch_size(self.cond, batch_size).to(device))
|
||||||
|
|
||||||
|
def can_concat(self, other):
|
||||||
|
if self.cond.shape != other.cond.shape:
|
||||||
|
return False
|
||||||
|
return True
|
||||||
|
|
||||||
|
def concat(self, others):
|
||||||
|
conds = [self.cond]
|
||||||
|
for x in others:
|
||||||
|
conds.append(x.cond)
|
||||||
|
return torch.cat(conds)
|
||||||
|
|
||||||
|
class CONDNoiseShape(CONDRegular):
|
||||||
|
def process_cond(self, batch_size, device, area, **kwargs):
|
||||||
|
data = self.cond[:,:,area[2]:area[0] + area[2],area[3]:area[1] + area[3]]
|
||||||
|
return self._copy_with(ldm_patched.modules.utils.repeat_to_batch_size(data, batch_size).to(device))
|
||||||
|
|
||||||
|
|
||||||
|
class CONDCrossAttn(CONDRegular):
|
||||||
|
def can_concat(self, other):
|
||||||
|
s1 = self.cond.shape
|
||||||
|
s2 = other.cond.shape
|
||||||
|
if s1 != s2:
|
||||||
|
if s1[0] != s2[0] or s1[2] != s2[2]: #these 2 cases should not happen
|
||||||
|
return False
|
||||||
|
|
||||||
|
mult_min = lcm(s1[1], s2[1])
|
||||||
|
diff = mult_min // min(s1[1], s2[1])
|
||||||
|
if diff > 4: #arbitrary limit on the padding because it's probably going to impact performance negatively if it's too much
|
||||||
|
return False
|
||||||
|
return True
|
||||||
|
|
||||||
|
def concat(self, others):
|
||||||
|
conds = [self.cond]
|
||||||
|
crossattn_max_len = self.cond.shape[1]
|
||||||
|
for x in others:
|
||||||
|
c = x.cond
|
||||||
|
crossattn_max_len = lcm(crossattn_max_len, c.shape[1])
|
||||||
|
conds.append(c)
|
||||||
|
|
||||||
|
out = []
|
||||||
|
for c in conds:
|
||||||
|
if c.shape[1] < crossattn_max_len:
|
||||||
|
c = c.repeat(1, crossattn_max_len // c.shape[1], 1) #padding with repeat doesn't change result
|
||||||
|
out.append(c)
|
||||||
|
return torch.cat(out)
|
||||||
|
|
||||||
|
class CONDConstant(CONDRegular):
|
||||||
|
def __init__(self, cond):
|
||||||
|
self.cond = cond
|
||||||
|
|
||||||
|
def process_cond(self, batch_size, device, **kwargs):
|
||||||
|
return self._copy_with(self.cond)
|
||||||
|
|
||||||
|
def can_concat(self, other):
|
||||||
|
if self.cond != other.cond:
|
||||||
|
return False
|
||||||
|
return True
|
||||||
|
|
||||||
|
def concat(self, others):
|
||||||
|
return self.cond
|
517
ldm_patched/modules/controlnet.py
Normal file
517
ldm_patched/modules/controlnet.py
Normal file
@ -0,0 +1,517 @@
|
|||||||
|
import torch
|
||||||
|
import math
|
||||||
|
import os
|
||||||
|
import contextlib
|
||||||
|
import ldm_patched.modules.utils
|
||||||
|
import ldm_patched.modules.model_management
|
||||||
|
import ldm_patched.modules.model_detection
|
||||||
|
import ldm_patched.modules.model_patcher
|
||||||
|
import ldm_patched.modules.ops
|
||||||
|
|
||||||
|
import ldm_patched.controlnet.cldm
|
||||||
|
import ldm_patched.t2ia.adapter
|
||||||
|
|
||||||
|
|
||||||
|
def broadcast_image_to(tensor, target_batch_size, batched_number):
|
||||||
|
current_batch_size = tensor.shape[0]
|
||||||
|
#print(current_batch_size, target_batch_size)
|
||||||
|
if current_batch_size == 1:
|
||||||
|
return tensor
|
||||||
|
|
||||||
|
per_batch = target_batch_size // batched_number
|
||||||
|
tensor = tensor[:per_batch]
|
||||||
|
|
||||||
|
if per_batch > tensor.shape[0]:
|
||||||
|
tensor = torch.cat([tensor] * (per_batch // tensor.shape[0]) + [tensor[:(per_batch % tensor.shape[0])]], dim=0)
|
||||||
|
|
||||||
|
current_batch_size = tensor.shape[0]
|
||||||
|
if current_batch_size == target_batch_size:
|
||||||
|
return tensor
|
||||||
|
else:
|
||||||
|
return torch.cat([tensor] * batched_number, dim=0)
|
||||||
|
|
||||||
|
class ControlBase:
|
||||||
|
def __init__(self, device=None):
|
||||||
|
self.cond_hint_original = None
|
||||||
|
self.cond_hint = None
|
||||||
|
self.strength = 1.0
|
||||||
|
self.timestep_percent_range = (0.0, 1.0)
|
||||||
|
self.global_average_pooling = False
|
||||||
|
self.timestep_range = None
|
||||||
|
|
||||||
|
if device is None:
|
||||||
|
device = ldm_patched.modules.model_management.get_torch_device()
|
||||||
|
self.device = device
|
||||||
|
self.previous_controlnet = None
|
||||||
|
|
||||||
|
def set_cond_hint(self, cond_hint, strength=1.0, timestep_percent_range=(0.0, 1.0)):
|
||||||
|
self.cond_hint_original = cond_hint
|
||||||
|
self.strength = strength
|
||||||
|
self.timestep_percent_range = timestep_percent_range
|
||||||
|
return self
|
||||||
|
|
||||||
|
def pre_run(self, model, percent_to_timestep_function):
|
||||||
|
self.timestep_range = (percent_to_timestep_function(self.timestep_percent_range[0]), percent_to_timestep_function(self.timestep_percent_range[1]))
|
||||||
|
if self.previous_controlnet is not None:
|
||||||
|
self.previous_controlnet.pre_run(model, percent_to_timestep_function)
|
||||||
|
|
||||||
|
def set_previous_controlnet(self, controlnet):
|
||||||
|
self.previous_controlnet = controlnet
|
||||||
|
return self
|
||||||
|
|
||||||
|
def cleanup(self):
|
||||||
|
if self.previous_controlnet is not None:
|
||||||
|
self.previous_controlnet.cleanup()
|
||||||
|
if self.cond_hint is not None:
|
||||||
|
del self.cond_hint
|
||||||
|
self.cond_hint = None
|
||||||
|
self.timestep_range = None
|
||||||
|
|
||||||
|
def get_models(self):
|
||||||
|
out = []
|
||||||
|
if self.previous_controlnet is not None:
|
||||||
|
out += self.previous_controlnet.get_models()
|
||||||
|
return out
|
||||||
|
|
||||||
|
def copy_to(self, c):
|
||||||
|
c.cond_hint_original = self.cond_hint_original
|
||||||
|
c.strength = self.strength
|
||||||
|
c.timestep_percent_range = self.timestep_percent_range
|
||||||
|
c.global_average_pooling = self.global_average_pooling
|
||||||
|
|
||||||
|
def inference_memory_requirements(self, dtype):
|
||||||
|
if self.previous_controlnet is not None:
|
||||||
|
return self.previous_controlnet.inference_memory_requirements(dtype)
|
||||||
|
return 0
|
||||||
|
|
||||||
|
def control_merge(self, control_input, control_output, control_prev, output_dtype):
|
||||||
|
out = {'input':[], 'middle':[], 'output': []}
|
||||||
|
|
||||||
|
if control_input is not None:
|
||||||
|
for i in range(len(control_input)):
|
||||||
|
key = 'input'
|
||||||
|
x = control_input[i]
|
||||||
|
if x is not None:
|
||||||
|
x *= self.strength
|
||||||
|
if x.dtype != output_dtype:
|
||||||
|
x = x.to(output_dtype)
|
||||||
|
out[key].insert(0, x)
|
||||||
|
|
||||||
|
if control_output is not None:
|
||||||
|
for i in range(len(control_output)):
|
||||||
|
if i == (len(control_output) - 1):
|
||||||
|
key = 'middle'
|
||||||
|
index = 0
|
||||||
|
else:
|
||||||
|
key = 'output'
|
||||||
|
index = i
|
||||||
|
x = control_output[i]
|
||||||
|
if x is not None:
|
||||||
|
if self.global_average_pooling:
|
||||||
|
x = torch.mean(x, dim=(2, 3), keepdim=True).repeat(1, 1, x.shape[2], x.shape[3])
|
||||||
|
|
||||||
|
x *= self.strength
|
||||||
|
if x.dtype != output_dtype:
|
||||||
|
x = x.to(output_dtype)
|
||||||
|
|
||||||
|
out[key].append(x)
|
||||||
|
if control_prev is not None:
|
||||||
|
for x in ['input', 'middle', 'output']:
|
||||||
|
o = out[x]
|
||||||
|
for i in range(len(control_prev[x])):
|
||||||
|
prev_val = control_prev[x][i]
|
||||||
|
if i >= len(o):
|
||||||
|
o.append(prev_val)
|
||||||
|
elif prev_val is not None:
|
||||||
|
if o[i] is None:
|
||||||
|
o[i] = prev_val
|
||||||
|
else:
|
||||||
|
if o[i].shape[0] < prev_val.shape[0]:
|
||||||
|
o[i] = prev_val + o[i]
|
||||||
|
else:
|
||||||
|
o[i] += prev_val
|
||||||
|
return out
|
||||||
|
|
||||||
|
class ControlNet(ControlBase):
|
||||||
|
def __init__(self, control_model, global_average_pooling=False, device=None, load_device=None, manual_cast_dtype=None):
|
||||||
|
super().__init__(device)
|
||||||
|
self.control_model = control_model
|
||||||
|
self.load_device = load_device
|
||||||
|
self.control_model_wrapped = ldm_patched.modules.model_patcher.ModelPatcher(self.control_model, load_device=load_device, offload_device=ldm_patched.modules.model_management.unet_offload_device())
|
||||||
|
self.global_average_pooling = global_average_pooling
|
||||||
|
self.model_sampling_current = None
|
||||||
|
self.manual_cast_dtype = manual_cast_dtype
|
||||||
|
|
||||||
|
def get_control(self, x_noisy, t, cond, batched_number):
|
||||||
|
control_prev = None
|
||||||
|
if self.previous_controlnet is not None:
|
||||||
|
control_prev = self.previous_controlnet.get_control(x_noisy, t, cond, batched_number)
|
||||||
|
|
||||||
|
if self.timestep_range is not None:
|
||||||
|
if t[0] > self.timestep_range[0] or t[0] < self.timestep_range[1]:
|
||||||
|
if control_prev is not None:
|
||||||
|
return control_prev
|
||||||
|
else:
|
||||||
|
return None
|
||||||
|
|
||||||
|
dtype = self.control_model.dtype
|
||||||
|
if self.manual_cast_dtype is not None:
|
||||||
|
dtype = self.manual_cast_dtype
|
||||||
|
|
||||||
|
output_dtype = x_noisy.dtype
|
||||||
|
if self.cond_hint is None or x_noisy.shape[2] * 8 != self.cond_hint.shape[2] or x_noisy.shape[3] * 8 != self.cond_hint.shape[3]:
|
||||||
|
if self.cond_hint is not None:
|
||||||
|
del self.cond_hint
|
||||||
|
self.cond_hint = None
|
||||||
|
self.cond_hint = ldm_patched.modules.utils.common_upscale(self.cond_hint_original, x_noisy.shape[3] * 8, x_noisy.shape[2] * 8, 'nearest-exact', "center").to(dtype).to(self.device)
|
||||||
|
if x_noisy.shape[0] != self.cond_hint.shape[0]:
|
||||||
|
self.cond_hint = broadcast_image_to(self.cond_hint, x_noisy.shape[0], batched_number)
|
||||||
|
|
||||||
|
context = cond['c_crossattn']
|
||||||
|
y = cond.get('y', None)
|
||||||
|
if y is not None:
|
||||||
|
y = y.to(dtype)
|
||||||
|
timestep = self.model_sampling_current.timestep(t)
|
||||||
|
x_noisy = self.model_sampling_current.calculate_input(t, x_noisy)
|
||||||
|
|
||||||
|
control = self.control_model(x=x_noisy.to(dtype), hint=self.cond_hint, timesteps=timestep.float(), context=context.to(dtype), y=y)
|
||||||
|
return self.control_merge(None, control, control_prev, output_dtype)
|
||||||
|
|
||||||
|
def copy(self):
|
||||||
|
c = ControlNet(self.control_model, global_average_pooling=self.global_average_pooling, load_device=self.load_device, manual_cast_dtype=self.manual_cast_dtype)
|
||||||
|
self.copy_to(c)
|
||||||
|
return c
|
||||||
|
|
||||||
|
def get_models(self):
|
||||||
|
out = super().get_models()
|
||||||
|
out.append(self.control_model_wrapped)
|
||||||
|
return out
|
||||||
|
|
||||||
|
def pre_run(self, model, percent_to_timestep_function):
|
||||||
|
super().pre_run(model, percent_to_timestep_function)
|
||||||
|
self.model_sampling_current = model.model_sampling
|
||||||
|
|
||||||
|
def cleanup(self):
|
||||||
|
self.model_sampling_current = None
|
||||||
|
super().cleanup()
|
||||||
|
|
||||||
|
class ControlLoraOps:
|
||||||
|
class Linear(torch.nn.Module):
|
||||||
|
def __init__(self, in_features: int, out_features: int, bias: bool = True,
|
||||||
|
device=None, dtype=None) -> None:
|
||||||
|
factory_kwargs = {'device': device, 'dtype': dtype}
|
||||||
|
super().__init__()
|
||||||
|
self.in_features = in_features
|
||||||
|
self.out_features = out_features
|
||||||
|
self.weight = None
|
||||||
|
self.up = None
|
||||||
|
self.down = None
|
||||||
|
self.bias = None
|
||||||
|
|
||||||
|
def forward(self, input):
|
||||||
|
weight, bias = ldm_patched.modules.ops.cast_bias_weight(self, input)
|
||||||
|
if self.up is not None:
|
||||||
|
return torch.nn.functional.linear(input, weight + (torch.mm(self.up.flatten(start_dim=1), self.down.flatten(start_dim=1))).reshape(self.weight.shape).type(input.dtype), bias)
|
||||||
|
else:
|
||||||
|
return torch.nn.functional.linear(input, weight, bias)
|
||||||
|
|
||||||
|
class Conv2d(torch.nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
in_channels,
|
||||||
|
out_channels,
|
||||||
|
kernel_size,
|
||||||
|
stride=1,
|
||||||
|
padding=0,
|
||||||
|
dilation=1,
|
||||||
|
groups=1,
|
||||||
|
bias=True,
|
||||||
|
padding_mode='zeros',
|
||||||
|
device=None,
|
||||||
|
dtype=None
|
||||||
|
):
|
||||||
|
super().__init__()
|
||||||
|
self.in_channels = in_channels
|
||||||
|
self.out_channels = out_channels
|
||||||
|
self.kernel_size = kernel_size
|
||||||
|
self.stride = stride
|
||||||
|
self.padding = padding
|
||||||
|
self.dilation = dilation
|
||||||
|
self.transposed = False
|
||||||
|
self.output_padding = 0
|
||||||
|
self.groups = groups
|
||||||
|
self.padding_mode = padding_mode
|
||||||
|
|
||||||
|
self.weight = None
|
||||||
|
self.bias = None
|
||||||
|
self.up = None
|
||||||
|
self.down = None
|
||||||
|
|
||||||
|
|
||||||
|
def forward(self, input):
|
||||||
|
weight, bias = ldm_patched.modules.ops.cast_bias_weight(self, input)
|
||||||
|
if self.up is not None:
|
||||||
|
return torch.nn.functional.conv2d(input, weight + (torch.mm(self.up.flatten(start_dim=1), self.down.flatten(start_dim=1))).reshape(self.weight.shape).type(input.dtype), bias, self.stride, self.padding, self.dilation, self.groups)
|
||||||
|
else:
|
||||||
|
return torch.nn.functional.conv2d(input, weight, bias, self.stride, self.padding, self.dilation, self.groups)
|
||||||
|
|
||||||
|
|
||||||
|
class ControlLora(ControlNet):
|
||||||
|
def __init__(self, control_weights, global_average_pooling=False, device=None):
|
||||||
|
ControlBase.__init__(self, device)
|
||||||
|
self.control_weights = control_weights
|
||||||
|
self.global_average_pooling = global_average_pooling
|
||||||
|
|
||||||
|
def pre_run(self, model, percent_to_timestep_function):
|
||||||
|
super().pre_run(model, percent_to_timestep_function)
|
||||||
|
controlnet_config = model.model_config.unet_config.copy()
|
||||||
|
controlnet_config.pop("out_channels")
|
||||||
|
controlnet_config["hint_channels"] = self.control_weights["input_hint_block.0.weight"].shape[1]
|
||||||
|
self.manual_cast_dtype = model.manual_cast_dtype
|
||||||
|
dtype = model.get_dtype()
|
||||||
|
if self.manual_cast_dtype is None:
|
||||||
|
class control_lora_ops(ControlLoraOps, ldm_patched.modules.ops.disable_weight_init):
|
||||||
|
pass
|
||||||
|
else:
|
||||||
|
class control_lora_ops(ControlLoraOps, ldm_patched.modules.ops.manual_cast):
|
||||||
|
pass
|
||||||
|
dtype = self.manual_cast_dtype
|
||||||
|
|
||||||
|
controlnet_config["operations"] = control_lora_ops
|
||||||
|
controlnet_config["dtype"] = dtype
|
||||||
|
self.control_model = ldm_patched.controlnet.cldm.ControlNet(**controlnet_config)
|
||||||
|
self.control_model.to(ldm_patched.modules.model_management.get_torch_device())
|
||||||
|
diffusion_model = model.diffusion_model
|
||||||
|
sd = diffusion_model.state_dict()
|
||||||
|
cm = self.control_model.state_dict()
|
||||||
|
|
||||||
|
for k in sd:
|
||||||
|
weight = sd[k]
|
||||||
|
try:
|
||||||
|
ldm_patched.modules.utils.set_attr(self.control_model, k, weight)
|
||||||
|
except:
|
||||||
|
pass
|
||||||
|
|
||||||
|
for k in self.control_weights:
|
||||||
|
if k not in {"lora_controlnet"}:
|
||||||
|
ldm_patched.modules.utils.set_attr(self.control_model, k, self.control_weights[k].to(dtype).to(ldm_patched.modules.model_management.get_torch_device()))
|
||||||
|
|
||||||
|
def copy(self):
|
||||||
|
c = ControlLora(self.control_weights, global_average_pooling=self.global_average_pooling)
|
||||||
|
self.copy_to(c)
|
||||||
|
return c
|
||||||
|
|
||||||
|
def cleanup(self):
|
||||||
|
del self.control_model
|
||||||
|
self.control_model = None
|
||||||
|
super().cleanup()
|
||||||
|
|
||||||
|
def get_models(self):
|
||||||
|
out = ControlBase.get_models(self)
|
||||||
|
return out
|
||||||
|
|
||||||
|
def inference_memory_requirements(self, dtype):
|
||||||
|
return ldm_patched.modules.utils.calculate_parameters(self.control_weights) * ldm_patched.modules.model_management.dtype_size(dtype) + ControlBase.inference_memory_requirements(self, dtype)
|
||||||
|
|
||||||
|
def load_controlnet(ckpt_path, model=None):
|
||||||
|
controlnet_data = ldm_patched.modules.utils.load_torch_file(ckpt_path, safe_load=True)
|
||||||
|
if "lora_controlnet" in controlnet_data:
|
||||||
|
return ControlLora(controlnet_data)
|
||||||
|
|
||||||
|
controlnet_config = None
|
||||||
|
if "controlnet_cond_embedding.conv_in.weight" in controlnet_data: #diffusers format
|
||||||
|
unet_dtype = ldm_patched.modules.model_management.unet_dtype()
|
||||||
|
controlnet_config = ldm_patched.modules.model_detection.unet_config_from_diffusers_unet(controlnet_data, unet_dtype)
|
||||||
|
diffusers_keys = ldm_patched.modules.utils.unet_to_diffusers(controlnet_config)
|
||||||
|
diffusers_keys["controlnet_mid_block.weight"] = "middle_block_out.0.weight"
|
||||||
|
diffusers_keys["controlnet_mid_block.bias"] = "middle_block_out.0.bias"
|
||||||
|
|
||||||
|
count = 0
|
||||||
|
loop = True
|
||||||
|
while loop:
|
||||||
|
suffix = [".weight", ".bias"]
|
||||||
|
for s in suffix:
|
||||||
|
k_in = "controlnet_down_blocks.{}{}".format(count, s)
|
||||||
|
k_out = "zero_convs.{}.0{}".format(count, s)
|
||||||
|
if k_in not in controlnet_data:
|
||||||
|
loop = False
|
||||||
|
break
|
||||||
|
diffusers_keys[k_in] = k_out
|
||||||
|
count += 1
|
||||||
|
|
||||||
|
count = 0
|
||||||
|
loop = True
|
||||||
|
while loop:
|
||||||
|
suffix = [".weight", ".bias"]
|
||||||
|
for s in suffix:
|
||||||
|
if count == 0:
|
||||||
|
k_in = "controlnet_cond_embedding.conv_in{}".format(s)
|
||||||
|
else:
|
||||||
|
k_in = "controlnet_cond_embedding.blocks.{}{}".format(count - 1, s)
|
||||||
|
k_out = "input_hint_block.{}{}".format(count * 2, s)
|
||||||
|
if k_in not in controlnet_data:
|
||||||
|
k_in = "controlnet_cond_embedding.conv_out{}".format(s)
|
||||||
|
loop = False
|
||||||
|
diffusers_keys[k_in] = k_out
|
||||||
|
count += 1
|
||||||
|
|
||||||
|
new_sd = {}
|
||||||
|
for k in diffusers_keys:
|
||||||
|
if k in controlnet_data:
|
||||||
|
new_sd[diffusers_keys[k]] = controlnet_data.pop(k)
|
||||||
|
|
||||||
|
leftover_keys = controlnet_data.keys()
|
||||||
|
if len(leftover_keys) > 0:
|
||||||
|
print("leftover keys:", leftover_keys)
|
||||||
|
controlnet_data = new_sd
|
||||||
|
|
||||||
|
pth_key = 'control_model.zero_convs.0.0.weight'
|
||||||
|
pth = False
|
||||||
|
key = 'zero_convs.0.0.weight'
|
||||||
|
if pth_key in controlnet_data:
|
||||||
|
pth = True
|
||||||
|
key = pth_key
|
||||||
|
prefix = "control_model."
|
||||||
|
elif key in controlnet_data:
|
||||||
|
prefix = ""
|
||||||
|
else:
|
||||||
|
net = load_t2i_adapter(controlnet_data)
|
||||||
|
if net is None:
|
||||||
|
print("error checkpoint does not contain controlnet or t2i adapter data", ckpt_path)
|
||||||
|
return net
|
||||||
|
|
||||||
|
if controlnet_config is None:
|
||||||
|
unet_dtype = ldm_patched.modules.model_management.unet_dtype()
|
||||||
|
controlnet_config = ldm_patched.modules.model_detection.model_config_from_unet(controlnet_data, prefix, unet_dtype, True).unet_config
|
||||||
|
load_device = ldm_patched.modules.model_management.get_torch_device()
|
||||||
|
manual_cast_dtype = ldm_patched.modules.model_management.unet_manual_cast(unet_dtype, load_device)
|
||||||
|
if manual_cast_dtype is not None:
|
||||||
|
controlnet_config["operations"] = ldm_patched.modules.ops.manual_cast
|
||||||
|
controlnet_config.pop("out_channels")
|
||||||
|
controlnet_config["hint_channels"] = controlnet_data["{}input_hint_block.0.weight".format(prefix)].shape[1]
|
||||||
|
control_model = ldm_patched.controlnet.cldm.ControlNet(**controlnet_config)
|
||||||
|
|
||||||
|
if pth:
|
||||||
|
if 'difference' in controlnet_data:
|
||||||
|
if model is not None:
|
||||||
|
ldm_patched.modules.model_management.load_models_gpu([model])
|
||||||
|
model_sd = model.model_state_dict()
|
||||||
|
for x in controlnet_data:
|
||||||
|
c_m = "control_model."
|
||||||
|
if x.startswith(c_m):
|
||||||
|
sd_key = "diffusion_model.{}".format(x[len(c_m):])
|
||||||
|
if sd_key in model_sd:
|
||||||
|
cd = controlnet_data[x]
|
||||||
|
cd += model_sd[sd_key].type(cd.dtype).to(cd.device)
|
||||||
|
else:
|
||||||
|
print("WARNING: Loaded a diff controlnet without a model. It will very likely not work.")
|
||||||
|
|
||||||
|
class WeightsLoader(torch.nn.Module):
|
||||||
|
pass
|
||||||
|
w = WeightsLoader()
|
||||||
|
w.control_model = control_model
|
||||||
|
missing, unexpected = w.load_state_dict(controlnet_data, strict=False)
|
||||||
|
else:
|
||||||
|
missing, unexpected = control_model.load_state_dict(controlnet_data, strict=False)
|
||||||
|
print(missing, unexpected)
|
||||||
|
|
||||||
|
global_average_pooling = False
|
||||||
|
filename = os.path.splitext(ckpt_path)[0]
|
||||||
|
if filename.endswith("_shuffle") or filename.endswith("_shuffle_fp16"): #TODO: smarter way of enabling global_average_pooling
|
||||||
|
global_average_pooling = True
|
||||||
|
|
||||||
|
control = ControlNet(control_model, global_average_pooling=global_average_pooling, load_device=load_device, manual_cast_dtype=manual_cast_dtype)
|
||||||
|
return control
|
||||||
|
|
||||||
|
class T2IAdapter(ControlBase):
|
||||||
|
def __init__(self, t2i_model, channels_in, device=None):
|
||||||
|
super().__init__(device)
|
||||||
|
self.t2i_model = t2i_model
|
||||||
|
self.channels_in = channels_in
|
||||||
|
self.control_input = None
|
||||||
|
|
||||||
|
def scale_image_to(self, width, height):
|
||||||
|
unshuffle_amount = self.t2i_model.unshuffle_amount
|
||||||
|
width = math.ceil(width / unshuffle_amount) * unshuffle_amount
|
||||||
|
height = math.ceil(height / unshuffle_amount) * unshuffle_amount
|
||||||
|
return width, height
|
||||||
|
|
||||||
|
def get_control(self, x_noisy, t, cond, batched_number):
|
||||||
|
control_prev = None
|
||||||
|
if self.previous_controlnet is not None:
|
||||||
|
control_prev = self.previous_controlnet.get_control(x_noisy, t, cond, batched_number)
|
||||||
|
|
||||||
|
if self.timestep_range is not None:
|
||||||
|
if t[0] > self.timestep_range[0] or t[0] < self.timestep_range[1]:
|
||||||
|
if control_prev is not None:
|
||||||
|
return control_prev
|
||||||
|
else:
|
||||||
|
return None
|
||||||
|
|
||||||
|
if self.cond_hint is None or x_noisy.shape[2] * 8 != self.cond_hint.shape[2] or x_noisy.shape[3] * 8 != self.cond_hint.shape[3]:
|
||||||
|
if self.cond_hint is not None:
|
||||||
|
del self.cond_hint
|
||||||
|
self.control_input = None
|
||||||
|
self.cond_hint = None
|
||||||
|
width, height = self.scale_image_to(x_noisy.shape[3] * 8, x_noisy.shape[2] * 8)
|
||||||
|
self.cond_hint = ldm_patched.modules.utils.common_upscale(self.cond_hint_original, width, height, 'nearest-exact', "center").float().to(self.device)
|
||||||
|
if self.channels_in == 1 and self.cond_hint.shape[1] > 1:
|
||||||
|
self.cond_hint = torch.mean(self.cond_hint, 1, keepdim=True)
|
||||||
|
if x_noisy.shape[0] != self.cond_hint.shape[0]:
|
||||||
|
self.cond_hint = broadcast_image_to(self.cond_hint, x_noisy.shape[0], batched_number)
|
||||||
|
if self.control_input is None:
|
||||||
|
self.t2i_model.to(x_noisy.dtype)
|
||||||
|
self.t2i_model.to(self.device)
|
||||||
|
self.control_input = self.t2i_model(self.cond_hint.to(x_noisy.dtype))
|
||||||
|
self.t2i_model.cpu()
|
||||||
|
|
||||||
|
control_input = list(map(lambda a: None if a is None else a.clone(), self.control_input))
|
||||||
|
mid = None
|
||||||
|
if self.t2i_model.xl == True:
|
||||||
|
mid = control_input[-1:]
|
||||||
|
control_input = control_input[:-1]
|
||||||
|
return self.control_merge(control_input, mid, control_prev, x_noisy.dtype)
|
||||||
|
|
||||||
|
def copy(self):
|
||||||
|
c = T2IAdapter(self.t2i_model, self.channels_in)
|
||||||
|
self.copy_to(c)
|
||||||
|
return c
|
||||||
|
|
||||||
|
def load_t2i_adapter(t2i_data):
|
||||||
|
if 'adapter' in t2i_data:
|
||||||
|
t2i_data = t2i_data['adapter']
|
||||||
|
if 'adapter.body.0.resnets.0.block1.weight' in t2i_data: #diffusers format
|
||||||
|
prefix_replace = {}
|
||||||
|
for i in range(4):
|
||||||
|
for j in range(2):
|
||||||
|
prefix_replace["adapter.body.{}.resnets.{}.".format(i, j)] = "body.{}.".format(i * 2 + j)
|
||||||
|
prefix_replace["adapter.body.{}.".format(i, j)] = "body.{}.".format(i * 2)
|
||||||
|
prefix_replace["adapter."] = ""
|
||||||
|
t2i_data = ldm_patched.modules.utils.state_dict_prefix_replace(t2i_data, prefix_replace)
|
||||||
|
keys = t2i_data.keys()
|
||||||
|
|
||||||
|
if "body.0.in_conv.weight" in keys:
|
||||||
|
cin = t2i_data['body.0.in_conv.weight'].shape[1]
|
||||||
|
model_ad = ldm_patched.t2ia.adapter.Adapter_light(cin=cin, channels=[320, 640, 1280, 1280], nums_rb=4)
|
||||||
|
elif 'conv_in.weight' in keys:
|
||||||
|
cin = t2i_data['conv_in.weight'].shape[1]
|
||||||
|
channel = t2i_data['conv_in.weight'].shape[0]
|
||||||
|
ksize = t2i_data['body.0.block2.weight'].shape[2]
|
||||||
|
use_conv = False
|
||||||
|
down_opts = list(filter(lambda a: a.endswith("down_opt.op.weight"), keys))
|
||||||
|
if len(down_opts) > 0:
|
||||||
|
use_conv = True
|
||||||
|
xl = False
|
||||||
|
if cin == 256 or cin == 768:
|
||||||
|
xl = True
|
||||||
|
model_ad = ldm_patched.t2ia.adapter.Adapter(cin=cin, channels=[channel, channel*2, channel*4, channel*4][:4], nums_rb=2, ksize=ksize, sk=True, use_conv=use_conv, xl=xl)
|
||||||
|
else:
|
||||||
|
return None
|
||||||
|
missing, unexpected = model_ad.load_state_dict(t2i_data)
|
||||||
|
if len(missing) > 0:
|
||||||
|
print("t2i missing", missing)
|
||||||
|
|
||||||
|
if len(unexpected) > 0:
|
||||||
|
print("t2i unexpected", unexpected)
|
||||||
|
|
||||||
|
return T2IAdapter(model_ad, model_ad.input_channels)
|
261
ldm_patched/modules/diffusers_convert.py
Normal file
261
ldm_patched/modules/diffusers_convert.py
Normal file
@ -0,0 +1,261 @@
|
|||||||
|
import re
|
||||||
|
import torch
|
||||||
|
|
||||||
|
# conversion code from https://github.com/huggingface/diffusers/blob/main/scripts/convert_diffusers_to_original_stable_diffusion.py
|
||||||
|
|
||||||
|
# =================#
|
||||||
|
# UNet Conversion #
|
||||||
|
# =================#
|
||||||
|
|
||||||
|
unet_conversion_map = [
|
||||||
|
# (stable-diffusion, HF Diffusers)
|
||||||
|
("time_embed.0.weight", "time_embedding.linear_1.weight"),
|
||||||
|
("time_embed.0.bias", "time_embedding.linear_1.bias"),
|
||||||
|
("time_embed.2.weight", "time_embedding.linear_2.weight"),
|
||||||
|
("time_embed.2.bias", "time_embedding.linear_2.bias"),
|
||||||
|
("input_blocks.0.0.weight", "conv_in.weight"),
|
||||||
|
("input_blocks.0.0.bias", "conv_in.bias"),
|
||||||
|
("out.0.weight", "conv_norm_out.weight"),
|
||||||
|
("out.0.bias", "conv_norm_out.bias"),
|
||||||
|
("out.2.weight", "conv_out.weight"),
|
||||||
|
("out.2.bias", "conv_out.bias"),
|
||||||
|
]
|
||||||
|
|
||||||
|
unet_conversion_map_resnet = [
|
||||||
|
# (stable-diffusion, HF Diffusers)
|
||||||
|
("in_layers.0", "norm1"),
|
||||||
|
("in_layers.2", "conv1"),
|
||||||
|
("out_layers.0", "norm2"),
|
||||||
|
("out_layers.3", "conv2"),
|
||||||
|
("emb_layers.1", "time_emb_proj"),
|
||||||
|
("skip_connection", "conv_shortcut"),
|
||||||
|
]
|
||||||
|
|
||||||
|
unet_conversion_map_layer = []
|
||||||
|
# hardcoded number of downblocks and resnets/attentions...
|
||||||
|
# would need smarter logic for other networks.
|
||||||
|
for i in range(4):
|
||||||
|
# loop over downblocks/upblocks
|
||||||
|
|
||||||
|
for j in range(2):
|
||||||
|
# loop over resnets/attentions for downblocks
|
||||||
|
hf_down_res_prefix = f"down_blocks.{i}.resnets.{j}."
|
||||||
|
sd_down_res_prefix = f"input_blocks.{3 * i + j + 1}.0."
|
||||||
|
unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))
|
||||||
|
|
||||||
|
if i < 3:
|
||||||
|
# no attention layers in down_blocks.3
|
||||||
|
hf_down_atn_prefix = f"down_blocks.{i}.attentions.{j}."
|
||||||
|
sd_down_atn_prefix = f"input_blocks.{3 * i + j + 1}.1."
|
||||||
|
unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))
|
||||||
|
|
||||||
|
for j in range(3):
|
||||||
|
# loop over resnets/attentions for upblocks
|
||||||
|
hf_up_res_prefix = f"up_blocks.{i}.resnets.{j}."
|
||||||
|
sd_up_res_prefix = f"output_blocks.{3 * i + j}.0."
|
||||||
|
unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))
|
||||||
|
|
||||||
|
if i > 0:
|
||||||
|
# no attention layers in up_blocks.0
|
||||||
|
hf_up_atn_prefix = f"up_blocks.{i}.attentions.{j}."
|
||||||
|
sd_up_atn_prefix = f"output_blocks.{3 * i + j}.1."
|
||||||
|
unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))
|
||||||
|
|
||||||
|
if i < 3:
|
||||||
|
# no downsample in down_blocks.3
|
||||||
|
hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0.conv."
|
||||||
|
sd_downsample_prefix = f"input_blocks.{3 * (i + 1)}.0.op."
|
||||||
|
unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))
|
||||||
|
|
||||||
|
# no upsample in up_blocks.3
|
||||||
|
hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
|
||||||
|
sd_upsample_prefix = f"output_blocks.{3 * i + 2}.{1 if i == 0 else 2}."
|
||||||
|
unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))
|
||||||
|
|
||||||
|
hf_mid_atn_prefix = "mid_block.attentions.0."
|
||||||
|
sd_mid_atn_prefix = "middle_block.1."
|
||||||
|
unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))
|
||||||
|
|
||||||
|
for j in range(2):
|
||||||
|
hf_mid_res_prefix = f"mid_block.resnets.{j}."
|
||||||
|
sd_mid_res_prefix = f"middle_block.{2 * j}."
|
||||||
|
unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))
|
||||||
|
|
||||||
|
|
||||||
|
def convert_unet_state_dict(unet_state_dict):
|
||||||
|
# buyer beware: this is a *brittle* function,
|
||||||
|
# and correct output requires that all of these pieces interact in
|
||||||
|
# the exact order in which I have arranged them.
|
||||||
|
mapping = {k: k for k in unet_state_dict.keys()}
|
||||||
|
for sd_name, hf_name in unet_conversion_map:
|
||||||
|
mapping[hf_name] = sd_name
|
||||||
|
for k, v in mapping.items():
|
||||||
|
if "resnets" in k:
|
||||||
|
for sd_part, hf_part in unet_conversion_map_resnet:
|
||||||
|
v = v.replace(hf_part, sd_part)
|
||||||
|
mapping[k] = v
|
||||||
|
for k, v in mapping.items():
|
||||||
|
for sd_part, hf_part in unet_conversion_map_layer:
|
||||||
|
v = v.replace(hf_part, sd_part)
|
||||||
|
mapping[k] = v
|
||||||
|
new_state_dict = {v: unet_state_dict[k] for k, v in mapping.items()}
|
||||||
|
return new_state_dict
|
||||||
|
|
||||||
|
|
||||||
|
# ================#
|
||||||
|
# VAE Conversion #
|
||||||
|
# ================#
|
||||||
|
|
||||||
|
vae_conversion_map = [
|
||||||
|
# (stable-diffusion, HF Diffusers)
|
||||||
|
("nin_shortcut", "conv_shortcut"),
|
||||||
|
("norm_out", "conv_norm_out"),
|
||||||
|
("mid.attn_1.", "mid_block.attentions.0."),
|
||||||
|
]
|
||||||
|
|
||||||
|
for i in range(4):
|
||||||
|
# down_blocks have two resnets
|
||||||
|
for j in range(2):
|
||||||
|
hf_down_prefix = f"encoder.down_blocks.{i}.resnets.{j}."
|
||||||
|
sd_down_prefix = f"encoder.down.{i}.block.{j}."
|
||||||
|
vae_conversion_map.append((sd_down_prefix, hf_down_prefix))
|
||||||
|
|
||||||
|
if i < 3:
|
||||||
|
hf_downsample_prefix = f"down_blocks.{i}.downsamplers.0."
|
||||||
|
sd_downsample_prefix = f"down.{i}.downsample."
|
||||||
|
vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))
|
||||||
|
|
||||||
|
hf_upsample_prefix = f"up_blocks.{i}.upsamplers.0."
|
||||||
|
sd_upsample_prefix = f"up.{3 - i}.upsample."
|
||||||
|
vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix))
|
||||||
|
|
||||||
|
# up_blocks have three resnets
|
||||||
|
# also, up blocks in hf are numbered in reverse from sd
|
||||||
|
for j in range(3):
|
||||||
|
hf_up_prefix = f"decoder.up_blocks.{i}.resnets.{j}."
|
||||||
|
sd_up_prefix = f"decoder.up.{3 - i}.block.{j}."
|
||||||
|
vae_conversion_map.append((sd_up_prefix, hf_up_prefix))
|
||||||
|
|
||||||
|
# this part accounts for mid blocks in both the encoder and the decoder
|
||||||
|
for i in range(2):
|
||||||
|
hf_mid_res_prefix = f"mid_block.resnets.{i}."
|
||||||
|
sd_mid_res_prefix = f"mid.block_{i + 1}."
|
||||||
|
vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))
|
||||||
|
|
||||||
|
vae_conversion_map_attn = [
|
||||||
|
# (stable-diffusion, HF Diffusers)
|
||||||
|
("norm.", "group_norm."),
|
||||||
|
("q.", "query."),
|
||||||
|
("k.", "key."),
|
||||||
|
("v.", "value."),
|
||||||
|
("q.", "to_q."),
|
||||||
|
("k.", "to_k."),
|
||||||
|
("v.", "to_v."),
|
||||||
|
("proj_out.", "to_out.0."),
|
||||||
|
("proj_out.", "proj_attn."),
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
def reshape_weight_for_sd(w):
|
||||||
|
# convert HF linear weights to SD conv2d weights
|
||||||
|
return w.reshape(*w.shape, 1, 1)
|
||||||
|
|
||||||
|
|
||||||
|
def convert_vae_state_dict(vae_state_dict):
|
||||||
|
mapping = {k: k for k in vae_state_dict.keys()}
|
||||||
|
for k, v in mapping.items():
|
||||||
|
for sd_part, hf_part in vae_conversion_map:
|
||||||
|
v = v.replace(hf_part, sd_part)
|
||||||
|
mapping[k] = v
|
||||||
|
for k, v in mapping.items():
|
||||||
|
if "attentions" in k:
|
||||||
|
for sd_part, hf_part in vae_conversion_map_attn:
|
||||||
|
v = v.replace(hf_part, sd_part)
|
||||||
|
mapping[k] = v
|
||||||
|
new_state_dict = {v: vae_state_dict[k] for k, v in mapping.items()}
|
||||||
|
weights_to_convert = ["q", "k", "v", "proj_out"]
|
||||||
|
for k, v in new_state_dict.items():
|
||||||
|
for weight_name in weights_to_convert:
|
||||||
|
if f"mid.attn_1.{weight_name}.weight" in k:
|
||||||
|
print(f"Reshaping {k} for SD format")
|
||||||
|
new_state_dict[k] = reshape_weight_for_sd(v)
|
||||||
|
return new_state_dict
|
||||||
|
|
||||||
|
|
||||||
|
# =========================#
|
||||||
|
# Text Encoder Conversion #
|
||||||
|
# =========================#
|
||||||
|
|
||||||
|
|
||||||
|
textenc_conversion_lst = [
|
||||||
|
# (stable-diffusion, HF Diffusers)
|
||||||
|
("resblocks.", "text_model.encoder.layers."),
|
||||||
|
("ln_1", "layer_norm1"),
|
||||||
|
("ln_2", "layer_norm2"),
|
||||||
|
(".c_fc.", ".fc1."),
|
||||||
|
(".c_proj.", ".fc2."),
|
||||||
|
(".attn", ".self_attn"),
|
||||||
|
("ln_final.", "transformer.text_model.final_layer_norm."),
|
||||||
|
("token_embedding.weight", "transformer.text_model.embeddings.token_embedding.weight"),
|
||||||
|
("positional_embedding", "transformer.text_model.embeddings.position_embedding.weight"),
|
||||||
|
]
|
||||||
|
protected = {re.escape(x[1]): x[0] for x in textenc_conversion_lst}
|
||||||
|
textenc_pattern = re.compile("|".join(protected.keys()))
|
||||||
|
|
||||||
|
# Ordering is from https://github.com/pytorch/pytorch/blob/master/test/cpp/api/modules.cpp
|
||||||
|
code2idx = {"q": 0, "k": 1, "v": 2}
|
||||||
|
|
||||||
|
|
||||||
|
def convert_text_enc_state_dict_v20(text_enc_dict, prefix=""):
|
||||||
|
new_state_dict = {}
|
||||||
|
capture_qkv_weight = {}
|
||||||
|
capture_qkv_bias = {}
|
||||||
|
for k, v in text_enc_dict.items():
|
||||||
|
if not k.startswith(prefix):
|
||||||
|
continue
|
||||||
|
if (
|
||||||
|
k.endswith(".self_attn.q_proj.weight")
|
||||||
|
or k.endswith(".self_attn.k_proj.weight")
|
||||||
|
or k.endswith(".self_attn.v_proj.weight")
|
||||||
|
):
|
||||||
|
k_pre = k[: -len(".q_proj.weight")]
|
||||||
|
k_code = k[-len("q_proj.weight")]
|
||||||
|
if k_pre not in capture_qkv_weight:
|
||||||
|
capture_qkv_weight[k_pre] = [None, None, None]
|
||||||
|
capture_qkv_weight[k_pre][code2idx[k_code]] = v
|
||||||
|
continue
|
||||||
|
|
||||||
|
if (
|
||||||
|
k.endswith(".self_attn.q_proj.bias")
|
||||||
|
or k.endswith(".self_attn.k_proj.bias")
|
||||||
|
or k.endswith(".self_attn.v_proj.bias")
|
||||||
|
):
|
||||||
|
k_pre = k[: -len(".q_proj.bias")]
|
||||||
|
k_code = k[-len("q_proj.bias")]
|
||||||
|
if k_pre not in capture_qkv_bias:
|
||||||
|
capture_qkv_bias[k_pre] = [None, None, None]
|
||||||
|
capture_qkv_bias[k_pre][code2idx[k_code]] = v
|
||||||
|
continue
|
||||||
|
|
||||||
|
relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k)
|
||||||
|
new_state_dict[relabelled_key] = v
|
||||||
|
|
||||||
|
for k_pre, tensors in capture_qkv_weight.items():
|
||||||
|
if None in tensors:
|
||||||
|
raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing")
|
||||||
|
relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k_pre)
|
||||||
|
new_state_dict[relabelled_key + ".in_proj_weight"] = torch.cat(tensors)
|
||||||
|
|
||||||
|
for k_pre, tensors in capture_qkv_bias.items():
|
||||||
|
if None in tensors:
|
||||||
|
raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing")
|
||||||
|
relabelled_key = textenc_pattern.sub(lambda m: protected[re.escape(m.group(0))], k_pre)
|
||||||
|
new_state_dict[relabelled_key + ".in_proj_bias"] = torch.cat(tensors)
|
||||||
|
|
||||||
|
return new_state_dict
|
||||||
|
|
||||||
|
|
||||||
|
def convert_text_enc_state_dict(text_enc_dict):
|
||||||
|
return text_enc_dict
|
||||||
|
|
||||||
|
|
37
ldm_patched/modules/diffusers_load.py
Normal file
37
ldm_patched/modules/diffusers_load.py
Normal file
@ -0,0 +1,37 @@
|
|||||||
|
import json
|
||||||
|
import os
|
||||||
|
|
||||||
|
import ldm_patched.modules.sd
|
||||||
|
|
||||||
|
def first_file(path, filenames):
|
||||||
|
for f in filenames:
|
||||||
|
p = os.path.join(path, f)
|
||||||
|
if os.path.exists(p):
|
||||||
|
return p
|
||||||
|
return None
|
||||||
|
|
||||||
|
def load_diffusers(model_path, output_vae=True, output_clip=True, embedding_directory=None):
|
||||||
|
diffusion_model_names = ["diffusion_pytorch_model.fp16.safetensors", "diffusion_pytorch_model.safetensors", "diffusion_pytorch_model.fp16.bin", "diffusion_pytorch_model.bin"]
|
||||||
|
unet_path = first_file(os.path.join(model_path, "unet"), diffusion_model_names)
|
||||||
|
vae_path = first_file(os.path.join(model_path, "vae"), diffusion_model_names)
|
||||||
|
|
||||||
|
text_encoder_model_names = ["model.fp16.safetensors", "model.safetensors", "pytorch_model.fp16.bin", "pytorch_model.bin"]
|
||||||
|
text_encoder1_path = first_file(os.path.join(model_path, "text_encoder"), text_encoder_model_names)
|
||||||
|
text_encoder2_path = first_file(os.path.join(model_path, "text_encoder_2"), text_encoder_model_names)
|
||||||
|
|
||||||
|
text_encoder_paths = [text_encoder1_path]
|
||||||
|
if text_encoder2_path is not None:
|
||||||
|
text_encoder_paths.append(text_encoder2_path)
|
||||||
|
|
||||||
|
unet = ldm_patched.modules.sd.load_unet(unet_path)
|
||||||
|
|
||||||
|
clip = None
|
||||||
|
if output_clip:
|
||||||
|
clip = ldm_patched.modules.sd.load_clip(text_encoder_paths, embedding_directory=embedding_directory)
|
||||||
|
|
||||||
|
vae = None
|
||||||
|
if output_vae:
|
||||||
|
sd = ldm_patched.modules.utils.load_torch_file(vae_path)
|
||||||
|
vae = ldm_patched.modules.sd.VAE(sd=sd)
|
||||||
|
|
||||||
|
return (unet, clip, vae)
|
341
ldm_patched/modules/gligen.py
Normal file
341
ldm_patched/modules/gligen.py
Normal file
@ -0,0 +1,341 @@
|
|||||||
|
import torch
|
||||||
|
from torch import nn, einsum
|
||||||
|
from ldm_patched.ldm.modules.attention import CrossAttention
|
||||||
|
from inspect import isfunction
|
||||||
|
|
||||||
|
|
||||||
|
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() if isfunction(d) else d
|
||||||
|
|
||||||
|
|
||||||
|
# feedforward
|
||||||
|
class GEGLU(nn.Module):
|
||||||
|
def __init__(self, dim_in, dim_out):
|
||||||
|
super().__init__()
|
||||||
|
self.proj = nn.Linear(dim_in, dim_out * 2)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
x, gate = self.proj(x).chunk(2, dim=-1)
|
||||||
|
return x * torch.nn.functional.gelu(gate)
|
||||||
|
|
||||||
|
|
||||||
|
class FeedForward(nn.Module):
|
||||||
|
def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.):
|
||||||
|
super().__init__()
|
||||||
|
inner_dim = int(dim * mult)
|
||||||
|
dim_out = default(dim_out, dim)
|
||||||
|
project_in = nn.Sequential(
|
||||||
|
nn.Linear(dim, inner_dim),
|
||||||
|
nn.GELU()
|
||||||
|
) if not glu else GEGLU(dim, inner_dim)
|
||||||
|
|
||||||
|
self.net = nn.Sequential(
|
||||||
|
project_in,
|
||||||
|
nn.Dropout(dropout),
|
||||||
|
nn.Linear(inner_dim, dim_out)
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
return self.net(x)
|
||||||
|
|
||||||
|
|
||||||
|
class GatedCrossAttentionDense(nn.Module):
|
||||||
|
def __init__(self, query_dim, context_dim, n_heads, d_head):
|
||||||
|
super().__init__()
|
||||||
|
|
||||||
|
self.attn = CrossAttention(
|
||||||
|
query_dim=query_dim,
|
||||||
|
context_dim=context_dim,
|
||||||
|
heads=n_heads,
|
||||||
|
dim_head=d_head)
|
||||||
|
self.ff = FeedForward(query_dim, glu=True)
|
||||||
|
|
||||||
|
self.norm1 = nn.LayerNorm(query_dim)
|
||||||
|
self.norm2 = nn.LayerNorm(query_dim)
|
||||||
|
|
||||||
|
self.register_parameter('alpha_attn', nn.Parameter(torch.tensor(0.)))
|
||||||
|
self.register_parameter('alpha_dense', nn.Parameter(torch.tensor(0.)))
|
||||||
|
|
||||||
|
# this can be useful: we can externally change magnitude of tanh(alpha)
|
||||||
|
# for example, when it is set to 0, then the entire model is same as
|
||||||
|
# original one
|
||||||
|
self.scale = 1
|
||||||
|
|
||||||
|
def forward(self, x, objs):
|
||||||
|
|
||||||
|
x = x + self.scale * \
|
||||||
|
torch.tanh(self.alpha_attn) * self.attn(self.norm1(x), objs, objs)
|
||||||
|
x = x + self.scale * \
|
||||||
|
torch.tanh(self.alpha_dense) * self.ff(self.norm2(x))
|
||||||
|
|
||||||
|
return x
|
||||||
|
|
||||||
|
|
||||||
|
class GatedSelfAttentionDense(nn.Module):
|
||||||
|
def __init__(self, query_dim, context_dim, n_heads, d_head):
|
||||||
|
super().__init__()
|
||||||
|
|
||||||
|
# we need a linear projection since we need cat visual feature and obj
|
||||||
|
# feature
|
||||||
|
self.linear = nn.Linear(context_dim, query_dim)
|
||||||
|
|
||||||
|
self.attn = CrossAttention(
|
||||||
|
query_dim=query_dim,
|
||||||
|
context_dim=query_dim,
|
||||||
|
heads=n_heads,
|
||||||
|
dim_head=d_head)
|
||||||
|
self.ff = FeedForward(query_dim, glu=True)
|
||||||
|
|
||||||
|
self.norm1 = nn.LayerNorm(query_dim)
|
||||||
|
self.norm2 = nn.LayerNorm(query_dim)
|
||||||
|
|
||||||
|
self.register_parameter('alpha_attn', nn.Parameter(torch.tensor(0.)))
|
||||||
|
self.register_parameter('alpha_dense', nn.Parameter(torch.tensor(0.)))
|
||||||
|
|
||||||
|
# this can be useful: we can externally change magnitude of tanh(alpha)
|
||||||
|
# for example, when it is set to 0, then the entire model is same as
|
||||||
|
# original one
|
||||||
|
self.scale = 1
|
||||||
|
|
||||||
|
def forward(self, x, objs):
|
||||||
|
|
||||||
|
N_visual = x.shape[1]
|
||||||
|
objs = self.linear(objs)
|
||||||
|
|
||||||
|
x = x + self.scale * torch.tanh(self.alpha_attn) * self.attn(
|
||||||
|
self.norm1(torch.cat([x, objs], dim=1)))[:, 0:N_visual, :]
|
||||||
|
x = x + self.scale * \
|
||||||
|
torch.tanh(self.alpha_dense) * self.ff(self.norm2(x))
|
||||||
|
|
||||||
|
return x
|
||||||
|
|
||||||
|
|
||||||
|
class GatedSelfAttentionDense2(nn.Module):
|
||||||
|
def __init__(self, query_dim, context_dim, n_heads, d_head):
|
||||||
|
super().__init__()
|
||||||
|
|
||||||
|
# we need a linear projection since we need cat visual feature and obj
|
||||||
|
# feature
|
||||||
|
self.linear = nn.Linear(context_dim, query_dim)
|
||||||
|
|
||||||
|
self.attn = CrossAttention(
|
||||||
|
query_dim=query_dim, context_dim=query_dim, dim_head=d_head)
|
||||||
|
self.ff = FeedForward(query_dim, glu=True)
|
||||||
|
|
||||||
|
self.norm1 = nn.LayerNorm(query_dim)
|
||||||
|
self.norm2 = nn.LayerNorm(query_dim)
|
||||||
|
|
||||||
|
self.register_parameter('alpha_attn', nn.Parameter(torch.tensor(0.)))
|
||||||
|
self.register_parameter('alpha_dense', nn.Parameter(torch.tensor(0.)))
|
||||||
|
|
||||||
|
# this can be useful: we can externally change magnitude of tanh(alpha)
|
||||||
|
# for example, when it is set to 0, then the entire model is same as
|
||||||
|
# original one
|
||||||
|
self.scale = 1
|
||||||
|
|
||||||
|
def forward(self, x, objs):
|
||||||
|
|
||||||
|
B, N_visual, _ = x.shape
|
||||||
|
B, N_ground, _ = objs.shape
|
||||||
|
|
||||||
|
objs = self.linear(objs)
|
||||||
|
|
||||||
|
# sanity check
|
||||||
|
size_v = math.sqrt(N_visual)
|
||||||
|
size_g = math.sqrt(N_ground)
|
||||||
|
assert int(size_v) == size_v, "Visual tokens must be square rootable"
|
||||||
|
assert int(size_g) == size_g, "Grounding tokens must be square rootable"
|
||||||
|
size_v = int(size_v)
|
||||||
|
size_g = int(size_g)
|
||||||
|
|
||||||
|
# select grounding token and resize it to visual token size as residual
|
||||||
|
out = self.attn(self.norm1(torch.cat([x, objs], dim=1)))[
|
||||||
|
:, N_visual:, :]
|
||||||
|
out = out.permute(0, 2, 1).reshape(B, -1, size_g, size_g)
|
||||||
|
out = torch.nn.functional.interpolate(
|
||||||
|
out, (size_v, size_v), mode='bicubic')
|
||||||
|
residual = out.reshape(B, -1, N_visual).permute(0, 2, 1)
|
||||||
|
|
||||||
|
# add residual to visual feature
|
||||||
|
x = x + self.scale * torch.tanh(self.alpha_attn) * residual
|
||||||
|
x = x + self.scale * \
|
||||||
|
torch.tanh(self.alpha_dense) * self.ff(self.norm2(x))
|
||||||
|
|
||||||
|
return x
|
||||||
|
|
||||||
|
|
||||||
|
class FourierEmbedder():
|
||||||
|
def __init__(self, num_freqs=64, temperature=100):
|
||||||
|
|
||||||
|
self.num_freqs = num_freqs
|
||||||
|
self.temperature = temperature
|
||||||
|
self.freq_bands = temperature ** (torch.arange(num_freqs) / num_freqs)
|
||||||
|
|
||||||
|
@torch.no_grad()
|
||||||
|
def __call__(self, x, cat_dim=-1):
|
||||||
|
"x: arbitrary shape of tensor. dim: cat dim"
|
||||||
|
out = []
|
||||||
|
for freq in self.freq_bands:
|
||||||
|
out.append(torch.sin(freq * x))
|
||||||
|
out.append(torch.cos(freq * x))
|
||||||
|
return torch.cat(out, cat_dim)
|
||||||
|
|
||||||
|
|
||||||
|
class PositionNet(nn.Module):
|
||||||
|
def __init__(self, in_dim, out_dim, fourier_freqs=8):
|
||||||
|
super().__init__()
|
||||||
|
self.in_dim = in_dim
|
||||||
|
self.out_dim = out_dim
|
||||||
|
|
||||||
|
self.fourier_embedder = FourierEmbedder(num_freqs=fourier_freqs)
|
||||||
|
self.position_dim = fourier_freqs * 2 * 4 # 2 is sin&cos, 4 is xyxy
|
||||||
|
|
||||||
|
self.linears = nn.Sequential(
|
||||||
|
nn.Linear(self.in_dim + self.position_dim, 512),
|
||||||
|
nn.SiLU(),
|
||||||
|
nn.Linear(512, 512),
|
||||||
|
nn.SiLU(),
|
||||||
|
nn.Linear(512, out_dim),
|
||||||
|
)
|
||||||
|
|
||||||
|
self.null_positive_feature = torch.nn.Parameter(
|
||||||
|
torch.zeros([self.in_dim]))
|
||||||
|
self.null_position_feature = torch.nn.Parameter(
|
||||||
|
torch.zeros([self.position_dim]))
|
||||||
|
|
||||||
|
def forward(self, boxes, masks, positive_embeddings):
|
||||||
|
B, N, _ = boxes.shape
|
||||||
|
dtype = self.linears[0].weight.dtype
|
||||||
|
masks = masks.unsqueeze(-1).to(dtype)
|
||||||
|
positive_embeddings = positive_embeddings.to(dtype)
|
||||||
|
|
||||||
|
# embedding position (it may includes padding as placeholder)
|
||||||
|
xyxy_embedding = self.fourier_embedder(boxes.to(dtype)) # B*N*4 --> B*N*C
|
||||||
|
|
||||||
|
# learnable null embedding
|
||||||
|
positive_null = self.null_positive_feature.view(1, 1, -1)
|
||||||
|
xyxy_null = self.null_position_feature.view(1, 1, -1)
|
||||||
|
|
||||||
|
# replace padding with learnable null embedding
|
||||||
|
positive_embeddings = positive_embeddings * \
|
||||||
|
masks + (1 - masks) * positive_null
|
||||||
|
xyxy_embedding = xyxy_embedding * masks + (1 - masks) * xyxy_null
|
||||||
|
|
||||||
|
objs = self.linears(
|
||||||
|
torch.cat([positive_embeddings, xyxy_embedding], dim=-1))
|
||||||
|
assert objs.shape == torch.Size([B, N, self.out_dim])
|
||||||
|
return objs
|
||||||
|
|
||||||
|
|
||||||
|
class Gligen(nn.Module):
|
||||||
|
def __init__(self, modules, position_net, key_dim):
|
||||||
|
super().__init__()
|
||||||
|
self.module_list = nn.ModuleList(modules)
|
||||||
|
self.position_net = position_net
|
||||||
|
self.key_dim = key_dim
|
||||||
|
self.max_objs = 30
|
||||||
|
self.current_device = torch.device("cpu")
|
||||||
|
|
||||||
|
def _set_position(self, boxes, masks, positive_embeddings):
|
||||||
|
objs = self.position_net(boxes, masks, positive_embeddings)
|
||||||
|
def func(x, extra_options):
|
||||||
|
key = extra_options["transformer_index"]
|
||||||
|
module = self.module_list[key]
|
||||||
|
return module(x, objs)
|
||||||
|
return func
|
||||||
|
|
||||||
|
def set_position(self, latent_image_shape, position_params, device):
|
||||||
|
batch, c, h, w = latent_image_shape
|
||||||
|
masks = torch.zeros([self.max_objs], device="cpu")
|
||||||
|
boxes = []
|
||||||
|
positive_embeddings = []
|
||||||
|
for p in position_params:
|
||||||
|
x1 = (p[4]) / w
|
||||||
|
y1 = (p[3]) / h
|
||||||
|
x2 = (p[4] + p[2]) / w
|
||||||
|
y2 = (p[3] + p[1]) / h
|
||||||
|
masks[len(boxes)] = 1.0
|
||||||
|
boxes += [torch.tensor((x1, y1, x2, y2)).unsqueeze(0)]
|
||||||
|
positive_embeddings += [p[0]]
|
||||||
|
append_boxes = []
|
||||||
|
append_conds = []
|
||||||
|
if len(boxes) < self.max_objs:
|
||||||
|
append_boxes = [torch.zeros(
|
||||||
|
[self.max_objs - len(boxes), 4], device="cpu")]
|
||||||
|
append_conds = [torch.zeros(
|
||||||
|
[self.max_objs - len(boxes), self.key_dim], device="cpu")]
|
||||||
|
|
||||||
|
box_out = torch.cat(
|
||||||
|
boxes + append_boxes).unsqueeze(0).repeat(batch, 1, 1)
|
||||||
|
masks = masks.unsqueeze(0).repeat(batch, 1)
|
||||||
|
conds = torch.cat(positive_embeddings +
|
||||||
|
append_conds).unsqueeze(0).repeat(batch, 1, 1)
|
||||||
|
return self._set_position(
|
||||||
|
box_out.to(device),
|
||||||
|
masks.to(device),
|
||||||
|
conds.to(device))
|
||||||
|
|
||||||
|
def set_empty(self, latent_image_shape, device):
|
||||||
|
batch, c, h, w = latent_image_shape
|
||||||
|
masks = torch.zeros([self.max_objs], device="cpu").repeat(batch, 1)
|
||||||
|
box_out = torch.zeros([self.max_objs, 4],
|
||||||
|
device="cpu").repeat(batch, 1, 1)
|
||||||
|
conds = torch.zeros([self.max_objs, self.key_dim],
|
||||||
|
device="cpu").repeat(batch, 1, 1)
|
||||||
|
return self._set_position(
|
||||||
|
box_out.to(device),
|
||||||
|
masks.to(device),
|
||||||
|
conds.to(device))
|
||||||
|
|
||||||
|
|
||||||
|
def load_gligen(sd):
|
||||||
|
sd_k = sd.keys()
|
||||||
|
output_list = []
|
||||||
|
key_dim = 768
|
||||||
|
for a in ["input_blocks", "middle_block", "output_blocks"]:
|
||||||
|
for b in range(20):
|
||||||
|
k_temp = filter(lambda k: "{}.{}.".format(a, b)
|
||||||
|
in k and ".fuser." in k, sd_k)
|
||||||
|
k_temp = map(lambda k: (k, k.split(".fuser.")[-1]), k_temp)
|
||||||
|
|
||||||
|
n_sd = {}
|
||||||
|
for k in k_temp:
|
||||||
|
n_sd[k[1]] = sd[k[0]]
|
||||||
|
if len(n_sd) > 0:
|
||||||
|
query_dim = n_sd["linear.weight"].shape[0]
|
||||||
|
key_dim = n_sd["linear.weight"].shape[1]
|
||||||
|
|
||||||
|
if key_dim == 768: # SD1.x
|
||||||
|
n_heads = 8
|
||||||
|
d_head = query_dim // n_heads
|
||||||
|
else:
|
||||||
|
d_head = 64
|
||||||
|
n_heads = query_dim // d_head
|
||||||
|
|
||||||
|
gated = GatedSelfAttentionDense(
|
||||||
|
query_dim, key_dim, n_heads, d_head)
|
||||||
|
gated.load_state_dict(n_sd, strict=False)
|
||||||
|
output_list.append(gated)
|
||||||
|
|
||||||
|
if "position_net.null_positive_feature" in sd_k:
|
||||||
|
in_dim = sd["position_net.null_positive_feature"].shape[0]
|
||||||
|
out_dim = sd["position_net.linears.4.weight"].shape[0]
|
||||||
|
|
||||||
|
class WeightsLoader(torch.nn.Module):
|
||||||
|
pass
|
||||||
|
w = WeightsLoader()
|
||||||
|
w.position_net = PositionNet(in_dim, out_dim)
|
||||||
|
w.load_state_dict(sd, strict=False)
|
||||||
|
|
||||||
|
gligen = Gligen(output_list, w.position_net, key_dim)
|
||||||
|
return gligen
|
39
ldm_patched/modules/latent_formats.py
Normal file
39
ldm_patched/modules/latent_formats.py
Normal file
@ -0,0 +1,39 @@
|
|||||||
|
|
||||||
|
class LatentFormat:
|
||||||
|
scale_factor = 1.0
|
||||||
|
latent_rgb_factors = None
|
||||||
|
taesd_decoder_name = None
|
||||||
|
|
||||||
|
def process_in(self, latent):
|
||||||
|
return latent * self.scale_factor
|
||||||
|
|
||||||
|
def process_out(self, latent):
|
||||||
|
return latent / self.scale_factor
|
||||||
|
|
||||||
|
class SD15(LatentFormat):
|
||||||
|
def __init__(self, scale_factor=0.18215):
|
||||||
|
self.scale_factor = scale_factor
|
||||||
|
self.latent_rgb_factors = [
|
||||||
|
# R G B
|
||||||
|
[ 0.3512, 0.2297, 0.3227],
|
||||||
|
[ 0.3250, 0.4974, 0.2350],
|
||||||
|
[-0.2829, 0.1762, 0.2721],
|
||||||
|
[-0.2120, -0.2616, -0.7177]
|
||||||
|
]
|
||||||
|
self.taesd_decoder_name = "taesd_decoder"
|
||||||
|
|
||||||
|
class SDXL(LatentFormat):
|
||||||
|
def __init__(self):
|
||||||
|
self.scale_factor = 0.13025
|
||||||
|
self.latent_rgb_factors = [
|
||||||
|
# R G B
|
||||||
|
[ 0.3920, 0.4054, 0.4549],
|
||||||
|
[-0.2634, -0.0196, 0.0653],
|
||||||
|
[ 0.0568, 0.1687, -0.0755],
|
||||||
|
[-0.3112, -0.2359, -0.2076]
|
||||||
|
]
|
||||||
|
self.taesd_decoder_name = "taesdxl_decoder"
|
||||||
|
|
||||||
|
class SD_X4(LatentFormat):
|
||||||
|
def __init__(self):
|
||||||
|
self.scale_factor = 0.08333
|
224
ldm_patched/modules/lora.py
Normal file
224
ldm_patched/modules/lora.py
Normal file
@ -0,0 +1,224 @@
|
|||||||
|
import ldm_patched.modules.utils
|
||||||
|
|
||||||
|
LORA_CLIP_MAP = {
|
||||||
|
"mlp.fc1": "mlp_fc1",
|
||||||
|
"mlp.fc2": "mlp_fc2",
|
||||||
|
"self_attn.k_proj": "self_attn_k_proj",
|
||||||
|
"self_attn.q_proj": "self_attn_q_proj",
|
||||||
|
"self_attn.v_proj": "self_attn_v_proj",
|
||||||
|
"self_attn.out_proj": "self_attn_out_proj",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def load_lora(lora, to_load):
|
||||||
|
patch_dict = {}
|
||||||
|
loaded_keys = set()
|
||||||
|
for x in to_load:
|
||||||
|
alpha_name = "{}.alpha".format(x)
|
||||||
|
alpha = None
|
||||||
|
if alpha_name in lora.keys():
|
||||||
|
alpha = lora[alpha_name].item()
|
||||||
|
loaded_keys.add(alpha_name)
|
||||||
|
|
||||||
|
regular_lora = "{}.lora_up.weight".format(x)
|
||||||
|
diffusers_lora = "{}_lora.up.weight".format(x)
|
||||||
|
transformers_lora = "{}.lora_linear_layer.up.weight".format(x)
|
||||||
|
A_name = None
|
||||||
|
|
||||||
|
if regular_lora in lora.keys():
|
||||||
|
A_name = regular_lora
|
||||||
|
B_name = "{}.lora_down.weight".format(x)
|
||||||
|
mid_name = "{}.lora_mid.weight".format(x)
|
||||||
|
elif diffusers_lora in lora.keys():
|
||||||
|
A_name = diffusers_lora
|
||||||
|
B_name = "{}_lora.down.weight".format(x)
|
||||||
|
mid_name = None
|
||||||
|
elif transformers_lora in lora.keys():
|
||||||
|
A_name = transformers_lora
|
||||||
|
B_name ="{}.lora_linear_layer.down.weight".format(x)
|
||||||
|
mid_name = None
|
||||||
|
|
||||||
|
if A_name is not None:
|
||||||
|
mid = None
|
||||||
|
if mid_name is not None and mid_name in lora.keys():
|
||||||
|
mid = lora[mid_name]
|
||||||
|
loaded_keys.add(mid_name)
|
||||||
|
patch_dict[to_load[x]] = ("lora", (lora[A_name], lora[B_name], alpha, mid))
|
||||||
|
loaded_keys.add(A_name)
|
||||||
|
loaded_keys.add(B_name)
|
||||||
|
|
||||||
|
|
||||||
|
######## loha
|
||||||
|
hada_w1_a_name = "{}.hada_w1_a".format(x)
|
||||||
|
hada_w1_b_name = "{}.hada_w1_b".format(x)
|
||||||
|
hada_w2_a_name = "{}.hada_w2_a".format(x)
|
||||||
|
hada_w2_b_name = "{}.hada_w2_b".format(x)
|
||||||
|
hada_t1_name = "{}.hada_t1".format(x)
|
||||||
|
hada_t2_name = "{}.hada_t2".format(x)
|
||||||
|
if hada_w1_a_name in lora.keys():
|
||||||
|
hada_t1 = None
|
||||||
|
hada_t2 = None
|
||||||
|
if hada_t1_name in lora.keys():
|
||||||
|
hada_t1 = lora[hada_t1_name]
|
||||||
|
hada_t2 = lora[hada_t2_name]
|
||||||
|
loaded_keys.add(hada_t1_name)
|
||||||
|
loaded_keys.add(hada_t2_name)
|
||||||
|
|
||||||
|
patch_dict[to_load[x]] = ("loha", (lora[hada_w1_a_name], lora[hada_w1_b_name], alpha, lora[hada_w2_a_name], lora[hada_w2_b_name], hada_t1, hada_t2))
|
||||||
|
loaded_keys.add(hada_w1_a_name)
|
||||||
|
loaded_keys.add(hada_w1_b_name)
|
||||||
|
loaded_keys.add(hada_w2_a_name)
|
||||||
|
loaded_keys.add(hada_w2_b_name)
|
||||||
|
|
||||||
|
|
||||||
|
######## lokr
|
||||||
|
lokr_w1_name = "{}.lokr_w1".format(x)
|
||||||
|
lokr_w2_name = "{}.lokr_w2".format(x)
|
||||||
|
lokr_w1_a_name = "{}.lokr_w1_a".format(x)
|
||||||
|
lokr_w1_b_name = "{}.lokr_w1_b".format(x)
|
||||||
|
lokr_t2_name = "{}.lokr_t2".format(x)
|
||||||
|
lokr_w2_a_name = "{}.lokr_w2_a".format(x)
|
||||||
|
lokr_w2_b_name = "{}.lokr_w2_b".format(x)
|
||||||
|
|
||||||
|
lokr_w1 = None
|
||||||
|
if lokr_w1_name in lora.keys():
|
||||||
|
lokr_w1 = lora[lokr_w1_name]
|
||||||
|
loaded_keys.add(lokr_w1_name)
|
||||||
|
|
||||||
|
lokr_w2 = None
|
||||||
|
if lokr_w2_name in lora.keys():
|
||||||
|
lokr_w2 = lora[lokr_w2_name]
|
||||||
|
loaded_keys.add(lokr_w2_name)
|
||||||
|
|
||||||
|
lokr_w1_a = None
|
||||||
|
if lokr_w1_a_name in lora.keys():
|
||||||
|
lokr_w1_a = lora[lokr_w1_a_name]
|
||||||
|
loaded_keys.add(lokr_w1_a_name)
|
||||||
|
|
||||||
|
lokr_w1_b = None
|
||||||
|
if lokr_w1_b_name in lora.keys():
|
||||||
|
lokr_w1_b = lora[lokr_w1_b_name]
|
||||||
|
loaded_keys.add(lokr_w1_b_name)
|
||||||
|
|
||||||
|
lokr_w2_a = None
|
||||||
|
if lokr_w2_a_name in lora.keys():
|
||||||
|
lokr_w2_a = lora[lokr_w2_a_name]
|
||||||
|
loaded_keys.add(lokr_w2_a_name)
|
||||||
|
|
||||||
|
lokr_w2_b = None
|
||||||
|
if lokr_w2_b_name in lora.keys():
|
||||||
|
lokr_w2_b = lora[lokr_w2_b_name]
|
||||||
|
loaded_keys.add(lokr_w2_b_name)
|
||||||
|
|
||||||
|
lokr_t2 = None
|
||||||
|
if lokr_t2_name in lora.keys():
|
||||||
|
lokr_t2 = lora[lokr_t2_name]
|
||||||
|
loaded_keys.add(lokr_t2_name)
|
||||||
|
|
||||||
|
if (lokr_w1 is not None) or (lokr_w2 is not None) or (lokr_w1_a is not None) or (lokr_w2_a is not None):
|
||||||
|
patch_dict[to_load[x]] = ("lokr", (lokr_w1, lokr_w2, alpha, lokr_w1_a, lokr_w1_b, lokr_w2_a, lokr_w2_b, lokr_t2))
|
||||||
|
|
||||||
|
#glora
|
||||||
|
a1_name = "{}.a1.weight".format(x)
|
||||||
|
a2_name = "{}.a2.weight".format(x)
|
||||||
|
b1_name = "{}.b1.weight".format(x)
|
||||||
|
b2_name = "{}.b2.weight".format(x)
|
||||||
|
if a1_name in lora:
|
||||||
|
patch_dict[to_load[x]] = ("glora", (lora[a1_name], lora[a2_name], lora[b1_name], lora[b2_name], alpha))
|
||||||
|
loaded_keys.add(a1_name)
|
||||||
|
loaded_keys.add(a2_name)
|
||||||
|
loaded_keys.add(b1_name)
|
||||||
|
loaded_keys.add(b2_name)
|
||||||
|
|
||||||
|
w_norm_name = "{}.w_norm".format(x)
|
||||||
|
b_norm_name = "{}.b_norm".format(x)
|
||||||
|
w_norm = lora.get(w_norm_name, None)
|
||||||
|
b_norm = lora.get(b_norm_name, None)
|
||||||
|
|
||||||
|
if w_norm is not None:
|
||||||
|
loaded_keys.add(w_norm_name)
|
||||||
|
patch_dict[to_load[x]] = ("diff", (w_norm,))
|
||||||
|
if b_norm is not None:
|
||||||
|
loaded_keys.add(b_norm_name)
|
||||||
|
patch_dict["{}.bias".format(to_load[x][:-len(".weight")])] = ("diff", (b_norm,))
|
||||||
|
|
||||||
|
diff_name = "{}.diff".format(x)
|
||||||
|
diff_weight = lora.get(diff_name, None)
|
||||||
|
if diff_weight is not None:
|
||||||
|
patch_dict[to_load[x]] = ("diff", (diff_weight,))
|
||||||
|
loaded_keys.add(diff_name)
|
||||||
|
|
||||||
|
diff_bias_name = "{}.diff_b".format(x)
|
||||||
|
diff_bias = lora.get(diff_bias_name, None)
|
||||||
|
if diff_bias is not None:
|
||||||
|
patch_dict["{}.bias".format(to_load[x][:-len(".weight")])] = ("diff", (diff_bias,))
|
||||||
|
loaded_keys.add(diff_bias_name)
|
||||||
|
|
||||||
|
for x in lora.keys():
|
||||||
|
if x not in loaded_keys:
|
||||||
|
print("lora key not loaded", x)
|
||||||
|
return patch_dict
|
||||||
|
|
||||||
|
def model_lora_keys_clip(model, key_map={}):
|
||||||
|
sdk = model.state_dict().keys()
|
||||||
|
|
||||||
|
text_model_lora_key = "lora_te_text_model_encoder_layers_{}_{}"
|
||||||
|
clip_l_present = False
|
||||||
|
for b in range(32): #TODO: clean up
|
||||||
|
for c in LORA_CLIP_MAP:
|
||||||
|
k = "clip_h.transformer.text_model.encoder.layers.{}.{}.weight".format(b, c)
|
||||||
|
if k in sdk:
|
||||||
|
lora_key = text_model_lora_key.format(b, LORA_CLIP_MAP[c])
|
||||||
|
key_map[lora_key] = k
|
||||||
|
lora_key = "lora_te1_text_model_encoder_layers_{}_{}".format(b, LORA_CLIP_MAP[c])
|
||||||
|
key_map[lora_key] = k
|
||||||
|
lora_key = "text_encoder.text_model.encoder.layers.{}.{}".format(b, c) #diffusers lora
|
||||||
|
key_map[lora_key] = k
|
||||||
|
|
||||||
|
k = "clip_l.transformer.text_model.encoder.layers.{}.{}.weight".format(b, c)
|
||||||
|
if k in sdk:
|
||||||
|
lora_key = text_model_lora_key.format(b, LORA_CLIP_MAP[c])
|
||||||
|
key_map[lora_key] = k
|
||||||
|
lora_key = "lora_te1_text_model_encoder_layers_{}_{}".format(b, LORA_CLIP_MAP[c]) #SDXL base
|
||||||
|
key_map[lora_key] = k
|
||||||
|
clip_l_present = True
|
||||||
|
lora_key = "text_encoder.text_model.encoder.layers.{}.{}".format(b, c) #diffusers lora
|
||||||
|
key_map[lora_key] = k
|
||||||
|
|
||||||
|
k = "clip_g.transformer.text_model.encoder.layers.{}.{}.weight".format(b, c)
|
||||||
|
if k in sdk:
|
||||||
|
if clip_l_present:
|
||||||
|
lora_key = "lora_te2_text_model_encoder_layers_{}_{}".format(b, LORA_CLIP_MAP[c]) #SDXL base
|
||||||
|
key_map[lora_key] = k
|
||||||
|
lora_key = "text_encoder_2.text_model.encoder.layers.{}.{}".format(b, c) #diffusers lora
|
||||||
|
key_map[lora_key] = k
|
||||||
|
else:
|
||||||
|
lora_key = "lora_te_text_model_encoder_layers_{}_{}".format(b, LORA_CLIP_MAP[c]) #TODO: test if this is correct for SDXL-Refiner
|
||||||
|
key_map[lora_key] = k
|
||||||
|
lora_key = "text_encoder.text_model.encoder.layers.{}.{}".format(b, c) #diffusers lora
|
||||||
|
key_map[lora_key] = k
|
||||||
|
|
||||||
|
return key_map
|
||||||
|
|
||||||
|
def model_lora_keys_unet(model, key_map={}):
|
||||||
|
sdk = model.state_dict().keys()
|
||||||
|
|
||||||
|
for k in sdk:
|
||||||
|
if k.startswith("diffusion_model.") and k.endswith(".weight"):
|
||||||
|
key_lora = k[len("diffusion_model."):-len(".weight")].replace(".", "_")
|
||||||
|
key_map["lora_unet_{}".format(key_lora)] = k
|
||||||
|
|
||||||
|
diffusers_keys = ldm_patched.modules.utils.unet_to_diffusers(model.model_config.unet_config)
|
||||||
|
for k in diffusers_keys:
|
||||||
|
if k.endswith(".weight"):
|
||||||
|
unet_key = "diffusion_model.{}".format(diffusers_keys[k])
|
||||||
|
key_lora = k[:-len(".weight")].replace(".", "_")
|
||||||
|
key_map["lora_unet_{}".format(key_lora)] = unet_key
|
||||||
|
|
||||||
|
diffusers_lora_prefix = ["", "unet."]
|
||||||
|
for p in diffusers_lora_prefix:
|
||||||
|
diffusers_lora_key = "{}{}".format(p, k[:-len(".weight")].replace(".to_", ".processor.to_"))
|
||||||
|
if diffusers_lora_key.endswith(".to_out.0"):
|
||||||
|
diffusers_lora_key = diffusers_lora_key[:-2]
|
||||||
|
key_map[diffusers_lora_key] = unet_key
|
||||||
|
return key_map
|
417
ldm_patched/modules/model_base.py
Normal file
417
ldm_patched/modules/model_base.py
Normal file
@ -0,0 +1,417 @@
|
|||||||
|
import torch
|
||||||
|
from ldm_patched.ldm.modules.diffusionmodules.openaimodel import UNetModel, Timestep
|
||||||
|
from ldm_patched.ldm.modules.encoders.noise_aug_modules import CLIPEmbeddingNoiseAugmentation
|
||||||
|
from ldm_patched.ldm.modules.diffusionmodules.upscaling import ImageConcatWithNoiseAugmentation
|
||||||
|
import ldm_patched.modules.model_management
|
||||||
|
import ldm_patched.modules.conds
|
||||||
|
import ldm_patched.modules.ops
|
||||||
|
from enum import Enum
|
||||||
|
import contextlib
|
||||||
|
from . import utils
|
||||||
|
|
||||||
|
class ModelType(Enum):
|
||||||
|
EPS = 1
|
||||||
|
V_PREDICTION = 2
|
||||||
|
V_PREDICTION_EDM = 3
|
||||||
|
|
||||||
|
|
||||||
|
from ldm_patched.modules.model_sampling import EPS, V_PREDICTION, ModelSamplingDiscrete, ModelSamplingContinuousEDM
|
||||||
|
|
||||||
|
|
||||||
|
def model_sampling(model_config, model_type):
|
||||||
|
s = ModelSamplingDiscrete
|
||||||
|
|
||||||
|
if model_type == ModelType.EPS:
|
||||||
|
c = EPS
|
||||||
|
elif model_type == ModelType.V_PREDICTION:
|
||||||
|
c = V_PREDICTION
|
||||||
|
elif model_type == ModelType.V_PREDICTION_EDM:
|
||||||
|
c = V_PREDICTION
|
||||||
|
s = ModelSamplingContinuousEDM
|
||||||
|
|
||||||
|
class ModelSampling(s, c):
|
||||||
|
pass
|
||||||
|
|
||||||
|
return ModelSampling(model_config)
|
||||||
|
|
||||||
|
|
||||||
|
class BaseModel(torch.nn.Module):
|
||||||
|
def __init__(self, model_config, model_type=ModelType.EPS, device=None):
|
||||||
|
super().__init__()
|
||||||
|
|
||||||
|
unet_config = model_config.unet_config
|
||||||
|
self.latent_format = model_config.latent_format
|
||||||
|
self.model_config = model_config
|
||||||
|
self.manual_cast_dtype = model_config.manual_cast_dtype
|
||||||
|
|
||||||
|
if not unet_config.get("disable_unet_model_creation", False):
|
||||||
|
if self.manual_cast_dtype is not None:
|
||||||
|
operations = ldm_patched.modules.ops.manual_cast
|
||||||
|
else:
|
||||||
|
operations = ldm_patched.modules.ops.disable_weight_init
|
||||||
|
self.diffusion_model = UNetModel(**unet_config, device=device, operations=operations)
|
||||||
|
self.model_type = model_type
|
||||||
|
self.model_sampling = model_sampling(model_config, model_type)
|
||||||
|
|
||||||
|
self.adm_channels = unet_config.get("adm_in_channels", None)
|
||||||
|
if self.adm_channels is None:
|
||||||
|
self.adm_channels = 0
|
||||||
|
self.inpaint_model = False
|
||||||
|
print("model_type", model_type.name)
|
||||||
|
print("UNet ADM Dimension", self.adm_channels)
|
||||||
|
|
||||||
|
def apply_model(self, x, t, c_concat=None, c_crossattn=None, control=None, transformer_options={}, **kwargs):
|
||||||
|
sigma = t
|
||||||
|
xc = self.model_sampling.calculate_input(sigma, x)
|
||||||
|
if c_concat is not None:
|
||||||
|
xc = torch.cat([xc] + [c_concat], dim=1)
|
||||||
|
|
||||||
|
context = c_crossattn
|
||||||
|
dtype = self.get_dtype()
|
||||||
|
|
||||||
|
if self.manual_cast_dtype is not None:
|
||||||
|
dtype = self.manual_cast_dtype
|
||||||
|
|
||||||
|
xc = xc.to(dtype)
|
||||||
|
t = self.model_sampling.timestep(t).float()
|
||||||
|
context = context.to(dtype)
|
||||||
|
extra_conds = {}
|
||||||
|
for o in kwargs:
|
||||||
|
extra = kwargs[o]
|
||||||
|
if hasattr(extra, "dtype"):
|
||||||
|
if extra.dtype != torch.int and extra.dtype != torch.long:
|
||||||
|
extra = extra.to(dtype)
|
||||||
|
extra_conds[o] = extra
|
||||||
|
|
||||||
|
model_output = self.diffusion_model(xc, t, context=context, control=control, transformer_options=transformer_options, **extra_conds).float()
|
||||||
|
return self.model_sampling.calculate_denoised(sigma, model_output, x)
|
||||||
|
|
||||||
|
def get_dtype(self):
|
||||||
|
return self.diffusion_model.dtype
|
||||||
|
|
||||||
|
def is_adm(self):
|
||||||
|
return self.adm_channels > 0
|
||||||
|
|
||||||
|
def encode_adm(self, **kwargs):
|
||||||
|
return None
|
||||||
|
|
||||||
|
def extra_conds(self, **kwargs):
|
||||||
|
out = {}
|
||||||
|
if self.inpaint_model:
|
||||||
|
concat_keys = ("mask", "masked_image")
|
||||||
|
cond_concat = []
|
||||||
|
denoise_mask = kwargs.get("concat_mask", kwargs.get("denoise_mask", None))
|
||||||
|
concat_latent_image = kwargs.get("concat_latent_image", None)
|
||||||
|
if concat_latent_image is None:
|
||||||
|
concat_latent_image = kwargs.get("latent_image", None)
|
||||||
|
else:
|
||||||
|
concat_latent_image = self.process_latent_in(concat_latent_image)
|
||||||
|
|
||||||
|
noise = kwargs.get("noise", None)
|
||||||
|
device = kwargs["device"]
|
||||||
|
|
||||||
|
if concat_latent_image.shape[1:] != noise.shape[1:]:
|
||||||
|
concat_latent_image = utils.common_upscale(concat_latent_image, noise.shape[-1], noise.shape[-2], "bilinear", "center")
|
||||||
|
|
||||||
|
concat_latent_image = utils.resize_to_batch_size(concat_latent_image, noise.shape[0])
|
||||||
|
|
||||||
|
if len(denoise_mask.shape) == len(noise.shape):
|
||||||
|
denoise_mask = denoise_mask[:,:1]
|
||||||
|
|
||||||
|
denoise_mask = denoise_mask.reshape((-1, 1, denoise_mask.shape[-2], denoise_mask.shape[-1]))
|
||||||
|
if denoise_mask.shape[-2:] != noise.shape[-2:]:
|
||||||
|
denoise_mask = utils.common_upscale(denoise_mask, noise.shape[-1], noise.shape[-2], "bilinear", "center")
|
||||||
|
denoise_mask = utils.resize_to_batch_size(denoise_mask.round(), noise.shape[0])
|
||||||
|
|
||||||
|
def blank_inpaint_image_like(latent_image):
|
||||||
|
blank_image = torch.ones_like(latent_image)
|
||||||
|
# these are the values for "zero" in pixel space translated to latent space
|
||||||
|
blank_image[:,0] *= 0.8223
|
||||||
|
blank_image[:,1] *= -0.6876
|
||||||
|
blank_image[:,2] *= 0.6364
|
||||||
|
blank_image[:,3] *= 0.1380
|
||||||
|
return blank_image
|
||||||
|
|
||||||
|
for ck in concat_keys:
|
||||||
|
if denoise_mask is not None:
|
||||||
|
if ck == "mask":
|
||||||
|
cond_concat.append(denoise_mask.to(device))
|
||||||
|
elif ck == "masked_image":
|
||||||
|
cond_concat.append(concat_latent_image.to(device)) #NOTE: the latent_image should be masked by the mask in pixel space
|
||||||
|
else:
|
||||||
|
if ck == "mask":
|
||||||
|
cond_concat.append(torch.ones_like(noise)[:,:1])
|
||||||
|
elif ck == "masked_image":
|
||||||
|
cond_concat.append(blank_inpaint_image_like(noise))
|
||||||
|
data = torch.cat(cond_concat, dim=1)
|
||||||
|
out['c_concat'] = ldm_patched.modules.conds.CONDNoiseShape(data)
|
||||||
|
|
||||||
|
adm = self.encode_adm(**kwargs)
|
||||||
|
if adm is not None:
|
||||||
|
out['y'] = ldm_patched.modules.conds.CONDRegular(adm)
|
||||||
|
|
||||||
|
cross_attn = kwargs.get("cross_attn", None)
|
||||||
|
if cross_attn is not None:
|
||||||
|
out['c_crossattn'] = ldm_patched.modules.conds.CONDCrossAttn(cross_attn)
|
||||||
|
|
||||||
|
return out
|
||||||
|
|
||||||
|
def load_model_weights(self, sd, unet_prefix=""):
|
||||||
|
to_load = {}
|
||||||
|
keys = list(sd.keys())
|
||||||
|
for k in keys:
|
||||||
|
if k.startswith(unet_prefix):
|
||||||
|
to_load[k[len(unet_prefix):]] = sd.pop(k)
|
||||||
|
|
||||||
|
to_load = self.model_config.process_unet_state_dict(to_load)
|
||||||
|
m, u = self.diffusion_model.load_state_dict(to_load, strict=False)
|
||||||
|
if len(m) > 0:
|
||||||
|
print("unet missing:", m)
|
||||||
|
|
||||||
|
if len(u) > 0:
|
||||||
|
print("unet unexpected:", u)
|
||||||
|
del to_load
|
||||||
|
return self
|
||||||
|
|
||||||
|
def process_latent_in(self, latent):
|
||||||
|
return self.latent_format.process_in(latent)
|
||||||
|
|
||||||
|
def process_latent_out(self, latent):
|
||||||
|
return self.latent_format.process_out(latent)
|
||||||
|
|
||||||
|
def state_dict_for_saving(self, clip_state_dict, vae_state_dict):
|
||||||
|
clip_state_dict = self.model_config.process_clip_state_dict_for_saving(clip_state_dict)
|
||||||
|
unet_state_dict = self.diffusion_model.state_dict()
|
||||||
|
unet_state_dict = self.model_config.process_unet_state_dict_for_saving(unet_state_dict)
|
||||||
|
vae_state_dict = self.model_config.process_vae_state_dict_for_saving(vae_state_dict)
|
||||||
|
if self.get_dtype() == torch.float16:
|
||||||
|
clip_state_dict = utils.convert_sd_to(clip_state_dict, torch.float16)
|
||||||
|
vae_state_dict = utils.convert_sd_to(vae_state_dict, torch.float16)
|
||||||
|
|
||||||
|
if self.model_type == ModelType.V_PREDICTION:
|
||||||
|
unet_state_dict["v_pred"] = torch.tensor([])
|
||||||
|
|
||||||
|
return {**unet_state_dict, **vae_state_dict, **clip_state_dict}
|
||||||
|
|
||||||
|
def set_inpaint(self):
|
||||||
|
self.inpaint_model = True
|
||||||
|
|
||||||
|
def memory_required(self, input_shape):
|
||||||
|
if ldm_patched.modules.model_management.xformers_enabled() or ldm_patched.modules.model_management.pytorch_attention_flash_attention():
|
||||||
|
dtype = self.get_dtype()
|
||||||
|
if self.manual_cast_dtype is not None:
|
||||||
|
dtype = self.manual_cast_dtype
|
||||||
|
#TODO: this needs to be tweaked
|
||||||
|
area = input_shape[0] * input_shape[2] * input_shape[3]
|
||||||
|
return (area * ldm_patched.modules.model_management.dtype_size(dtype) / 50) * (1024 * 1024)
|
||||||
|
else:
|
||||||
|
#TODO: this formula might be too aggressive since I tweaked the sub-quad and split algorithms to use less memory.
|
||||||
|
area = input_shape[0] * input_shape[2] * input_shape[3]
|
||||||
|
return (((area * 0.6) / 0.9) + 1024) * (1024 * 1024)
|
||||||
|
|
||||||
|
|
||||||
|
def unclip_adm(unclip_conditioning, device, noise_augmentor, noise_augment_merge=0.0):
|
||||||
|
adm_inputs = []
|
||||||
|
weights = []
|
||||||
|
noise_aug = []
|
||||||
|
for unclip_cond in unclip_conditioning:
|
||||||
|
for adm_cond in unclip_cond["clip_vision_output"].image_embeds:
|
||||||
|
weight = unclip_cond["strength"]
|
||||||
|
noise_augment = unclip_cond["noise_augmentation"]
|
||||||
|
noise_level = round((noise_augmentor.max_noise_level - 1) * noise_augment)
|
||||||
|
c_adm, noise_level_emb = noise_augmentor(adm_cond.to(device), noise_level=torch.tensor([noise_level], device=device))
|
||||||
|
adm_out = torch.cat((c_adm, noise_level_emb), 1) * weight
|
||||||
|
weights.append(weight)
|
||||||
|
noise_aug.append(noise_augment)
|
||||||
|
adm_inputs.append(adm_out)
|
||||||
|
|
||||||
|
if len(noise_aug) > 1:
|
||||||
|
adm_out = torch.stack(adm_inputs).sum(0)
|
||||||
|
noise_augment = noise_augment_merge
|
||||||
|
noise_level = round((noise_augmentor.max_noise_level - 1) * noise_augment)
|
||||||
|
c_adm, noise_level_emb = noise_augmentor(adm_out[:, :noise_augmentor.time_embed.dim], noise_level=torch.tensor([noise_level], device=device))
|
||||||
|
adm_out = torch.cat((c_adm, noise_level_emb), 1)
|
||||||
|
|
||||||
|
return adm_out
|
||||||
|
|
||||||
|
class SD21UNCLIP(BaseModel):
|
||||||
|
def __init__(self, model_config, noise_aug_config, model_type=ModelType.V_PREDICTION, device=None):
|
||||||
|
super().__init__(model_config, model_type, device=device)
|
||||||
|
self.noise_augmentor = CLIPEmbeddingNoiseAugmentation(**noise_aug_config)
|
||||||
|
|
||||||
|
def encode_adm(self, **kwargs):
|
||||||
|
unclip_conditioning = kwargs.get("unclip_conditioning", None)
|
||||||
|
device = kwargs["device"]
|
||||||
|
if unclip_conditioning is None:
|
||||||
|
return torch.zeros((1, self.adm_channels))
|
||||||
|
else:
|
||||||
|
return unclip_adm(unclip_conditioning, device, self.noise_augmentor, kwargs.get("unclip_noise_augment_merge", 0.05))
|
||||||
|
|
||||||
|
def sdxl_pooled(args, noise_augmentor):
|
||||||
|
if "unclip_conditioning" in args:
|
||||||
|
return unclip_adm(args.get("unclip_conditioning", None), args["device"], noise_augmentor)[:,:1280]
|
||||||
|
else:
|
||||||
|
return args["pooled_output"]
|
||||||
|
|
||||||
|
class SDXLRefiner(BaseModel):
|
||||||
|
def __init__(self, model_config, model_type=ModelType.EPS, device=None):
|
||||||
|
super().__init__(model_config, model_type, device=device)
|
||||||
|
self.embedder = Timestep(256)
|
||||||
|
self.noise_augmentor = CLIPEmbeddingNoiseAugmentation(**{"noise_schedule_config": {"timesteps": 1000, "beta_schedule": "squaredcos_cap_v2"}, "timestep_dim": 1280})
|
||||||
|
|
||||||
|
def encode_adm(self, **kwargs):
|
||||||
|
clip_pooled = sdxl_pooled(kwargs, self.noise_augmentor)
|
||||||
|
width = kwargs.get("width", 768)
|
||||||
|
height = kwargs.get("height", 768)
|
||||||
|
crop_w = kwargs.get("crop_w", 0)
|
||||||
|
crop_h = kwargs.get("crop_h", 0)
|
||||||
|
|
||||||
|
if kwargs.get("prompt_type", "") == "negative":
|
||||||
|
aesthetic_score = kwargs.get("aesthetic_score", 2.5)
|
||||||
|
else:
|
||||||
|
aesthetic_score = kwargs.get("aesthetic_score", 6)
|
||||||
|
|
||||||
|
out = []
|
||||||
|
out.append(self.embedder(torch.Tensor([height])))
|
||||||
|
out.append(self.embedder(torch.Tensor([width])))
|
||||||
|
out.append(self.embedder(torch.Tensor([crop_h])))
|
||||||
|
out.append(self.embedder(torch.Tensor([crop_w])))
|
||||||
|
out.append(self.embedder(torch.Tensor([aesthetic_score])))
|
||||||
|
flat = torch.flatten(torch.cat(out)).unsqueeze(dim=0).repeat(clip_pooled.shape[0], 1)
|
||||||
|
return torch.cat((clip_pooled.to(flat.device), flat), dim=1)
|
||||||
|
|
||||||
|
class SDXL(BaseModel):
|
||||||
|
def __init__(self, model_config, model_type=ModelType.EPS, device=None):
|
||||||
|
super().__init__(model_config, model_type, device=device)
|
||||||
|
self.embedder = Timestep(256)
|
||||||
|
self.noise_augmentor = CLIPEmbeddingNoiseAugmentation(**{"noise_schedule_config": {"timesteps": 1000, "beta_schedule": "squaredcos_cap_v2"}, "timestep_dim": 1280})
|
||||||
|
|
||||||
|
def encode_adm(self, **kwargs):
|
||||||
|
clip_pooled = sdxl_pooled(kwargs, self.noise_augmentor)
|
||||||
|
width = kwargs.get("width", 768)
|
||||||
|
height = kwargs.get("height", 768)
|
||||||
|
crop_w = kwargs.get("crop_w", 0)
|
||||||
|
crop_h = kwargs.get("crop_h", 0)
|
||||||
|
target_width = kwargs.get("target_width", width)
|
||||||
|
target_height = kwargs.get("target_height", height)
|
||||||
|
|
||||||
|
out = []
|
||||||
|
out.append(self.embedder(torch.Tensor([height])))
|
||||||
|
out.append(self.embedder(torch.Tensor([width])))
|
||||||
|
out.append(self.embedder(torch.Tensor([crop_h])))
|
||||||
|
out.append(self.embedder(torch.Tensor([crop_w])))
|
||||||
|
out.append(self.embedder(torch.Tensor([target_height])))
|
||||||
|
out.append(self.embedder(torch.Tensor([target_width])))
|
||||||
|
flat = torch.flatten(torch.cat(out)).unsqueeze(dim=0).repeat(clip_pooled.shape[0], 1)
|
||||||
|
return torch.cat((clip_pooled.to(flat.device), flat), dim=1)
|
||||||
|
|
||||||
|
class SVD_img2vid(BaseModel):
|
||||||
|
def __init__(self, model_config, model_type=ModelType.V_PREDICTION_EDM, device=None):
|
||||||
|
super().__init__(model_config, model_type, device=device)
|
||||||
|
self.embedder = Timestep(256)
|
||||||
|
|
||||||
|
def encode_adm(self, **kwargs):
|
||||||
|
fps_id = kwargs.get("fps", 6) - 1
|
||||||
|
motion_bucket_id = kwargs.get("motion_bucket_id", 127)
|
||||||
|
augmentation = kwargs.get("augmentation_level", 0)
|
||||||
|
|
||||||
|
out = []
|
||||||
|
out.append(self.embedder(torch.Tensor([fps_id])))
|
||||||
|
out.append(self.embedder(torch.Tensor([motion_bucket_id])))
|
||||||
|
out.append(self.embedder(torch.Tensor([augmentation])))
|
||||||
|
|
||||||
|
flat = torch.flatten(torch.cat(out)).unsqueeze(dim=0)
|
||||||
|
return flat
|
||||||
|
|
||||||
|
def extra_conds(self, **kwargs):
|
||||||
|
out = {}
|
||||||
|
adm = self.encode_adm(**kwargs)
|
||||||
|
if adm is not None:
|
||||||
|
out['y'] = ldm_patched.modules.conds.CONDRegular(adm)
|
||||||
|
|
||||||
|
latent_image = kwargs.get("concat_latent_image", None)
|
||||||
|
noise = kwargs.get("noise", None)
|
||||||
|
device = kwargs["device"]
|
||||||
|
|
||||||
|
if latent_image is None:
|
||||||
|
latent_image = torch.zeros_like(noise)
|
||||||
|
|
||||||
|
if latent_image.shape[1:] != noise.shape[1:]:
|
||||||
|
latent_image = utils.common_upscale(latent_image, noise.shape[-1], noise.shape[-2], "bilinear", "center")
|
||||||
|
|
||||||
|
latent_image = utils.resize_to_batch_size(latent_image, noise.shape[0])
|
||||||
|
|
||||||
|
out['c_concat'] = ldm_patched.modules.conds.CONDNoiseShape(latent_image)
|
||||||
|
|
||||||
|
cross_attn = kwargs.get("cross_attn", None)
|
||||||
|
if cross_attn is not None:
|
||||||
|
out['c_crossattn'] = ldm_patched.modules.conds.CONDCrossAttn(cross_attn)
|
||||||
|
|
||||||
|
if "time_conditioning" in kwargs:
|
||||||
|
out["time_context"] = ldm_patched.modules.conds.CONDCrossAttn(kwargs["time_conditioning"])
|
||||||
|
|
||||||
|
out['image_only_indicator'] = ldm_patched.modules.conds.CONDConstant(torch.zeros((1,), device=device))
|
||||||
|
out['num_video_frames'] = ldm_patched.modules.conds.CONDConstant(noise.shape[0])
|
||||||
|
return out
|
||||||
|
|
||||||
|
class Stable_Zero123(BaseModel):
|
||||||
|
def __init__(self, model_config, model_type=ModelType.EPS, device=None, cc_projection_weight=None, cc_projection_bias=None):
|
||||||
|
super().__init__(model_config, model_type, device=device)
|
||||||
|
self.cc_projection = ldm_patched.modules.ops.manual_cast.Linear(cc_projection_weight.shape[1], cc_projection_weight.shape[0], dtype=self.get_dtype(), device=device)
|
||||||
|
self.cc_projection.weight.copy_(cc_projection_weight)
|
||||||
|
self.cc_projection.bias.copy_(cc_projection_bias)
|
||||||
|
|
||||||
|
def extra_conds(self, **kwargs):
|
||||||
|
out = {}
|
||||||
|
|
||||||
|
latent_image = kwargs.get("concat_latent_image", None)
|
||||||
|
noise = kwargs.get("noise", None)
|
||||||
|
|
||||||
|
if latent_image is None:
|
||||||
|
latent_image = torch.zeros_like(noise)
|
||||||
|
|
||||||
|
if latent_image.shape[1:] != noise.shape[1:]:
|
||||||
|
latent_image = utils.common_upscale(latent_image, noise.shape[-1], noise.shape[-2], "bilinear", "center")
|
||||||
|
|
||||||
|
latent_image = utils.resize_to_batch_size(latent_image, noise.shape[0])
|
||||||
|
|
||||||
|
out['c_concat'] = ldm_patched.modules.conds.CONDNoiseShape(latent_image)
|
||||||
|
|
||||||
|
cross_attn = kwargs.get("cross_attn", None)
|
||||||
|
if cross_attn is not None:
|
||||||
|
if cross_attn.shape[-1] != 768:
|
||||||
|
cross_attn = self.cc_projection(cross_attn)
|
||||||
|
out['c_crossattn'] = ldm_patched.modules.conds.CONDCrossAttn(cross_attn)
|
||||||
|
return out
|
||||||
|
|
||||||
|
class SD_X4Upscaler(BaseModel):
|
||||||
|
def __init__(self, model_config, model_type=ModelType.V_PREDICTION, device=None):
|
||||||
|
super().__init__(model_config, model_type, device=device)
|
||||||
|
self.noise_augmentor = ImageConcatWithNoiseAugmentation(noise_schedule_config={"linear_start": 0.0001, "linear_end": 0.02}, max_noise_level=350)
|
||||||
|
|
||||||
|
def extra_conds(self, **kwargs):
|
||||||
|
out = {}
|
||||||
|
|
||||||
|
image = kwargs.get("concat_image", None)
|
||||||
|
noise = kwargs.get("noise", None)
|
||||||
|
noise_augment = kwargs.get("noise_augmentation", 0.0)
|
||||||
|
device = kwargs["device"]
|
||||||
|
seed = kwargs["seed"] - 10
|
||||||
|
|
||||||
|
noise_level = round((self.noise_augmentor.max_noise_level) * noise_augment)
|
||||||
|
|
||||||
|
if image is None:
|
||||||
|
image = torch.zeros_like(noise)[:,:3]
|
||||||
|
|
||||||
|
if image.shape[1:] != noise.shape[1:]:
|
||||||
|
image = utils.common_upscale(image.to(device), noise.shape[-1], noise.shape[-2], "bilinear", "center")
|
||||||
|
|
||||||
|
noise_level = torch.tensor([noise_level], device=device)
|
||||||
|
if noise_augment > 0:
|
||||||
|
image, noise_level = self.noise_augmentor(image.to(device), noise_level=noise_level, seed=seed)
|
||||||
|
|
||||||
|
image = utils.resize_to_batch_size(image, noise.shape[0])
|
||||||
|
|
||||||
|
out['c_concat'] = ldm_patched.modules.conds.CONDNoiseShape(image)
|
||||||
|
out['y'] = ldm_patched.modules.conds.CONDRegular(noise_level)
|
||||||
|
return out
|
320
ldm_patched/modules/model_detection.py
Normal file
320
ldm_patched/modules/model_detection.py
Normal file
@ -0,0 +1,320 @@
|
|||||||
|
import ldm_patched.modules.supported_models
|
||||||
|
import ldm_patched.modules.supported_models_base
|
||||||
|
|
||||||
|
def count_blocks(state_dict_keys, prefix_string):
|
||||||
|
count = 0
|
||||||
|
while True:
|
||||||
|
c = False
|
||||||
|
for k in state_dict_keys:
|
||||||
|
if k.startswith(prefix_string.format(count)):
|
||||||
|
c = True
|
||||||
|
break
|
||||||
|
if c == False:
|
||||||
|
break
|
||||||
|
count += 1
|
||||||
|
return count
|
||||||
|
|
||||||
|
def calculate_transformer_depth(prefix, state_dict_keys, state_dict):
|
||||||
|
context_dim = None
|
||||||
|
use_linear_in_transformer = False
|
||||||
|
|
||||||
|
transformer_prefix = prefix + "1.transformer_blocks."
|
||||||
|
transformer_keys = sorted(list(filter(lambda a: a.startswith(transformer_prefix), state_dict_keys)))
|
||||||
|
if len(transformer_keys) > 0:
|
||||||
|
last_transformer_depth = count_blocks(state_dict_keys, transformer_prefix + '{}')
|
||||||
|
context_dim = state_dict['{}0.attn2.to_k.weight'.format(transformer_prefix)].shape[1]
|
||||||
|
use_linear_in_transformer = len(state_dict['{}1.proj_in.weight'.format(prefix)].shape) == 2
|
||||||
|
time_stack = '{}1.time_stack.0.attn1.to_q.weight'.format(prefix) in state_dict or '{}1.time_mix_blocks.0.attn1.to_q.weight'.format(prefix) in state_dict
|
||||||
|
return last_transformer_depth, context_dim, use_linear_in_transformer, time_stack
|
||||||
|
return None
|
||||||
|
|
||||||
|
def detect_unet_config(state_dict, key_prefix, dtype):
|
||||||
|
state_dict_keys = list(state_dict.keys())
|
||||||
|
|
||||||
|
unet_config = {
|
||||||
|
"use_checkpoint": False,
|
||||||
|
"image_size": 32,
|
||||||
|
"use_spatial_transformer": True,
|
||||||
|
"legacy": False
|
||||||
|
}
|
||||||
|
|
||||||
|
y_input = '{}label_emb.0.0.weight'.format(key_prefix)
|
||||||
|
if y_input in state_dict_keys:
|
||||||
|
unet_config["num_classes"] = "sequential"
|
||||||
|
unet_config["adm_in_channels"] = state_dict[y_input].shape[1]
|
||||||
|
else:
|
||||||
|
unet_config["adm_in_channels"] = None
|
||||||
|
|
||||||
|
unet_config["dtype"] = dtype
|
||||||
|
model_channels = state_dict['{}input_blocks.0.0.weight'.format(key_prefix)].shape[0]
|
||||||
|
in_channels = state_dict['{}input_blocks.0.0.weight'.format(key_prefix)].shape[1]
|
||||||
|
|
||||||
|
out_key = '{}out.2.weight'.format(key_prefix)
|
||||||
|
if out_key in state_dict:
|
||||||
|
out_channels = state_dict[out_key].shape[0]
|
||||||
|
else:
|
||||||
|
out_channels = 4
|
||||||
|
|
||||||
|
num_res_blocks = []
|
||||||
|
channel_mult = []
|
||||||
|
attention_resolutions = []
|
||||||
|
transformer_depth = []
|
||||||
|
transformer_depth_output = []
|
||||||
|
context_dim = None
|
||||||
|
use_linear_in_transformer = False
|
||||||
|
|
||||||
|
video_model = False
|
||||||
|
|
||||||
|
current_res = 1
|
||||||
|
count = 0
|
||||||
|
|
||||||
|
last_res_blocks = 0
|
||||||
|
last_channel_mult = 0
|
||||||
|
|
||||||
|
input_block_count = count_blocks(state_dict_keys, '{}input_blocks'.format(key_prefix) + '.{}.')
|
||||||
|
for count in range(input_block_count):
|
||||||
|
prefix = '{}input_blocks.{}.'.format(key_prefix, count)
|
||||||
|
prefix_output = '{}output_blocks.{}.'.format(key_prefix, input_block_count - count - 1)
|
||||||
|
|
||||||
|
block_keys = sorted(list(filter(lambda a: a.startswith(prefix), state_dict_keys)))
|
||||||
|
if len(block_keys) == 0:
|
||||||
|
break
|
||||||
|
|
||||||
|
block_keys_output = sorted(list(filter(lambda a: a.startswith(prefix_output), state_dict_keys)))
|
||||||
|
|
||||||
|
if "{}0.op.weight".format(prefix) in block_keys: #new layer
|
||||||
|
num_res_blocks.append(last_res_blocks)
|
||||||
|
channel_mult.append(last_channel_mult)
|
||||||
|
|
||||||
|
current_res *= 2
|
||||||
|
last_res_blocks = 0
|
||||||
|
last_channel_mult = 0
|
||||||
|
out = calculate_transformer_depth(prefix_output, state_dict_keys, state_dict)
|
||||||
|
if out is not None:
|
||||||
|
transformer_depth_output.append(out[0])
|
||||||
|
else:
|
||||||
|
transformer_depth_output.append(0)
|
||||||
|
else:
|
||||||
|
res_block_prefix = "{}0.in_layers.0.weight".format(prefix)
|
||||||
|
if res_block_prefix in block_keys:
|
||||||
|
last_res_blocks += 1
|
||||||
|
last_channel_mult = state_dict["{}0.out_layers.3.weight".format(prefix)].shape[0] // model_channels
|
||||||
|
|
||||||
|
out = calculate_transformer_depth(prefix, state_dict_keys, state_dict)
|
||||||
|
if out is not None:
|
||||||
|
transformer_depth.append(out[0])
|
||||||
|
if context_dim is None:
|
||||||
|
context_dim = out[1]
|
||||||
|
use_linear_in_transformer = out[2]
|
||||||
|
video_model = out[3]
|
||||||
|
else:
|
||||||
|
transformer_depth.append(0)
|
||||||
|
|
||||||
|
res_block_prefix = "{}0.in_layers.0.weight".format(prefix_output)
|
||||||
|
if res_block_prefix in block_keys_output:
|
||||||
|
out = calculate_transformer_depth(prefix_output, state_dict_keys, state_dict)
|
||||||
|
if out is not None:
|
||||||
|
transformer_depth_output.append(out[0])
|
||||||
|
else:
|
||||||
|
transformer_depth_output.append(0)
|
||||||
|
|
||||||
|
|
||||||
|
num_res_blocks.append(last_res_blocks)
|
||||||
|
channel_mult.append(last_channel_mult)
|
||||||
|
if "{}middle_block.1.proj_in.weight".format(key_prefix) in state_dict_keys:
|
||||||
|
transformer_depth_middle = count_blocks(state_dict_keys, '{}middle_block.1.transformer_blocks.'.format(key_prefix) + '{}')
|
||||||
|
else:
|
||||||
|
transformer_depth_middle = -1
|
||||||
|
|
||||||
|
unet_config["in_channels"] = in_channels
|
||||||
|
unet_config["out_channels"] = out_channels
|
||||||
|
unet_config["model_channels"] = model_channels
|
||||||
|
unet_config["num_res_blocks"] = num_res_blocks
|
||||||
|
unet_config["transformer_depth"] = transformer_depth
|
||||||
|
unet_config["transformer_depth_output"] = transformer_depth_output
|
||||||
|
unet_config["channel_mult"] = channel_mult
|
||||||
|
unet_config["transformer_depth_middle"] = transformer_depth_middle
|
||||||
|
unet_config['use_linear_in_transformer'] = use_linear_in_transformer
|
||||||
|
unet_config["context_dim"] = context_dim
|
||||||
|
|
||||||
|
if video_model:
|
||||||
|
unet_config["extra_ff_mix_layer"] = True
|
||||||
|
unet_config["use_spatial_context"] = True
|
||||||
|
unet_config["merge_strategy"] = "learned_with_images"
|
||||||
|
unet_config["merge_factor"] = 0.0
|
||||||
|
unet_config["video_kernel_size"] = [3, 1, 1]
|
||||||
|
unet_config["use_temporal_resblock"] = True
|
||||||
|
unet_config["use_temporal_attention"] = True
|
||||||
|
else:
|
||||||
|
unet_config["use_temporal_resblock"] = False
|
||||||
|
unet_config["use_temporal_attention"] = False
|
||||||
|
|
||||||
|
return unet_config
|
||||||
|
|
||||||
|
def model_config_from_unet_config(unet_config):
|
||||||
|
for model_config in ldm_patched.modules.supported_models.models:
|
||||||
|
if model_config.matches(unet_config):
|
||||||
|
return model_config(unet_config)
|
||||||
|
|
||||||
|
print("no match", unet_config)
|
||||||
|
return None
|
||||||
|
|
||||||
|
def model_config_from_unet(state_dict, unet_key_prefix, dtype, use_base_if_no_match=False):
|
||||||
|
unet_config = detect_unet_config(state_dict, unet_key_prefix, dtype)
|
||||||
|
model_config = model_config_from_unet_config(unet_config)
|
||||||
|
if model_config is None and use_base_if_no_match:
|
||||||
|
return ldm_patched.modules.supported_models_base.BASE(unet_config)
|
||||||
|
else:
|
||||||
|
return model_config
|
||||||
|
|
||||||
|
def convert_config(unet_config):
|
||||||
|
new_config = unet_config.copy()
|
||||||
|
num_res_blocks = new_config.get("num_res_blocks", None)
|
||||||
|
channel_mult = new_config.get("channel_mult", None)
|
||||||
|
|
||||||
|
if isinstance(num_res_blocks, int):
|
||||||
|
num_res_blocks = len(channel_mult) * [num_res_blocks]
|
||||||
|
|
||||||
|
if "attention_resolutions" in new_config:
|
||||||
|
attention_resolutions = new_config.pop("attention_resolutions")
|
||||||
|
transformer_depth = new_config.get("transformer_depth", None)
|
||||||
|
transformer_depth_middle = new_config.get("transformer_depth_middle", None)
|
||||||
|
|
||||||
|
if isinstance(transformer_depth, int):
|
||||||
|
transformer_depth = len(channel_mult) * [transformer_depth]
|
||||||
|
if transformer_depth_middle is None:
|
||||||
|
transformer_depth_middle = transformer_depth[-1]
|
||||||
|
t_in = []
|
||||||
|
t_out = []
|
||||||
|
s = 1
|
||||||
|
for i in range(len(num_res_blocks)):
|
||||||
|
res = num_res_blocks[i]
|
||||||
|
d = 0
|
||||||
|
if s in attention_resolutions:
|
||||||
|
d = transformer_depth[i]
|
||||||
|
|
||||||
|
t_in += [d] * res
|
||||||
|
t_out += [d] * (res + 1)
|
||||||
|
s *= 2
|
||||||
|
transformer_depth = t_in
|
||||||
|
transformer_depth_output = t_out
|
||||||
|
new_config["transformer_depth"] = t_in
|
||||||
|
new_config["transformer_depth_output"] = t_out
|
||||||
|
new_config["transformer_depth_middle"] = transformer_depth_middle
|
||||||
|
|
||||||
|
new_config["num_res_blocks"] = num_res_blocks
|
||||||
|
return new_config
|
||||||
|
|
||||||
|
|
||||||
|
def unet_config_from_diffusers_unet(state_dict, dtype):
|
||||||
|
match = {}
|
||||||
|
transformer_depth = []
|
||||||
|
|
||||||
|
attn_res = 1
|
||||||
|
down_blocks = count_blocks(state_dict, "down_blocks.{}")
|
||||||
|
for i in range(down_blocks):
|
||||||
|
attn_blocks = count_blocks(state_dict, "down_blocks.{}.attentions.".format(i) + '{}')
|
||||||
|
for ab in range(attn_blocks):
|
||||||
|
transformer_count = count_blocks(state_dict, "down_blocks.{}.attentions.{}.transformer_blocks.".format(i, ab) + '{}')
|
||||||
|
transformer_depth.append(transformer_count)
|
||||||
|
if transformer_count > 0:
|
||||||
|
match["context_dim"] = state_dict["down_blocks.{}.attentions.{}.transformer_blocks.0.attn2.to_k.weight".format(i, ab)].shape[1]
|
||||||
|
|
||||||
|
attn_res *= 2
|
||||||
|
if attn_blocks == 0:
|
||||||
|
transformer_depth.append(0)
|
||||||
|
transformer_depth.append(0)
|
||||||
|
|
||||||
|
match["transformer_depth"] = transformer_depth
|
||||||
|
|
||||||
|
match["model_channels"] = state_dict["conv_in.weight"].shape[0]
|
||||||
|
match["in_channels"] = state_dict["conv_in.weight"].shape[1]
|
||||||
|
match["adm_in_channels"] = None
|
||||||
|
if "class_embedding.linear_1.weight" in state_dict:
|
||||||
|
match["adm_in_channels"] = state_dict["class_embedding.linear_1.weight"].shape[1]
|
||||||
|
elif "add_embedding.linear_1.weight" in state_dict:
|
||||||
|
match["adm_in_channels"] = state_dict["add_embedding.linear_1.weight"].shape[1]
|
||||||
|
|
||||||
|
SDXL = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
|
||||||
|
'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320,
|
||||||
|
'num_res_blocks': [2, 2, 2], 'transformer_depth': [0, 0, 2, 2, 10, 10], 'channel_mult': [1, 2, 4], 'transformer_depth_middle': 10,
|
||||||
|
'use_linear_in_transformer': True, 'context_dim': 2048, 'num_head_channels': 64, 'transformer_depth_output': [0, 0, 0, 2, 2, 2, 10, 10, 10],
|
||||||
|
'use_temporal_attention': False, 'use_temporal_resblock': False}
|
||||||
|
|
||||||
|
SDXL_refiner = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
|
||||||
|
'num_classes': 'sequential', 'adm_in_channels': 2560, 'dtype': dtype, 'in_channels': 4, 'model_channels': 384,
|
||||||
|
'num_res_blocks': [2, 2, 2, 2], 'transformer_depth': [0, 0, 4, 4, 4, 4, 0, 0], 'channel_mult': [1, 2, 4, 4], 'transformer_depth_middle': 4,
|
||||||
|
'use_linear_in_transformer': True, 'context_dim': 1280, 'num_head_channels': 64, 'transformer_depth_output': [0, 0, 0, 4, 4, 4, 4, 4, 4, 0, 0, 0],
|
||||||
|
'use_temporal_attention': False, 'use_temporal_resblock': False}
|
||||||
|
|
||||||
|
SD21 = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
|
||||||
|
'adm_in_channels': None, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320, 'num_res_blocks': [2, 2, 2, 2],
|
||||||
|
'transformer_depth': [1, 1, 1, 1, 1, 1, 0, 0], 'channel_mult': [1, 2, 4, 4], 'transformer_depth_middle': 1, 'use_linear_in_transformer': True,
|
||||||
|
'context_dim': 1024, 'num_head_channels': 64, 'transformer_depth_output': [1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
|
||||||
|
'use_temporal_attention': False, 'use_temporal_resblock': False}
|
||||||
|
|
||||||
|
SD21_uncliph = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
|
||||||
|
'num_classes': 'sequential', 'adm_in_channels': 2048, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320,
|
||||||
|
'num_res_blocks': [2, 2, 2, 2], 'transformer_depth': [1, 1, 1, 1, 1, 1, 0, 0], 'channel_mult': [1, 2, 4, 4], 'transformer_depth_middle': 1,
|
||||||
|
'use_linear_in_transformer': True, 'context_dim': 1024, 'num_head_channels': 64, 'transformer_depth_output': [1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
|
||||||
|
'use_temporal_attention': False, 'use_temporal_resblock': False}
|
||||||
|
|
||||||
|
SD21_unclipl = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
|
||||||
|
'num_classes': 'sequential', 'adm_in_channels': 1536, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320,
|
||||||
|
'num_res_blocks': [2, 2, 2, 2], 'transformer_depth': [1, 1, 1, 1, 1, 1, 0, 0], 'channel_mult': [1, 2, 4, 4], 'transformer_depth_middle': 1,
|
||||||
|
'use_linear_in_transformer': True, 'context_dim': 1024, 'num_head_channels': 64, 'transformer_depth_output': [1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
|
||||||
|
'use_temporal_attention': False, 'use_temporal_resblock': False}
|
||||||
|
|
||||||
|
SD15 = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False, 'adm_in_channels': None,
|
||||||
|
'dtype': dtype, 'in_channels': 4, 'model_channels': 320, 'num_res_blocks': [2, 2, 2, 2], 'transformer_depth': [1, 1, 1, 1, 1, 1, 0, 0],
|
||||||
|
'channel_mult': [1, 2, 4, 4], 'transformer_depth_middle': 1, 'use_linear_in_transformer': False, 'context_dim': 768, 'num_heads': 8,
|
||||||
|
'transformer_depth_output': [1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
|
||||||
|
'use_temporal_attention': False, 'use_temporal_resblock': False}
|
||||||
|
|
||||||
|
SDXL_mid_cnet = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
|
||||||
|
'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320,
|
||||||
|
'num_res_blocks': [2, 2, 2], 'transformer_depth': [0, 0, 0, 0, 1, 1], 'channel_mult': [1, 2, 4], 'transformer_depth_middle': 1,
|
||||||
|
'use_linear_in_transformer': True, 'context_dim': 2048, 'num_head_channels': 64, 'transformer_depth_output': [0, 0, 0, 0, 0, 0, 1, 1, 1],
|
||||||
|
'use_temporal_attention': False, 'use_temporal_resblock': False}
|
||||||
|
|
||||||
|
SDXL_small_cnet = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
|
||||||
|
'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320,
|
||||||
|
'num_res_blocks': [2, 2, 2], 'transformer_depth': [0, 0, 0, 0, 0, 0], 'channel_mult': [1, 2, 4], 'transformer_depth_middle': 0,
|
||||||
|
'use_linear_in_transformer': True, 'num_head_channels': 64, 'context_dim': 1, 'transformer_depth_output': [0, 0, 0, 0, 0, 0, 0, 0, 0],
|
||||||
|
'use_temporal_attention': False, 'use_temporal_resblock': False}
|
||||||
|
|
||||||
|
SDXL_diffusers_inpaint = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
|
||||||
|
'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 9, 'model_channels': 320,
|
||||||
|
'num_res_blocks': [2, 2, 2], 'transformer_depth': [0, 0, 2, 2, 10, 10], 'channel_mult': [1, 2, 4], 'transformer_depth_middle': 10,
|
||||||
|
'use_linear_in_transformer': True, 'context_dim': 2048, 'num_head_channels': 64, 'transformer_depth_output': [0, 0, 0, 2, 2, 2, 10, 10, 10],
|
||||||
|
'use_temporal_attention': False, 'use_temporal_resblock': False}
|
||||||
|
|
||||||
|
SSD_1B = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
|
||||||
|
'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320,
|
||||||
|
'num_res_blocks': [2, 2, 2], 'transformer_depth': [0, 0, 2, 2, 4, 4], 'transformer_depth_output': [0, 0, 0, 1, 1, 2, 10, 4, 4],
|
||||||
|
'channel_mult': [1, 2, 4], 'transformer_depth_middle': -1, 'use_linear_in_transformer': True, 'context_dim': 2048, 'num_head_channels': 64,
|
||||||
|
'use_temporal_attention': False, 'use_temporal_resblock': False}
|
||||||
|
|
||||||
|
Segmind_Vega = {'use_checkpoint': False, 'image_size': 32, 'out_channels': 4, 'use_spatial_transformer': True, 'legacy': False,
|
||||||
|
'num_classes': 'sequential', 'adm_in_channels': 2816, 'dtype': dtype, 'in_channels': 4, 'model_channels': 320,
|
||||||
|
'num_res_blocks': [2, 2, 2], 'transformer_depth': [0, 0, 1, 1, 2, 2], 'transformer_depth_output': [0, 0, 0, 1, 1, 1, 2, 2, 2],
|
||||||
|
'channel_mult': [1, 2, 4], 'transformer_depth_middle': -1, 'use_linear_in_transformer': True, 'context_dim': 2048, 'num_head_channels': 64,
|
||||||
|
'use_temporal_attention': False, 'use_temporal_resblock': False}
|
||||||
|
|
||||||
|
supported_models = [SDXL, SDXL_refiner, SD21, SD15, SD21_uncliph, SD21_unclipl, SDXL_mid_cnet, SDXL_small_cnet, SDXL_diffusers_inpaint, SSD_1B, Segmind_Vega]
|
||||||
|
|
||||||
|
for unet_config in supported_models:
|
||||||
|
matches = True
|
||||||
|
for k in match:
|
||||||
|
if match[k] != unet_config[k]:
|
||||||
|
matches = False
|
||||||
|
break
|
||||||
|
if matches:
|
||||||
|
return convert_config(unet_config)
|
||||||
|
return None
|
||||||
|
|
||||||
|
def model_config_from_diffusers_unet(state_dict, dtype):
|
||||||
|
unet_config = unet_config_from_diffusers_unet(state_dict, dtype)
|
||||||
|
if unet_config is not None:
|
||||||
|
return model_config_from_unet_config(unet_config)
|
||||||
|
return None
|
804
ldm_patched/modules/model_management.py
Normal file
804
ldm_patched/modules/model_management.py
Normal file
@ -0,0 +1,804 @@
|
|||||||
|
import psutil
|
||||||
|
from enum import Enum
|
||||||
|
from ldm_patched.modules.args_parser import args
|
||||||
|
import ldm_patched.modules.utils
|
||||||
|
import torch
|
||||||
|
import sys
|
||||||
|
|
||||||
|
class VRAMState(Enum):
|
||||||
|
DISABLED = 0 #No vram present: no need to move models to vram
|
||||||
|
NO_VRAM = 1 #Very low vram: enable all the options to save vram
|
||||||
|
LOW_VRAM = 2
|
||||||
|
NORMAL_VRAM = 3
|
||||||
|
HIGH_VRAM = 4
|
||||||
|
SHARED = 5 #No dedicated vram: memory shared between CPU and GPU but models still need to be moved between both.
|
||||||
|
|
||||||
|
class CPUState(Enum):
|
||||||
|
GPU = 0
|
||||||
|
CPU = 1
|
||||||
|
MPS = 2
|
||||||
|
|
||||||
|
# Determine VRAM State
|
||||||
|
vram_state = VRAMState.NORMAL_VRAM
|
||||||
|
set_vram_to = VRAMState.NORMAL_VRAM
|
||||||
|
cpu_state = CPUState.GPU
|
||||||
|
|
||||||
|
total_vram = 0
|
||||||
|
|
||||||
|
lowvram_available = True
|
||||||
|
xpu_available = False
|
||||||
|
|
||||||
|
if args.pytorch_deterministic:
|
||||||
|
print("Using deterministic algorithms for pytorch")
|
||||||
|
torch.use_deterministic_algorithms(True, warn_only=True)
|
||||||
|
|
||||||
|
directml_enabled = False
|
||||||
|
if args.directml is not None:
|
||||||
|
import torch_directml
|
||||||
|
directml_enabled = True
|
||||||
|
device_index = args.directml
|
||||||
|
if device_index < 0:
|
||||||
|
directml_device = torch_directml.device()
|
||||||
|
else:
|
||||||
|
directml_device = torch_directml.device(device_index)
|
||||||
|
print("Using directml with device:", torch_directml.device_name(device_index))
|
||||||
|
# torch_directml.disable_tiled_resources(True)
|
||||||
|
lowvram_available = False #TODO: need to find a way to get free memory in directml before this can be enabled by default.
|
||||||
|
|
||||||
|
try:
|
||||||
|
import intel_extension_for_pytorch as ipex
|
||||||
|
if torch.xpu.is_available():
|
||||||
|
xpu_available = True
|
||||||
|
except:
|
||||||
|
pass
|
||||||
|
|
||||||
|
try:
|
||||||
|
if torch.backends.mps.is_available():
|
||||||
|
cpu_state = CPUState.MPS
|
||||||
|
import torch.mps
|
||||||
|
except:
|
||||||
|
pass
|
||||||
|
|
||||||
|
if args.always_cpu:
|
||||||
|
cpu_state = CPUState.CPU
|
||||||
|
|
||||||
|
def is_intel_xpu():
|
||||||
|
global cpu_state
|
||||||
|
global xpu_available
|
||||||
|
if cpu_state == CPUState.GPU:
|
||||||
|
if xpu_available:
|
||||||
|
return True
|
||||||
|
return False
|
||||||
|
|
||||||
|
def get_torch_device():
|
||||||
|
global directml_enabled
|
||||||
|
global cpu_state
|
||||||
|
if directml_enabled:
|
||||||
|
global directml_device
|
||||||
|
return directml_device
|
||||||
|
if cpu_state == CPUState.MPS:
|
||||||
|
return torch.device("mps")
|
||||||
|
if cpu_state == CPUState.CPU:
|
||||||
|
return torch.device("cpu")
|
||||||
|
else:
|
||||||
|
if is_intel_xpu():
|
||||||
|
return torch.device("xpu")
|
||||||
|
else:
|
||||||
|
return torch.device(torch.cuda.current_device())
|
||||||
|
|
||||||
|
def get_total_memory(dev=None, torch_total_too=False):
|
||||||
|
global directml_enabled
|
||||||
|
if dev is None:
|
||||||
|
dev = get_torch_device()
|
||||||
|
|
||||||
|
if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):
|
||||||
|
mem_total = psutil.virtual_memory().total
|
||||||
|
mem_total_torch = mem_total
|
||||||
|
else:
|
||||||
|
if directml_enabled:
|
||||||
|
mem_total = 1024 * 1024 * 1024 #TODO
|
||||||
|
mem_total_torch = mem_total
|
||||||
|
elif is_intel_xpu():
|
||||||
|
stats = torch.xpu.memory_stats(dev)
|
||||||
|
mem_reserved = stats['reserved_bytes.all.current']
|
||||||
|
mem_total = torch.xpu.get_device_properties(dev).total_memory
|
||||||
|
mem_total_torch = mem_reserved
|
||||||
|
else:
|
||||||
|
stats = torch.cuda.memory_stats(dev)
|
||||||
|
mem_reserved = stats['reserved_bytes.all.current']
|
||||||
|
_, mem_total_cuda = torch.cuda.mem_get_info(dev)
|
||||||
|
mem_total_torch = mem_reserved
|
||||||
|
mem_total = mem_total_cuda
|
||||||
|
|
||||||
|
if torch_total_too:
|
||||||
|
return (mem_total, mem_total_torch)
|
||||||
|
else:
|
||||||
|
return mem_total
|
||||||
|
|
||||||
|
total_vram = get_total_memory(get_torch_device()) / (1024 * 1024)
|
||||||
|
total_ram = psutil.virtual_memory().total / (1024 * 1024)
|
||||||
|
print("Total VRAM {:0.0f} MB, total RAM {:0.0f} MB".format(total_vram, total_ram))
|
||||||
|
if not args.always_normal_vram and not args.always_cpu:
|
||||||
|
if lowvram_available and total_vram <= 4096:
|
||||||
|
print("Trying to enable lowvram mode because your GPU seems to have 4GB or less. If you don't want this use: --always-normal-vram")
|
||||||
|
set_vram_to = VRAMState.LOW_VRAM
|
||||||
|
|
||||||
|
try:
|
||||||
|
OOM_EXCEPTION = torch.cuda.OutOfMemoryError
|
||||||
|
except:
|
||||||
|
OOM_EXCEPTION = Exception
|
||||||
|
|
||||||
|
XFORMERS_VERSION = ""
|
||||||
|
XFORMERS_ENABLED_VAE = True
|
||||||
|
if args.disable_xformers:
|
||||||
|
XFORMERS_IS_AVAILABLE = False
|
||||||
|
else:
|
||||||
|
try:
|
||||||
|
import xformers
|
||||||
|
import xformers.ops
|
||||||
|
XFORMERS_IS_AVAILABLE = True
|
||||||
|
try:
|
||||||
|
XFORMERS_IS_AVAILABLE = xformers._has_cpp_library
|
||||||
|
except:
|
||||||
|
pass
|
||||||
|
try:
|
||||||
|
XFORMERS_VERSION = xformers.version.__version__
|
||||||
|
print("xformers version:", XFORMERS_VERSION)
|
||||||
|
if XFORMERS_VERSION.startswith("0.0.18"):
|
||||||
|
print()
|
||||||
|
print("WARNING: This version of xformers has a major bug where you will get black images when generating high resolution images.")
|
||||||
|
print("Please downgrade or upgrade xformers to a different version.")
|
||||||
|
print()
|
||||||
|
XFORMERS_ENABLED_VAE = False
|
||||||
|
except:
|
||||||
|
pass
|
||||||
|
except:
|
||||||
|
XFORMERS_IS_AVAILABLE = False
|
||||||
|
|
||||||
|
def is_nvidia():
|
||||||
|
global cpu_state
|
||||||
|
if cpu_state == CPUState.GPU:
|
||||||
|
if torch.version.cuda:
|
||||||
|
return True
|
||||||
|
return False
|
||||||
|
|
||||||
|
ENABLE_PYTORCH_ATTENTION = False
|
||||||
|
if args.attention_pytorch:
|
||||||
|
ENABLE_PYTORCH_ATTENTION = True
|
||||||
|
XFORMERS_IS_AVAILABLE = False
|
||||||
|
|
||||||
|
VAE_DTYPE = torch.float32
|
||||||
|
|
||||||
|
try:
|
||||||
|
if is_nvidia():
|
||||||
|
torch_version = torch.version.__version__
|
||||||
|
if int(torch_version[0]) >= 2:
|
||||||
|
if ENABLE_PYTORCH_ATTENTION == False and args.attention_split == False and args.attention_quad == False:
|
||||||
|
ENABLE_PYTORCH_ATTENTION = True
|
||||||
|
if torch.cuda.is_bf16_supported():
|
||||||
|
VAE_DTYPE = torch.bfloat16
|
||||||
|
if is_intel_xpu():
|
||||||
|
if args.attention_split == False and args.attention_quad == False:
|
||||||
|
ENABLE_PYTORCH_ATTENTION = True
|
||||||
|
except:
|
||||||
|
pass
|
||||||
|
|
||||||
|
if is_intel_xpu():
|
||||||
|
VAE_DTYPE = torch.bfloat16
|
||||||
|
|
||||||
|
if args.vae_in_cpu:
|
||||||
|
VAE_DTYPE = torch.float32
|
||||||
|
|
||||||
|
if args.vae_in_fp16:
|
||||||
|
VAE_DTYPE = torch.float16
|
||||||
|
elif args.vae_in_bf16:
|
||||||
|
VAE_DTYPE = torch.bfloat16
|
||||||
|
elif args.vae_in_fp32:
|
||||||
|
VAE_DTYPE = torch.float32
|
||||||
|
|
||||||
|
|
||||||
|
if ENABLE_PYTORCH_ATTENTION:
|
||||||
|
torch.backends.cuda.enable_math_sdp(True)
|
||||||
|
torch.backends.cuda.enable_flash_sdp(True)
|
||||||
|
torch.backends.cuda.enable_mem_efficient_sdp(True)
|
||||||
|
|
||||||
|
if args.always_low_vram:
|
||||||
|
set_vram_to = VRAMState.LOW_VRAM
|
||||||
|
lowvram_available = True
|
||||||
|
elif args.always_no_vram:
|
||||||
|
set_vram_to = VRAMState.NO_VRAM
|
||||||
|
elif args.always_high_vram or args.always_gpu:
|
||||||
|
vram_state = VRAMState.HIGH_VRAM
|
||||||
|
|
||||||
|
FORCE_FP32 = False
|
||||||
|
FORCE_FP16 = False
|
||||||
|
if args.all_in_fp32:
|
||||||
|
print("Forcing FP32, if this improves things please report it.")
|
||||||
|
FORCE_FP32 = True
|
||||||
|
|
||||||
|
if args.all_in_fp16:
|
||||||
|
print("Forcing FP16.")
|
||||||
|
FORCE_FP16 = True
|
||||||
|
|
||||||
|
if lowvram_available:
|
||||||
|
if set_vram_to in (VRAMState.LOW_VRAM, VRAMState.NO_VRAM):
|
||||||
|
vram_state = set_vram_to
|
||||||
|
|
||||||
|
|
||||||
|
if cpu_state != CPUState.GPU:
|
||||||
|
vram_state = VRAMState.DISABLED
|
||||||
|
|
||||||
|
if cpu_state == CPUState.MPS:
|
||||||
|
vram_state = VRAMState.SHARED
|
||||||
|
|
||||||
|
print(f"Set vram state to: {vram_state.name}")
|
||||||
|
|
||||||
|
ALWAYS_VRAM_OFFLOAD = args.always_offload_from_vram
|
||||||
|
|
||||||
|
if ALWAYS_VRAM_OFFLOAD:
|
||||||
|
print("Always offload VRAM")
|
||||||
|
|
||||||
|
def get_torch_device_name(device):
|
||||||
|
if hasattr(device, 'type'):
|
||||||
|
if device.type == "cuda":
|
||||||
|
try:
|
||||||
|
allocator_backend = torch.cuda.get_allocator_backend()
|
||||||
|
except:
|
||||||
|
allocator_backend = ""
|
||||||
|
return "{} {} : {}".format(device, torch.cuda.get_device_name(device), allocator_backend)
|
||||||
|
else:
|
||||||
|
return "{}".format(device.type)
|
||||||
|
elif is_intel_xpu():
|
||||||
|
return "{} {}".format(device, torch.xpu.get_device_name(device))
|
||||||
|
else:
|
||||||
|
return "CUDA {}: {}".format(device, torch.cuda.get_device_name(device))
|
||||||
|
|
||||||
|
try:
|
||||||
|
print("Device:", get_torch_device_name(get_torch_device()))
|
||||||
|
except:
|
||||||
|
print("Could not pick default device.")
|
||||||
|
|
||||||
|
print("VAE dtype:", VAE_DTYPE)
|
||||||
|
|
||||||
|
current_loaded_models = []
|
||||||
|
|
||||||
|
def module_size(module):
|
||||||
|
module_mem = 0
|
||||||
|
sd = module.state_dict()
|
||||||
|
for k in sd:
|
||||||
|
t = sd[k]
|
||||||
|
module_mem += t.nelement() * t.element_size()
|
||||||
|
return module_mem
|
||||||
|
|
||||||
|
class LoadedModel:
|
||||||
|
def __init__(self, model):
|
||||||
|
self.model = model
|
||||||
|
self.model_accelerated = False
|
||||||
|
self.device = model.load_device
|
||||||
|
|
||||||
|
def model_memory(self):
|
||||||
|
return self.model.model_size()
|
||||||
|
|
||||||
|
def model_memory_required(self, device):
|
||||||
|
if device == self.model.current_device:
|
||||||
|
return 0
|
||||||
|
else:
|
||||||
|
return self.model_memory()
|
||||||
|
|
||||||
|
def model_load(self, lowvram_model_memory=0):
|
||||||
|
patch_model_to = None
|
||||||
|
if lowvram_model_memory == 0:
|
||||||
|
patch_model_to = self.device
|
||||||
|
|
||||||
|
self.model.model_patches_to(self.device)
|
||||||
|
self.model.model_patches_to(self.model.model_dtype())
|
||||||
|
|
||||||
|
try:
|
||||||
|
self.real_model = self.model.patch_model(device_to=patch_model_to) #TODO: do something with loras and offloading to CPU
|
||||||
|
except Exception as e:
|
||||||
|
self.model.unpatch_model(self.model.offload_device)
|
||||||
|
self.model_unload()
|
||||||
|
raise e
|
||||||
|
|
||||||
|
if lowvram_model_memory > 0:
|
||||||
|
print("loading in lowvram mode", lowvram_model_memory/(1024 * 1024))
|
||||||
|
mem_counter = 0
|
||||||
|
for m in self.real_model.modules():
|
||||||
|
if hasattr(m, "ldm_patched_cast_weights"):
|
||||||
|
m.prev_ldm_patched_cast_weights = m.ldm_patched_cast_weights
|
||||||
|
m.ldm_patched_cast_weights = True
|
||||||
|
module_mem = module_size(m)
|
||||||
|
if mem_counter + module_mem < lowvram_model_memory:
|
||||||
|
m.to(self.device)
|
||||||
|
mem_counter += module_mem
|
||||||
|
elif hasattr(m, "weight"): #only modules with ldm_patched_cast_weights can be set to lowvram mode
|
||||||
|
m.to(self.device)
|
||||||
|
mem_counter += module_size(m)
|
||||||
|
print("lowvram: loaded module regularly", m)
|
||||||
|
|
||||||
|
self.model_accelerated = True
|
||||||
|
|
||||||
|
if is_intel_xpu() and not args.disable_ipex_hijack:
|
||||||
|
self.real_model = torch.xpu.optimize(self.real_model.eval(), inplace=True, auto_kernel_selection=True, graph_mode=True)
|
||||||
|
|
||||||
|
return self.real_model
|
||||||
|
|
||||||
|
def model_unload(self):
|
||||||
|
if self.model_accelerated:
|
||||||
|
for m in self.real_model.modules():
|
||||||
|
if hasattr(m, "prev_ldm_patched_cast_weights"):
|
||||||
|
m.ldm_patched_cast_weights = m.prev_ldm_patched_cast_weights
|
||||||
|
del m.prev_ldm_patched_cast_weights
|
||||||
|
|
||||||
|
self.model_accelerated = False
|
||||||
|
|
||||||
|
self.model.unpatch_model(self.model.offload_device)
|
||||||
|
self.model.model_patches_to(self.model.offload_device)
|
||||||
|
|
||||||
|
def __eq__(self, other):
|
||||||
|
return self.model is other.model
|
||||||
|
|
||||||
|
def minimum_inference_memory():
|
||||||
|
return (1024 * 1024 * 1024)
|
||||||
|
|
||||||
|
def unload_model_clones(model):
|
||||||
|
to_unload = []
|
||||||
|
for i in range(len(current_loaded_models)):
|
||||||
|
if model.is_clone(current_loaded_models[i].model):
|
||||||
|
to_unload = [i] + to_unload
|
||||||
|
|
||||||
|
for i in to_unload:
|
||||||
|
print("unload clone", i)
|
||||||
|
current_loaded_models.pop(i).model_unload()
|
||||||
|
|
||||||
|
def free_memory(memory_required, device, keep_loaded=[]):
|
||||||
|
unloaded_model = False
|
||||||
|
for i in range(len(current_loaded_models) -1, -1, -1):
|
||||||
|
if not ALWAYS_VRAM_OFFLOAD:
|
||||||
|
if get_free_memory(device) > memory_required:
|
||||||
|
break
|
||||||
|
shift_model = current_loaded_models[i]
|
||||||
|
if shift_model.device == device:
|
||||||
|
if shift_model not in keep_loaded:
|
||||||
|
m = current_loaded_models.pop(i)
|
||||||
|
m.model_unload()
|
||||||
|
del m
|
||||||
|
unloaded_model = True
|
||||||
|
|
||||||
|
if unloaded_model:
|
||||||
|
soft_empty_cache()
|
||||||
|
else:
|
||||||
|
if vram_state != VRAMState.HIGH_VRAM:
|
||||||
|
mem_free_total, mem_free_torch = get_free_memory(device, torch_free_too=True)
|
||||||
|
if mem_free_torch > mem_free_total * 0.25:
|
||||||
|
soft_empty_cache()
|
||||||
|
|
||||||
|
def load_models_gpu(models, memory_required=0):
|
||||||
|
global vram_state
|
||||||
|
|
||||||
|
inference_memory = minimum_inference_memory()
|
||||||
|
extra_mem = max(inference_memory, memory_required)
|
||||||
|
|
||||||
|
models_to_load = []
|
||||||
|
models_already_loaded = []
|
||||||
|
for x in models:
|
||||||
|
loaded_model = LoadedModel(x)
|
||||||
|
|
||||||
|
if loaded_model in current_loaded_models:
|
||||||
|
index = current_loaded_models.index(loaded_model)
|
||||||
|
current_loaded_models.insert(0, current_loaded_models.pop(index))
|
||||||
|
models_already_loaded.append(loaded_model)
|
||||||
|
else:
|
||||||
|
if hasattr(x, "model"):
|
||||||
|
print(f"Requested to load {x.model.__class__.__name__}")
|
||||||
|
models_to_load.append(loaded_model)
|
||||||
|
|
||||||
|
if len(models_to_load) == 0:
|
||||||
|
devs = set(map(lambda a: a.device, models_already_loaded))
|
||||||
|
for d in devs:
|
||||||
|
if d != torch.device("cpu"):
|
||||||
|
free_memory(extra_mem, d, models_already_loaded)
|
||||||
|
return
|
||||||
|
|
||||||
|
print(f"Loading {len(models_to_load)} new model{'s' if len(models_to_load) > 1 else ''}")
|
||||||
|
|
||||||
|
total_memory_required = {}
|
||||||
|
for loaded_model in models_to_load:
|
||||||
|
unload_model_clones(loaded_model.model)
|
||||||
|
total_memory_required[loaded_model.device] = total_memory_required.get(loaded_model.device, 0) + loaded_model.model_memory_required(loaded_model.device)
|
||||||
|
|
||||||
|
for device in total_memory_required:
|
||||||
|
if device != torch.device("cpu"):
|
||||||
|
free_memory(total_memory_required[device] * 1.3 + extra_mem, device, models_already_loaded)
|
||||||
|
|
||||||
|
for loaded_model in models_to_load:
|
||||||
|
model = loaded_model.model
|
||||||
|
torch_dev = model.load_device
|
||||||
|
if is_device_cpu(torch_dev):
|
||||||
|
vram_set_state = VRAMState.DISABLED
|
||||||
|
else:
|
||||||
|
vram_set_state = vram_state
|
||||||
|
lowvram_model_memory = 0
|
||||||
|
if lowvram_available and (vram_set_state == VRAMState.LOW_VRAM or vram_set_state == VRAMState.NORMAL_VRAM):
|
||||||
|
model_size = loaded_model.model_memory_required(torch_dev)
|
||||||
|
current_free_mem = get_free_memory(torch_dev)
|
||||||
|
lowvram_model_memory = int(max(64 * (1024 * 1024), (current_free_mem - 1024 * (1024 * 1024)) / 1.3 ))
|
||||||
|
if model_size > (current_free_mem - inference_memory): #only switch to lowvram if really necessary
|
||||||
|
vram_set_state = VRAMState.LOW_VRAM
|
||||||
|
else:
|
||||||
|
lowvram_model_memory = 0
|
||||||
|
|
||||||
|
if vram_set_state == VRAMState.NO_VRAM:
|
||||||
|
lowvram_model_memory = 64 * 1024 * 1024
|
||||||
|
|
||||||
|
cur_loaded_model = loaded_model.model_load(lowvram_model_memory)
|
||||||
|
current_loaded_models.insert(0, loaded_model)
|
||||||
|
return
|
||||||
|
|
||||||
|
|
||||||
|
def load_model_gpu(model):
|
||||||
|
return load_models_gpu([model])
|
||||||
|
|
||||||
|
def cleanup_models():
|
||||||
|
to_delete = []
|
||||||
|
for i in range(len(current_loaded_models)):
|
||||||
|
if sys.getrefcount(current_loaded_models[i].model) <= 2:
|
||||||
|
to_delete = [i] + to_delete
|
||||||
|
|
||||||
|
for i in to_delete:
|
||||||
|
x = current_loaded_models.pop(i)
|
||||||
|
x.model_unload()
|
||||||
|
del x
|
||||||
|
|
||||||
|
def dtype_size(dtype):
|
||||||
|
dtype_size = 4
|
||||||
|
if dtype == torch.float16 or dtype == torch.bfloat16:
|
||||||
|
dtype_size = 2
|
||||||
|
elif dtype == torch.float32:
|
||||||
|
dtype_size = 4
|
||||||
|
else:
|
||||||
|
try:
|
||||||
|
dtype_size = dtype.itemsize
|
||||||
|
except: #Old pytorch doesn't have .itemsize
|
||||||
|
pass
|
||||||
|
return dtype_size
|
||||||
|
|
||||||
|
def unet_offload_device():
|
||||||
|
if vram_state == VRAMState.HIGH_VRAM:
|
||||||
|
return get_torch_device()
|
||||||
|
else:
|
||||||
|
return torch.device("cpu")
|
||||||
|
|
||||||
|
def unet_inital_load_device(parameters, dtype):
|
||||||
|
torch_dev = get_torch_device()
|
||||||
|
if vram_state == VRAMState.HIGH_VRAM:
|
||||||
|
return torch_dev
|
||||||
|
|
||||||
|
cpu_dev = torch.device("cpu")
|
||||||
|
if ALWAYS_VRAM_OFFLOAD:
|
||||||
|
return cpu_dev
|
||||||
|
|
||||||
|
model_size = dtype_size(dtype) * parameters
|
||||||
|
|
||||||
|
mem_dev = get_free_memory(torch_dev)
|
||||||
|
mem_cpu = get_free_memory(cpu_dev)
|
||||||
|
if mem_dev > mem_cpu and model_size < mem_dev:
|
||||||
|
return torch_dev
|
||||||
|
else:
|
||||||
|
return cpu_dev
|
||||||
|
|
||||||
|
def unet_dtype(device=None, model_params=0):
|
||||||
|
if args.unet_in_bf16:
|
||||||
|
return torch.bfloat16
|
||||||
|
if args.unet_in_fp16:
|
||||||
|
return torch.float16
|
||||||
|
if args.unet_in_fp8_e4m3fn:
|
||||||
|
return torch.float8_e4m3fn
|
||||||
|
if args.unet_in_fp8_e5m2:
|
||||||
|
return torch.float8_e5m2
|
||||||
|
if should_use_fp16(device=device, model_params=model_params):
|
||||||
|
return torch.float16
|
||||||
|
return torch.float32
|
||||||
|
|
||||||
|
# None means no manual cast
|
||||||
|
def unet_manual_cast(weight_dtype, inference_device):
|
||||||
|
if weight_dtype == torch.float32:
|
||||||
|
return None
|
||||||
|
|
||||||
|
fp16_supported = ldm_patched.modules.model_management.should_use_fp16(inference_device, prioritize_performance=False)
|
||||||
|
if fp16_supported and weight_dtype == torch.float16:
|
||||||
|
return None
|
||||||
|
|
||||||
|
if fp16_supported:
|
||||||
|
return torch.float16
|
||||||
|
else:
|
||||||
|
return torch.float32
|
||||||
|
|
||||||
|
def text_encoder_offload_device():
|
||||||
|
if args.always_gpu:
|
||||||
|
return get_torch_device()
|
||||||
|
else:
|
||||||
|
return torch.device("cpu")
|
||||||
|
|
||||||
|
def text_encoder_device():
|
||||||
|
if args.always_gpu:
|
||||||
|
return get_torch_device()
|
||||||
|
elif vram_state == VRAMState.HIGH_VRAM or vram_state == VRAMState.NORMAL_VRAM:
|
||||||
|
if is_intel_xpu():
|
||||||
|
return torch.device("cpu")
|
||||||
|
if should_use_fp16(prioritize_performance=False):
|
||||||
|
return get_torch_device()
|
||||||
|
else:
|
||||||
|
return torch.device("cpu")
|
||||||
|
else:
|
||||||
|
return torch.device("cpu")
|
||||||
|
|
||||||
|
def text_encoder_dtype(device=None):
|
||||||
|
if args.clip_in_fp8_e4m3fn:
|
||||||
|
return torch.float8_e4m3fn
|
||||||
|
elif args.clip_in_fp8_e5m2:
|
||||||
|
return torch.float8_e5m2
|
||||||
|
elif args.clip_in_fp16:
|
||||||
|
return torch.float16
|
||||||
|
elif args.clip_in_fp32:
|
||||||
|
return torch.float32
|
||||||
|
|
||||||
|
if is_device_cpu(device):
|
||||||
|
return torch.float16
|
||||||
|
|
||||||
|
if should_use_fp16(device, prioritize_performance=False):
|
||||||
|
return torch.float16
|
||||||
|
else:
|
||||||
|
return torch.float32
|
||||||
|
|
||||||
|
def intermediate_device():
|
||||||
|
if args.always_gpu:
|
||||||
|
return get_torch_device()
|
||||||
|
else:
|
||||||
|
return torch.device("cpu")
|
||||||
|
|
||||||
|
def vae_device():
|
||||||
|
if args.vae_in_cpu:
|
||||||
|
return torch.device("cpu")
|
||||||
|
return get_torch_device()
|
||||||
|
|
||||||
|
def vae_offload_device():
|
||||||
|
if args.always_gpu:
|
||||||
|
return get_torch_device()
|
||||||
|
else:
|
||||||
|
return torch.device("cpu")
|
||||||
|
|
||||||
|
def vae_dtype():
|
||||||
|
global VAE_DTYPE
|
||||||
|
return VAE_DTYPE
|
||||||
|
|
||||||
|
def get_autocast_device(dev):
|
||||||
|
if hasattr(dev, 'type'):
|
||||||
|
return dev.type
|
||||||
|
return "cuda"
|
||||||
|
|
||||||
|
def supports_dtype(device, dtype): #TODO
|
||||||
|
if dtype == torch.float32:
|
||||||
|
return True
|
||||||
|
if is_device_cpu(device):
|
||||||
|
return False
|
||||||
|
if dtype == torch.float16:
|
||||||
|
return True
|
||||||
|
if dtype == torch.bfloat16:
|
||||||
|
return True
|
||||||
|
return False
|
||||||
|
|
||||||
|
def device_supports_non_blocking(device):
|
||||||
|
if is_device_mps(device):
|
||||||
|
return False #pytorch bug? mps doesn't support non blocking
|
||||||
|
return True
|
||||||
|
|
||||||
|
def cast_to_device(tensor, device, dtype, copy=False):
|
||||||
|
device_supports_cast = False
|
||||||
|
if tensor.dtype == torch.float32 or tensor.dtype == torch.float16:
|
||||||
|
device_supports_cast = True
|
||||||
|
elif tensor.dtype == torch.bfloat16:
|
||||||
|
if hasattr(device, 'type') and device.type.startswith("cuda"):
|
||||||
|
device_supports_cast = True
|
||||||
|
elif is_intel_xpu():
|
||||||
|
device_supports_cast = True
|
||||||
|
|
||||||
|
non_blocking = device_supports_non_blocking(device)
|
||||||
|
|
||||||
|
if device_supports_cast:
|
||||||
|
if copy:
|
||||||
|
if tensor.device == device:
|
||||||
|
return tensor.to(dtype, copy=copy, non_blocking=non_blocking)
|
||||||
|
return tensor.to(device, copy=copy, non_blocking=non_blocking).to(dtype, non_blocking=non_blocking)
|
||||||
|
else:
|
||||||
|
return tensor.to(device, non_blocking=non_blocking).to(dtype, non_blocking=non_blocking)
|
||||||
|
else:
|
||||||
|
return tensor.to(device, dtype, copy=copy, non_blocking=non_blocking)
|
||||||
|
|
||||||
|
def xformers_enabled():
|
||||||
|
global directml_enabled
|
||||||
|
global cpu_state
|
||||||
|
if cpu_state != CPUState.GPU:
|
||||||
|
return False
|
||||||
|
if is_intel_xpu():
|
||||||
|
return False
|
||||||
|
if directml_enabled:
|
||||||
|
return False
|
||||||
|
return XFORMERS_IS_AVAILABLE
|
||||||
|
|
||||||
|
|
||||||
|
def xformers_enabled_vae():
|
||||||
|
enabled = xformers_enabled()
|
||||||
|
if not enabled:
|
||||||
|
return False
|
||||||
|
|
||||||
|
return XFORMERS_ENABLED_VAE
|
||||||
|
|
||||||
|
def pytorch_attention_enabled():
|
||||||
|
global ENABLE_PYTORCH_ATTENTION
|
||||||
|
return ENABLE_PYTORCH_ATTENTION
|
||||||
|
|
||||||
|
def pytorch_attention_flash_attention():
|
||||||
|
global ENABLE_PYTORCH_ATTENTION
|
||||||
|
if ENABLE_PYTORCH_ATTENTION:
|
||||||
|
#TODO: more reliable way of checking for flash attention?
|
||||||
|
if is_nvidia(): #pytorch flash attention only works on Nvidia
|
||||||
|
return True
|
||||||
|
return False
|
||||||
|
|
||||||
|
def get_free_memory(dev=None, torch_free_too=False):
|
||||||
|
global directml_enabled
|
||||||
|
if dev is None:
|
||||||
|
dev = get_torch_device()
|
||||||
|
|
||||||
|
if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):
|
||||||
|
mem_free_total = psutil.virtual_memory().available
|
||||||
|
mem_free_torch = mem_free_total
|
||||||
|
else:
|
||||||
|
if directml_enabled:
|
||||||
|
mem_free_total = 1024 * 1024 * 1024 #TODO
|
||||||
|
mem_free_torch = mem_free_total
|
||||||
|
elif is_intel_xpu():
|
||||||
|
stats = torch.xpu.memory_stats(dev)
|
||||||
|
mem_active = stats['active_bytes.all.current']
|
||||||
|
mem_allocated = stats['allocated_bytes.all.current']
|
||||||
|
mem_reserved = stats['reserved_bytes.all.current']
|
||||||
|
mem_free_torch = mem_reserved - mem_active
|
||||||
|
mem_free_total = torch.xpu.get_device_properties(dev).total_memory - mem_allocated
|
||||||
|
else:
|
||||||
|
stats = torch.cuda.memory_stats(dev)
|
||||||
|
mem_active = stats['active_bytes.all.current']
|
||||||
|
mem_reserved = stats['reserved_bytes.all.current']
|
||||||
|
mem_free_cuda, _ = torch.cuda.mem_get_info(dev)
|
||||||
|
mem_free_torch = mem_reserved - mem_active
|
||||||
|
mem_free_total = mem_free_cuda + mem_free_torch
|
||||||
|
|
||||||
|
if torch_free_too:
|
||||||
|
return (mem_free_total, mem_free_torch)
|
||||||
|
else:
|
||||||
|
return mem_free_total
|
||||||
|
|
||||||
|
def cpu_mode():
|
||||||
|
global cpu_state
|
||||||
|
return cpu_state == CPUState.CPU
|
||||||
|
|
||||||
|
def mps_mode():
|
||||||
|
global cpu_state
|
||||||
|
return cpu_state == CPUState.MPS
|
||||||
|
|
||||||
|
def is_device_cpu(device):
|
||||||
|
if hasattr(device, 'type'):
|
||||||
|
if (device.type == 'cpu'):
|
||||||
|
return True
|
||||||
|
return False
|
||||||
|
|
||||||
|
def is_device_mps(device):
|
||||||
|
if hasattr(device, 'type'):
|
||||||
|
if (device.type == 'mps'):
|
||||||
|
return True
|
||||||
|
return False
|
||||||
|
|
||||||
|
def should_use_fp16(device=None, model_params=0, prioritize_performance=True):
|
||||||
|
global directml_enabled
|
||||||
|
|
||||||
|
if device is not None:
|
||||||
|
if is_device_cpu(device):
|
||||||
|
return False
|
||||||
|
|
||||||
|
if FORCE_FP16:
|
||||||
|
return True
|
||||||
|
|
||||||
|
if device is not None: #TODO
|
||||||
|
if is_device_mps(device):
|
||||||
|
return False
|
||||||
|
|
||||||
|
if FORCE_FP32:
|
||||||
|
return False
|
||||||
|
|
||||||
|
if directml_enabled:
|
||||||
|
return False
|
||||||
|
|
||||||
|
if cpu_mode() or mps_mode():
|
||||||
|
return False #TODO ?
|
||||||
|
|
||||||
|
if is_intel_xpu():
|
||||||
|
return True
|
||||||
|
|
||||||
|
if torch.cuda.is_bf16_supported():
|
||||||
|
return True
|
||||||
|
|
||||||
|
props = torch.cuda.get_device_properties("cuda")
|
||||||
|
if props.major < 6:
|
||||||
|
return False
|
||||||
|
|
||||||
|
fp16_works = False
|
||||||
|
#FP16 is confirmed working on a 1080 (GP104) but it's a bit slower than FP32 so it should only be enabled
|
||||||
|
#when the model doesn't actually fit on the card
|
||||||
|
#TODO: actually test if GP106 and others have the same type of behavior
|
||||||
|
nvidia_10_series = ["1080", "1070", "titan x", "p3000", "p3200", "p4000", "p4200", "p5000", "p5200", "p6000", "1060", "1050"]
|
||||||
|
for x in nvidia_10_series:
|
||||||
|
if x in props.name.lower():
|
||||||
|
fp16_works = True
|
||||||
|
|
||||||
|
if fp16_works:
|
||||||
|
free_model_memory = (get_free_memory() * 0.9 - minimum_inference_memory())
|
||||||
|
if (not prioritize_performance) or model_params * 4 > free_model_memory:
|
||||||
|
return True
|
||||||
|
|
||||||
|
if props.major < 7:
|
||||||
|
return False
|
||||||
|
|
||||||
|
#FP16 is just broken on these cards
|
||||||
|
nvidia_16_series = ["1660", "1650", "1630", "T500", "T550", "T600", "MX550", "MX450", "CMP 30HX", "T2000", "T1000", "T1200"]
|
||||||
|
for x in nvidia_16_series:
|
||||||
|
if x in props.name:
|
||||||
|
return False
|
||||||
|
|
||||||
|
return True
|
||||||
|
|
||||||
|
def soft_empty_cache(force=False):
|
||||||
|
global cpu_state
|
||||||
|
if cpu_state == CPUState.MPS:
|
||||||
|
torch.mps.empty_cache()
|
||||||
|
elif is_intel_xpu():
|
||||||
|
torch.xpu.empty_cache()
|
||||||
|
elif torch.cuda.is_available():
|
||||||
|
if force or is_nvidia(): #This seems to make things worse on ROCm so I only do it for cuda
|
||||||
|
torch.cuda.empty_cache()
|
||||||
|
torch.cuda.ipc_collect()
|
||||||
|
|
||||||
|
def unload_all_models():
|
||||||
|
free_memory(1e30, get_torch_device())
|
||||||
|
|
||||||
|
|
||||||
|
def resolve_lowvram_weight(weight, model, key): #TODO: remove
|
||||||
|
return weight
|
||||||
|
|
||||||
|
#TODO: might be cleaner to put this somewhere else
|
||||||
|
import threading
|
||||||
|
|
||||||
|
class InterruptProcessingException(Exception):
|
||||||
|
pass
|
||||||
|
|
||||||
|
interrupt_processing_mutex = threading.RLock()
|
||||||
|
|
||||||
|
interrupt_processing = False
|
||||||
|
def interrupt_current_processing(value=True):
|
||||||
|
global interrupt_processing
|
||||||
|
global interrupt_processing_mutex
|
||||||
|
with interrupt_processing_mutex:
|
||||||
|
interrupt_processing = value
|
||||||
|
|
||||||
|
def processing_interrupted():
|
||||||
|
global interrupt_processing
|
||||||
|
global interrupt_processing_mutex
|
||||||
|
with interrupt_processing_mutex:
|
||||||
|
return interrupt_processing
|
||||||
|
|
||||||
|
def throw_exception_if_processing_interrupted():
|
||||||
|
global interrupt_processing
|
||||||
|
global interrupt_processing_mutex
|
||||||
|
with interrupt_processing_mutex:
|
||||||
|
if interrupt_processing:
|
||||||
|
interrupt_processing = False
|
||||||
|
raise InterruptProcessingException()
|
357
ldm_patched/modules/model_patcher.py
Normal file
357
ldm_patched/modules/model_patcher.py
Normal file
@ -0,0 +1,357 @@
|
|||||||
|
import torch
|
||||||
|
import copy
|
||||||
|
import inspect
|
||||||
|
|
||||||
|
import ldm_patched.modules.utils
|
||||||
|
import ldm_patched.modules.model_management
|
||||||
|
|
||||||
|
class ModelPatcher:
|
||||||
|
def __init__(self, model, load_device, offload_device, size=0, current_device=None, weight_inplace_update=False):
|
||||||
|
self.size = size
|
||||||
|
self.model = model
|
||||||
|
self.patches = {}
|
||||||
|
self.backup = {}
|
||||||
|
self.object_patches = {}
|
||||||
|
self.object_patches_backup = {}
|
||||||
|
self.model_options = {"transformer_options":{}}
|
||||||
|
self.model_size()
|
||||||
|
self.load_device = load_device
|
||||||
|
self.offload_device = offload_device
|
||||||
|
if current_device is None:
|
||||||
|
self.current_device = self.offload_device
|
||||||
|
else:
|
||||||
|
self.current_device = current_device
|
||||||
|
|
||||||
|
self.weight_inplace_update = weight_inplace_update
|
||||||
|
|
||||||
|
def model_size(self):
|
||||||
|
if self.size > 0:
|
||||||
|
return self.size
|
||||||
|
model_sd = self.model.state_dict()
|
||||||
|
self.size = ldm_patched.modules.model_management.module_size(self.model)
|
||||||
|
self.model_keys = set(model_sd.keys())
|
||||||
|
return self.size
|
||||||
|
|
||||||
|
def clone(self):
|
||||||
|
n = ModelPatcher(self.model, self.load_device, self.offload_device, self.size, self.current_device, weight_inplace_update=self.weight_inplace_update)
|
||||||
|
n.patches = {}
|
||||||
|
for k in self.patches:
|
||||||
|
n.patches[k] = self.patches[k][:]
|
||||||
|
|
||||||
|
n.object_patches = self.object_patches.copy()
|
||||||
|
n.model_options = copy.deepcopy(self.model_options)
|
||||||
|
n.model_keys = self.model_keys
|
||||||
|
return n
|
||||||
|
|
||||||
|
def is_clone(self, other):
|
||||||
|
if hasattr(other, 'model') and self.model is other.model:
|
||||||
|
return True
|
||||||
|
return False
|
||||||
|
|
||||||
|
def memory_required(self, input_shape):
|
||||||
|
return self.model.memory_required(input_shape=input_shape)
|
||||||
|
|
||||||
|
def set_model_sampler_cfg_function(self, sampler_cfg_function, disable_cfg1_optimization=False):
|
||||||
|
if len(inspect.signature(sampler_cfg_function).parameters) == 3:
|
||||||
|
self.model_options["sampler_cfg_function"] = lambda args: sampler_cfg_function(args["cond"], args["uncond"], args["cond_scale"]) #Old way
|
||||||
|
else:
|
||||||
|
self.model_options["sampler_cfg_function"] = sampler_cfg_function
|
||||||
|
if disable_cfg1_optimization:
|
||||||
|
self.model_options["disable_cfg1_optimization"] = True
|
||||||
|
|
||||||
|
def set_model_sampler_post_cfg_function(self, post_cfg_function, disable_cfg1_optimization=False):
|
||||||
|
self.model_options["sampler_post_cfg_function"] = self.model_options.get("sampler_post_cfg_function", []) + [post_cfg_function]
|
||||||
|
if disable_cfg1_optimization:
|
||||||
|
self.model_options["disable_cfg1_optimization"] = True
|
||||||
|
|
||||||
|
def set_model_unet_function_wrapper(self, unet_wrapper_function):
|
||||||
|
self.model_options["model_function_wrapper"] = unet_wrapper_function
|
||||||
|
|
||||||
|
def set_model_patch(self, patch, name):
|
||||||
|
to = self.model_options["transformer_options"]
|
||||||
|
if "patches" not in to:
|
||||||
|
to["patches"] = {}
|
||||||
|
to["patches"][name] = to["patches"].get(name, []) + [patch]
|
||||||
|
|
||||||
|
def set_model_patch_replace(self, patch, name, block_name, number, transformer_index=None):
|
||||||
|
to = self.model_options["transformer_options"]
|
||||||
|
if "patches_replace" not in to:
|
||||||
|
to["patches_replace"] = {}
|
||||||
|
if name not in to["patches_replace"]:
|
||||||
|
to["patches_replace"][name] = {}
|
||||||
|
if transformer_index is not None:
|
||||||
|
block = (block_name, number, transformer_index)
|
||||||
|
else:
|
||||||
|
block = (block_name, number)
|
||||||
|
to["patches_replace"][name][block] = patch
|
||||||
|
|
||||||
|
def set_model_attn1_patch(self, patch):
|
||||||
|
self.set_model_patch(patch, "attn1_patch")
|
||||||
|
|
||||||
|
def set_model_attn2_patch(self, patch):
|
||||||
|
self.set_model_patch(patch, "attn2_patch")
|
||||||
|
|
||||||
|
def set_model_attn1_replace(self, patch, block_name, number, transformer_index=None):
|
||||||
|
self.set_model_patch_replace(patch, "attn1", block_name, number, transformer_index)
|
||||||
|
|
||||||
|
def set_model_attn2_replace(self, patch, block_name, number, transformer_index=None):
|
||||||
|
self.set_model_patch_replace(patch, "attn2", block_name, number, transformer_index)
|
||||||
|
|
||||||
|
def set_model_attn1_output_patch(self, patch):
|
||||||
|
self.set_model_patch(patch, "attn1_output_patch")
|
||||||
|
|
||||||
|
def set_model_attn2_output_patch(self, patch):
|
||||||
|
self.set_model_patch(patch, "attn2_output_patch")
|
||||||
|
|
||||||
|
def set_model_input_block_patch(self, patch):
|
||||||
|
self.set_model_patch(patch, "input_block_patch")
|
||||||
|
|
||||||
|
def set_model_input_block_patch_after_skip(self, patch):
|
||||||
|
self.set_model_patch(patch, "input_block_patch_after_skip")
|
||||||
|
|
||||||
|
def set_model_output_block_patch(self, patch):
|
||||||
|
self.set_model_patch(patch, "output_block_patch")
|
||||||
|
|
||||||
|
def add_object_patch(self, name, obj):
|
||||||
|
self.object_patches[name] = obj
|
||||||
|
|
||||||
|
def model_patches_to(self, device):
|
||||||
|
to = self.model_options["transformer_options"]
|
||||||
|
if "patches" in to:
|
||||||
|
patches = to["patches"]
|
||||||
|
for name in patches:
|
||||||
|
patch_list = patches[name]
|
||||||
|
for i in range(len(patch_list)):
|
||||||
|
if hasattr(patch_list[i], "to"):
|
||||||
|
patch_list[i] = patch_list[i].to(device)
|
||||||
|
if "patches_replace" in to:
|
||||||
|
patches = to["patches_replace"]
|
||||||
|
for name in patches:
|
||||||
|
patch_list = patches[name]
|
||||||
|
for k in patch_list:
|
||||||
|
if hasattr(patch_list[k], "to"):
|
||||||
|
patch_list[k] = patch_list[k].to(device)
|
||||||
|
if "model_function_wrapper" in self.model_options:
|
||||||
|
wrap_func = self.model_options["model_function_wrapper"]
|
||||||
|
if hasattr(wrap_func, "to"):
|
||||||
|
self.model_options["model_function_wrapper"] = wrap_func.to(device)
|
||||||
|
|
||||||
|
def model_dtype(self):
|
||||||
|
if hasattr(self.model, "get_dtype"):
|
||||||
|
return self.model.get_dtype()
|
||||||
|
|
||||||
|
def add_patches(self, patches, strength_patch=1.0, strength_model=1.0):
|
||||||
|
p = set()
|
||||||
|
for k in patches:
|
||||||
|
if k in self.model_keys:
|
||||||
|
p.add(k)
|
||||||
|
current_patches = self.patches.get(k, [])
|
||||||
|
current_patches.append((strength_patch, patches[k], strength_model))
|
||||||
|
self.patches[k] = current_patches
|
||||||
|
|
||||||
|
return list(p)
|
||||||
|
|
||||||
|
def get_key_patches(self, filter_prefix=None):
|
||||||
|
ldm_patched.modules.model_management.unload_model_clones(self)
|
||||||
|
model_sd = self.model_state_dict()
|
||||||
|
p = {}
|
||||||
|
for k in model_sd:
|
||||||
|
if filter_prefix is not None:
|
||||||
|
if not k.startswith(filter_prefix):
|
||||||
|
continue
|
||||||
|
if k in self.patches:
|
||||||
|
p[k] = [model_sd[k]] + self.patches[k]
|
||||||
|
else:
|
||||||
|
p[k] = (model_sd[k],)
|
||||||
|
return p
|
||||||
|
|
||||||
|
def model_state_dict(self, filter_prefix=None):
|
||||||
|
sd = self.model.state_dict()
|
||||||
|
keys = list(sd.keys())
|
||||||
|
if filter_prefix is not None:
|
||||||
|
for k in keys:
|
||||||
|
if not k.startswith(filter_prefix):
|
||||||
|
sd.pop(k)
|
||||||
|
return sd
|
||||||
|
|
||||||
|
def patch_model(self, device_to=None, patch_weights=True):
|
||||||
|
for k in self.object_patches:
|
||||||
|
old = getattr(self.model, k)
|
||||||
|
if k not in self.object_patches_backup:
|
||||||
|
self.object_patches_backup[k] = old
|
||||||
|
setattr(self.model, k, self.object_patches[k])
|
||||||
|
|
||||||
|
if patch_weights:
|
||||||
|
model_sd = self.model_state_dict()
|
||||||
|
for key in self.patches:
|
||||||
|
if key not in model_sd:
|
||||||
|
print("could not patch. key doesn't exist in model:", key)
|
||||||
|
continue
|
||||||
|
|
||||||
|
weight = model_sd[key]
|
||||||
|
|
||||||
|
inplace_update = self.weight_inplace_update
|
||||||
|
|
||||||
|
if key not in self.backup:
|
||||||
|
self.backup[key] = weight.to(device=self.offload_device, copy=inplace_update)
|
||||||
|
|
||||||
|
if device_to is not None:
|
||||||
|
temp_weight = ldm_patched.modules.model_management.cast_to_device(weight, device_to, torch.float32, copy=True)
|
||||||
|
else:
|
||||||
|
temp_weight = weight.to(torch.float32, copy=True)
|
||||||
|
out_weight = self.calculate_weight(self.patches[key], temp_weight, key).to(weight.dtype)
|
||||||
|
if inplace_update:
|
||||||
|
ldm_patched.modules.utils.copy_to_param(self.model, key, out_weight)
|
||||||
|
else:
|
||||||
|
ldm_patched.modules.utils.set_attr(self.model, key, out_weight)
|
||||||
|
del temp_weight
|
||||||
|
|
||||||
|
if device_to is not None:
|
||||||
|
self.model.to(device_to)
|
||||||
|
self.current_device = device_to
|
||||||
|
|
||||||
|
return self.model
|
||||||
|
|
||||||
|
def calculate_weight(self, patches, weight, key):
|
||||||
|
for p in patches:
|
||||||
|
alpha = p[0]
|
||||||
|
v = p[1]
|
||||||
|
strength_model = p[2]
|
||||||
|
|
||||||
|
if strength_model != 1.0:
|
||||||
|
weight *= strength_model
|
||||||
|
|
||||||
|
if isinstance(v, list):
|
||||||
|
v = (self.calculate_weight(v[1:], v[0].clone(), key), )
|
||||||
|
|
||||||
|
if len(v) == 1:
|
||||||
|
patch_type = "diff"
|
||||||
|
elif len(v) == 2:
|
||||||
|
patch_type = v[0]
|
||||||
|
v = v[1]
|
||||||
|
|
||||||
|
if patch_type == "diff":
|
||||||
|
w1 = v[0]
|
||||||
|
if alpha != 0.0:
|
||||||
|
if w1.shape != weight.shape:
|
||||||
|
print("WARNING SHAPE MISMATCH {} WEIGHT NOT MERGED {} != {}".format(key, w1.shape, weight.shape))
|
||||||
|
else:
|
||||||
|
weight += alpha * ldm_patched.modules.model_management.cast_to_device(w1, weight.device, weight.dtype)
|
||||||
|
elif patch_type == "lora": #lora/locon
|
||||||
|
mat1 = ldm_patched.modules.model_management.cast_to_device(v[0], weight.device, torch.float32)
|
||||||
|
mat2 = ldm_patched.modules.model_management.cast_to_device(v[1], weight.device, torch.float32)
|
||||||
|
if v[2] is not None:
|
||||||
|
alpha *= v[2] / mat2.shape[0]
|
||||||
|
if v[3] is not None:
|
||||||
|
#locon mid weights, hopefully the math is fine because I didn't properly test it
|
||||||
|
mat3 = ldm_patched.modules.model_management.cast_to_device(v[3], weight.device, torch.float32)
|
||||||
|
final_shape = [mat2.shape[1], mat2.shape[0], mat3.shape[2], mat3.shape[3]]
|
||||||
|
mat2 = torch.mm(mat2.transpose(0, 1).flatten(start_dim=1), mat3.transpose(0, 1).flatten(start_dim=1)).reshape(final_shape).transpose(0, 1)
|
||||||
|
try:
|
||||||
|
weight += (alpha * torch.mm(mat1.flatten(start_dim=1), mat2.flatten(start_dim=1))).reshape(weight.shape).type(weight.dtype)
|
||||||
|
except Exception as e:
|
||||||
|
print("ERROR", key, e)
|
||||||
|
elif patch_type == "lokr":
|
||||||
|
w1 = v[0]
|
||||||
|
w2 = v[1]
|
||||||
|
w1_a = v[3]
|
||||||
|
w1_b = v[4]
|
||||||
|
w2_a = v[5]
|
||||||
|
w2_b = v[6]
|
||||||
|
t2 = v[7]
|
||||||
|
dim = None
|
||||||
|
|
||||||
|
if w1 is None:
|
||||||
|
dim = w1_b.shape[0]
|
||||||
|
w1 = torch.mm(ldm_patched.modules.model_management.cast_to_device(w1_a, weight.device, torch.float32),
|
||||||
|
ldm_patched.modules.model_management.cast_to_device(w1_b, weight.device, torch.float32))
|
||||||
|
else:
|
||||||
|
w1 = ldm_patched.modules.model_management.cast_to_device(w1, weight.device, torch.float32)
|
||||||
|
|
||||||
|
if w2 is None:
|
||||||
|
dim = w2_b.shape[0]
|
||||||
|
if t2 is None:
|
||||||
|
w2 = torch.mm(ldm_patched.modules.model_management.cast_to_device(w2_a, weight.device, torch.float32),
|
||||||
|
ldm_patched.modules.model_management.cast_to_device(w2_b, weight.device, torch.float32))
|
||||||
|
else:
|
||||||
|
w2 = torch.einsum('i j k l, j r, i p -> p r k l',
|
||||||
|
ldm_patched.modules.model_management.cast_to_device(t2, weight.device, torch.float32),
|
||||||
|
ldm_patched.modules.model_management.cast_to_device(w2_b, weight.device, torch.float32),
|
||||||
|
ldm_patched.modules.model_management.cast_to_device(w2_a, weight.device, torch.float32))
|
||||||
|
else:
|
||||||
|
w2 = ldm_patched.modules.model_management.cast_to_device(w2, weight.device, torch.float32)
|
||||||
|
|
||||||
|
if len(w2.shape) == 4:
|
||||||
|
w1 = w1.unsqueeze(2).unsqueeze(2)
|
||||||
|
if v[2] is not None and dim is not None:
|
||||||
|
alpha *= v[2] / dim
|
||||||
|
|
||||||
|
try:
|
||||||
|
weight += alpha * torch.kron(w1, w2).reshape(weight.shape).type(weight.dtype)
|
||||||
|
except Exception as e:
|
||||||
|
print("ERROR", key, e)
|
||||||
|
elif patch_type == "loha":
|
||||||
|
w1a = v[0]
|
||||||
|
w1b = v[1]
|
||||||
|
if v[2] is not None:
|
||||||
|
alpha *= v[2] / w1b.shape[0]
|
||||||
|
w2a = v[3]
|
||||||
|
w2b = v[4]
|
||||||
|
if v[5] is not None: #cp decomposition
|
||||||
|
t1 = v[5]
|
||||||
|
t2 = v[6]
|
||||||
|
m1 = torch.einsum('i j k l, j r, i p -> p r k l',
|
||||||
|
ldm_patched.modules.model_management.cast_to_device(t1, weight.device, torch.float32),
|
||||||
|
ldm_patched.modules.model_management.cast_to_device(w1b, weight.device, torch.float32),
|
||||||
|
ldm_patched.modules.model_management.cast_to_device(w1a, weight.device, torch.float32))
|
||||||
|
|
||||||
|
m2 = torch.einsum('i j k l, j r, i p -> p r k l',
|
||||||
|
ldm_patched.modules.model_management.cast_to_device(t2, weight.device, torch.float32),
|
||||||
|
ldm_patched.modules.model_management.cast_to_device(w2b, weight.device, torch.float32),
|
||||||
|
ldm_patched.modules.model_management.cast_to_device(w2a, weight.device, torch.float32))
|
||||||
|
else:
|
||||||
|
m1 = torch.mm(ldm_patched.modules.model_management.cast_to_device(w1a, weight.device, torch.float32),
|
||||||
|
ldm_patched.modules.model_management.cast_to_device(w1b, weight.device, torch.float32))
|
||||||
|
m2 = torch.mm(ldm_patched.modules.model_management.cast_to_device(w2a, weight.device, torch.float32),
|
||||||
|
ldm_patched.modules.model_management.cast_to_device(w2b, weight.device, torch.float32))
|
||||||
|
|
||||||
|
try:
|
||||||
|
weight += (alpha * m1 * m2).reshape(weight.shape).type(weight.dtype)
|
||||||
|
except Exception as e:
|
||||||
|
print("ERROR", key, e)
|
||||||
|
elif patch_type == "glora":
|
||||||
|
if v[4] is not None:
|
||||||
|
alpha *= v[4] / v[0].shape[0]
|
||||||
|
|
||||||
|
a1 = ldm_patched.modules.model_management.cast_to_device(v[0].flatten(start_dim=1), weight.device, torch.float32)
|
||||||
|
a2 = ldm_patched.modules.model_management.cast_to_device(v[1].flatten(start_dim=1), weight.device, torch.float32)
|
||||||
|
b1 = ldm_patched.modules.model_management.cast_to_device(v[2].flatten(start_dim=1), weight.device, torch.float32)
|
||||||
|
b2 = ldm_patched.modules.model_management.cast_to_device(v[3].flatten(start_dim=1), weight.device, torch.float32)
|
||||||
|
|
||||||
|
weight += ((torch.mm(b2, b1) + torch.mm(torch.mm(weight.flatten(start_dim=1), a2), a1)) * alpha).reshape(weight.shape).type(weight.dtype)
|
||||||
|
else:
|
||||||
|
print("patch type not recognized", patch_type, key)
|
||||||
|
|
||||||
|
return weight
|
||||||
|
|
||||||
|
def unpatch_model(self, device_to=None):
|
||||||
|
keys = list(self.backup.keys())
|
||||||
|
|
||||||
|
if self.weight_inplace_update:
|
||||||
|
for k in keys:
|
||||||
|
ldm_patched.modules.utils.copy_to_param(self.model, k, self.backup[k])
|
||||||
|
else:
|
||||||
|
for k in keys:
|
||||||
|
ldm_patched.modules.utils.set_attr(self.model, k, self.backup[k])
|
||||||
|
|
||||||
|
self.backup = {}
|
||||||
|
|
||||||
|
if device_to is not None:
|
||||||
|
self.model.to(device_to)
|
||||||
|
self.current_device = device_to
|
||||||
|
|
||||||
|
keys = list(self.object_patches_backup.keys())
|
||||||
|
for k in keys:
|
||||||
|
setattr(self.model, k, self.object_patches_backup[k])
|
||||||
|
|
||||||
|
self.object_patches_backup = {}
|
136
ldm_patched/modules/model_sampling.py
Normal file
136
ldm_patched/modules/model_sampling.py
Normal file
@ -0,0 +1,136 @@
|
|||||||
|
import torch
|
||||||
|
import numpy as np
|
||||||
|
from ldm_patched.ldm.modules.diffusionmodules.util import make_beta_schedule
|
||||||
|
import math
|
||||||
|
|
||||||
|
class EPS:
|
||||||
|
def calculate_input(self, sigma, noise):
|
||||||
|
sigma = sigma.view(sigma.shape[:1] + (1,) * (noise.ndim - 1))
|
||||||
|
return noise / (sigma ** 2 + self.sigma_data ** 2) ** 0.5
|
||||||
|
|
||||||
|
def calculate_denoised(self, sigma, model_output, model_input):
|
||||||
|
sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1))
|
||||||
|
return model_input - model_output * sigma
|
||||||
|
|
||||||
|
|
||||||
|
class V_PREDICTION(EPS):
|
||||||
|
def calculate_denoised(self, sigma, model_output, model_input):
|
||||||
|
sigma = sigma.view(sigma.shape[:1] + (1,) * (model_output.ndim - 1))
|
||||||
|
return model_input * self.sigma_data ** 2 / (sigma ** 2 + self.sigma_data ** 2) - model_output * sigma * self.sigma_data / (sigma ** 2 + self.sigma_data ** 2) ** 0.5
|
||||||
|
|
||||||
|
|
||||||
|
class ModelSamplingDiscrete(torch.nn.Module):
|
||||||
|
def __init__(self, model_config=None):
|
||||||
|
super().__init__()
|
||||||
|
|
||||||
|
if model_config is not None:
|
||||||
|
sampling_settings = model_config.sampling_settings
|
||||||
|
else:
|
||||||
|
sampling_settings = {}
|
||||||
|
|
||||||
|
beta_schedule = sampling_settings.get("beta_schedule", "linear")
|
||||||
|
linear_start = sampling_settings.get("linear_start", 0.00085)
|
||||||
|
linear_end = sampling_settings.get("linear_end", 0.012)
|
||||||
|
|
||||||
|
self._register_schedule(given_betas=None, beta_schedule=beta_schedule, timesteps=1000, linear_start=linear_start, linear_end=linear_end, cosine_s=8e-3)
|
||||||
|
self.sigma_data = 1.0
|
||||||
|
|
||||||
|
def _register_schedule(self, given_betas=None, beta_schedule="linear", timesteps=1000,
|
||||||
|
linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
|
||||||
|
if given_betas is not None:
|
||||||
|
betas = given_betas
|
||||||
|
else:
|
||||||
|
betas = make_beta_schedule(beta_schedule, timesteps, linear_start=linear_start, linear_end=linear_end, cosine_s=cosine_s)
|
||||||
|
alphas = 1. - betas
|
||||||
|
alphas_cumprod = torch.tensor(np.cumprod(alphas, axis=0), dtype=torch.float32)
|
||||||
|
# alphas_cumprod_prev = np.append(1., alphas_cumprod[:-1])
|
||||||
|
|
||||||
|
timesteps, = betas.shape
|
||||||
|
self.num_timesteps = int(timesteps)
|
||||||
|
self.linear_start = linear_start
|
||||||
|
self.linear_end = linear_end
|
||||||
|
|
||||||
|
# self.register_buffer('betas', torch.tensor(betas, dtype=torch.float32))
|
||||||
|
# self.register_buffer('alphas_cumprod', torch.tensor(alphas_cumprod, dtype=torch.float32))
|
||||||
|
# self.register_buffer('alphas_cumprod_prev', torch.tensor(alphas_cumprod_prev, dtype=torch.float32))
|
||||||
|
|
||||||
|
sigmas = ((1 - alphas_cumprod) / alphas_cumprod) ** 0.5
|
||||||
|
self.set_sigmas(sigmas)
|
||||||
|
|
||||||
|
def set_sigmas(self, sigmas):
|
||||||
|
self.register_buffer('sigmas', sigmas)
|
||||||
|
self.register_buffer('log_sigmas', sigmas.log())
|
||||||
|
|
||||||
|
@property
|
||||||
|
def sigma_min(self):
|
||||||
|
return self.sigmas[0]
|
||||||
|
|
||||||
|
@property
|
||||||
|
def sigma_max(self):
|
||||||
|
return self.sigmas[-1]
|
||||||
|
|
||||||
|
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).to(sigma.device)
|
||||||
|
|
||||||
|
def sigma(self, timestep):
|
||||||
|
t = torch.clamp(timestep.float().to(self.log_sigmas.device), 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 percent_to_sigma(self, percent):
|
||||||
|
if percent <= 0.0:
|
||||||
|
return 999999999.9
|
||||||
|
if percent >= 1.0:
|
||||||
|
return 0.0
|
||||||
|
percent = 1.0 - percent
|
||||||
|
return self.sigma(torch.tensor(percent * 999.0)).item()
|
||||||
|
|
||||||
|
|
||||||
|
class ModelSamplingContinuousEDM(torch.nn.Module):
|
||||||
|
def __init__(self, model_config=None):
|
||||||
|
super().__init__()
|
||||||
|
self.sigma_data = 1.0
|
||||||
|
|
||||||
|
if model_config is not None:
|
||||||
|
sampling_settings = model_config.sampling_settings
|
||||||
|
else:
|
||||||
|
sampling_settings = {}
|
||||||
|
|
||||||
|
sigma_min = sampling_settings.get("sigma_min", 0.002)
|
||||||
|
sigma_max = sampling_settings.get("sigma_max", 120.0)
|
||||||
|
self.set_sigma_range(sigma_min, sigma_max)
|
||||||
|
|
||||||
|
def set_sigma_range(self, sigma_min, sigma_max):
|
||||||
|
sigmas = torch.linspace(math.log(sigma_min), math.log(sigma_max), 1000).exp()
|
||||||
|
|
||||||
|
self.register_buffer('sigmas', sigmas) #for compatibility with some schedulers
|
||||||
|
self.register_buffer('log_sigmas', sigmas.log())
|
||||||
|
|
||||||
|
@property
|
||||||
|
def sigma_min(self):
|
||||||
|
return self.sigmas[0]
|
||||||
|
|
||||||
|
@property
|
||||||
|
def sigma_max(self):
|
||||||
|
return self.sigmas[-1]
|
||||||
|
|
||||||
|
def timestep(self, sigma):
|
||||||
|
return 0.25 * sigma.log()
|
||||||
|
|
||||||
|
def sigma(self, timestep):
|
||||||
|
return (timestep / 0.25).exp()
|
||||||
|
|
||||||
|
def percent_to_sigma(self, percent):
|
||||||
|
if percent <= 0.0:
|
||||||
|
return 999999999.9
|
||||||
|
if percent >= 1.0:
|
||||||
|
return 0.0
|
||||||
|
percent = 1.0 - percent
|
||||||
|
|
||||||
|
log_sigma_min = math.log(self.sigma_min)
|
||||||
|
return math.exp((math.log(self.sigma_max) - log_sigma_min) * percent + log_sigma_min)
|
115
ldm_patched/modules/ops.py
Normal file
115
ldm_patched/modules/ops.py
Normal file
@ -0,0 +1,115 @@
|
|||||||
|
import torch
|
||||||
|
from contextlib import contextmanager
|
||||||
|
import ldm_patched.modules.model_management
|
||||||
|
|
||||||
|
def cast_bias_weight(s, input):
|
||||||
|
bias = None
|
||||||
|
non_blocking = ldm_patched.modules.model_management.device_supports_non_blocking(input.device)
|
||||||
|
if s.bias is not None:
|
||||||
|
bias = s.bias.to(device=input.device, dtype=input.dtype, non_blocking=non_blocking)
|
||||||
|
weight = s.weight.to(device=input.device, dtype=input.dtype, non_blocking=non_blocking)
|
||||||
|
return weight, bias
|
||||||
|
|
||||||
|
|
||||||
|
class disable_weight_init:
|
||||||
|
class Linear(torch.nn.Linear):
|
||||||
|
ldm_patched_cast_weights = False
|
||||||
|
def reset_parameters(self):
|
||||||
|
return None
|
||||||
|
|
||||||
|
def forward_ldm_patched_cast_weights(self, input):
|
||||||
|
weight, bias = cast_bias_weight(self, input)
|
||||||
|
return torch.nn.functional.linear(input, weight, bias)
|
||||||
|
|
||||||
|
def forward(self, *args, **kwargs):
|
||||||
|
if self.ldm_patched_cast_weights:
|
||||||
|
return self.forward_ldm_patched_cast_weights(*args, **kwargs)
|
||||||
|
else:
|
||||||
|
return super().forward(*args, **kwargs)
|
||||||
|
|
||||||
|
class Conv2d(torch.nn.Conv2d):
|
||||||
|
ldm_patched_cast_weights = False
|
||||||
|
def reset_parameters(self):
|
||||||
|
return None
|
||||||
|
|
||||||
|
def forward_ldm_patched_cast_weights(self, input):
|
||||||
|
weight, bias = cast_bias_weight(self, input)
|
||||||
|
return self._conv_forward(input, weight, bias)
|
||||||
|
|
||||||
|
def forward(self, *args, **kwargs):
|
||||||
|
if self.ldm_patched_cast_weights:
|
||||||
|
return self.forward_ldm_patched_cast_weights(*args, **kwargs)
|
||||||
|
else:
|
||||||
|
return super().forward(*args, **kwargs)
|
||||||
|
|
||||||
|
class Conv3d(torch.nn.Conv3d):
|
||||||
|
ldm_patched_cast_weights = False
|
||||||
|
def reset_parameters(self):
|
||||||
|
return None
|
||||||
|
|
||||||
|
def forward_ldm_patched_cast_weights(self, input):
|
||||||
|
weight, bias = cast_bias_weight(self, input)
|
||||||
|
return self._conv_forward(input, weight, bias)
|
||||||
|
|
||||||
|
def forward(self, *args, **kwargs):
|
||||||
|
if self.ldm_patched_cast_weights:
|
||||||
|
return self.forward_ldm_patched_cast_weights(*args, **kwargs)
|
||||||
|
else:
|
||||||
|
return super().forward(*args, **kwargs)
|
||||||
|
|
||||||
|
class GroupNorm(torch.nn.GroupNorm):
|
||||||
|
ldm_patched_cast_weights = False
|
||||||
|
def reset_parameters(self):
|
||||||
|
return None
|
||||||
|
|
||||||
|
def forward_ldm_patched_cast_weights(self, input):
|
||||||
|
weight, bias = cast_bias_weight(self, input)
|
||||||
|
return torch.nn.functional.group_norm(input, self.num_groups, weight, bias, self.eps)
|
||||||
|
|
||||||
|
def forward(self, *args, **kwargs):
|
||||||
|
if self.ldm_patched_cast_weights:
|
||||||
|
return self.forward_ldm_patched_cast_weights(*args, **kwargs)
|
||||||
|
else:
|
||||||
|
return super().forward(*args, **kwargs)
|
||||||
|
|
||||||
|
|
||||||
|
class LayerNorm(torch.nn.LayerNorm):
|
||||||
|
ldm_patched_cast_weights = False
|
||||||
|
def reset_parameters(self):
|
||||||
|
return None
|
||||||
|
|
||||||
|
def forward_ldm_patched_cast_weights(self, input):
|
||||||
|
weight, bias = cast_bias_weight(self, input)
|
||||||
|
return torch.nn.functional.layer_norm(input, self.normalized_shape, weight, bias, self.eps)
|
||||||
|
|
||||||
|
def forward(self, *args, **kwargs):
|
||||||
|
if self.ldm_patched_cast_weights:
|
||||||
|
return self.forward_ldm_patched_cast_weights(*args, **kwargs)
|
||||||
|
else:
|
||||||
|
return super().forward(*args, **kwargs)
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def conv_nd(s, dims, *args, **kwargs):
|
||||||
|
if dims == 2:
|
||||||
|
return s.Conv2d(*args, **kwargs)
|
||||||
|
elif dims == 3:
|
||||||
|
return s.Conv3d(*args, **kwargs)
|
||||||
|
else:
|
||||||
|
raise ValueError(f"unsupported dimensions: {dims}")
|
||||||
|
|
||||||
|
|
||||||
|
class manual_cast(disable_weight_init):
|
||||||
|
class Linear(disable_weight_init.Linear):
|
||||||
|
ldm_patched_cast_weights = True
|
||||||
|
|
||||||
|
class Conv2d(disable_weight_init.Conv2d):
|
||||||
|
ldm_patched_cast_weights = True
|
||||||
|
|
||||||
|
class Conv3d(disable_weight_init.Conv3d):
|
||||||
|
ldm_patched_cast_weights = True
|
||||||
|
|
||||||
|
class GroupNorm(disable_weight_init.GroupNorm):
|
||||||
|
ldm_patched_cast_weights = True
|
||||||
|
|
||||||
|
class LayerNorm(disable_weight_init.LayerNorm):
|
||||||
|
ldm_patched_cast_weights = True
|
6
ldm_patched/modules/options.py
Normal file
6
ldm_patched/modules/options.py
Normal file
@ -0,0 +1,6 @@
|
|||||||
|
|
||||||
|
args_parsing = False
|
||||||
|
|
||||||
|
def enable_args_parsing(enable=True):
|
||||||
|
global args_parsing
|
||||||
|
args_parsing = enable
|
118
ldm_patched/modules/sample.py
Normal file
118
ldm_patched/modules/sample.py
Normal file
@ -0,0 +1,118 @@
|
|||||||
|
import torch
|
||||||
|
import ldm_patched.modules.model_management
|
||||||
|
import ldm_patched.modules.samplers
|
||||||
|
import ldm_patched.modules.conds
|
||||||
|
import ldm_patched.modules.utils
|
||||||
|
import math
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
def prepare_noise(latent_image, seed, noise_inds=None):
|
||||||
|
"""
|
||||||
|
creates random noise given a latent image and a seed.
|
||||||
|
optional arg skip can be used to skip and discard x number of noise generations for a given seed
|
||||||
|
"""
|
||||||
|
generator = torch.manual_seed(seed)
|
||||||
|
if noise_inds is None:
|
||||||
|
return torch.randn(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, generator=generator, device="cpu")
|
||||||
|
|
||||||
|
unique_inds, inverse = np.unique(noise_inds, return_inverse=True)
|
||||||
|
noises = []
|
||||||
|
for i in range(unique_inds[-1]+1):
|
||||||
|
noise = torch.randn([1] + list(latent_image.size())[1:], dtype=latent_image.dtype, layout=latent_image.layout, generator=generator, device="cpu")
|
||||||
|
if i in unique_inds:
|
||||||
|
noises.append(noise)
|
||||||
|
noises = [noises[i] for i in inverse]
|
||||||
|
noises = torch.cat(noises, axis=0)
|
||||||
|
return noises
|
||||||
|
|
||||||
|
def prepare_mask(noise_mask, shape, device):
|
||||||
|
"""ensures noise mask is of proper dimensions"""
|
||||||
|
noise_mask = torch.nn.functional.interpolate(noise_mask.reshape((-1, 1, noise_mask.shape[-2], noise_mask.shape[-1])), size=(shape[2], shape[3]), mode="bilinear")
|
||||||
|
noise_mask = torch.cat([noise_mask] * shape[1], dim=1)
|
||||||
|
noise_mask = ldm_patched.modules.utils.repeat_to_batch_size(noise_mask, shape[0])
|
||||||
|
noise_mask = noise_mask.to(device)
|
||||||
|
return noise_mask
|
||||||
|
|
||||||
|
def get_models_from_cond(cond, model_type):
|
||||||
|
models = []
|
||||||
|
for c in cond:
|
||||||
|
if model_type in c:
|
||||||
|
models += [c[model_type]]
|
||||||
|
return models
|
||||||
|
|
||||||
|
def convert_cond(cond):
|
||||||
|
out = []
|
||||||
|
for c in cond:
|
||||||
|
temp = c[1].copy()
|
||||||
|
model_conds = temp.get("model_conds", {})
|
||||||
|
if c[0] is not None:
|
||||||
|
model_conds["c_crossattn"] = ldm_patched.modules.conds.CONDCrossAttn(c[0]) #TODO: remove
|
||||||
|
temp["cross_attn"] = c[0]
|
||||||
|
temp["model_conds"] = model_conds
|
||||||
|
out.append(temp)
|
||||||
|
return out
|
||||||
|
|
||||||
|
def get_additional_models(positive, negative, dtype):
|
||||||
|
"""loads additional models in positive and negative conditioning"""
|
||||||
|
control_nets = set(get_models_from_cond(positive, "control") + get_models_from_cond(negative, "control"))
|
||||||
|
|
||||||
|
inference_memory = 0
|
||||||
|
control_models = []
|
||||||
|
for m in control_nets:
|
||||||
|
control_models += m.get_models()
|
||||||
|
inference_memory += m.inference_memory_requirements(dtype)
|
||||||
|
|
||||||
|
gligen = get_models_from_cond(positive, "gligen") + get_models_from_cond(negative, "gligen")
|
||||||
|
gligen = [x[1] for x in gligen]
|
||||||
|
models = control_models + gligen
|
||||||
|
return models, inference_memory
|
||||||
|
|
||||||
|
def cleanup_additional_models(models):
|
||||||
|
"""cleanup additional models that were loaded"""
|
||||||
|
for m in models:
|
||||||
|
if hasattr(m, 'cleanup'):
|
||||||
|
m.cleanup()
|
||||||
|
|
||||||
|
def prepare_sampling(model, noise_shape, positive, negative, noise_mask):
|
||||||
|
device = model.load_device
|
||||||
|
positive = convert_cond(positive)
|
||||||
|
negative = convert_cond(negative)
|
||||||
|
|
||||||
|
if noise_mask is not None:
|
||||||
|
noise_mask = prepare_mask(noise_mask, noise_shape, device)
|
||||||
|
|
||||||
|
real_model = None
|
||||||
|
models, inference_memory = get_additional_models(positive, negative, model.model_dtype())
|
||||||
|
ldm_patched.modules.model_management.load_models_gpu([model] + models, model.memory_required([noise_shape[0] * 2] + list(noise_shape[1:])) + inference_memory)
|
||||||
|
real_model = model.model
|
||||||
|
|
||||||
|
return real_model, positive, negative, noise_mask, models
|
||||||
|
|
||||||
|
|
||||||
|
def sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0, disable_noise=False, start_step=None, last_step=None, force_full_denoise=False, noise_mask=None, sigmas=None, callback=None, disable_pbar=False, seed=None):
|
||||||
|
real_model, positive_copy, negative_copy, noise_mask, models = prepare_sampling(model, noise.shape, positive, negative, noise_mask)
|
||||||
|
|
||||||
|
noise = noise.to(model.load_device)
|
||||||
|
latent_image = latent_image.to(model.load_device)
|
||||||
|
|
||||||
|
sampler = ldm_patched.modules.samplers.KSampler(real_model, steps=steps, device=model.load_device, sampler=sampler_name, scheduler=scheduler, denoise=denoise, model_options=model.model_options)
|
||||||
|
|
||||||
|
samples = sampler.sample(noise, positive_copy, negative_copy, cfg=cfg, latent_image=latent_image, start_step=start_step, last_step=last_step, force_full_denoise=force_full_denoise, denoise_mask=noise_mask, sigmas=sigmas, callback=callback, disable_pbar=disable_pbar, seed=seed)
|
||||||
|
samples = samples.to(ldm_patched.modules.model_management.intermediate_device())
|
||||||
|
|
||||||
|
cleanup_additional_models(models)
|
||||||
|
cleanup_additional_models(set(get_models_from_cond(positive_copy, "control") + get_models_from_cond(negative_copy, "control")))
|
||||||
|
return samples
|
||||||
|
|
||||||
|
def sample_custom(model, noise, cfg, sampler, sigmas, positive, negative, latent_image, noise_mask=None, callback=None, disable_pbar=False, seed=None):
|
||||||
|
real_model, positive_copy, negative_copy, noise_mask, models = prepare_sampling(model, noise.shape, positive, negative, noise_mask)
|
||||||
|
noise = noise.to(model.load_device)
|
||||||
|
latent_image = latent_image.to(model.load_device)
|
||||||
|
sigmas = sigmas.to(model.load_device)
|
||||||
|
|
||||||
|
samples = ldm_patched.modules.samplers.sample(real_model, noise, positive_copy, negative_copy, cfg, model.load_device, sampler, sigmas, model_options=model.model_options, latent_image=latent_image, denoise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed)
|
||||||
|
samples = samples.to(ldm_patched.modules.model_management.intermediate_device())
|
||||||
|
cleanup_additional_models(models)
|
||||||
|
cleanup_additional_models(set(get_models_from_cond(positive_copy, "control") + get_models_from_cond(negative_copy, "control")))
|
||||||
|
return samples
|
||||||
|
|
716
ldm_patched/modules/samplers.py
Normal file
716
ldm_patched/modules/samplers.py
Normal file
@ -0,0 +1,716 @@
|
|||||||
|
from ldm_patched.k_diffusion import sampling as k_diffusion_sampling
|
||||||
|
from ldm_patched.unipc import uni_pc
|
||||||
|
import torch
|
||||||
|
import enum
|
||||||
|
import collections
|
||||||
|
from ldm_patched.modules import model_management
|
||||||
|
import math
|
||||||
|
from ldm_patched.modules import model_base
|
||||||
|
import ldm_patched.modules.utils
|
||||||
|
import ldm_patched.modules.conds
|
||||||
|
|
||||||
|
def get_area_and_mult(conds, x_in, timestep_in):
|
||||||
|
area = (x_in.shape[2], x_in.shape[3], 0, 0)
|
||||||
|
strength = 1.0
|
||||||
|
|
||||||
|
if 'timestep_start' in conds:
|
||||||
|
timestep_start = conds['timestep_start']
|
||||||
|
if timestep_in[0] > timestep_start:
|
||||||
|
return None
|
||||||
|
if 'timestep_end' in conds:
|
||||||
|
timestep_end = conds['timestep_end']
|
||||||
|
if timestep_in[0] < timestep_end:
|
||||||
|
return None
|
||||||
|
if 'area' in conds:
|
||||||
|
area = conds['area']
|
||||||
|
if 'strength' in conds:
|
||||||
|
strength = conds['strength']
|
||||||
|
|
||||||
|
input_x = x_in[:,:,area[2]:area[0] + area[2],area[3]:area[1] + area[3]]
|
||||||
|
if 'mask' in conds:
|
||||||
|
# Scale the mask to the size of the input
|
||||||
|
# The mask should have been resized as we began the sampling process
|
||||||
|
mask_strength = 1.0
|
||||||
|
if "mask_strength" in conds:
|
||||||
|
mask_strength = conds["mask_strength"]
|
||||||
|
mask = conds['mask']
|
||||||
|
assert(mask.shape[1] == x_in.shape[2])
|
||||||
|
assert(mask.shape[2] == x_in.shape[3])
|
||||||
|
mask = mask[:,area[2]:area[0] + area[2],area[3]:area[1] + area[3]] * mask_strength
|
||||||
|
mask = mask.unsqueeze(1).repeat(input_x.shape[0] // mask.shape[0], input_x.shape[1], 1, 1)
|
||||||
|
else:
|
||||||
|
mask = torch.ones_like(input_x)
|
||||||
|
mult = mask * strength
|
||||||
|
|
||||||
|
if 'mask' not in conds:
|
||||||
|
rr = 8
|
||||||
|
if area[2] != 0:
|
||||||
|
for t in range(rr):
|
||||||
|
mult[:,:,t:1+t,:] *= ((1.0/rr) * (t + 1))
|
||||||
|
if (area[0] + area[2]) < x_in.shape[2]:
|
||||||
|
for t in range(rr):
|
||||||
|
mult[:,:,area[0] - 1 - t:area[0] - t,:] *= ((1.0/rr) * (t + 1))
|
||||||
|
if area[3] != 0:
|
||||||
|
for t in range(rr):
|
||||||
|
mult[:,:,:,t:1+t] *= ((1.0/rr) * (t + 1))
|
||||||
|
if (area[1] + area[3]) < x_in.shape[3]:
|
||||||
|
for t in range(rr):
|
||||||
|
mult[:,:,:,area[1] - 1 - t:area[1] - t] *= ((1.0/rr) * (t + 1))
|
||||||
|
|
||||||
|
conditioning = {}
|
||||||
|
model_conds = conds["model_conds"]
|
||||||
|
for c in model_conds:
|
||||||
|
conditioning[c] = model_conds[c].process_cond(batch_size=x_in.shape[0], device=x_in.device, area=area)
|
||||||
|
|
||||||
|
control = conds.get('control', None)
|
||||||
|
|
||||||
|
patches = None
|
||||||
|
if 'gligen' in conds:
|
||||||
|
gligen = conds['gligen']
|
||||||
|
patches = {}
|
||||||
|
gligen_type = gligen[0]
|
||||||
|
gligen_model = gligen[1]
|
||||||
|
if gligen_type == "position":
|
||||||
|
gligen_patch = gligen_model.model.set_position(input_x.shape, gligen[2], input_x.device)
|
||||||
|
else:
|
||||||
|
gligen_patch = gligen_model.model.set_empty(input_x.shape, input_x.device)
|
||||||
|
|
||||||
|
patches['middle_patch'] = [gligen_patch]
|
||||||
|
|
||||||
|
cond_obj = collections.namedtuple('cond_obj', ['input_x', 'mult', 'conditioning', 'area', 'control', 'patches'])
|
||||||
|
return cond_obj(input_x, mult, conditioning, area, control, patches)
|
||||||
|
|
||||||
|
def cond_equal_size(c1, c2):
|
||||||
|
if c1 is c2:
|
||||||
|
return True
|
||||||
|
if c1.keys() != c2.keys():
|
||||||
|
return False
|
||||||
|
for k in c1:
|
||||||
|
if not c1[k].can_concat(c2[k]):
|
||||||
|
return False
|
||||||
|
return True
|
||||||
|
|
||||||
|
def can_concat_cond(c1, c2):
|
||||||
|
if c1.input_x.shape != c2.input_x.shape:
|
||||||
|
return False
|
||||||
|
|
||||||
|
def objects_concatable(obj1, obj2):
|
||||||
|
if (obj1 is None) != (obj2 is None):
|
||||||
|
return False
|
||||||
|
if obj1 is not None:
|
||||||
|
if obj1 is not obj2:
|
||||||
|
return False
|
||||||
|
return True
|
||||||
|
|
||||||
|
if not objects_concatable(c1.control, c2.control):
|
||||||
|
return False
|
||||||
|
|
||||||
|
if not objects_concatable(c1.patches, c2.patches):
|
||||||
|
return False
|
||||||
|
|
||||||
|
return cond_equal_size(c1.conditioning, c2.conditioning)
|
||||||
|
|
||||||
|
def cond_cat(c_list):
|
||||||
|
c_crossattn = []
|
||||||
|
c_concat = []
|
||||||
|
c_adm = []
|
||||||
|
crossattn_max_len = 0
|
||||||
|
|
||||||
|
temp = {}
|
||||||
|
for x in c_list:
|
||||||
|
for k in x:
|
||||||
|
cur = temp.get(k, [])
|
||||||
|
cur.append(x[k])
|
||||||
|
temp[k] = cur
|
||||||
|
|
||||||
|
out = {}
|
||||||
|
for k in temp:
|
||||||
|
conds = temp[k]
|
||||||
|
out[k] = conds[0].concat(conds[1:])
|
||||||
|
|
||||||
|
return out
|
||||||
|
|
||||||
|
def calc_cond_uncond_batch(model, cond, uncond, x_in, timestep, model_options):
|
||||||
|
out_cond = torch.zeros_like(x_in)
|
||||||
|
out_count = torch.ones_like(x_in) * 1e-37
|
||||||
|
|
||||||
|
out_uncond = torch.zeros_like(x_in)
|
||||||
|
out_uncond_count = torch.ones_like(x_in) * 1e-37
|
||||||
|
|
||||||
|
COND = 0
|
||||||
|
UNCOND = 1
|
||||||
|
|
||||||
|
to_run = []
|
||||||
|
for x in cond:
|
||||||
|
p = get_area_and_mult(x, x_in, timestep)
|
||||||
|
if p is None:
|
||||||
|
continue
|
||||||
|
|
||||||
|
to_run += [(p, COND)]
|
||||||
|
if uncond is not None:
|
||||||
|
for x in uncond:
|
||||||
|
p = get_area_and_mult(x, x_in, timestep)
|
||||||
|
if p is None:
|
||||||
|
continue
|
||||||
|
|
||||||
|
to_run += [(p, UNCOND)]
|
||||||
|
|
||||||
|
while len(to_run) > 0:
|
||||||
|
first = to_run[0]
|
||||||
|
first_shape = first[0][0].shape
|
||||||
|
to_batch_temp = []
|
||||||
|
for x in range(len(to_run)):
|
||||||
|
if can_concat_cond(to_run[x][0], first[0]):
|
||||||
|
to_batch_temp += [x]
|
||||||
|
|
||||||
|
to_batch_temp.reverse()
|
||||||
|
to_batch = to_batch_temp[:1]
|
||||||
|
|
||||||
|
free_memory = model_management.get_free_memory(x_in.device)
|
||||||
|
for i in range(1, len(to_batch_temp) + 1):
|
||||||
|
batch_amount = to_batch_temp[:len(to_batch_temp)//i]
|
||||||
|
input_shape = [len(batch_amount) * first_shape[0]] + list(first_shape)[1:]
|
||||||
|
if model.memory_required(input_shape) < free_memory:
|
||||||
|
to_batch = batch_amount
|
||||||
|
break
|
||||||
|
|
||||||
|
input_x = []
|
||||||
|
mult = []
|
||||||
|
c = []
|
||||||
|
cond_or_uncond = []
|
||||||
|
area = []
|
||||||
|
control = None
|
||||||
|
patches = None
|
||||||
|
for x in to_batch:
|
||||||
|
o = to_run.pop(x)
|
||||||
|
p = o[0]
|
||||||
|
input_x.append(p.input_x)
|
||||||
|
mult.append(p.mult)
|
||||||
|
c.append(p.conditioning)
|
||||||
|
area.append(p.area)
|
||||||
|
cond_or_uncond.append(o[1])
|
||||||
|
control = p.control
|
||||||
|
patches = p.patches
|
||||||
|
|
||||||
|
batch_chunks = len(cond_or_uncond)
|
||||||
|
input_x = torch.cat(input_x)
|
||||||
|
c = cond_cat(c)
|
||||||
|
timestep_ = torch.cat([timestep] * batch_chunks)
|
||||||
|
|
||||||
|
if control is not None:
|
||||||
|
c['control'] = control.get_control(input_x, timestep_, c, len(cond_or_uncond))
|
||||||
|
|
||||||
|
transformer_options = {}
|
||||||
|
if 'transformer_options' in model_options:
|
||||||
|
transformer_options = model_options['transformer_options'].copy()
|
||||||
|
|
||||||
|
if patches is not None:
|
||||||
|
if "patches" in transformer_options:
|
||||||
|
cur_patches = transformer_options["patches"].copy()
|
||||||
|
for p in patches:
|
||||||
|
if p in cur_patches:
|
||||||
|
cur_patches[p] = cur_patches[p] + patches[p]
|
||||||
|
else:
|
||||||
|
cur_patches[p] = patches[p]
|
||||||
|
else:
|
||||||
|
transformer_options["patches"] = patches
|
||||||
|
|
||||||
|
transformer_options["cond_or_uncond"] = cond_or_uncond[:]
|
||||||
|
transformer_options["sigmas"] = timestep
|
||||||
|
|
||||||
|
c['transformer_options'] = transformer_options
|
||||||
|
|
||||||
|
if 'model_function_wrapper' in model_options:
|
||||||
|
output = model_options['model_function_wrapper'](model.apply_model, {"input": input_x, "timestep": timestep_, "c": c, "cond_or_uncond": cond_or_uncond}).chunk(batch_chunks)
|
||||||
|
else:
|
||||||
|
output = model.apply_model(input_x, timestep_, **c).chunk(batch_chunks)
|
||||||
|
del input_x
|
||||||
|
|
||||||
|
for o in range(batch_chunks):
|
||||||
|
if cond_or_uncond[o] == COND:
|
||||||
|
out_cond[:,:,area[o][2]:area[o][0] + area[o][2],area[o][3]:area[o][1] + area[o][3]] += output[o] * mult[o]
|
||||||
|
out_count[:,:,area[o][2]:area[o][0] + area[o][2],area[o][3]:area[o][1] + area[o][3]] += mult[o]
|
||||||
|
else:
|
||||||
|
out_uncond[:,:,area[o][2]:area[o][0] + area[o][2],area[o][3]:area[o][1] + area[o][3]] += output[o] * mult[o]
|
||||||
|
out_uncond_count[:,:,area[o][2]:area[o][0] + area[o][2],area[o][3]:area[o][1] + area[o][3]] += mult[o]
|
||||||
|
del mult
|
||||||
|
|
||||||
|
out_cond /= out_count
|
||||||
|
del out_count
|
||||||
|
out_uncond /= out_uncond_count
|
||||||
|
del out_uncond_count
|
||||||
|
return out_cond, out_uncond
|
||||||
|
|
||||||
|
#The main sampling function shared by all the samplers
|
||||||
|
#Returns denoised
|
||||||
|
def sampling_function(model, x, timestep, uncond, cond, cond_scale, model_options={}, seed=None):
|
||||||
|
if math.isclose(cond_scale, 1.0) and model_options.get("disable_cfg1_optimization", False) == False:
|
||||||
|
uncond_ = None
|
||||||
|
else:
|
||||||
|
uncond_ = uncond
|
||||||
|
|
||||||
|
cond_pred, uncond_pred = calc_cond_uncond_batch(model, cond, uncond_, x, timestep, model_options)
|
||||||
|
if "sampler_cfg_function" in model_options:
|
||||||
|
args = {"cond": x - cond_pred, "uncond": x - uncond_pred, "cond_scale": cond_scale, "timestep": timestep, "input": x, "sigma": timestep,
|
||||||
|
"cond_denoised": cond_pred, "uncond_denoised": uncond_pred, "model": model, "model_options": model_options}
|
||||||
|
cfg_result = x - model_options["sampler_cfg_function"](args)
|
||||||
|
else:
|
||||||
|
cfg_result = uncond_pred + (cond_pred - uncond_pred) * cond_scale
|
||||||
|
|
||||||
|
for fn in model_options.get("sampler_post_cfg_function", []):
|
||||||
|
args = {"denoised": cfg_result, "cond": cond, "uncond": uncond, "model": model, "uncond_denoised": uncond_pred, "cond_denoised": cond_pred,
|
||||||
|
"sigma": timestep, "model_options": model_options, "input": x}
|
||||||
|
cfg_result = fn(args)
|
||||||
|
|
||||||
|
return cfg_result
|
||||||
|
|
||||||
|
class CFGNoisePredictor(torch.nn.Module):
|
||||||
|
def __init__(self, model):
|
||||||
|
super().__init__()
|
||||||
|
self.inner_model = model
|
||||||
|
def apply_model(self, x, timestep, cond, uncond, cond_scale, model_options={}, seed=None):
|
||||||
|
out = sampling_function(self.inner_model, x, timestep, uncond, cond, cond_scale, model_options=model_options, seed=seed)
|
||||||
|
return out
|
||||||
|
def forward(self, *args, **kwargs):
|
||||||
|
return self.apply_model(*args, **kwargs)
|
||||||
|
|
||||||
|
class KSamplerX0Inpaint(torch.nn.Module):
|
||||||
|
def __init__(self, model):
|
||||||
|
super().__init__()
|
||||||
|
self.inner_model = model
|
||||||
|
def forward(self, x, sigma, uncond, cond, cond_scale, denoise_mask, model_options={}, seed=None):
|
||||||
|
if denoise_mask is not None:
|
||||||
|
latent_mask = 1. - denoise_mask
|
||||||
|
x = x * denoise_mask + (self.latent_image + self.noise * sigma.reshape([sigma.shape[0]] + [1] * (len(self.noise.shape) - 1))) * latent_mask
|
||||||
|
out = self.inner_model(x, sigma, cond=cond, uncond=uncond, cond_scale=cond_scale, model_options=model_options, seed=seed)
|
||||||
|
if denoise_mask is not None:
|
||||||
|
out = out * denoise_mask + self.latent_image * latent_mask
|
||||||
|
return out
|
||||||
|
|
||||||
|
def simple_scheduler(model, steps):
|
||||||
|
s = model.model_sampling
|
||||||
|
sigs = []
|
||||||
|
ss = len(s.sigmas) / steps
|
||||||
|
for x in range(steps):
|
||||||
|
sigs += [float(s.sigmas[-(1 + int(x * ss))])]
|
||||||
|
sigs += [0.0]
|
||||||
|
return torch.FloatTensor(sigs)
|
||||||
|
|
||||||
|
def ddim_scheduler(model, steps):
|
||||||
|
s = model.model_sampling
|
||||||
|
sigs = []
|
||||||
|
ss = len(s.sigmas) // steps
|
||||||
|
x = 1
|
||||||
|
while x < len(s.sigmas):
|
||||||
|
sigs += [float(s.sigmas[x])]
|
||||||
|
x += ss
|
||||||
|
sigs = sigs[::-1]
|
||||||
|
sigs += [0.0]
|
||||||
|
return torch.FloatTensor(sigs)
|
||||||
|
|
||||||
|
def normal_scheduler(model, steps, sgm=False, floor=False):
|
||||||
|
s = model.model_sampling
|
||||||
|
start = s.timestep(s.sigma_max)
|
||||||
|
end = s.timestep(s.sigma_min)
|
||||||
|
|
||||||
|
if sgm:
|
||||||
|
timesteps = torch.linspace(start, end, steps + 1)[:-1]
|
||||||
|
else:
|
||||||
|
timesteps = torch.linspace(start, end, steps)
|
||||||
|
|
||||||
|
sigs = []
|
||||||
|
for x in range(len(timesteps)):
|
||||||
|
ts = timesteps[x]
|
||||||
|
sigs.append(s.sigma(ts))
|
||||||
|
sigs += [0.0]
|
||||||
|
return torch.FloatTensor(sigs)
|
||||||
|
|
||||||
|
def get_mask_aabb(masks):
|
||||||
|
if masks.numel() == 0:
|
||||||
|
return torch.zeros((0, 4), device=masks.device, dtype=torch.int)
|
||||||
|
|
||||||
|
b = masks.shape[0]
|
||||||
|
|
||||||
|
bounding_boxes = torch.zeros((b, 4), device=masks.device, dtype=torch.int)
|
||||||
|
is_empty = torch.zeros((b), device=masks.device, dtype=torch.bool)
|
||||||
|
for i in range(b):
|
||||||
|
mask = masks[i]
|
||||||
|
if mask.numel() == 0:
|
||||||
|
continue
|
||||||
|
if torch.max(mask != 0) == False:
|
||||||
|
is_empty[i] = True
|
||||||
|
continue
|
||||||
|
y, x = torch.where(mask)
|
||||||
|
bounding_boxes[i, 0] = torch.min(x)
|
||||||
|
bounding_boxes[i, 1] = torch.min(y)
|
||||||
|
bounding_boxes[i, 2] = torch.max(x)
|
||||||
|
bounding_boxes[i, 3] = torch.max(y)
|
||||||
|
|
||||||
|
return bounding_boxes, is_empty
|
||||||
|
|
||||||
|
def resolve_areas_and_cond_masks(conditions, h, w, device):
|
||||||
|
# We need to decide on an area outside the sampling loop in order to properly generate opposite areas of equal sizes.
|
||||||
|
# While we're doing this, we can also resolve the mask device and scaling for performance reasons
|
||||||
|
for i in range(len(conditions)):
|
||||||
|
c = conditions[i]
|
||||||
|
if 'area' in c:
|
||||||
|
area = c['area']
|
||||||
|
if area[0] == "percentage":
|
||||||
|
modified = c.copy()
|
||||||
|
area = (max(1, round(area[1] * h)), max(1, round(area[2] * w)), round(area[3] * h), round(area[4] * w))
|
||||||
|
modified['area'] = area
|
||||||
|
c = modified
|
||||||
|
conditions[i] = c
|
||||||
|
|
||||||
|
if 'mask' in c:
|
||||||
|
mask = c['mask']
|
||||||
|
mask = mask.to(device=device)
|
||||||
|
modified = c.copy()
|
||||||
|
if len(mask.shape) == 2:
|
||||||
|
mask = mask.unsqueeze(0)
|
||||||
|
if mask.shape[1] != h or mask.shape[2] != w:
|
||||||
|
mask = torch.nn.functional.interpolate(mask.unsqueeze(1), size=(h, w), mode='bilinear', align_corners=False).squeeze(1)
|
||||||
|
|
||||||
|
if modified.get("set_area_to_bounds", False):
|
||||||
|
bounds = torch.max(torch.abs(mask),dim=0).values.unsqueeze(0)
|
||||||
|
boxes, is_empty = get_mask_aabb(bounds)
|
||||||
|
if is_empty[0]:
|
||||||
|
# Use the minimum possible size for efficiency reasons. (Since the mask is all-0, this becomes a noop anyway)
|
||||||
|
modified['area'] = (8, 8, 0, 0)
|
||||||
|
else:
|
||||||
|
box = boxes[0]
|
||||||
|
H, W, Y, X = (box[3] - box[1] + 1, box[2] - box[0] + 1, box[1], box[0])
|
||||||
|
H = max(8, H)
|
||||||
|
W = max(8, W)
|
||||||
|
area = (int(H), int(W), int(Y), int(X))
|
||||||
|
modified['area'] = area
|
||||||
|
|
||||||
|
modified['mask'] = mask
|
||||||
|
conditions[i] = modified
|
||||||
|
|
||||||
|
def create_cond_with_same_area_if_none(conds, c):
|
||||||
|
if 'area' not in c:
|
||||||
|
return
|
||||||
|
|
||||||
|
c_area = c['area']
|
||||||
|
smallest = None
|
||||||
|
for x in conds:
|
||||||
|
if 'area' in x:
|
||||||
|
a = x['area']
|
||||||
|
if c_area[2] >= a[2] and c_area[3] >= a[3]:
|
||||||
|
if a[0] + a[2] >= c_area[0] + c_area[2]:
|
||||||
|
if a[1] + a[3] >= c_area[1] + c_area[3]:
|
||||||
|
if smallest is None:
|
||||||
|
smallest = x
|
||||||
|
elif 'area' not in smallest:
|
||||||
|
smallest = x
|
||||||
|
else:
|
||||||
|
if smallest['area'][0] * smallest['area'][1] > a[0] * a[1]:
|
||||||
|
smallest = x
|
||||||
|
else:
|
||||||
|
if smallest is None:
|
||||||
|
smallest = x
|
||||||
|
if smallest is None:
|
||||||
|
return
|
||||||
|
if 'area' in smallest:
|
||||||
|
if smallest['area'] == c_area:
|
||||||
|
return
|
||||||
|
|
||||||
|
out = c.copy()
|
||||||
|
out['model_conds'] = smallest['model_conds'].copy() #TODO: which fields should be copied?
|
||||||
|
conds += [out]
|
||||||
|
|
||||||
|
def calculate_start_end_timesteps(model, conds):
|
||||||
|
s = model.model_sampling
|
||||||
|
for t in range(len(conds)):
|
||||||
|
x = conds[t]
|
||||||
|
|
||||||
|
timestep_start = None
|
||||||
|
timestep_end = None
|
||||||
|
if 'start_percent' in x:
|
||||||
|
timestep_start = s.percent_to_sigma(x['start_percent'])
|
||||||
|
if 'end_percent' in x:
|
||||||
|
timestep_end = s.percent_to_sigma(x['end_percent'])
|
||||||
|
|
||||||
|
if (timestep_start is not None) or (timestep_end is not None):
|
||||||
|
n = x.copy()
|
||||||
|
if (timestep_start is not None):
|
||||||
|
n['timestep_start'] = timestep_start
|
||||||
|
if (timestep_end is not None):
|
||||||
|
n['timestep_end'] = timestep_end
|
||||||
|
conds[t] = n
|
||||||
|
|
||||||
|
def pre_run_control(model, conds):
|
||||||
|
s = model.model_sampling
|
||||||
|
for t in range(len(conds)):
|
||||||
|
x = conds[t]
|
||||||
|
|
||||||
|
timestep_start = None
|
||||||
|
timestep_end = None
|
||||||
|
percent_to_timestep_function = lambda a: s.percent_to_sigma(a)
|
||||||
|
if 'control' in x:
|
||||||
|
x['control'].pre_run(model, percent_to_timestep_function)
|
||||||
|
|
||||||
|
def apply_empty_x_to_equal_area(conds, uncond, name, uncond_fill_func):
|
||||||
|
cond_cnets = []
|
||||||
|
cond_other = []
|
||||||
|
uncond_cnets = []
|
||||||
|
uncond_other = []
|
||||||
|
for t in range(len(conds)):
|
||||||
|
x = conds[t]
|
||||||
|
if 'area' not in x:
|
||||||
|
if name in x and x[name] is not None:
|
||||||
|
cond_cnets.append(x[name])
|
||||||
|
else:
|
||||||
|
cond_other.append((x, t))
|
||||||
|
for t in range(len(uncond)):
|
||||||
|
x = uncond[t]
|
||||||
|
if 'area' not in x:
|
||||||
|
if name in x and x[name] is not None:
|
||||||
|
uncond_cnets.append(x[name])
|
||||||
|
else:
|
||||||
|
uncond_other.append((x, t))
|
||||||
|
|
||||||
|
if len(uncond_cnets) > 0:
|
||||||
|
return
|
||||||
|
|
||||||
|
for x in range(len(cond_cnets)):
|
||||||
|
temp = uncond_other[x % len(uncond_other)]
|
||||||
|
o = temp[0]
|
||||||
|
if name in o and o[name] is not None:
|
||||||
|
n = o.copy()
|
||||||
|
n[name] = uncond_fill_func(cond_cnets, x)
|
||||||
|
uncond += [n]
|
||||||
|
else:
|
||||||
|
n = o.copy()
|
||||||
|
n[name] = uncond_fill_func(cond_cnets, x)
|
||||||
|
uncond[temp[1]] = n
|
||||||
|
|
||||||
|
def encode_model_conds(model_function, conds, noise, device, prompt_type, **kwargs):
|
||||||
|
for t in range(len(conds)):
|
||||||
|
x = conds[t]
|
||||||
|
params = x.copy()
|
||||||
|
params["device"] = device
|
||||||
|
params["noise"] = noise
|
||||||
|
params["width"] = params.get("width", noise.shape[3] * 8)
|
||||||
|
params["height"] = params.get("height", noise.shape[2] * 8)
|
||||||
|
params["prompt_type"] = params.get("prompt_type", prompt_type)
|
||||||
|
for k in kwargs:
|
||||||
|
if k not in params:
|
||||||
|
params[k] = kwargs[k]
|
||||||
|
|
||||||
|
out = model_function(**params)
|
||||||
|
x = x.copy()
|
||||||
|
model_conds = x['model_conds'].copy()
|
||||||
|
for k in out:
|
||||||
|
model_conds[k] = out[k]
|
||||||
|
x['model_conds'] = model_conds
|
||||||
|
conds[t] = x
|
||||||
|
return conds
|
||||||
|
|
||||||
|
class Sampler:
|
||||||
|
def sample(self):
|
||||||
|
pass
|
||||||
|
|
||||||
|
def max_denoise(self, model_wrap, sigmas):
|
||||||
|
max_sigma = float(model_wrap.inner_model.model_sampling.sigma_max)
|
||||||
|
sigma = float(sigmas[0])
|
||||||
|
return math.isclose(max_sigma, sigma, rel_tol=1e-05) or sigma > max_sigma
|
||||||
|
|
||||||
|
class UNIPC(Sampler):
|
||||||
|
def sample(self, model_wrap, sigmas, extra_args, callback, noise, latent_image=None, denoise_mask=None, disable_pbar=False):
|
||||||
|
return uni_pc.sample_unipc(model_wrap, noise, latent_image, sigmas, max_denoise=self.max_denoise(model_wrap, sigmas), extra_args=extra_args, noise_mask=denoise_mask, callback=callback, disable=disable_pbar)
|
||||||
|
|
||||||
|
class UNIPCBH2(Sampler):
|
||||||
|
def sample(self, model_wrap, sigmas, extra_args, callback, noise, latent_image=None, denoise_mask=None, disable_pbar=False):
|
||||||
|
return uni_pc.sample_unipc(model_wrap, noise, latent_image, sigmas, max_denoise=self.max_denoise(model_wrap, sigmas), extra_args=extra_args, noise_mask=denoise_mask, callback=callback, variant='bh2', disable=disable_pbar)
|
||||||
|
|
||||||
|
KSAMPLER_NAMES = ["euler", "euler_ancestral", "heun", "heunpp2","dpm_2", "dpm_2_ancestral",
|
||||||
|
"lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_sde_gpu",
|
||||||
|
"dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm"]
|
||||||
|
|
||||||
|
class KSAMPLER(Sampler):
|
||||||
|
def __init__(self, sampler_function, extra_options={}, inpaint_options={}):
|
||||||
|
self.sampler_function = sampler_function
|
||||||
|
self.extra_options = extra_options
|
||||||
|
self.inpaint_options = inpaint_options
|
||||||
|
|
||||||
|
def sample(self, model_wrap, sigmas, extra_args, callback, noise, latent_image=None, denoise_mask=None, disable_pbar=False):
|
||||||
|
extra_args["denoise_mask"] = denoise_mask
|
||||||
|
model_k = KSamplerX0Inpaint(model_wrap)
|
||||||
|
model_k.latent_image = latent_image
|
||||||
|
if self.inpaint_options.get("random", False): #TODO: Should this be the default?
|
||||||
|
generator = torch.manual_seed(extra_args.get("seed", 41) + 1)
|
||||||
|
model_k.noise = torch.randn(noise.shape, generator=generator, device="cpu").to(noise.dtype).to(noise.device)
|
||||||
|
else:
|
||||||
|
model_k.noise = noise
|
||||||
|
|
||||||
|
if self.max_denoise(model_wrap, sigmas):
|
||||||
|
noise = noise * torch.sqrt(1.0 + sigmas[0] ** 2.0)
|
||||||
|
else:
|
||||||
|
noise = noise * sigmas[0]
|
||||||
|
|
||||||
|
k_callback = None
|
||||||
|
total_steps = len(sigmas) - 1
|
||||||
|
if callback is not None:
|
||||||
|
k_callback = lambda x: callback(x["i"], x["denoised"], x["x"], total_steps)
|
||||||
|
|
||||||
|
if latent_image is not None:
|
||||||
|
noise += latent_image
|
||||||
|
|
||||||
|
samples = self.sampler_function(model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar, **self.extra_options)
|
||||||
|
return samples
|
||||||
|
|
||||||
|
|
||||||
|
def ksampler(sampler_name, extra_options={}, inpaint_options={}):
|
||||||
|
if sampler_name == "dpm_fast":
|
||||||
|
def dpm_fast_function(model, noise, sigmas, extra_args, callback, disable):
|
||||||
|
sigma_min = sigmas[-1]
|
||||||
|
if sigma_min == 0:
|
||||||
|
sigma_min = sigmas[-2]
|
||||||
|
total_steps = len(sigmas) - 1
|
||||||
|
return k_diffusion_sampling.sample_dpm_fast(model, noise, sigma_min, sigmas[0], total_steps, extra_args=extra_args, callback=callback, disable=disable)
|
||||||
|
sampler_function = dpm_fast_function
|
||||||
|
elif sampler_name == "dpm_adaptive":
|
||||||
|
def dpm_adaptive_function(model, noise, sigmas, extra_args, callback, disable):
|
||||||
|
sigma_min = sigmas[-1]
|
||||||
|
if sigma_min == 0:
|
||||||
|
sigma_min = sigmas[-2]
|
||||||
|
return k_diffusion_sampling.sample_dpm_adaptive(model, noise, sigma_min, sigmas[0], extra_args=extra_args, callback=callback, disable=disable)
|
||||||
|
sampler_function = dpm_adaptive_function
|
||||||
|
else:
|
||||||
|
sampler_function = getattr(k_diffusion_sampling, "sample_{}".format(sampler_name))
|
||||||
|
|
||||||
|
return KSAMPLER(sampler_function, extra_options, inpaint_options)
|
||||||
|
|
||||||
|
def wrap_model(model):
|
||||||
|
model_denoise = CFGNoisePredictor(model)
|
||||||
|
return model_denoise
|
||||||
|
|
||||||
|
def sample(model, noise, positive, negative, cfg, device, sampler, sigmas, model_options={}, latent_image=None, denoise_mask=None, callback=None, disable_pbar=False, seed=None):
|
||||||
|
positive = positive[:]
|
||||||
|
negative = negative[:]
|
||||||
|
|
||||||
|
resolve_areas_and_cond_masks(positive, noise.shape[2], noise.shape[3], device)
|
||||||
|
resolve_areas_and_cond_masks(negative, noise.shape[2], noise.shape[3], device)
|
||||||
|
|
||||||
|
model_wrap = wrap_model(model)
|
||||||
|
|
||||||
|
calculate_start_end_timesteps(model, negative)
|
||||||
|
calculate_start_end_timesteps(model, positive)
|
||||||
|
|
||||||
|
if latent_image is not None:
|
||||||
|
latent_image = model.process_latent_in(latent_image)
|
||||||
|
|
||||||
|
if hasattr(model, 'extra_conds'):
|
||||||
|
positive = encode_model_conds(model.extra_conds, positive, noise, device, "positive", latent_image=latent_image, denoise_mask=denoise_mask, seed=seed)
|
||||||
|
negative = encode_model_conds(model.extra_conds, negative, noise, device, "negative", latent_image=latent_image, denoise_mask=denoise_mask, seed=seed)
|
||||||
|
|
||||||
|
#make sure each cond area has an opposite one with the same area
|
||||||
|
for c in positive:
|
||||||
|
create_cond_with_same_area_if_none(negative, c)
|
||||||
|
for c in negative:
|
||||||
|
create_cond_with_same_area_if_none(positive, c)
|
||||||
|
|
||||||
|
pre_run_control(model, negative + positive)
|
||||||
|
|
||||||
|
apply_empty_x_to_equal_area(list(filter(lambda c: c.get('control_apply_to_uncond', False) == True, positive)), negative, 'control', lambda cond_cnets, x: cond_cnets[x])
|
||||||
|
apply_empty_x_to_equal_area(positive, negative, 'gligen', lambda cond_cnets, x: cond_cnets[x])
|
||||||
|
|
||||||
|
extra_args = {"cond":positive, "uncond":negative, "cond_scale": cfg, "model_options": model_options, "seed":seed}
|
||||||
|
|
||||||
|
samples = sampler.sample(model_wrap, sigmas, extra_args, callback, noise, latent_image, denoise_mask, disable_pbar)
|
||||||
|
return model.process_latent_out(samples.to(torch.float32))
|
||||||
|
|
||||||
|
SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform"]
|
||||||
|
SAMPLER_NAMES = KSAMPLER_NAMES + ["ddim", "uni_pc", "uni_pc_bh2"]
|
||||||
|
|
||||||
|
def calculate_sigmas_scheduler(model, scheduler_name, steps):
|
||||||
|
if scheduler_name == "karras":
|
||||||
|
sigmas = k_diffusion_sampling.get_sigmas_karras(n=steps, sigma_min=float(model.model_sampling.sigma_min), sigma_max=float(model.model_sampling.sigma_max))
|
||||||
|
elif scheduler_name == "exponential":
|
||||||
|
sigmas = k_diffusion_sampling.get_sigmas_exponential(n=steps, sigma_min=float(model.model_sampling.sigma_min), sigma_max=float(model.model_sampling.sigma_max))
|
||||||
|
elif scheduler_name == "normal":
|
||||||
|
sigmas = normal_scheduler(model, steps)
|
||||||
|
elif scheduler_name == "simple":
|
||||||
|
sigmas = simple_scheduler(model, steps)
|
||||||
|
elif scheduler_name == "ddim_uniform":
|
||||||
|
sigmas = ddim_scheduler(model, steps)
|
||||||
|
elif scheduler_name == "sgm_uniform":
|
||||||
|
sigmas = normal_scheduler(model, steps, sgm=True)
|
||||||
|
else:
|
||||||
|
print("error invalid scheduler", self.scheduler)
|
||||||
|
return sigmas
|
||||||
|
|
||||||
|
def sampler_object(name):
|
||||||
|
if name == "uni_pc":
|
||||||
|
sampler = UNIPC()
|
||||||
|
elif name == "uni_pc_bh2":
|
||||||
|
sampler = UNIPCBH2()
|
||||||
|
elif name == "ddim":
|
||||||
|
sampler = ksampler("euler", inpaint_options={"random": True})
|
||||||
|
else:
|
||||||
|
sampler = ksampler(name)
|
||||||
|
return sampler
|
||||||
|
|
||||||
|
class KSampler:
|
||||||
|
SCHEDULERS = SCHEDULER_NAMES
|
||||||
|
SAMPLERS = SAMPLER_NAMES
|
||||||
|
|
||||||
|
def __init__(self, model, steps, device, sampler=None, scheduler=None, denoise=None, model_options={}):
|
||||||
|
self.model = model
|
||||||
|
self.device = device
|
||||||
|
if scheduler not in self.SCHEDULERS:
|
||||||
|
scheduler = self.SCHEDULERS[0]
|
||||||
|
if sampler not in self.SAMPLERS:
|
||||||
|
sampler = self.SAMPLERS[0]
|
||||||
|
self.scheduler = scheduler
|
||||||
|
self.sampler = sampler
|
||||||
|
self.set_steps(steps, denoise)
|
||||||
|
self.denoise = denoise
|
||||||
|
self.model_options = model_options
|
||||||
|
|
||||||
|
def calculate_sigmas(self, steps):
|
||||||
|
sigmas = None
|
||||||
|
|
||||||
|
discard_penultimate_sigma = False
|
||||||
|
if self.sampler in ['dpm_2', 'dpm_2_ancestral', 'uni_pc', 'uni_pc_bh2']:
|
||||||
|
steps += 1
|
||||||
|
discard_penultimate_sigma = True
|
||||||
|
|
||||||
|
sigmas = calculate_sigmas_scheduler(self.model, self.scheduler, steps)
|
||||||
|
|
||||||
|
if discard_penultimate_sigma:
|
||||||
|
sigmas = torch.cat([sigmas[:-2], sigmas[-1:]])
|
||||||
|
return sigmas
|
||||||
|
|
||||||
|
def set_steps(self, steps, denoise=None):
|
||||||
|
self.steps = steps
|
||||||
|
if denoise is None or denoise > 0.9999:
|
||||||
|
self.sigmas = self.calculate_sigmas(steps).to(self.device)
|
||||||
|
else:
|
||||||
|
new_steps = int(steps/denoise)
|
||||||
|
sigmas = self.calculate_sigmas(new_steps).to(self.device)
|
||||||
|
self.sigmas = sigmas[-(steps + 1):]
|
||||||
|
|
||||||
|
def sample(self, noise, positive, negative, cfg, latent_image=None, start_step=None, last_step=None, force_full_denoise=False, denoise_mask=None, sigmas=None, callback=None, disable_pbar=False, seed=None):
|
||||||
|
if sigmas is None:
|
||||||
|
sigmas = self.sigmas
|
||||||
|
|
||||||
|
if last_step is not None and last_step < (len(sigmas) - 1):
|
||||||
|
sigmas = sigmas[:last_step + 1]
|
||||||
|
if force_full_denoise:
|
||||||
|
sigmas[-1] = 0
|
||||||
|
|
||||||
|
if start_step is not None:
|
||||||
|
if start_step < (len(sigmas) - 1):
|
||||||
|
sigmas = sigmas[start_step:]
|
||||||
|
else:
|
||||||
|
if latent_image is not None:
|
||||||
|
return latent_image
|
||||||
|
else:
|
||||||
|
return torch.zeros_like(noise)
|
||||||
|
|
||||||
|
sampler = sampler_object(self.sampler)
|
||||||
|
|
||||||
|
return sample(self.model, noise, positive, negative, cfg, self.device, sampler, sigmas, self.model_options, latent_image=latent_image, denoise_mask=denoise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed)
|
540
ldm_patched/modules/sd.py
Normal file
540
ldm_patched/modules/sd.py
Normal file
@ -0,0 +1,540 @@
|
|||||||
|
import torch
|
||||||
|
import contextlib
|
||||||
|
import math
|
||||||
|
|
||||||
|
from ldm_patched.modules import model_management
|
||||||
|
from ldm_patched.ldm.util import instantiate_from_config
|
||||||
|
from ldm_patched.ldm.models.autoencoder import AutoencoderKL, AutoencodingEngine
|
||||||
|
import yaml
|
||||||
|
|
||||||
|
import ldm_patched.modules.utils
|
||||||
|
|
||||||
|
from . import clip_vision
|
||||||
|
from . import gligen
|
||||||
|
from . import diffusers_convert
|
||||||
|
from . import model_base
|
||||||
|
from . import model_detection
|
||||||
|
|
||||||
|
from . import sd1_clip
|
||||||
|
from . import sd2_clip
|
||||||
|
from . import sdxl_clip
|
||||||
|
|
||||||
|
import ldm_patched.modules.model_patcher
|
||||||
|
import ldm_patched.modules.lora
|
||||||
|
import ldm_patched.t2ia.adapter
|
||||||
|
import ldm_patched.modules.supported_models_base
|
||||||
|
import ldm_patched.taesd.taesd
|
||||||
|
|
||||||
|
def load_model_weights(model, sd):
|
||||||
|
m, u = model.load_state_dict(sd, strict=False)
|
||||||
|
m = set(m)
|
||||||
|
unexpected_keys = set(u)
|
||||||
|
|
||||||
|
k = list(sd.keys())
|
||||||
|
for x in k:
|
||||||
|
if x not in unexpected_keys:
|
||||||
|
w = sd.pop(x)
|
||||||
|
del w
|
||||||
|
if len(m) > 0:
|
||||||
|
print("extra", m)
|
||||||
|
return model
|
||||||
|
|
||||||
|
def load_clip_weights(model, sd):
|
||||||
|
k = list(sd.keys())
|
||||||
|
for x in k:
|
||||||
|
if x.startswith("cond_stage_model.transformer.") and not x.startswith("cond_stage_model.transformer.text_model."):
|
||||||
|
y = x.replace("cond_stage_model.transformer.", "cond_stage_model.transformer.text_model.")
|
||||||
|
sd[y] = sd.pop(x)
|
||||||
|
|
||||||
|
if 'cond_stage_model.transformer.text_model.embeddings.position_ids' in sd:
|
||||||
|
ids = sd['cond_stage_model.transformer.text_model.embeddings.position_ids']
|
||||||
|
if ids.dtype == torch.float32:
|
||||||
|
sd['cond_stage_model.transformer.text_model.embeddings.position_ids'] = ids.round()
|
||||||
|
|
||||||
|
sd = ldm_patched.modules.utils.transformers_convert(sd, "cond_stage_model.model.", "cond_stage_model.transformer.text_model.", 24)
|
||||||
|
return load_model_weights(model, sd)
|
||||||
|
|
||||||
|
|
||||||
|
def load_lora_for_models(model, clip, lora, strength_model, strength_clip):
|
||||||
|
key_map = {}
|
||||||
|
if model is not None:
|
||||||
|
key_map = ldm_patched.modules.lora.model_lora_keys_unet(model.model, key_map)
|
||||||
|
if clip is not None:
|
||||||
|
key_map = ldm_patched.modules.lora.model_lora_keys_clip(clip.cond_stage_model, key_map)
|
||||||
|
|
||||||
|
loaded = ldm_patched.modules.lora.load_lora(lora, key_map)
|
||||||
|
if model is not None:
|
||||||
|
new_modelpatcher = model.clone()
|
||||||
|
k = new_modelpatcher.add_patches(loaded, strength_model)
|
||||||
|
else:
|
||||||
|
k = ()
|
||||||
|
new_modelpatcher = None
|
||||||
|
|
||||||
|
if clip is not None:
|
||||||
|
new_clip = clip.clone()
|
||||||
|
k1 = new_clip.add_patches(loaded, strength_clip)
|
||||||
|
else:
|
||||||
|
k1 = ()
|
||||||
|
new_clip = None
|
||||||
|
k = set(k)
|
||||||
|
k1 = set(k1)
|
||||||
|
for x in loaded:
|
||||||
|
if (x not in k) and (x not in k1):
|
||||||
|
print("NOT LOADED", x)
|
||||||
|
|
||||||
|
return (new_modelpatcher, new_clip)
|
||||||
|
|
||||||
|
|
||||||
|
class CLIP:
|
||||||
|
def __init__(self, target=None, embedding_directory=None, no_init=False):
|
||||||
|
if no_init:
|
||||||
|
return
|
||||||
|
params = target.params.copy()
|
||||||
|
clip = target.clip
|
||||||
|
tokenizer = target.tokenizer
|
||||||
|
|
||||||
|
load_device = model_management.text_encoder_device()
|
||||||
|
offload_device = model_management.text_encoder_offload_device()
|
||||||
|
params['device'] = offload_device
|
||||||
|
params['dtype'] = model_management.text_encoder_dtype(load_device)
|
||||||
|
|
||||||
|
self.cond_stage_model = clip(**(params))
|
||||||
|
|
||||||
|
self.tokenizer = tokenizer(embedding_directory=embedding_directory)
|
||||||
|
self.patcher = ldm_patched.modules.model_patcher.ModelPatcher(self.cond_stage_model, load_device=load_device, offload_device=offload_device)
|
||||||
|
self.layer_idx = None
|
||||||
|
|
||||||
|
def clone(self):
|
||||||
|
n = CLIP(no_init=True)
|
||||||
|
n.patcher = self.patcher.clone()
|
||||||
|
n.cond_stage_model = self.cond_stage_model
|
||||||
|
n.tokenizer = self.tokenizer
|
||||||
|
n.layer_idx = self.layer_idx
|
||||||
|
return n
|
||||||
|
|
||||||
|
def add_patches(self, patches, strength_patch=1.0, strength_model=1.0):
|
||||||
|
return self.patcher.add_patches(patches, strength_patch, strength_model)
|
||||||
|
|
||||||
|
def clip_layer(self, layer_idx):
|
||||||
|
self.layer_idx = layer_idx
|
||||||
|
|
||||||
|
def tokenize(self, text, return_word_ids=False):
|
||||||
|
return self.tokenizer.tokenize_with_weights(text, return_word_ids)
|
||||||
|
|
||||||
|
def encode_from_tokens(self, tokens, return_pooled=False):
|
||||||
|
if self.layer_idx is not None:
|
||||||
|
self.cond_stage_model.clip_layer(self.layer_idx)
|
||||||
|
else:
|
||||||
|
self.cond_stage_model.reset_clip_layer()
|
||||||
|
|
||||||
|
self.load_model()
|
||||||
|
cond, pooled = self.cond_stage_model.encode_token_weights(tokens)
|
||||||
|
if return_pooled:
|
||||||
|
return cond, pooled
|
||||||
|
return cond
|
||||||
|
|
||||||
|
def encode(self, text):
|
||||||
|
tokens = self.tokenize(text)
|
||||||
|
return self.encode_from_tokens(tokens)
|
||||||
|
|
||||||
|
def load_sd(self, sd):
|
||||||
|
return self.cond_stage_model.load_sd(sd)
|
||||||
|
|
||||||
|
def get_sd(self):
|
||||||
|
return self.cond_stage_model.state_dict()
|
||||||
|
|
||||||
|
def load_model(self):
|
||||||
|
model_management.load_model_gpu(self.patcher)
|
||||||
|
return self.patcher
|
||||||
|
|
||||||
|
def get_key_patches(self):
|
||||||
|
return self.patcher.get_key_patches()
|
||||||
|
|
||||||
|
class VAE:
|
||||||
|
def __init__(self, sd=None, device=None, config=None, dtype=None):
|
||||||
|
if 'decoder.up_blocks.0.resnets.0.norm1.weight' in sd.keys(): #diffusers format
|
||||||
|
sd = diffusers_convert.convert_vae_state_dict(sd)
|
||||||
|
|
||||||
|
self.memory_used_encode = lambda shape, dtype: (1767 * shape[2] * shape[3]) * model_management.dtype_size(dtype) #These are for AutoencoderKL and need tweaking (should be lower)
|
||||||
|
self.memory_used_decode = lambda shape, dtype: (2178 * shape[2] * shape[3] * 64) * model_management.dtype_size(dtype)
|
||||||
|
self.downscale_ratio = 8
|
||||||
|
self.latent_channels = 4
|
||||||
|
|
||||||
|
if config is None:
|
||||||
|
if "decoder.mid.block_1.mix_factor" in sd:
|
||||||
|
encoder_config = {'double_z': True, 'z_channels': 4, 'resolution': 256, 'in_channels': 3, 'out_ch': 3, 'ch': 128, 'ch_mult': [1, 2, 4, 4], 'num_res_blocks': 2, 'attn_resolutions': [], 'dropout': 0.0}
|
||||||
|
decoder_config = encoder_config.copy()
|
||||||
|
decoder_config["video_kernel_size"] = [3, 1, 1]
|
||||||
|
decoder_config["alpha"] = 0.0
|
||||||
|
self.first_stage_model = AutoencodingEngine(regularizer_config={'target': "ldm_patched.ldm.models.autoencoder.DiagonalGaussianRegularizer"},
|
||||||
|
encoder_config={'target': "ldm_patched.ldm.modules.diffusionmodules.model.Encoder", 'params': encoder_config},
|
||||||
|
decoder_config={'target': "ldm_patched.ldm.modules.temporal_ae.VideoDecoder", 'params': decoder_config})
|
||||||
|
elif "taesd_decoder.1.weight" in sd:
|
||||||
|
self.first_stage_model = ldm_patched.taesd.taesd.TAESD()
|
||||||
|
else:
|
||||||
|
#default SD1.x/SD2.x VAE parameters
|
||||||
|
ddconfig = {'double_z': True, 'z_channels': 4, 'resolution': 256, 'in_channels': 3, 'out_ch': 3, 'ch': 128, 'ch_mult': [1, 2, 4, 4], 'num_res_blocks': 2, 'attn_resolutions': [], 'dropout': 0.0}
|
||||||
|
|
||||||
|
if 'encoder.down.2.downsample.conv.weight' not in sd: #Stable diffusion x4 upscaler VAE
|
||||||
|
ddconfig['ch_mult'] = [1, 2, 4]
|
||||||
|
self.downscale_ratio = 4
|
||||||
|
|
||||||
|
self.first_stage_model = AutoencoderKL(ddconfig=ddconfig, embed_dim=4)
|
||||||
|
else:
|
||||||
|
self.first_stage_model = AutoencoderKL(**(config['params']))
|
||||||
|
self.first_stage_model = self.first_stage_model.eval()
|
||||||
|
|
||||||
|
m, u = self.first_stage_model.load_state_dict(sd, strict=False)
|
||||||
|
if len(m) > 0:
|
||||||
|
print("Missing VAE keys", m)
|
||||||
|
|
||||||
|
if len(u) > 0:
|
||||||
|
print("Leftover VAE keys", u)
|
||||||
|
|
||||||
|
if device is None:
|
||||||
|
device = model_management.vae_device()
|
||||||
|
self.device = device
|
||||||
|
offload_device = model_management.vae_offload_device()
|
||||||
|
if dtype is None:
|
||||||
|
dtype = model_management.vae_dtype()
|
||||||
|
self.vae_dtype = dtype
|
||||||
|
self.first_stage_model.to(self.vae_dtype)
|
||||||
|
self.output_device = model_management.intermediate_device()
|
||||||
|
|
||||||
|
self.patcher = ldm_patched.modules.model_patcher.ModelPatcher(self.first_stage_model, load_device=self.device, offload_device=offload_device)
|
||||||
|
|
||||||
|
def decode_tiled_(self, samples, tile_x=64, tile_y=64, overlap = 16):
|
||||||
|
steps = samples.shape[0] * ldm_patched.modules.utils.get_tiled_scale_steps(samples.shape[3], samples.shape[2], tile_x, tile_y, overlap)
|
||||||
|
steps += samples.shape[0] * ldm_patched.modules.utils.get_tiled_scale_steps(samples.shape[3], samples.shape[2], tile_x // 2, tile_y * 2, overlap)
|
||||||
|
steps += samples.shape[0] * ldm_patched.modules.utils.get_tiled_scale_steps(samples.shape[3], samples.shape[2], tile_x * 2, tile_y // 2, overlap)
|
||||||
|
pbar = ldm_patched.modules.utils.ProgressBar(steps)
|
||||||
|
|
||||||
|
decode_fn = lambda a: (self.first_stage_model.decode(a.to(self.vae_dtype).to(self.device)) + 1.0).float()
|
||||||
|
output = torch.clamp((
|
||||||
|
(ldm_patched.modules.utils.tiled_scale(samples, decode_fn, tile_x // 2, tile_y * 2, overlap, upscale_amount = self.downscale_ratio, output_device=self.output_device, pbar = pbar) +
|
||||||
|
ldm_patched.modules.utils.tiled_scale(samples, decode_fn, tile_x * 2, tile_y // 2, overlap, upscale_amount = self.downscale_ratio, output_device=self.output_device, pbar = pbar) +
|
||||||
|
ldm_patched.modules.utils.tiled_scale(samples, decode_fn, tile_x, tile_y, overlap, upscale_amount = self.downscale_ratio, output_device=self.output_device, pbar = pbar))
|
||||||
|
/ 3.0) / 2.0, min=0.0, max=1.0)
|
||||||
|
return output
|
||||||
|
|
||||||
|
def encode_tiled_(self, pixel_samples, tile_x=512, tile_y=512, overlap = 64):
|
||||||
|
steps = pixel_samples.shape[0] * ldm_patched.modules.utils.get_tiled_scale_steps(pixel_samples.shape[3], pixel_samples.shape[2], tile_x, tile_y, overlap)
|
||||||
|
steps += pixel_samples.shape[0] * ldm_patched.modules.utils.get_tiled_scale_steps(pixel_samples.shape[3], pixel_samples.shape[2], tile_x // 2, tile_y * 2, overlap)
|
||||||
|
steps += pixel_samples.shape[0] * ldm_patched.modules.utils.get_tiled_scale_steps(pixel_samples.shape[3], pixel_samples.shape[2], tile_x * 2, tile_y // 2, overlap)
|
||||||
|
pbar = ldm_patched.modules.utils.ProgressBar(steps)
|
||||||
|
|
||||||
|
encode_fn = lambda a: self.first_stage_model.encode((2. * a - 1.).to(self.vae_dtype).to(self.device)).float()
|
||||||
|
samples = ldm_patched.modules.utils.tiled_scale(pixel_samples, encode_fn, tile_x, tile_y, overlap, upscale_amount = (1/self.downscale_ratio), out_channels=self.latent_channels, output_device=self.output_device, pbar=pbar)
|
||||||
|
samples += ldm_patched.modules.utils.tiled_scale(pixel_samples, encode_fn, tile_x * 2, tile_y // 2, overlap, upscale_amount = (1/self.downscale_ratio), out_channels=self.latent_channels, output_device=self.output_device, pbar=pbar)
|
||||||
|
samples += ldm_patched.modules.utils.tiled_scale(pixel_samples, encode_fn, tile_x // 2, tile_y * 2, overlap, upscale_amount = (1/self.downscale_ratio), out_channels=self.latent_channels, output_device=self.output_device, pbar=pbar)
|
||||||
|
samples /= 3.0
|
||||||
|
return samples
|
||||||
|
|
||||||
|
def decode(self, samples_in):
|
||||||
|
try:
|
||||||
|
memory_used = self.memory_used_decode(samples_in.shape, self.vae_dtype)
|
||||||
|
model_management.load_models_gpu([self.patcher], memory_required=memory_used)
|
||||||
|
free_memory = model_management.get_free_memory(self.device)
|
||||||
|
batch_number = int(free_memory / memory_used)
|
||||||
|
batch_number = max(1, batch_number)
|
||||||
|
|
||||||
|
pixel_samples = torch.empty((samples_in.shape[0], 3, round(samples_in.shape[2] * self.downscale_ratio), round(samples_in.shape[3] * self.downscale_ratio)), device=self.output_device)
|
||||||
|
for x in range(0, samples_in.shape[0], batch_number):
|
||||||
|
samples = samples_in[x:x+batch_number].to(self.vae_dtype).to(self.device)
|
||||||
|
pixel_samples[x:x+batch_number] = torch.clamp((self.first_stage_model.decode(samples).to(self.output_device).float() + 1.0) / 2.0, min=0.0, max=1.0)
|
||||||
|
except model_management.OOM_EXCEPTION as e:
|
||||||
|
print("Warning: Ran out of memory when regular VAE decoding, retrying with tiled VAE decoding.")
|
||||||
|
pixel_samples = self.decode_tiled_(samples_in)
|
||||||
|
|
||||||
|
pixel_samples = pixel_samples.to(self.output_device).movedim(1,-1)
|
||||||
|
return pixel_samples
|
||||||
|
|
||||||
|
def decode_tiled(self, samples, tile_x=64, tile_y=64, overlap = 16):
|
||||||
|
model_management.load_model_gpu(self.patcher)
|
||||||
|
output = self.decode_tiled_(samples, tile_x, tile_y, overlap)
|
||||||
|
return output.movedim(1,-1)
|
||||||
|
|
||||||
|
def encode(self, pixel_samples):
|
||||||
|
pixel_samples = pixel_samples.movedim(-1,1)
|
||||||
|
try:
|
||||||
|
memory_used = self.memory_used_encode(pixel_samples.shape, self.vae_dtype)
|
||||||
|
model_management.load_models_gpu([self.patcher], memory_required=memory_used)
|
||||||
|
free_memory = model_management.get_free_memory(self.device)
|
||||||
|
batch_number = int(free_memory / memory_used)
|
||||||
|
batch_number = max(1, batch_number)
|
||||||
|
samples = torch.empty((pixel_samples.shape[0], self.latent_channels, round(pixel_samples.shape[2] // self.downscale_ratio), round(pixel_samples.shape[3] // self.downscale_ratio)), device=self.output_device)
|
||||||
|
for x in range(0, pixel_samples.shape[0], batch_number):
|
||||||
|
pixels_in = (2. * pixel_samples[x:x+batch_number] - 1.).to(self.vae_dtype).to(self.device)
|
||||||
|
samples[x:x+batch_number] = self.first_stage_model.encode(pixels_in).to(self.output_device).float()
|
||||||
|
|
||||||
|
except model_management.OOM_EXCEPTION as e:
|
||||||
|
print("Warning: Ran out of memory when regular VAE encoding, retrying with tiled VAE encoding.")
|
||||||
|
samples = self.encode_tiled_(pixel_samples)
|
||||||
|
|
||||||
|
return samples
|
||||||
|
|
||||||
|
def encode_tiled(self, pixel_samples, tile_x=512, tile_y=512, overlap = 64):
|
||||||
|
model_management.load_model_gpu(self.patcher)
|
||||||
|
pixel_samples = pixel_samples.movedim(-1,1)
|
||||||
|
samples = self.encode_tiled_(pixel_samples, tile_x=tile_x, tile_y=tile_y, overlap=overlap)
|
||||||
|
return samples
|
||||||
|
|
||||||
|
def get_sd(self):
|
||||||
|
return self.first_stage_model.state_dict()
|
||||||
|
|
||||||
|
class StyleModel:
|
||||||
|
def __init__(self, model, device="cpu"):
|
||||||
|
self.model = model
|
||||||
|
|
||||||
|
def get_cond(self, input):
|
||||||
|
return self.model(input.last_hidden_state)
|
||||||
|
|
||||||
|
|
||||||
|
def load_style_model(ckpt_path):
|
||||||
|
model_data = ldm_patched.modules.utils.load_torch_file(ckpt_path, safe_load=True)
|
||||||
|
keys = model_data.keys()
|
||||||
|
if "style_embedding" in keys:
|
||||||
|
model = ldm_patched.t2ia.adapter.StyleAdapter(width=1024, context_dim=768, num_head=8, n_layes=3, num_token=8)
|
||||||
|
else:
|
||||||
|
raise Exception("invalid style model {}".format(ckpt_path))
|
||||||
|
model.load_state_dict(model_data)
|
||||||
|
return StyleModel(model)
|
||||||
|
|
||||||
|
|
||||||
|
def load_clip(ckpt_paths, embedding_directory=None):
|
||||||
|
clip_data = []
|
||||||
|
for p in ckpt_paths:
|
||||||
|
clip_data.append(ldm_patched.modules.utils.load_torch_file(p, safe_load=True))
|
||||||
|
|
||||||
|
class EmptyClass:
|
||||||
|
pass
|
||||||
|
|
||||||
|
for i in range(len(clip_data)):
|
||||||
|
if "transformer.resblocks.0.ln_1.weight" in clip_data[i]:
|
||||||
|
clip_data[i] = ldm_patched.modules.utils.transformers_convert(clip_data[i], "", "text_model.", 32)
|
||||||
|
|
||||||
|
clip_target = EmptyClass()
|
||||||
|
clip_target.params = {}
|
||||||
|
if len(clip_data) == 1:
|
||||||
|
if "text_model.encoder.layers.30.mlp.fc1.weight" in clip_data[0]:
|
||||||
|
clip_target.clip = sdxl_clip.SDXLRefinerClipModel
|
||||||
|
clip_target.tokenizer = sdxl_clip.SDXLTokenizer
|
||||||
|
elif "text_model.encoder.layers.22.mlp.fc1.weight" in clip_data[0]:
|
||||||
|
clip_target.clip = sd2_clip.SD2ClipModel
|
||||||
|
clip_target.tokenizer = sd2_clip.SD2Tokenizer
|
||||||
|
else:
|
||||||
|
clip_target.clip = sd1_clip.SD1ClipModel
|
||||||
|
clip_target.tokenizer = sd1_clip.SD1Tokenizer
|
||||||
|
else:
|
||||||
|
clip_target.clip = sdxl_clip.SDXLClipModel
|
||||||
|
clip_target.tokenizer = sdxl_clip.SDXLTokenizer
|
||||||
|
|
||||||
|
clip = CLIP(clip_target, embedding_directory=embedding_directory)
|
||||||
|
for c in clip_data:
|
||||||
|
m, u = clip.load_sd(c)
|
||||||
|
if len(m) > 0:
|
||||||
|
print("clip missing:", m)
|
||||||
|
|
||||||
|
if len(u) > 0:
|
||||||
|
print("clip unexpected:", u)
|
||||||
|
return clip
|
||||||
|
|
||||||
|
def load_gligen(ckpt_path):
|
||||||
|
data = ldm_patched.modules.utils.load_torch_file(ckpt_path, safe_load=True)
|
||||||
|
model = gligen.load_gligen(data)
|
||||||
|
if model_management.should_use_fp16():
|
||||||
|
model = model.half()
|
||||||
|
return ldm_patched.modules.model_patcher.ModelPatcher(model, load_device=model_management.get_torch_device(), offload_device=model_management.unet_offload_device())
|
||||||
|
|
||||||
|
def load_checkpoint(config_path=None, ckpt_path=None, output_vae=True, output_clip=True, embedding_directory=None, state_dict=None, config=None):
|
||||||
|
#TODO: this function is a mess and should be removed eventually
|
||||||
|
if config is None:
|
||||||
|
with open(config_path, 'r') as stream:
|
||||||
|
config = yaml.safe_load(stream)
|
||||||
|
model_config_params = config['model']['params']
|
||||||
|
clip_config = model_config_params['cond_stage_config']
|
||||||
|
scale_factor = model_config_params['scale_factor']
|
||||||
|
vae_config = model_config_params['first_stage_config']
|
||||||
|
|
||||||
|
fp16 = False
|
||||||
|
if "unet_config" in model_config_params:
|
||||||
|
if "params" in model_config_params["unet_config"]:
|
||||||
|
unet_config = model_config_params["unet_config"]["params"]
|
||||||
|
if "use_fp16" in unet_config:
|
||||||
|
fp16 = unet_config.pop("use_fp16")
|
||||||
|
if fp16:
|
||||||
|
unet_config["dtype"] = torch.float16
|
||||||
|
|
||||||
|
noise_aug_config = None
|
||||||
|
if "noise_aug_config" in model_config_params:
|
||||||
|
noise_aug_config = model_config_params["noise_aug_config"]
|
||||||
|
|
||||||
|
model_type = model_base.ModelType.EPS
|
||||||
|
|
||||||
|
if "parameterization" in model_config_params:
|
||||||
|
if model_config_params["parameterization"] == "v":
|
||||||
|
model_type = model_base.ModelType.V_PREDICTION
|
||||||
|
|
||||||
|
clip = None
|
||||||
|
vae = None
|
||||||
|
|
||||||
|
class WeightsLoader(torch.nn.Module):
|
||||||
|
pass
|
||||||
|
|
||||||
|
if state_dict is None:
|
||||||
|
state_dict = ldm_patched.modules.utils.load_torch_file(ckpt_path)
|
||||||
|
|
||||||
|
class EmptyClass:
|
||||||
|
pass
|
||||||
|
|
||||||
|
model_config = ldm_patched.modules.supported_models_base.BASE({})
|
||||||
|
|
||||||
|
from . import latent_formats
|
||||||
|
model_config.latent_format = latent_formats.SD15(scale_factor=scale_factor)
|
||||||
|
model_config.unet_config = model_detection.convert_config(unet_config)
|
||||||
|
|
||||||
|
if config['model']["target"].endswith("ImageEmbeddingConditionedLatentDiffusion"):
|
||||||
|
model = model_base.SD21UNCLIP(model_config, noise_aug_config["params"], model_type=model_type)
|
||||||
|
else:
|
||||||
|
model = model_base.BaseModel(model_config, model_type=model_type)
|
||||||
|
|
||||||
|
if config['model']["target"].endswith("LatentInpaintDiffusion"):
|
||||||
|
model.set_inpaint()
|
||||||
|
|
||||||
|
if fp16:
|
||||||
|
model = model.half()
|
||||||
|
|
||||||
|
offload_device = model_management.unet_offload_device()
|
||||||
|
model = model.to(offload_device)
|
||||||
|
model.load_model_weights(state_dict, "model.diffusion_model.")
|
||||||
|
|
||||||
|
if output_vae:
|
||||||
|
vae_sd = ldm_patched.modules.utils.state_dict_prefix_replace(state_dict, {"first_stage_model.": ""}, filter_keys=True)
|
||||||
|
vae = VAE(sd=vae_sd, config=vae_config)
|
||||||
|
|
||||||
|
if output_clip:
|
||||||
|
w = WeightsLoader()
|
||||||
|
clip_target = EmptyClass()
|
||||||
|
clip_target.params = clip_config.get("params", {})
|
||||||
|
if clip_config["target"].endswith("FrozenOpenCLIPEmbedder"):
|
||||||
|
clip_target.clip = sd2_clip.SD2ClipModel
|
||||||
|
clip_target.tokenizer = sd2_clip.SD2Tokenizer
|
||||||
|
clip = CLIP(clip_target, embedding_directory=embedding_directory)
|
||||||
|
w.cond_stage_model = clip.cond_stage_model.clip_h
|
||||||
|
elif clip_config["target"].endswith("FrozenCLIPEmbedder"):
|
||||||
|
clip_target.clip = sd1_clip.SD1ClipModel
|
||||||
|
clip_target.tokenizer = sd1_clip.SD1Tokenizer
|
||||||
|
clip = CLIP(clip_target, embedding_directory=embedding_directory)
|
||||||
|
w.cond_stage_model = clip.cond_stage_model.clip_l
|
||||||
|
load_clip_weights(w, state_dict)
|
||||||
|
|
||||||
|
return (ldm_patched.modules.model_patcher.ModelPatcher(model, load_device=model_management.get_torch_device(), offload_device=offload_device), clip, vae)
|
||||||
|
|
||||||
|
def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, output_clipvision=False, embedding_directory=None, output_model=True):
|
||||||
|
sd = ldm_patched.modules.utils.load_torch_file(ckpt_path)
|
||||||
|
sd_keys = sd.keys()
|
||||||
|
clip = None
|
||||||
|
clipvision = None
|
||||||
|
vae = None
|
||||||
|
model = None
|
||||||
|
model_patcher = None
|
||||||
|
clip_target = None
|
||||||
|
|
||||||
|
parameters = ldm_patched.modules.utils.calculate_parameters(sd, "model.diffusion_model.")
|
||||||
|
unet_dtype = model_management.unet_dtype(model_params=parameters)
|
||||||
|
load_device = model_management.get_torch_device()
|
||||||
|
manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device)
|
||||||
|
|
||||||
|
class WeightsLoader(torch.nn.Module):
|
||||||
|
pass
|
||||||
|
|
||||||
|
model_config = model_detection.model_config_from_unet(sd, "model.diffusion_model.", unet_dtype)
|
||||||
|
model_config.set_manual_cast(manual_cast_dtype)
|
||||||
|
|
||||||
|
if model_config is None:
|
||||||
|
raise RuntimeError("ERROR: Could not detect model type of: {}".format(ckpt_path))
|
||||||
|
|
||||||
|
if model_config.clip_vision_prefix is not None:
|
||||||
|
if output_clipvision:
|
||||||
|
clipvision = clip_vision.load_clipvision_from_sd(sd, model_config.clip_vision_prefix, True)
|
||||||
|
|
||||||
|
if output_model:
|
||||||
|
inital_load_device = model_management.unet_inital_load_device(parameters, unet_dtype)
|
||||||
|
offload_device = model_management.unet_offload_device()
|
||||||
|
model = model_config.get_model(sd, "model.diffusion_model.", device=inital_load_device)
|
||||||
|
model.load_model_weights(sd, "model.diffusion_model.")
|
||||||
|
|
||||||
|
if output_vae:
|
||||||
|
vae_sd = ldm_patched.modules.utils.state_dict_prefix_replace(sd, {"first_stage_model.": ""}, filter_keys=True)
|
||||||
|
vae_sd = model_config.process_vae_state_dict(vae_sd)
|
||||||
|
vae = VAE(sd=vae_sd)
|
||||||
|
|
||||||
|
if output_clip:
|
||||||
|
w = WeightsLoader()
|
||||||
|
clip_target = model_config.clip_target()
|
||||||
|
if clip_target is not None:
|
||||||
|
clip = CLIP(clip_target, embedding_directory=embedding_directory)
|
||||||
|
w.cond_stage_model = clip.cond_stage_model
|
||||||
|
sd = model_config.process_clip_state_dict(sd)
|
||||||
|
load_model_weights(w, sd)
|
||||||
|
|
||||||
|
left_over = sd.keys()
|
||||||
|
if len(left_over) > 0:
|
||||||
|
print("left over keys:", left_over)
|
||||||
|
|
||||||
|
if output_model:
|
||||||
|
model_patcher = ldm_patched.modules.model_patcher.ModelPatcher(model, load_device=load_device, offload_device=model_management.unet_offload_device(), current_device=inital_load_device)
|
||||||
|
if inital_load_device != torch.device("cpu"):
|
||||||
|
print("loaded straight to GPU")
|
||||||
|
model_management.load_model_gpu(model_patcher)
|
||||||
|
|
||||||
|
return (model_patcher, clip, vae, clipvision)
|
||||||
|
|
||||||
|
|
||||||
|
def load_unet_state_dict(sd): #load unet in diffusers format
|
||||||
|
parameters = ldm_patched.modules.utils.calculate_parameters(sd)
|
||||||
|
unet_dtype = model_management.unet_dtype(model_params=parameters)
|
||||||
|
load_device = model_management.get_torch_device()
|
||||||
|
manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device)
|
||||||
|
|
||||||
|
if "input_blocks.0.0.weight" in sd: #ldm
|
||||||
|
model_config = model_detection.model_config_from_unet(sd, "", unet_dtype)
|
||||||
|
if model_config is None:
|
||||||
|
return None
|
||||||
|
new_sd = sd
|
||||||
|
|
||||||
|
else: #diffusers
|
||||||
|
model_config = model_detection.model_config_from_diffusers_unet(sd, unet_dtype)
|
||||||
|
if model_config is None:
|
||||||
|
return None
|
||||||
|
|
||||||
|
diffusers_keys = ldm_patched.modules.utils.unet_to_diffusers(model_config.unet_config)
|
||||||
|
|
||||||
|
new_sd = {}
|
||||||
|
for k in diffusers_keys:
|
||||||
|
if k in sd:
|
||||||
|
new_sd[diffusers_keys[k]] = sd.pop(k)
|
||||||
|
else:
|
||||||
|
print(diffusers_keys[k], k)
|
||||||
|
offload_device = model_management.unet_offload_device()
|
||||||
|
model_config.set_manual_cast(manual_cast_dtype)
|
||||||
|
model = model_config.get_model(new_sd, "")
|
||||||
|
model = model.to(offload_device)
|
||||||
|
model.load_model_weights(new_sd, "")
|
||||||
|
left_over = sd.keys()
|
||||||
|
if len(left_over) > 0:
|
||||||
|
print("left over keys in unet:", left_over)
|
||||||
|
return ldm_patched.modules.model_patcher.ModelPatcher(model, load_device=load_device, offload_device=offload_device)
|
||||||
|
|
||||||
|
def load_unet(unet_path):
|
||||||
|
sd = ldm_patched.modules.utils.load_torch_file(unet_path)
|
||||||
|
model = load_unet_state_dict(sd)
|
||||||
|
if model is None:
|
||||||
|
print("ERROR UNSUPPORTED UNET", unet_path)
|
||||||
|
raise RuntimeError("ERROR: Could not detect model type of: {}".format(unet_path))
|
||||||
|
return model
|
||||||
|
|
||||||
|
def save_checkpoint(output_path, model, clip, vae, metadata=None):
|
||||||
|
model_management.load_models_gpu([model, clip.load_model()])
|
||||||
|
sd = model.model.state_dict_for_saving(clip.get_sd(), vae.get_sd())
|
||||||
|
ldm_patched.modules.utils.save_torch_file(sd, output_path, metadata=metadata)
|
519
ldm_patched/modules/sd1_clip.py
Normal file
519
ldm_patched/modules/sd1_clip.py
Normal file
@ -0,0 +1,519 @@
|
|||||||
|
import os
|
||||||
|
|
||||||
|
from transformers import CLIPTokenizer
|
||||||
|
import ldm_patched.modules.ops
|
||||||
|
import torch
|
||||||
|
import traceback
|
||||||
|
import zipfile
|
||||||
|
from . import model_management
|
||||||
|
import contextlib
|
||||||
|
import ldm_patched.modules.clip_model
|
||||||
|
import json
|
||||||
|
|
||||||
|
def gen_empty_tokens(special_tokens, length):
|
||||||
|
start_token = special_tokens.get("start", None)
|
||||||
|
end_token = special_tokens.get("end", None)
|
||||||
|
pad_token = special_tokens.get("pad")
|
||||||
|
output = []
|
||||||
|
if start_token is not None:
|
||||||
|
output.append(start_token)
|
||||||
|
if end_token is not None:
|
||||||
|
output.append(end_token)
|
||||||
|
output += [pad_token] * (length - len(output))
|
||||||
|
return output
|
||||||
|
|
||||||
|
class ClipTokenWeightEncoder:
|
||||||
|
def encode_token_weights(self, token_weight_pairs):
|
||||||
|
to_encode = list()
|
||||||
|
max_token_len = 0
|
||||||
|
has_weights = False
|
||||||
|
for x in token_weight_pairs:
|
||||||
|
tokens = list(map(lambda a: a[0], x))
|
||||||
|
max_token_len = max(len(tokens), max_token_len)
|
||||||
|
has_weights = has_weights or not all(map(lambda a: a[1] == 1.0, x))
|
||||||
|
to_encode.append(tokens)
|
||||||
|
|
||||||
|
sections = len(to_encode)
|
||||||
|
if has_weights or sections == 0:
|
||||||
|
to_encode.append(gen_empty_tokens(self.special_tokens, max_token_len))
|
||||||
|
|
||||||
|
out, pooled = self.encode(to_encode)
|
||||||
|
if pooled is not None:
|
||||||
|
first_pooled = pooled[0:1].to(model_management.intermediate_device())
|
||||||
|
else:
|
||||||
|
first_pooled = pooled
|
||||||
|
|
||||||
|
output = []
|
||||||
|
for k in range(0, sections):
|
||||||
|
z = out[k:k+1]
|
||||||
|
if has_weights:
|
||||||
|
z_empty = out[-1]
|
||||||
|
for i in range(len(z)):
|
||||||
|
for j in range(len(z[i])):
|
||||||
|
weight = token_weight_pairs[k][j][1]
|
||||||
|
if weight != 1.0:
|
||||||
|
z[i][j] = (z[i][j] - z_empty[j]) * weight + z_empty[j]
|
||||||
|
output.append(z)
|
||||||
|
|
||||||
|
if (len(output) == 0):
|
||||||
|
return out[-1:].to(model_management.intermediate_device()), first_pooled
|
||||||
|
return torch.cat(output, dim=-2).to(model_management.intermediate_device()), first_pooled
|
||||||
|
|
||||||
|
class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
|
||||||
|
"""Uses the CLIP transformer encoder for text (from huggingface)"""
|
||||||
|
LAYERS = [
|
||||||
|
"last",
|
||||||
|
"pooled",
|
||||||
|
"hidden"
|
||||||
|
]
|
||||||
|
def __init__(self, version="openai/clip-vit-large-patch14", device="cpu", max_length=77,
|
||||||
|
freeze=True, layer="last", layer_idx=None, textmodel_json_config=None, dtype=None, model_class=ldm_patched.modules.clip_model.CLIPTextModel,
|
||||||
|
special_tokens={"start": 49406, "end": 49407, "pad": 49407}, layer_norm_hidden_state=True): # clip-vit-base-patch32
|
||||||
|
super().__init__()
|
||||||
|
assert layer in self.LAYERS
|
||||||
|
|
||||||
|
if textmodel_json_config is None:
|
||||||
|
textmodel_json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "sd1_clip_config.json")
|
||||||
|
|
||||||
|
with open(textmodel_json_config) as f:
|
||||||
|
config = json.load(f)
|
||||||
|
|
||||||
|
self.transformer = model_class(config, dtype, device, ldm_patched.modules.ops.manual_cast)
|
||||||
|
self.num_layers = self.transformer.num_layers
|
||||||
|
|
||||||
|
self.max_length = max_length
|
||||||
|
if freeze:
|
||||||
|
self.freeze()
|
||||||
|
self.layer = layer
|
||||||
|
self.layer_idx = None
|
||||||
|
self.special_tokens = special_tokens
|
||||||
|
self.text_projection = torch.nn.Parameter(torch.eye(self.transformer.get_input_embeddings().weight.shape[1]))
|
||||||
|
self.logit_scale = torch.nn.Parameter(torch.tensor(4.6055))
|
||||||
|
self.enable_attention_masks = False
|
||||||
|
|
||||||
|
self.layer_norm_hidden_state = layer_norm_hidden_state
|
||||||
|
if layer == "hidden":
|
||||||
|
assert layer_idx is not None
|
||||||
|
assert abs(layer_idx) < self.num_layers
|
||||||
|
self.clip_layer(layer_idx)
|
||||||
|
self.layer_default = (self.layer, self.layer_idx)
|
||||||
|
|
||||||
|
def freeze(self):
|
||||||
|
self.transformer = self.transformer.eval()
|
||||||
|
#self.train = disabled_train
|
||||||
|
for param in self.parameters():
|
||||||
|
param.requires_grad = False
|
||||||
|
|
||||||
|
def clip_layer(self, layer_idx):
|
||||||
|
if abs(layer_idx) > self.num_layers:
|
||||||
|
self.layer = "last"
|
||||||
|
else:
|
||||||
|
self.layer = "hidden"
|
||||||
|
self.layer_idx = layer_idx
|
||||||
|
|
||||||
|
def reset_clip_layer(self):
|
||||||
|
self.layer = self.layer_default[0]
|
||||||
|
self.layer_idx = self.layer_default[1]
|
||||||
|
|
||||||
|
def set_up_textual_embeddings(self, tokens, current_embeds):
|
||||||
|
out_tokens = []
|
||||||
|
next_new_token = token_dict_size = current_embeds.weight.shape[0] - 1
|
||||||
|
embedding_weights = []
|
||||||
|
|
||||||
|
for x in tokens:
|
||||||
|
tokens_temp = []
|
||||||
|
for y in x:
|
||||||
|
if isinstance(y, int):
|
||||||
|
if y == token_dict_size: #EOS token
|
||||||
|
y = -1
|
||||||
|
tokens_temp += [y]
|
||||||
|
else:
|
||||||
|
if y.shape[0] == current_embeds.weight.shape[1]:
|
||||||
|
embedding_weights += [y]
|
||||||
|
tokens_temp += [next_new_token]
|
||||||
|
next_new_token += 1
|
||||||
|
else:
|
||||||
|
print("WARNING: shape mismatch when trying to apply embedding, embedding will be ignored", y.shape[0], current_embeds.weight.shape[1])
|
||||||
|
while len(tokens_temp) < len(x):
|
||||||
|
tokens_temp += [self.special_tokens["pad"]]
|
||||||
|
out_tokens += [tokens_temp]
|
||||||
|
|
||||||
|
n = token_dict_size
|
||||||
|
if len(embedding_weights) > 0:
|
||||||
|
new_embedding = torch.nn.Embedding(next_new_token + 1, current_embeds.weight.shape[1], device=current_embeds.weight.device, dtype=current_embeds.weight.dtype)
|
||||||
|
new_embedding.weight[:token_dict_size] = current_embeds.weight[:-1]
|
||||||
|
for x in embedding_weights:
|
||||||
|
new_embedding.weight[n] = x
|
||||||
|
n += 1
|
||||||
|
new_embedding.weight[n] = current_embeds.weight[-1] #EOS embedding
|
||||||
|
self.transformer.set_input_embeddings(new_embedding)
|
||||||
|
|
||||||
|
processed_tokens = []
|
||||||
|
for x in out_tokens:
|
||||||
|
processed_tokens += [list(map(lambda a: n if a == -1 else a, x))] #The EOS token should always be the largest one
|
||||||
|
|
||||||
|
return processed_tokens
|
||||||
|
|
||||||
|
def forward(self, tokens):
|
||||||
|
backup_embeds = self.transformer.get_input_embeddings()
|
||||||
|
device = backup_embeds.weight.device
|
||||||
|
tokens = self.set_up_textual_embeddings(tokens, backup_embeds)
|
||||||
|
tokens = torch.LongTensor(tokens).to(device)
|
||||||
|
|
||||||
|
attention_mask = None
|
||||||
|
if self.enable_attention_masks:
|
||||||
|
attention_mask = torch.zeros_like(tokens)
|
||||||
|
max_token = self.transformer.get_input_embeddings().weight.shape[0] - 1
|
||||||
|
for x in range(attention_mask.shape[0]):
|
||||||
|
for y in range(attention_mask.shape[1]):
|
||||||
|
attention_mask[x, y] = 1
|
||||||
|
if tokens[x, y] == max_token:
|
||||||
|
break
|
||||||
|
|
||||||
|
outputs = self.transformer(tokens, attention_mask, intermediate_output=self.layer_idx, final_layer_norm_intermediate=self.layer_norm_hidden_state)
|
||||||
|
self.transformer.set_input_embeddings(backup_embeds)
|
||||||
|
|
||||||
|
if self.layer == "last":
|
||||||
|
z = outputs[0]
|
||||||
|
else:
|
||||||
|
z = outputs[1]
|
||||||
|
|
||||||
|
if outputs[2] is not None:
|
||||||
|
pooled_output = outputs[2].float()
|
||||||
|
else:
|
||||||
|
pooled_output = None
|
||||||
|
|
||||||
|
if self.text_projection is not None and pooled_output is not None:
|
||||||
|
pooled_output = pooled_output.float().to(self.text_projection.device) @ self.text_projection.float()
|
||||||
|
return z.float(), pooled_output
|
||||||
|
|
||||||
|
def encode(self, tokens):
|
||||||
|
return self(tokens)
|
||||||
|
|
||||||
|
def load_sd(self, sd):
|
||||||
|
if "text_projection" in sd:
|
||||||
|
self.text_projection[:] = sd.pop("text_projection")
|
||||||
|
if "text_projection.weight" in sd:
|
||||||
|
self.text_projection[:] = sd.pop("text_projection.weight").transpose(0, 1)
|
||||||
|
return self.transformer.load_state_dict(sd, strict=False)
|
||||||
|
|
||||||
|
def parse_parentheses(string):
|
||||||
|
result = []
|
||||||
|
current_item = ""
|
||||||
|
nesting_level = 0
|
||||||
|
for char in string:
|
||||||
|
if char == "(":
|
||||||
|
if nesting_level == 0:
|
||||||
|
if current_item:
|
||||||
|
result.append(current_item)
|
||||||
|
current_item = "("
|
||||||
|
else:
|
||||||
|
current_item = "("
|
||||||
|
else:
|
||||||
|
current_item += char
|
||||||
|
nesting_level += 1
|
||||||
|
elif char == ")":
|
||||||
|
nesting_level -= 1
|
||||||
|
if nesting_level == 0:
|
||||||
|
result.append(current_item + ")")
|
||||||
|
current_item = ""
|
||||||
|
else:
|
||||||
|
current_item += char
|
||||||
|
else:
|
||||||
|
current_item += char
|
||||||
|
if current_item:
|
||||||
|
result.append(current_item)
|
||||||
|
return result
|
||||||
|
|
||||||
|
def token_weights(string, current_weight):
|
||||||
|
a = parse_parentheses(string)
|
||||||
|
out = []
|
||||||
|
for x in a:
|
||||||
|
weight = current_weight
|
||||||
|
if len(x) >= 2 and x[-1] == ')' and x[0] == '(':
|
||||||
|
x = x[1:-1]
|
||||||
|
xx = x.rfind(":")
|
||||||
|
weight *= 1.1
|
||||||
|
if xx > 0:
|
||||||
|
try:
|
||||||
|
weight = float(x[xx+1:])
|
||||||
|
x = x[:xx]
|
||||||
|
except:
|
||||||
|
pass
|
||||||
|
out += token_weights(x, weight)
|
||||||
|
else:
|
||||||
|
out += [(x, current_weight)]
|
||||||
|
return out
|
||||||
|
|
||||||
|
def escape_important(text):
|
||||||
|
text = text.replace("\\)", "\0\1")
|
||||||
|
text = text.replace("\\(", "\0\2")
|
||||||
|
return text
|
||||||
|
|
||||||
|
def unescape_important(text):
|
||||||
|
text = text.replace("\0\1", ")")
|
||||||
|
text = text.replace("\0\2", "(")
|
||||||
|
return text
|
||||||
|
|
||||||
|
def safe_load_embed_zip(embed_path):
|
||||||
|
with zipfile.ZipFile(embed_path) as myzip:
|
||||||
|
names = list(filter(lambda a: "data/" in a, myzip.namelist()))
|
||||||
|
names.reverse()
|
||||||
|
for n in names:
|
||||||
|
with myzip.open(n) as myfile:
|
||||||
|
data = myfile.read()
|
||||||
|
number = len(data) // 4
|
||||||
|
length_embed = 1024 #sd2.x
|
||||||
|
if number < 768:
|
||||||
|
continue
|
||||||
|
if number % 768 == 0:
|
||||||
|
length_embed = 768 #sd1.x
|
||||||
|
num_embeds = number // length_embed
|
||||||
|
embed = torch.frombuffer(data, dtype=torch.float)
|
||||||
|
out = embed.reshape((num_embeds, length_embed)).clone()
|
||||||
|
del embed
|
||||||
|
return out
|
||||||
|
|
||||||
|
def expand_directory_list(directories):
|
||||||
|
dirs = set()
|
||||||
|
for x in directories:
|
||||||
|
dirs.add(x)
|
||||||
|
for root, subdir, file in os.walk(x, followlinks=True):
|
||||||
|
dirs.add(root)
|
||||||
|
return list(dirs)
|
||||||
|
|
||||||
|
def load_embed(embedding_name, embedding_directory, embedding_size, embed_key=None):
|
||||||
|
if isinstance(embedding_directory, str):
|
||||||
|
embedding_directory = [embedding_directory]
|
||||||
|
|
||||||
|
embedding_directory = expand_directory_list(embedding_directory)
|
||||||
|
|
||||||
|
valid_file = None
|
||||||
|
for embed_dir in embedding_directory:
|
||||||
|
embed_path = os.path.abspath(os.path.join(embed_dir, embedding_name))
|
||||||
|
embed_dir = os.path.abspath(embed_dir)
|
||||||
|
try:
|
||||||
|
if os.path.commonpath((embed_dir, embed_path)) != embed_dir:
|
||||||
|
continue
|
||||||
|
except:
|
||||||
|
continue
|
||||||
|
if not os.path.isfile(embed_path):
|
||||||
|
extensions = ['.safetensors', '.pt', '.bin']
|
||||||
|
for x in extensions:
|
||||||
|
t = embed_path + x
|
||||||
|
if os.path.isfile(t):
|
||||||
|
valid_file = t
|
||||||
|
break
|
||||||
|
else:
|
||||||
|
valid_file = embed_path
|
||||||
|
if valid_file is not None:
|
||||||
|
break
|
||||||
|
|
||||||
|
if valid_file is None:
|
||||||
|
return None
|
||||||
|
|
||||||
|
embed_path = valid_file
|
||||||
|
|
||||||
|
embed_out = None
|
||||||
|
|
||||||
|
try:
|
||||||
|
if embed_path.lower().endswith(".safetensors"):
|
||||||
|
import safetensors.torch
|
||||||
|
embed = safetensors.torch.load_file(embed_path, device="cpu")
|
||||||
|
else:
|
||||||
|
if 'weights_only' in torch.load.__code__.co_varnames:
|
||||||
|
try:
|
||||||
|
embed = torch.load(embed_path, weights_only=True, map_location="cpu")
|
||||||
|
except:
|
||||||
|
embed_out = safe_load_embed_zip(embed_path)
|
||||||
|
else:
|
||||||
|
embed = torch.load(embed_path, map_location="cpu")
|
||||||
|
except Exception as e:
|
||||||
|
print(traceback.format_exc())
|
||||||
|
print()
|
||||||
|
print("error loading embedding, skipping loading:", embedding_name)
|
||||||
|
return None
|
||||||
|
|
||||||
|
if embed_out is None:
|
||||||
|
if 'string_to_param' in embed:
|
||||||
|
values = embed['string_to_param'].values()
|
||||||
|
embed_out = next(iter(values))
|
||||||
|
elif isinstance(embed, list):
|
||||||
|
out_list = []
|
||||||
|
for x in range(len(embed)):
|
||||||
|
for k in embed[x]:
|
||||||
|
t = embed[x][k]
|
||||||
|
if t.shape[-1] != embedding_size:
|
||||||
|
continue
|
||||||
|
out_list.append(t.reshape(-1, t.shape[-1]))
|
||||||
|
embed_out = torch.cat(out_list, dim=0)
|
||||||
|
elif embed_key is not None and embed_key in embed:
|
||||||
|
embed_out = embed[embed_key]
|
||||||
|
else:
|
||||||
|
values = embed.values()
|
||||||
|
embed_out = next(iter(values))
|
||||||
|
return embed_out
|
||||||
|
|
||||||
|
class SDTokenizer:
|
||||||
|
def __init__(self, tokenizer_path=None, max_length=77, pad_with_end=True, embedding_directory=None, embedding_size=768, embedding_key='clip_l', tokenizer_class=CLIPTokenizer, has_start_token=True, pad_to_max_length=True):
|
||||||
|
if tokenizer_path is None:
|
||||||
|
tokenizer_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "sd1_tokenizer")
|
||||||
|
self.tokenizer = tokenizer_class.from_pretrained(tokenizer_path)
|
||||||
|
self.max_length = max_length
|
||||||
|
|
||||||
|
empty = self.tokenizer('')["input_ids"]
|
||||||
|
if has_start_token:
|
||||||
|
self.tokens_start = 1
|
||||||
|
self.start_token = empty[0]
|
||||||
|
self.end_token = empty[1]
|
||||||
|
else:
|
||||||
|
self.tokens_start = 0
|
||||||
|
self.start_token = None
|
||||||
|
self.end_token = empty[0]
|
||||||
|
self.pad_with_end = pad_with_end
|
||||||
|
self.pad_to_max_length = pad_to_max_length
|
||||||
|
|
||||||
|
vocab = self.tokenizer.get_vocab()
|
||||||
|
self.inv_vocab = {v: k for k, v in vocab.items()}
|
||||||
|
self.embedding_directory = embedding_directory
|
||||||
|
self.max_word_length = 8
|
||||||
|
self.embedding_identifier = "embedding:"
|
||||||
|
self.embedding_size = embedding_size
|
||||||
|
self.embedding_key = embedding_key
|
||||||
|
|
||||||
|
def _try_get_embedding(self, embedding_name:str):
|
||||||
|
'''
|
||||||
|
Takes a potential embedding name and tries to retrieve it.
|
||||||
|
Returns a Tuple consisting of the embedding and any leftover string, embedding can be None.
|
||||||
|
'''
|
||||||
|
embed = load_embed(embedding_name, self.embedding_directory, self.embedding_size, self.embedding_key)
|
||||||
|
if embed is None:
|
||||||
|
stripped = embedding_name.strip(',')
|
||||||
|
if len(stripped) < len(embedding_name):
|
||||||
|
embed = load_embed(stripped, self.embedding_directory, self.embedding_size, self.embedding_key)
|
||||||
|
return (embed, embedding_name[len(stripped):])
|
||||||
|
return (embed, "")
|
||||||
|
|
||||||
|
|
||||||
|
def tokenize_with_weights(self, text:str, return_word_ids=False):
|
||||||
|
'''
|
||||||
|
Takes a prompt and converts it to a list of (token, weight, word id) elements.
|
||||||
|
Tokens can both be integer tokens and pre computed CLIP tensors.
|
||||||
|
Word id values are unique per word and embedding, where the id 0 is reserved for non word tokens.
|
||||||
|
Returned list has the dimensions NxM where M is the input size of CLIP
|
||||||
|
'''
|
||||||
|
if self.pad_with_end:
|
||||||
|
pad_token = self.end_token
|
||||||
|
else:
|
||||||
|
pad_token = 0
|
||||||
|
|
||||||
|
text = escape_important(text)
|
||||||
|
parsed_weights = token_weights(text, 1.0)
|
||||||
|
|
||||||
|
#tokenize words
|
||||||
|
tokens = []
|
||||||
|
for weighted_segment, weight in parsed_weights:
|
||||||
|
to_tokenize = unescape_important(weighted_segment).replace("\n", " ").split(' ')
|
||||||
|
to_tokenize = [x for x in to_tokenize if x != ""]
|
||||||
|
for word in to_tokenize:
|
||||||
|
#if we find an embedding, deal with the embedding
|
||||||
|
if word.startswith(self.embedding_identifier) and self.embedding_directory is not None:
|
||||||
|
embedding_name = word[len(self.embedding_identifier):].strip('\n')
|
||||||
|
embed, leftover = self._try_get_embedding(embedding_name)
|
||||||
|
if embed is None:
|
||||||
|
print(f"warning, embedding:{embedding_name} does not exist, ignoring")
|
||||||
|
else:
|
||||||
|
if len(embed.shape) == 1:
|
||||||
|
tokens.append([(embed, weight)])
|
||||||
|
else:
|
||||||
|
tokens.append([(embed[x], weight) for x in range(embed.shape[0])])
|
||||||
|
#if we accidentally have leftover text, continue parsing using leftover, else move on to next word
|
||||||
|
if leftover != "":
|
||||||
|
word = leftover
|
||||||
|
else:
|
||||||
|
continue
|
||||||
|
#parse word
|
||||||
|
tokens.append([(t, weight) for t in self.tokenizer(word)["input_ids"][self.tokens_start:-1]])
|
||||||
|
|
||||||
|
#reshape token array to CLIP input size
|
||||||
|
batched_tokens = []
|
||||||
|
batch = []
|
||||||
|
if self.start_token is not None:
|
||||||
|
batch.append((self.start_token, 1.0, 0))
|
||||||
|
batched_tokens.append(batch)
|
||||||
|
for i, t_group in enumerate(tokens):
|
||||||
|
#determine if we're going to try and keep the tokens in a single batch
|
||||||
|
is_large = len(t_group) >= self.max_word_length
|
||||||
|
|
||||||
|
while len(t_group) > 0:
|
||||||
|
if len(t_group) + len(batch) > self.max_length - 1:
|
||||||
|
remaining_length = self.max_length - len(batch) - 1
|
||||||
|
#break word in two and add end token
|
||||||
|
if is_large:
|
||||||
|
batch.extend([(t,w,i+1) for t,w in t_group[:remaining_length]])
|
||||||
|
batch.append((self.end_token, 1.0, 0))
|
||||||
|
t_group = t_group[remaining_length:]
|
||||||
|
#add end token and pad
|
||||||
|
else:
|
||||||
|
batch.append((self.end_token, 1.0, 0))
|
||||||
|
if self.pad_to_max_length:
|
||||||
|
batch.extend([(pad_token, 1.0, 0)] * (remaining_length))
|
||||||
|
#start new batch
|
||||||
|
batch = []
|
||||||
|
if self.start_token is not None:
|
||||||
|
batch.append((self.start_token, 1.0, 0))
|
||||||
|
batched_tokens.append(batch)
|
||||||
|
else:
|
||||||
|
batch.extend([(t,w,i+1) for t,w in t_group])
|
||||||
|
t_group = []
|
||||||
|
|
||||||
|
#fill last batch
|
||||||
|
batch.append((self.end_token, 1.0, 0))
|
||||||
|
if self.pad_to_max_length:
|
||||||
|
batch.extend([(pad_token, 1.0, 0)] * (self.max_length - len(batch)))
|
||||||
|
|
||||||
|
if not return_word_ids:
|
||||||
|
batched_tokens = [[(t, w) for t, w,_ in x] for x in batched_tokens]
|
||||||
|
|
||||||
|
return batched_tokens
|
||||||
|
|
||||||
|
|
||||||
|
def untokenize(self, token_weight_pair):
|
||||||
|
return list(map(lambda a: (a, self.inv_vocab[a[0]]), token_weight_pair))
|
||||||
|
|
||||||
|
|
||||||
|
class SD1Tokenizer:
|
||||||
|
def __init__(self, embedding_directory=None, clip_name="l", tokenizer=SDTokenizer):
|
||||||
|
self.clip_name = clip_name
|
||||||
|
self.clip = "clip_{}".format(self.clip_name)
|
||||||
|
setattr(self, self.clip, tokenizer(embedding_directory=embedding_directory))
|
||||||
|
|
||||||
|
def tokenize_with_weights(self, text:str, return_word_ids=False):
|
||||||
|
out = {}
|
||||||
|
out[self.clip_name] = getattr(self, self.clip).tokenize_with_weights(text, return_word_ids)
|
||||||
|
return out
|
||||||
|
|
||||||
|
def untokenize(self, token_weight_pair):
|
||||||
|
return getattr(self, self.clip).untokenize(token_weight_pair)
|
||||||
|
|
||||||
|
|
||||||
|
class SD1ClipModel(torch.nn.Module):
|
||||||
|
def __init__(self, device="cpu", dtype=None, clip_name="l", clip_model=SDClipModel, **kwargs):
|
||||||
|
super().__init__()
|
||||||
|
self.clip_name = clip_name
|
||||||
|
self.clip = "clip_{}".format(self.clip_name)
|
||||||
|
setattr(self, self.clip, clip_model(device=device, dtype=dtype, **kwargs))
|
||||||
|
|
||||||
|
def clip_layer(self, layer_idx):
|
||||||
|
getattr(self, self.clip).clip_layer(layer_idx)
|
||||||
|
|
||||||
|
def reset_clip_layer(self):
|
||||||
|
getattr(self, self.clip).reset_clip_layer()
|
||||||
|
|
||||||
|
def encode_token_weights(self, token_weight_pairs):
|
||||||
|
token_weight_pairs = token_weight_pairs[self.clip_name]
|
||||||
|
out, pooled = getattr(self, self.clip).encode_token_weights(token_weight_pairs)
|
||||||
|
return out, pooled
|
||||||
|
|
||||||
|
def load_sd(self, sd):
|
||||||
|
return getattr(self, self.clip).load_sd(sd)
|
25
ldm_patched/modules/sd1_clip_config.json
Normal file
25
ldm_patched/modules/sd1_clip_config.json
Normal file
@ -0,0 +1,25 @@
|
|||||||
|
{
|
||||||
|
"_name_or_path": "openai/clip-vit-large-patch14",
|
||||||
|
"architectures": [
|
||||||
|
"CLIPTextModel"
|
||||||
|
],
|
||||||
|
"attention_dropout": 0.0,
|
||||||
|
"bos_token_id": 0,
|
||||||
|
"dropout": 0.0,
|
||||||
|
"eos_token_id": 2,
|
||||||
|
"hidden_act": "quick_gelu",
|
||||||
|
"hidden_size": 768,
|
||||||
|
"initializer_factor": 1.0,
|
||||||
|
"initializer_range": 0.02,
|
||||||
|
"intermediate_size": 3072,
|
||||||
|
"layer_norm_eps": 1e-05,
|
||||||
|
"max_position_embeddings": 77,
|
||||||
|
"model_type": "clip_text_model",
|
||||||
|
"num_attention_heads": 12,
|
||||||
|
"num_hidden_layers": 12,
|
||||||
|
"pad_token_id": 1,
|
||||||
|
"projection_dim": 768,
|
||||||
|
"torch_dtype": "float32",
|
||||||
|
"transformers_version": "4.24.0",
|
||||||
|
"vocab_size": 49408
|
||||||
|
}
|
48895
ldm_patched/modules/sd1_tokenizer/merges.txt
Normal file
48895
ldm_patched/modules/sd1_tokenizer/merges.txt
Normal file
File diff suppressed because it is too large
Load Diff
24
ldm_patched/modules/sd1_tokenizer/special_tokens_map.json
Normal file
24
ldm_patched/modules/sd1_tokenizer/special_tokens_map.json
Normal file
@ -0,0 +1,24 @@
|
|||||||
|
{
|
||||||
|
"bos_token": {
|
||||||
|
"content": "<|startoftext|>",
|
||||||
|
"lstrip": false,
|
||||||
|
"normalized": true,
|
||||||
|
"rstrip": false,
|
||||||
|
"single_word": false
|
||||||
|
},
|
||||||
|
"eos_token": {
|
||||||
|
"content": "<|endoftext|>",
|
||||||
|
"lstrip": false,
|
||||||
|
"normalized": true,
|
||||||
|
"rstrip": false,
|
||||||
|
"single_word": false
|
||||||
|
},
|
||||||
|
"pad_token": "<|endoftext|>",
|
||||||
|
"unk_token": {
|
||||||
|
"content": "<|endoftext|>",
|
||||||
|
"lstrip": false,
|
||||||
|
"normalized": true,
|
||||||
|
"rstrip": false,
|
||||||
|
"single_word": false
|
||||||
|
}
|
||||||
|
}
|
34
ldm_patched/modules/sd1_tokenizer/tokenizer_config.json
Normal file
34
ldm_patched/modules/sd1_tokenizer/tokenizer_config.json
Normal file
@ -0,0 +1,34 @@
|
|||||||
|
{
|
||||||
|
"add_prefix_space": false,
|
||||||
|
"bos_token": {
|
||||||
|
"__type": "AddedToken",
|
||||||
|
"content": "<|startoftext|>",
|
||||||
|
"lstrip": false,
|
||||||
|
"normalized": true,
|
||||||
|
"rstrip": false,
|
||||||
|
"single_word": false
|
||||||
|
},
|
||||||
|
"do_lower_case": true,
|
||||||
|
"eos_token": {
|
||||||
|
"__type": "AddedToken",
|
||||||
|
"content": "<|endoftext|>",
|
||||||
|
"lstrip": false,
|
||||||
|
"normalized": true,
|
||||||
|
"rstrip": false,
|
||||||
|
"single_word": false
|
||||||
|
},
|
||||||
|
"errors": "replace",
|
||||||
|
"model_max_length": 77,
|
||||||
|
"name_or_path": "openai/clip-vit-large-patch14",
|
||||||
|
"pad_token": "<|endoftext|>",
|
||||||
|
"special_tokens_map_file": "./special_tokens_map.json",
|
||||||
|
"tokenizer_class": "CLIPTokenizer",
|
||||||
|
"unk_token": {
|
||||||
|
"__type": "AddedToken",
|
||||||
|
"content": "<|endoftext|>",
|
||||||
|
"lstrip": false,
|
||||||
|
"normalized": true,
|
||||||
|
"rstrip": false,
|
||||||
|
"single_word": false
|
||||||
|
}
|
||||||
|
}
|
49410
ldm_patched/modules/sd1_tokenizer/vocab.json
Normal file
49410
ldm_patched/modules/sd1_tokenizer/vocab.json
Normal file
File diff suppressed because it is too large
Load Diff
24
ldm_patched/modules/sd2_clip.py
Normal file
24
ldm_patched/modules/sd2_clip.py
Normal file
@ -0,0 +1,24 @@
|
|||||||
|
from ldm_patched.modules import sd1_clip
|
||||||
|
import torch
|
||||||
|
import os
|
||||||
|
|
||||||
|
class SD2ClipHModel(sd1_clip.SDClipModel):
|
||||||
|
def __init__(self, arch="ViT-H-14", device="cpu", max_length=77, freeze=True, layer="penultimate", layer_idx=None, dtype=None):
|
||||||
|
if layer == "penultimate":
|
||||||
|
layer="hidden"
|
||||||
|
layer_idx=-2
|
||||||
|
|
||||||
|
textmodel_json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "sd2_clip_config.json")
|
||||||
|
super().__init__(device=device, freeze=freeze, layer=layer, layer_idx=layer_idx, textmodel_json_config=textmodel_json_config, dtype=dtype, special_tokens={"start": 49406, "end": 49407, "pad": 0})
|
||||||
|
|
||||||
|
class SD2ClipHTokenizer(sd1_clip.SDTokenizer):
|
||||||
|
def __init__(self, tokenizer_path=None, embedding_directory=None):
|
||||||
|
super().__init__(tokenizer_path, pad_with_end=False, embedding_directory=embedding_directory, embedding_size=1024)
|
||||||
|
|
||||||
|
class SD2Tokenizer(sd1_clip.SD1Tokenizer):
|
||||||
|
def __init__(self, embedding_directory=None):
|
||||||
|
super().__init__(embedding_directory=embedding_directory, clip_name="h", tokenizer=SD2ClipHTokenizer)
|
||||||
|
|
||||||
|
class SD2ClipModel(sd1_clip.SD1ClipModel):
|
||||||
|
def __init__(self, device="cpu", dtype=None, **kwargs):
|
||||||
|
super().__init__(device=device, dtype=dtype, clip_name="h", clip_model=SD2ClipHModel, **kwargs)
|
23
ldm_patched/modules/sd2_clip_config.json
Normal file
23
ldm_patched/modules/sd2_clip_config.json
Normal file
@ -0,0 +1,23 @@
|
|||||||
|
{
|
||||||
|
"architectures": [
|
||||||
|
"CLIPTextModel"
|
||||||
|
],
|
||||||
|
"attention_dropout": 0.0,
|
||||||
|
"bos_token_id": 0,
|
||||||
|
"dropout": 0.0,
|
||||||
|
"eos_token_id": 2,
|
||||||
|
"hidden_act": "gelu",
|
||||||
|
"hidden_size": 1024,
|
||||||
|
"initializer_factor": 1.0,
|
||||||
|
"initializer_range": 0.02,
|
||||||
|
"intermediate_size": 4096,
|
||||||
|
"layer_norm_eps": 1e-05,
|
||||||
|
"max_position_embeddings": 77,
|
||||||
|
"model_type": "clip_text_model",
|
||||||
|
"num_attention_heads": 16,
|
||||||
|
"num_hidden_layers": 24,
|
||||||
|
"pad_token_id": 1,
|
||||||
|
"projection_dim": 1024,
|
||||||
|
"torch_dtype": "float32",
|
||||||
|
"vocab_size": 49408
|
||||||
|
}
|
66
ldm_patched/modules/sdxl_clip.py
Normal file
66
ldm_patched/modules/sdxl_clip.py
Normal file
@ -0,0 +1,66 @@
|
|||||||
|
from ldm_patched.modules import sd1_clip
|
||||||
|
import torch
|
||||||
|
import os
|
||||||
|
|
||||||
|
class SDXLClipG(sd1_clip.SDClipModel):
|
||||||
|
def __init__(self, device="cpu", max_length=77, freeze=True, layer="penultimate", layer_idx=None, dtype=None):
|
||||||
|
if layer == "penultimate":
|
||||||
|
layer="hidden"
|
||||||
|
layer_idx=-2
|
||||||
|
|
||||||
|
textmodel_json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "clip_config_bigg.json")
|
||||||
|
super().__init__(device=device, freeze=freeze, layer=layer, layer_idx=layer_idx, textmodel_json_config=textmodel_json_config, dtype=dtype,
|
||||||
|
special_tokens={"start": 49406, "end": 49407, "pad": 0}, layer_norm_hidden_state=False)
|
||||||
|
|
||||||
|
def load_sd(self, sd):
|
||||||
|
return super().load_sd(sd)
|
||||||
|
|
||||||
|
class SDXLClipGTokenizer(sd1_clip.SDTokenizer):
|
||||||
|
def __init__(self, tokenizer_path=None, embedding_directory=None):
|
||||||
|
super().__init__(tokenizer_path, pad_with_end=False, embedding_directory=embedding_directory, embedding_size=1280, embedding_key='clip_g')
|
||||||
|
|
||||||
|
|
||||||
|
class SDXLTokenizer:
|
||||||
|
def __init__(self, embedding_directory=None):
|
||||||
|
self.clip_l = sd1_clip.SDTokenizer(embedding_directory=embedding_directory)
|
||||||
|
self.clip_g = SDXLClipGTokenizer(embedding_directory=embedding_directory)
|
||||||
|
|
||||||
|
def tokenize_with_weights(self, text:str, return_word_ids=False):
|
||||||
|
out = {}
|
||||||
|
out["g"] = self.clip_g.tokenize_with_weights(text, return_word_ids)
|
||||||
|
out["l"] = self.clip_l.tokenize_with_weights(text, return_word_ids)
|
||||||
|
return out
|
||||||
|
|
||||||
|
def untokenize(self, token_weight_pair):
|
||||||
|
return self.clip_g.untokenize(token_weight_pair)
|
||||||
|
|
||||||
|
class SDXLClipModel(torch.nn.Module):
|
||||||
|
def __init__(self, device="cpu", dtype=None):
|
||||||
|
super().__init__()
|
||||||
|
self.clip_l = sd1_clip.SDClipModel(layer="hidden", layer_idx=-2, device=device, dtype=dtype, layer_norm_hidden_state=False)
|
||||||
|
self.clip_g = SDXLClipG(device=device, dtype=dtype)
|
||||||
|
|
||||||
|
def clip_layer(self, layer_idx):
|
||||||
|
self.clip_l.clip_layer(layer_idx)
|
||||||
|
self.clip_g.clip_layer(layer_idx)
|
||||||
|
|
||||||
|
def reset_clip_layer(self):
|
||||||
|
self.clip_g.reset_clip_layer()
|
||||||
|
self.clip_l.reset_clip_layer()
|
||||||
|
|
||||||
|
def encode_token_weights(self, token_weight_pairs):
|
||||||
|
token_weight_pairs_g = token_weight_pairs["g"]
|
||||||
|
token_weight_pairs_l = token_weight_pairs["l"]
|
||||||
|
g_out, g_pooled = self.clip_g.encode_token_weights(token_weight_pairs_g)
|
||||||
|
l_out, l_pooled = self.clip_l.encode_token_weights(token_weight_pairs_l)
|
||||||
|
return torch.cat([l_out, g_out], dim=-1), g_pooled
|
||||||
|
|
||||||
|
def load_sd(self, sd):
|
||||||
|
if "text_model.encoder.layers.30.mlp.fc1.weight" in sd:
|
||||||
|
return self.clip_g.load_sd(sd)
|
||||||
|
else:
|
||||||
|
return self.clip_l.load_sd(sd)
|
||||||
|
|
||||||
|
class SDXLRefinerClipModel(sd1_clip.SD1ClipModel):
|
||||||
|
def __init__(self, device="cpu", dtype=None):
|
||||||
|
super().__init__(device=device, dtype=dtype, clip_name="g", clip_model=SDXLClipG)
|
310
ldm_patched/modules/supported_models.py
Normal file
310
ldm_patched/modules/supported_models.py
Normal file
@ -0,0 +1,310 @@
|
|||||||
|
import torch
|
||||||
|
from . import model_base
|
||||||
|
from . import utils
|
||||||
|
|
||||||
|
from . import sd1_clip
|
||||||
|
from . import sd2_clip
|
||||||
|
from . import sdxl_clip
|
||||||
|
|
||||||
|
from . import supported_models_base
|
||||||
|
from . import latent_formats
|
||||||
|
|
||||||
|
from . import diffusers_convert
|
||||||
|
|
||||||
|
class SD15(supported_models_base.BASE):
|
||||||
|
unet_config = {
|
||||||
|
"context_dim": 768,
|
||||||
|
"model_channels": 320,
|
||||||
|
"use_linear_in_transformer": False,
|
||||||
|
"adm_in_channels": None,
|
||||||
|
"use_temporal_attention": False,
|
||||||
|
}
|
||||||
|
|
||||||
|
unet_extra_config = {
|
||||||
|
"num_heads": 8,
|
||||||
|
"num_head_channels": -1,
|
||||||
|
}
|
||||||
|
|
||||||
|
latent_format = latent_formats.SD15
|
||||||
|
|
||||||
|
def process_clip_state_dict(self, state_dict):
|
||||||
|
k = list(state_dict.keys())
|
||||||
|
for x in k:
|
||||||
|
if x.startswith("cond_stage_model.transformer.") and not x.startswith("cond_stage_model.transformer.text_model."):
|
||||||
|
y = x.replace("cond_stage_model.transformer.", "cond_stage_model.transformer.text_model.")
|
||||||
|
state_dict[y] = state_dict.pop(x)
|
||||||
|
|
||||||
|
if 'cond_stage_model.transformer.text_model.embeddings.position_ids' in state_dict:
|
||||||
|
ids = state_dict['cond_stage_model.transformer.text_model.embeddings.position_ids']
|
||||||
|
if ids.dtype == torch.float32:
|
||||||
|
state_dict['cond_stage_model.transformer.text_model.embeddings.position_ids'] = ids.round()
|
||||||
|
|
||||||
|
replace_prefix = {}
|
||||||
|
replace_prefix["cond_stage_model."] = "cond_stage_model.clip_l."
|
||||||
|
state_dict = utils.state_dict_prefix_replace(state_dict, replace_prefix)
|
||||||
|
return state_dict
|
||||||
|
|
||||||
|
def process_clip_state_dict_for_saving(self, state_dict):
|
||||||
|
replace_prefix = {"clip_l.": "cond_stage_model."}
|
||||||
|
return utils.state_dict_prefix_replace(state_dict, replace_prefix)
|
||||||
|
|
||||||
|
def clip_target(self):
|
||||||
|
return supported_models_base.ClipTarget(sd1_clip.SD1Tokenizer, sd1_clip.SD1ClipModel)
|
||||||
|
|
||||||
|
class SD20(supported_models_base.BASE):
|
||||||
|
unet_config = {
|
||||||
|
"context_dim": 1024,
|
||||||
|
"model_channels": 320,
|
||||||
|
"use_linear_in_transformer": True,
|
||||||
|
"adm_in_channels": None,
|
||||||
|
"use_temporal_attention": False,
|
||||||
|
}
|
||||||
|
|
||||||
|
latent_format = latent_formats.SD15
|
||||||
|
|
||||||
|
def model_type(self, state_dict, prefix=""):
|
||||||
|
if self.unet_config["in_channels"] == 4: #SD2.0 inpainting models are not v prediction
|
||||||
|
k = "{}output_blocks.11.1.transformer_blocks.0.norm1.bias".format(prefix)
|
||||||
|
out = state_dict[k]
|
||||||
|
if torch.std(out, unbiased=False) > 0.09: # not sure how well this will actually work. I guess we will find out.
|
||||||
|
return model_base.ModelType.V_PREDICTION
|
||||||
|
return model_base.ModelType.EPS
|
||||||
|
|
||||||
|
def process_clip_state_dict(self, state_dict):
|
||||||
|
replace_prefix = {}
|
||||||
|
replace_prefix["conditioner.embedders.0.model."] = "cond_stage_model.model." #SD2 in sgm format
|
||||||
|
state_dict = utils.state_dict_prefix_replace(state_dict, replace_prefix)
|
||||||
|
|
||||||
|
state_dict = utils.transformers_convert(state_dict, "cond_stage_model.model.", "cond_stage_model.clip_h.transformer.text_model.", 24)
|
||||||
|
return state_dict
|
||||||
|
|
||||||
|
def process_clip_state_dict_for_saving(self, state_dict):
|
||||||
|
replace_prefix = {}
|
||||||
|
replace_prefix["clip_h"] = "cond_stage_model.model"
|
||||||
|
state_dict = utils.state_dict_prefix_replace(state_dict, replace_prefix)
|
||||||
|
state_dict = diffusers_convert.convert_text_enc_state_dict_v20(state_dict)
|
||||||
|
return state_dict
|
||||||
|
|
||||||
|
def clip_target(self):
|
||||||
|
return supported_models_base.ClipTarget(sd2_clip.SD2Tokenizer, sd2_clip.SD2ClipModel)
|
||||||
|
|
||||||
|
class SD21UnclipL(SD20):
|
||||||
|
unet_config = {
|
||||||
|
"context_dim": 1024,
|
||||||
|
"model_channels": 320,
|
||||||
|
"use_linear_in_transformer": True,
|
||||||
|
"adm_in_channels": 1536,
|
||||||
|
"use_temporal_attention": False,
|
||||||
|
}
|
||||||
|
|
||||||
|
clip_vision_prefix = "embedder.model.visual."
|
||||||
|
noise_aug_config = {"noise_schedule_config": {"timesteps": 1000, "beta_schedule": "squaredcos_cap_v2"}, "timestep_dim": 768}
|
||||||
|
|
||||||
|
|
||||||
|
class SD21UnclipH(SD20):
|
||||||
|
unet_config = {
|
||||||
|
"context_dim": 1024,
|
||||||
|
"model_channels": 320,
|
||||||
|
"use_linear_in_transformer": True,
|
||||||
|
"adm_in_channels": 2048,
|
||||||
|
"use_temporal_attention": False,
|
||||||
|
}
|
||||||
|
|
||||||
|
clip_vision_prefix = "embedder.model.visual."
|
||||||
|
noise_aug_config = {"noise_schedule_config": {"timesteps": 1000, "beta_schedule": "squaredcos_cap_v2"}, "timestep_dim": 1024}
|
||||||
|
|
||||||
|
class SDXLRefiner(supported_models_base.BASE):
|
||||||
|
unet_config = {
|
||||||
|
"model_channels": 384,
|
||||||
|
"use_linear_in_transformer": True,
|
||||||
|
"context_dim": 1280,
|
||||||
|
"adm_in_channels": 2560,
|
||||||
|
"transformer_depth": [0, 0, 4, 4, 4, 4, 0, 0],
|
||||||
|
"use_temporal_attention": False,
|
||||||
|
}
|
||||||
|
|
||||||
|
latent_format = latent_formats.SDXL
|
||||||
|
|
||||||
|
def get_model(self, state_dict, prefix="", device=None):
|
||||||
|
return model_base.SDXLRefiner(self, device=device)
|
||||||
|
|
||||||
|
def process_clip_state_dict(self, state_dict):
|
||||||
|
keys_to_replace = {}
|
||||||
|
replace_prefix = {}
|
||||||
|
|
||||||
|
state_dict = utils.transformers_convert(state_dict, "conditioner.embedders.0.model.", "cond_stage_model.clip_g.transformer.text_model.", 32)
|
||||||
|
keys_to_replace["conditioner.embedders.0.model.text_projection"] = "cond_stage_model.clip_g.text_projection"
|
||||||
|
keys_to_replace["conditioner.embedders.0.model.logit_scale"] = "cond_stage_model.clip_g.logit_scale"
|
||||||
|
|
||||||
|
state_dict = utils.state_dict_key_replace(state_dict, keys_to_replace)
|
||||||
|
return state_dict
|
||||||
|
|
||||||
|
def process_clip_state_dict_for_saving(self, state_dict):
|
||||||
|
replace_prefix = {}
|
||||||
|
state_dict_g = diffusers_convert.convert_text_enc_state_dict_v20(state_dict, "clip_g")
|
||||||
|
if "clip_g.transformer.text_model.embeddings.position_ids" in state_dict_g:
|
||||||
|
state_dict_g.pop("clip_g.transformer.text_model.embeddings.position_ids")
|
||||||
|
replace_prefix["clip_g"] = "conditioner.embedders.0.model"
|
||||||
|
state_dict_g = utils.state_dict_prefix_replace(state_dict_g, replace_prefix)
|
||||||
|
return state_dict_g
|
||||||
|
|
||||||
|
def clip_target(self):
|
||||||
|
return supported_models_base.ClipTarget(sdxl_clip.SDXLTokenizer, sdxl_clip.SDXLRefinerClipModel)
|
||||||
|
|
||||||
|
class SDXL(supported_models_base.BASE):
|
||||||
|
unet_config = {
|
||||||
|
"model_channels": 320,
|
||||||
|
"use_linear_in_transformer": True,
|
||||||
|
"transformer_depth": [0, 0, 2, 2, 10, 10],
|
||||||
|
"context_dim": 2048,
|
||||||
|
"adm_in_channels": 2816,
|
||||||
|
"use_temporal_attention": False,
|
||||||
|
}
|
||||||
|
|
||||||
|
latent_format = latent_formats.SDXL
|
||||||
|
|
||||||
|
def model_type(self, state_dict, prefix=""):
|
||||||
|
if "v_pred" in state_dict:
|
||||||
|
return model_base.ModelType.V_PREDICTION
|
||||||
|
else:
|
||||||
|
return model_base.ModelType.EPS
|
||||||
|
|
||||||
|
def get_model(self, state_dict, prefix="", device=None):
|
||||||
|
out = model_base.SDXL(self, model_type=self.model_type(state_dict, prefix), device=device)
|
||||||
|
if self.inpaint_model():
|
||||||
|
out.set_inpaint()
|
||||||
|
return out
|
||||||
|
|
||||||
|
def process_clip_state_dict(self, state_dict):
|
||||||
|
keys_to_replace = {}
|
||||||
|
replace_prefix = {}
|
||||||
|
|
||||||
|
replace_prefix["conditioner.embedders.0.transformer.text_model"] = "cond_stage_model.clip_l.transformer.text_model"
|
||||||
|
state_dict = utils.transformers_convert(state_dict, "conditioner.embedders.1.model.", "cond_stage_model.clip_g.transformer.text_model.", 32)
|
||||||
|
keys_to_replace["conditioner.embedders.1.model.text_projection"] = "cond_stage_model.clip_g.text_projection"
|
||||||
|
keys_to_replace["conditioner.embedders.1.model.text_projection.weight"] = "cond_stage_model.clip_g.text_projection"
|
||||||
|
keys_to_replace["conditioner.embedders.1.model.logit_scale"] = "cond_stage_model.clip_g.logit_scale"
|
||||||
|
|
||||||
|
state_dict = utils.state_dict_prefix_replace(state_dict, replace_prefix)
|
||||||
|
state_dict = utils.state_dict_key_replace(state_dict, keys_to_replace)
|
||||||
|
return state_dict
|
||||||
|
|
||||||
|
def process_clip_state_dict_for_saving(self, state_dict):
|
||||||
|
replace_prefix = {}
|
||||||
|
keys_to_replace = {}
|
||||||
|
state_dict_g = diffusers_convert.convert_text_enc_state_dict_v20(state_dict, "clip_g")
|
||||||
|
if "clip_g.transformer.text_model.embeddings.position_ids" in state_dict_g:
|
||||||
|
state_dict_g.pop("clip_g.transformer.text_model.embeddings.position_ids")
|
||||||
|
for k in state_dict:
|
||||||
|
if k.startswith("clip_l"):
|
||||||
|
state_dict_g[k] = state_dict[k]
|
||||||
|
|
||||||
|
replace_prefix["clip_g"] = "conditioner.embedders.1.model"
|
||||||
|
replace_prefix["clip_l"] = "conditioner.embedders.0"
|
||||||
|
state_dict_g = utils.state_dict_prefix_replace(state_dict_g, replace_prefix)
|
||||||
|
return state_dict_g
|
||||||
|
|
||||||
|
def clip_target(self):
|
||||||
|
return supported_models_base.ClipTarget(sdxl_clip.SDXLTokenizer, sdxl_clip.SDXLClipModel)
|
||||||
|
|
||||||
|
class SSD1B(SDXL):
|
||||||
|
unet_config = {
|
||||||
|
"model_channels": 320,
|
||||||
|
"use_linear_in_transformer": True,
|
||||||
|
"transformer_depth": [0, 0, 2, 2, 4, 4],
|
||||||
|
"context_dim": 2048,
|
||||||
|
"adm_in_channels": 2816,
|
||||||
|
"use_temporal_attention": False,
|
||||||
|
}
|
||||||
|
|
||||||
|
class Segmind_Vega(SDXL):
|
||||||
|
unet_config = {
|
||||||
|
"model_channels": 320,
|
||||||
|
"use_linear_in_transformer": True,
|
||||||
|
"transformer_depth": [0, 0, 1, 1, 2, 2],
|
||||||
|
"context_dim": 2048,
|
||||||
|
"adm_in_channels": 2816,
|
||||||
|
"use_temporal_attention": False,
|
||||||
|
}
|
||||||
|
|
||||||
|
class SVD_img2vid(supported_models_base.BASE):
|
||||||
|
unet_config = {
|
||||||
|
"model_channels": 320,
|
||||||
|
"in_channels": 8,
|
||||||
|
"use_linear_in_transformer": True,
|
||||||
|
"transformer_depth": [1, 1, 1, 1, 1, 1, 0, 0],
|
||||||
|
"context_dim": 1024,
|
||||||
|
"adm_in_channels": 768,
|
||||||
|
"use_temporal_attention": True,
|
||||||
|
"use_temporal_resblock": True
|
||||||
|
}
|
||||||
|
|
||||||
|
clip_vision_prefix = "conditioner.embedders.0.open_clip.model.visual."
|
||||||
|
|
||||||
|
latent_format = latent_formats.SD15
|
||||||
|
|
||||||
|
sampling_settings = {"sigma_max": 700.0, "sigma_min": 0.002}
|
||||||
|
|
||||||
|
def get_model(self, state_dict, prefix="", device=None):
|
||||||
|
out = model_base.SVD_img2vid(self, device=device)
|
||||||
|
return out
|
||||||
|
|
||||||
|
def clip_target(self):
|
||||||
|
return None
|
||||||
|
|
||||||
|
class Stable_Zero123(supported_models_base.BASE):
|
||||||
|
unet_config = {
|
||||||
|
"context_dim": 768,
|
||||||
|
"model_channels": 320,
|
||||||
|
"use_linear_in_transformer": False,
|
||||||
|
"adm_in_channels": None,
|
||||||
|
"use_temporal_attention": False,
|
||||||
|
"in_channels": 8,
|
||||||
|
}
|
||||||
|
|
||||||
|
unet_extra_config = {
|
||||||
|
"num_heads": 8,
|
||||||
|
"num_head_channels": -1,
|
||||||
|
}
|
||||||
|
|
||||||
|
clip_vision_prefix = "cond_stage_model.model.visual."
|
||||||
|
|
||||||
|
latent_format = latent_formats.SD15
|
||||||
|
|
||||||
|
def get_model(self, state_dict, prefix="", device=None):
|
||||||
|
out = model_base.Stable_Zero123(self, device=device, cc_projection_weight=state_dict["cc_projection.weight"], cc_projection_bias=state_dict["cc_projection.bias"])
|
||||||
|
return out
|
||||||
|
|
||||||
|
def clip_target(self):
|
||||||
|
return None
|
||||||
|
|
||||||
|
class SD_X4Upscaler(SD20):
|
||||||
|
unet_config = {
|
||||||
|
"context_dim": 1024,
|
||||||
|
"model_channels": 256,
|
||||||
|
'in_channels': 7,
|
||||||
|
"use_linear_in_transformer": True,
|
||||||
|
"adm_in_channels": None,
|
||||||
|
"use_temporal_attention": False,
|
||||||
|
}
|
||||||
|
|
||||||
|
unet_extra_config = {
|
||||||
|
"disable_self_attentions": [True, True, True, False],
|
||||||
|
"num_classes": 1000,
|
||||||
|
"num_heads": 8,
|
||||||
|
"num_head_channels": -1,
|
||||||
|
}
|
||||||
|
|
||||||
|
latent_format = latent_formats.SD_X4
|
||||||
|
|
||||||
|
sampling_settings = {
|
||||||
|
"linear_start": 0.0001,
|
||||||
|
"linear_end": 0.02,
|
||||||
|
}
|
||||||
|
|
||||||
|
def get_model(self, state_dict, prefix="", device=None):
|
||||||
|
out = model_base.SD_X4Upscaler(self, device=device)
|
||||||
|
return out
|
||||||
|
|
||||||
|
models = [Stable_Zero123, SD15, SD20, SD21UnclipL, SD21UnclipH, SDXLRefiner, SDXL, SSD1B, Segmind_Vega, SD_X4Upscaler]
|
||||||
|
models += [SVD_img2vid]
|
77
ldm_patched/modules/supported_models_base.py
Normal file
77
ldm_patched/modules/supported_models_base.py
Normal file
@ -0,0 +1,77 @@
|
|||||||
|
import torch
|
||||||
|
from . import model_base
|
||||||
|
from . import utils
|
||||||
|
from . import latent_formats
|
||||||
|
|
||||||
|
class ClipTarget:
|
||||||
|
def __init__(self, tokenizer, clip):
|
||||||
|
self.clip = clip
|
||||||
|
self.tokenizer = tokenizer
|
||||||
|
self.params = {}
|
||||||
|
|
||||||
|
class BASE:
|
||||||
|
unet_config = {}
|
||||||
|
unet_extra_config = {
|
||||||
|
"num_heads": -1,
|
||||||
|
"num_head_channels": 64,
|
||||||
|
}
|
||||||
|
|
||||||
|
clip_prefix = []
|
||||||
|
clip_vision_prefix = None
|
||||||
|
noise_aug_config = None
|
||||||
|
sampling_settings = {}
|
||||||
|
latent_format = latent_formats.LatentFormat
|
||||||
|
|
||||||
|
manual_cast_dtype = None
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def matches(s, unet_config):
|
||||||
|
for k in s.unet_config:
|
||||||
|
if s.unet_config[k] != unet_config[k]:
|
||||||
|
return False
|
||||||
|
return True
|
||||||
|
|
||||||
|
def model_type(self, state_dict, prefix=""):
|
||||||
|
return model_base.ModelType.EPS
|
||||||
|
|
||||||
|
def inpaint_model(self):
|
||||||
|
return self.unet_config["in_channels"] > 4
|
||||||
|
|
||||||
|
def __init__(self, unet_config):
|
||||||
|
self.unet_config = unet_config
|
||||||
|
self.latent_format = self.latent_format()
|
||||||
|
for x in self.unet_extra_config:
|
||||||
|
self.unet_config[x] = self.unet_extra_config[x]
|
||||||
|
|
||||||
|
def get_model(self, state_dict, prefix="", device=None):
|
||||||
|
if self.noise_aug_config is not None:
|
||||||
|
out = model_base.SD21UNCLIP(self, self.noise_aug_config, model_type=self.model_type(state_dict, prefix), device=device)
|
||||||
|
else:
|
||||||
|
out = model_base.BaseModel(self, model_type=self.model_type(state_dict, prefix), device=device)
|
||||||
|
if self.inpaint_model():
|
||||||
|
out.set_inpaint()
|
||||||
|
return out
|
||||||
|
|
||||||
|
def process_clip_state_dict(self, state_dict):
|
||||||
|
return state_dict
|
||||||
|
|
||||||
|
def process_unet_state_dict(self, state_dict):
|
||||||
|
return state_dict
|
||||||
|
|
||||||
|
def process_vae_state_dict(self, state_dict):
|
||||||
|
return state_dict
|
||||||
|
|
||||||
|
def process_clip_state_dict_for_saving(self, state_dict):
|
||||||
|
replace_prefix = {"": "cond_stage_model."}
|
||||||
|
return utils.state_dict_prefix_replace(state_dict, replace_prefix)
|
||||||
|
|
||||||
|
def process_unet_state_dict_for_saving(self, state_dict):
|
||||||
|
replace_prefix = {"": "model.diffusion_model."}
|
||||||
|
return utils.state_dict_prefix_replace(state_dict, replace_prefix)
|
||||||
|
|
||||||
|
def process_vae_state_dict_for_saving(self, state_dict):
|
||||||
|
replace_prefix = {"": "first_stage_model."}
|
||||||
|
return utils.state_dict_prefix_replace(state_dict, replace_prefix)
|
||||||
|
|
||||||
|
def set_manual_cast(self, manual_cast_dtype):
|
||||||
|
self.manual_cast_dtype = manual_cast_dtype
|
461
ldm_patched/modules/utils.py
Normal file
461
ldm_patched/modules/utils.py
Normal file
@ -0,0 +1,461 @@
|
|||||||
|
import torch
|
||||||
|
import math
|
||||||
|
import struct
|
||||||
|
import ldm_patched.modules.checkpoint_pickle
|
||||||
|
import safetensors.torch
|
||||||
|
import numpy as np
|
||||||
|
from PIL import Image
|
||||||
|
|
||||||
|
def load_torch_file(ckpt, safe_load=False, device=None):
|
||||||
|
if device is None:
|
||||||
|
device = torch.device("cpu")
|
||||||
|
if ckpt.lower().endswith(".safetensors"):
|
||||||
|
sd = safetensors.torch.load_file(ckpt, device=device.type)
|
||||||
|
else:
|
||||||
|
if safe_load:
|
||||||
|
if not 'weights_only' in torch.load.__code__.co_varnames:
|
||||||
|
print("Warning torch.load doesn't support weights_only on this pytorch version, loading unsafely.")
|
||||||
|
safe_load = False
|
||||||
|
if safe_load:
|
||||||
|
pl_sd = torch.load(ckpt, map_location=device, weights_only=True)
|
||||||
|
else:
|
||||||
|
pl_sd = torch.load(ckpt, map_location=device, pickle_module=ldm_patched.modules.checkpoint_pickle)
|
||||||
|
if "global_step" in pl_sd:
|
||||||
|
print(f"Global Step: {pl_sd['global_step']}")
|
||||||
|
if "state_dict" in pl_sd:
|
||||||
|
sd = pl_sd["state_dict"]
|
||||||
|
else:
|
||||||
|
sd = pl_sd
|
||||||
|
return sd
|
||||||
|
|
||||||
|
def save_torch_file(sd, ckpt, metadata=None):
|
||||||
|
if metadata is not None:
|
||||||
|
safetensors.torch.save_file(sd, ckpt, metadata=metadata)
|
||||||
|
else:
|
||||||
|
safetensors.torch.save_file(sd, ckpt)
|
||||||
|
|
||||||
|
def calculate_parameters(sd, prefix=""):
|
||||||
|
params = 0
|
||||||
|
for k in sd.keys():
|
||||||
|
if k.startswith(prefix):
|
||||||
|
params += sd[k].nelement()
|
||||||
|
return params
|
||||||
|
|
||||||
|
def state_dict_key_replace(state_dict, keys_to_replace):
|
||||||
|
for x in keys_to_replace:
|
||||||
|
if x in state_dict:
|
||||||
|
state_dict[keys_to_replace[x]] = state_dict.pop(x)
|
||||||
|
return state_dict
|
||||||
|
|
||||||
|
def state_dict_prefix_replace(state_dict, replace_prefix, filter_keys=False):
|
||||||
|
if filter_keys:
|
||||||
|
out = {}
|
||||||
|
else:
|
||||||
|
out = state_dict
|
||||||
|
for rp in replace_prefix:
|
||||||
|
replace = list(map(lambda a: (a, "{}{}".format(replace_prefix[rp], a[len(rp):])), filter(lambda a: a.startswith(rp), state_dict.keys())))
|
||||||
|
for x in replace:
|
||||||
|
w = state_dict.pop(x[0])
|
||||||
|
out[x[1]] = w
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
def transformers_convert(sd, prefix_from, prefix_to, number):
|
||||||
|
keys_to_replace = {
|
||||||
|
"{}positional_embedding": "{}embeddings.position_embedding.weight",
|
||||||
|
"{}token_embedding.weight": "{}embeddings.token_embedding.weight",
|
||||||
|
"{}ln_final.weight": "{}final_layer_norm.weight",
|
||||||
|
"{}ln_final.bias": "{}final_layer_norm.bias",
|
||||||
|
}
|
||||||
|
|
||||||
|
for k in keys_to_replace:
|
||||||
|
x = k.format(prefix_from)
|
||||||
|
if x in sd:
|
||||||
|
sd[keys_to_replace[k].format(prefix_to)] = sd.pop(x)
|
||||||
|
|
||||||
|
resblock_to_replace = {
|
||||||
|
"ln_1": "layer_norm1",
|
||||||
|
"ln_2": "layer_norm2",
|
||||||
|
"mlp.c_fc": "mlp.fc1",
|
||||||
|
"mlp.c_proj": "mlp.fc2",
|
||||||
|
"attn.out_proj": "self_attn.out_proj",
|
||||||
|
}
|
||||||
|
|
||||||
|
for resblock in range(number):
|
||||||
|
for x in resblock_to_replace:
|
||||||
|
for y in ["weight", "bias"]:
|
||||||
|
k = "{}transformer.resblocks.{}.{}.{}".format(prefix_from, resblock, x, y)
|
||||||
|
k_to = "{}encoder.layers.{}.{}.{}".format(prefix_to, resblock, resblock_to_replace[x], y)
|
||||||
|
if k in sd:
|
||||||
|
sd[k_to] = sd.pop(k)
|
||||||
|
|
||||||
|
for y in ["weight", "bias"]:
|
||||||
|
k_from = "{}transformer.resblocks.{}.attn.in_proj_{}".format(prefix_from, resblock, y)
|
||||||
|
if k_from in sd:
|
||||||
|
weights = sd.pop(k_from)
|
||||||
|
shape_from = weights.shape[0] // 3
|
||||||
|
for x in range(3):
|
||||||
|
p = ["self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj"]
|
||||||
|
k_to = "{}encoder.layers.{}.{}.{}".format(prefix_to, resblock, p[x], y)
|
||||||
|
sd[k_to] = weights[shape_from*x:shape_from*(x + 1)]
|
||||||
|
return sd
|
||||||
|
|
||||||
|
UNET_MAP_ATTENTIONS = {
|
||||||
|
"proj_in.weight",
|
||||||
|
"proj_in.bias",
|
||||||
|
"proj_out.weight",
|
||||||
|
"proj_out.bias",
|
||||||
|
"norm.weight",
|
||||||
|
"norm.bias",
|
||||||
|
}
|
||||||
|
|
||||||
|
TRANSFORMER_BLOCKS = {
|
||||||
|
"norm1.weight",
|
||||||
|
"norm1.bias",
|
||||||
|
"norm2.weight",
|
||||||
|
"norm2.bias",
|
||||||
|
"norm3.weight",
|
||||||
|
"norm3.bias",
|
||||||
|
"attn1.to_q.weight",
|
||||||
|
"attn1.to_k.weight",
|
||||||
|
"attn1.to_v.weight",
|
||||||
|
"attn1.to_out.0.weight",
|
||||||
|
"attn1.to_out.0.bias",
|
||||||
|
"attn2.to_q.weight",
|
||||||
|
"attn2.to_k.weight",
|
||||||
|
"attn2.to_v.weight",
|
||||||
|
"attn2.to_out.0.weight",
|
||||||
|
"attn2.to_out.0.bias",
|
||||||
|
"ff.net.0.proj.weight",
|
||||||
|
"ff.net.0.proj.bias",
|
||||||
|
"ff.net.2.weight",
|
||||||
|
"ff.net.2.bias",
|
||||||
|
}
|
||||||
|
|
||||||
|
UNET_MAP_RESNET = {
|
||||||
|
"in_layers.2.weight": "conv1.weight",
|
||||||
|
"in_layers.2.bias": "conv1.bias",
|
||||||
|
"emb_layers.1.weight": "time_emb_proj.weight",
|
||||||
|
"emb_layers.1.bias": "time_emb_proj.bias",
|
||||||
|
"out_layers.3.weight": "conv2.weight",
|
||||||
|
"out_layers.3.bias": "conv2.bias",
|
||||||
|
"skip_connection.weight": "conv_shortcut.weight",
|
||||||
|
"skip_connection.bias": "conv_shortcut.bias",
|
||||||
|
"in_layers.0.weight": "norm1.weight",
|
||||||
|
"in_layers.0.bias": "norm1.bias",
|
||||||
|
"out_layers.0.weight": "norm2.weight",
|
||||||
|
"out_layers.0.bias": "norm2.bias",
|
||||||
|
}
|
||||||
|
|
||||||
|
UNET_MAP_BASIC = {
|
||||||
|
("label_emb.0.0.weight", "class_embedding.linear_1.weight"),
|
||||||
|
("label_emb.0.0.bias", "class_embedding.linear_1.bias"),
|
||||||
|
("label_emb.0.2.weight", "class_embedding.linear_2.weight"),
|
||||||
|
("label_emb.0.2.bias", "class_embedding.linear_2.bias"),
|
||||||
|
("label_emb.0.0.weight", "add_embedding.linear_1.weight"),
|
||||||
|
("label_emb.0.0.bias", "add_embedding.linear_1.bias"),
|
||||||
|
("label_emb.0.2.weight", "add_embedding.linear_2.weight"),
|
||||||
|
("label_emb.0.2.bias", "add_embedding.linear_2.bias"),
|
||||||
|
("input_blocks.0.0.weight", "conv_in.weight"),
|
||||||
|
("input_blocks.0.0.bias", "conv_in.bias"),
|
||||||
|
("out.0.weight", "conv_norm_out.weight"),
|
||||||
|
("out.0.bias", "conv_norm_out.bias"),
|
||||||
|
("out.2.weight", "conv_out.weight"),
|
||||||
|
("out.2.bias", "conv_out.bias"),
|
||||||
|
("time_embed.0.weight", "time_embedding.linear_1.weight"),
|
||||||
|
("time_embed.0.bias", "time_embedding.linear_1.bias"),
|
||||||
|
("time_embed.2.weight", "time_embedding.linear_2.weight"),
|
||||||
|
("time_embed.2.bias", "time_embedding.linear_2.bias")
|
||||||
|
}
|
||||||
|
|
||||||
|
def unet_to_diffusers(unet_config):
|
||||||
|
num_res_blocks = unet_config["num_res_blocks"]
|
||||||
|
channel_mult = unet_config["channel_mult"]
|
||||||
|
transformer_depth = unet_config["transformer_depth"][:]
|
||||||
|
transformer_depth_output = unet_config["transformer_depth_output"][:]
|
||||||
|
num_blocks = len(channel_mult)
|
||||||
|
|
||||||
|
transformers_mid = unet_config.get("transformer_depth_middle", None)
|
||||||
|
|
||||||
|
diffusers_unet_map = {}
|
||||||
|
for x in range(num_blocks):
|
||||||
|
n = 1 + (num_res_blocks[x] + 1) * x
|
||||||
|
for i in range(num_res_blocks[x]):
|
||||||
|
for b in UNET_MAP_RESNET:
|
||||||
|
diffusers_unet_map["down_blocks.{}.resnets.{}.{}".format(x, i, UNET_MAP_RESNET[b])] = "input_blocks.{}.0.{}".format(n, b)
|
||||||
|
num_transformers = transformer_depth.pop(0)
|
||||||
|
if num_transformers > 0:
|
||||||
|
for b in UNET_MAP_ATTENTIONS:
|
||||||
|
diffusers_unet_map["down_blocks.{}.attentions.{}.{}".format(x, i, b)] = "input_blocks.{}.1.{}".format(n, b)
|
||||||
|
for t in range(num_transformers):
|
||||||
|
for b in TRANSFORMER_BLOCKS:
|
||||||
|
diffusers_unet_map["down_blocks.{}.attentions.{}.transformer_blocks.{}.{}".format(x, i, t, b)] = "input_blocks.{}.1.transformer_blocks.{}.{}".format(n, t, b)
|
||||||
|
n += 1
|
||||||
|
for k in ["weight", "bias"]:
|
||||||
|
diffusers_unet_map["down_blocks.{}.downsamplers.0.conv.{}".format(x, k)] = "input_blocks.{}.0.op.{}".format(n, k)
|
||||||
|
|
||||||
|
i = 0
|
||||||
|
for b in UNET_MAP_ATTENTIONS:
|
||||||
|
diffusers_unet_map["mid_block.attentions.{}.{}".format(i, b)] = "middle_block.1.{}".format(b)
|
||||||
|
for t in range(transformers_mid):
|
||||||
|
for b in TRANSFORMER_BLOCKS:
|
||||||
|
diffusers_unet_map["mid_block.attentions.{}.transformer_blocks.{}.{}".format(i, t, b)] = "middle_block.1.transformer_blocks.{}.{}".format(t, b)
|
||||||
|
|
||||||
|
for i, n in enumerate([0, 2]):
|
||||||
|
for b in UNET_MAP_RESNET:
|
||||||
|
diffusers_unet_map["mid_block.resnets.{}.{}".format(i, UNET_MAP_RESNET[b])] = "middle_block.{}.{}".format(n, b)
|
||||||
|
|
||||||
|
num_res_blocks = list(reversed(num_res_blocks))
|
||||||
|
for x in range(num_blocks):
|
||||||
|
n = (num_res_blocks[x] + 1) * x
|
||||||
|
l = num_res_blocks[x] + 1
|
||||||
|
for i in range(l):
|
||||||
|
c = 0
|
||||||
|
for b in UNET_MAP_RESNET:
|
||||||
|
diffusers_unet_map["up_blocks.{}.resnets.{}.{}".format(x, i, UNET_MAP_RESNET[b])] = "output_blocks.{}.0.{}".format(n, b)
|
||||||
|
c += 1
|
||||||
|
num_transformers = transformer_depth_output.pop()
|
||||||
|
if num_transformers > 0:
|
||||||
|
c += 1
|
||||||
|
for b in UNET_MAP_ATTENTIONS:
|
||||||
|
diffusers_unet_map["up_blocks.{}.attentions.{}.{}".format(x, i, b)] = "output_blocks.{}.1.{}".format(n, b)
|
||||||
|
for t in range(num_transformers):
|
||||||
|
for b in TRANSFORMER_BLOCKS:
|
||||||
|
diffusers_unet_map["up_blocks.{}.attentions.{}.transformer_blocks.{}.{}".format(x, i, t, b)] = "output_blocks.{}.1.transformer_blocks.{}.{}".format(n, t, b)
|
||||||
|
if i == l - 1:
|
||||||
|
for k in ["weight", "bias"]:
|
||||||
|
diffusers_unet_map["up_blocks.{}.upsamplers.0.conv.{}".format(x, k)] = "output_blocks.{}.{}.conv.{}".format(n, c, k)
|
||||||
|
n += 1
|
||||||
|
|
||||||
|
for k in UNET_MAP_BASIC:
|
||||||
|
diffusers_unet_map[k[1]] = k[0]
|
||||||
|
|
||||||
|
return diffusers_unet_map
|
||||||
|
|
||||||
|
def repeat_to_batch_size(tensor, batch_size):
|
||||||
|
if tensor.shape[0] > batch_size:
|
||||||
|
return tensor[:batch_size]
|
||||||
|
elif tensor.shape[0] < batch_size:
|
||||||
|
return tensor.repeat([math.ceil(batch_size / tensor.shape[0])] + [1] * (len(tensor.shape) - 1))[:batch_size]
|
||||||
|
return tensor
|
||||||
|
|
||||||
|
def resize_to_batch_size(tensor, batch_size):
|
||||||
|
in_batch_size = tensor.shape[0]
|
||||||
|
if in_batch_size == batch_size:
|
||||||
|
return tensor
|
||||||
|
|
||||||
|
if batch_size <= 1:
|
||||||
|
return tensor[:batch_size]
|
||||||
|
|
||||||
|
output = torch.empty([batch_size] + list(tensor.shape)[1:], dtype=tensor.dtype, device=tensor.device)
|
||||||
|
if batch_size < in_batch_size:
|
||||||
|
scale = (in_batch_size - 1) / (batch_size - 1)
|
||||||
|
for i in range(batch_size):
|
||||||
|
output[i] = tensor[min(round(i * scale), in_batch_size - 1)]
|
||||||
|
else:
|
||||||
|
scale = in_batch_size / batch_size
|
||||||
|
for i in range(batch_size):
|
||||||
|
output[i] = tensor[min(math.floor((i + 0.5) * scale), in_batch_size - 1)]
|
||||||
|
|
||||||
|
return output
|
||||||
|
|
||||||
|
def convert_sd_to(state_dict, dtype):
|
||||||
|
keys = list(state_dict.keys())
|
||||||
|
for k in keys:
|
||||||
|
state_dict[k] = state_dict[k].to(dtype)
|
||||||
|
return state_dict
|
||||||
|
|
||||||
|
def safetensors_header(safetensors_path, max_size=100*1024*1024):
|
||||||
|
with open(safetensors_path, "rb") as f:
|
||||||
|
header = f.read(8)
|
||||||
|
length_of_header = struct.unpack('<Q', header)[0]
|
||||||
|
if length_of_header > max_size:
|
||||||
|
return None
|
||||||
|
return f.read(length_of_header)
|
||||||
|
|
||||||
|
def set_attr(obj, attr, value):
|
||||||
|
attrs = attr.split(".")
|
||||||
|
for name in attrs[:-1]:
|
||||||
|
obj = getattr(obj, name)
|
||||||
|
prev = getattr(obj, attrs[-1])
|
||||||
|
setattr(obj, attrs[-1], torch.nn.Parameter(value, requires_grad=False))
|
||||||
|
del prev
|
||||||
|
|
||||||
|
def copy_to_param(obj, attr, value):
|
||||||
|
# inplace update tensor instead of replacing it
|
||||||
|
attrs = attr.split(".")
|
||||||
|
for name in attrs[:-1]:
|
||||||
|
obj = getattr(obj, name)
|
||||||
|
prev = getattr(obj, attrs[-1])
|
||||||
|
prev.data.copy_(value)
|
||||||
|
|
||||||
|
def get_attr(obj, attr):
|
||||||
|
attrs = attr.split(".")
|
||||||
|
for name in attrs:
|
||||||
|
obj = getattr(obj, name)
|
||||||
|
return obj
|
||||||
|
|
||||||
|
def bislerp(samples, width, height):
|
||||||
|
def slerp(b1, b2, r):
|
||||||
|
'''slerps batches b1, b2 according to ratio r, batches should be flat e.g. NxC'''
|
||||||
|
|
||||||
|
c = b1.shape[-1]
|
||||||
|
|
||||||
|
#norms
|
||||||
|
b1_norms = torch.norm(b1, dim=-1, keepdim=True)
|
||||||
|
b2_norms = torch.norm(b2, dim=-1, keepdim=True)
|
||||||
|
|
||||||
|
#normalize
|
||||||
|
b1_normalized = b1 / b1_norms
|
||||||
|
b2_normalized = b2 / b2_norms
|
||||||
|
|
||||||
|
#zero when norms are zero
|
||||||
|
b1_normalized[b1_norms.expand(-1,c) == 0.0] = 0.0
|
||||||
|
b2_normalized[b2_norms.expand(-1,c) == 0.0] = 0.0
|
||||||
|
|
||||||
|
#slerp
|
||||||
|
dot = (b1_normalized*b2_normalized).sum(1)
|
||||||
|
omega = torch.acos(dot)
|
||||||
|
so = torch.sin(omega)
|
||||||
|
|
||||||
|
#technically not mathematically correct, but more pleasing?
|
||||||
|
res = (torch.sin((1.0-r.squeeze(1))*omega)/so).unsqueeze(1)*b1_normalized + (torch.sin(r.squeeze(1)*omega)/so).unsqueeze(1) * b2_normalized
|
||||||
|
res *= (b1_norms * (1.0-r) + b2_norms * r).expand(-1,c)
|
||||||
|
|
||||||
|
#edge cases for same or polar opposites
|
||||||
|
res[dot > 1 - 1e-5] = b1[dot > 1 - 1e-5]
|
||||||
|
res[dot < 1e-5 - 1] = (b1 * (1.0-r) + b2 * r)[dot < 1e-5 - 1]
|
||||||
|
return res
|
||||||
|
|
||||||
|
def generate_bilinear_data(length_old, length_new, device):
|
||||||
|
coords_1 = torch.arange(length_old, dtype=torch.float32, device=device).reshape((1,1,1,-1))
|
||||||
|
coords_1 = torch.nn.functional.interpolate(coords_1, size=(1, length_new), mode="bilinear")
|
||||||
|
ratios = coords_1 - coords_1.floor()
|
||||||
|
coords_1 = coords_1.to(torch.int64)
|
||||||
|
|
||||||
|
coords_2 = torch.arange(length_old, dtype=torch.float32, device=device).reshape((1,1,1,-1)) + 1
|
||||||
|
coords_2[:,:,:,-1] -= 1
|
||||||
|
coords_2 = torch.nn.functional.interpolate(coords_2, size=(1, length_new), mode="bilinear")
|
||||||
|
coords_2 = coords_2.to(torch.int64)
|
||||||
|
return ratios, coords_1, coords_2
|
||||||
|
|
||||||
|
orig_dtype = samples.dtype
|
||||||
|
samples = samples.float()
|
||||||
|
n,c,h,w = samples.shape
|
||||||
|
h_new, w_new = (height, width)
|
||||||
|
|
||||||
|
#linear w
|
||||||
|
ratios, coords_1, coords_2 = generate_bilinear_data(w, w_new, samples.device)
|
||||||
|
coords_1 = coords_1.expand((n, c, h, -1))
|
||||||
|
coords_2 = coords_2.expand((n, c, h, -1))
|
||||||
|
ratios = ratios.expand((n, 1, h, -1))
|
||||||
|
|
||||||
|
pass_1 = samples.gather(-1,coords_1).movedim(1, -1).reshape((-1,c))
|
||||||
|
pass_2 = samples.gather(-1,coords_2).movedim(1, -1).reshape((-1,c))
|
||||||
|
ratios = ratios.movedim(1, -1).reshape((-1,1))
|
||||||
|
|
||||||
|
result = slerp(pass_1, pass_2, ratios)
|
||||||
|
result = result.reshape(n, h, w_new, c).movedim(-1, 1)
|
||||||
|
|
||||||
|
#linear h
|
||||||
|
ratios, coords_1, coords_2 = generate_bilinear_data(h, h_new, samples.device)
|
||||||
|
coords_1 = coords_1.reshape((1,1,-1,1)).expand((n, c, -1, w_new))
|
||||||
|
coords_2 = coords_2.reshape((1,1,-1,1)).expand((n, c, -1, w_new))
|
||||||
|
ratios = ratios.reshape((1,1,-1,1)).expand((n, 1, -1, w_new))
|
||||||
|
|
||||||
|
pass_1 = result.gather(-2,coords_1).movedim(1, -1).reshape((-1,c))
|
||||||
|
pass_2 = result.gather(-2,coords_2).movedim(1, -1).reshape((-1,c))
|
||||||
|
ratios = ratios.movedim(1, -1).reshape((-1,1))
|
||||||
|
|
||||||
|
result = slerp(pass_1, pass_2, ratios)
|
||||||
|
result = result.reshape(n, h_new, w_new, c).movedim(-1, 1)
|
||||||
|
return result.to(orig_dtype)
|
||||||
|
|
||||||
|
def lanczos(samples, width, height):
|
||||||
|
images = [Image.fromarray(np.clip(255. * image.movedim(0, -1).cpu().numpy(), 0, 255).astype(np.uint8)) for image in samples]
|
||||||
|
images = [image.resize((width, height), resample=Image.Resampling.LANCZOS) for image in images]
|
||||||
|
images = [torch.from_numpy(np.array(image).astype(np.float32) / 255.0).movedim(-1, 0) for image in images]
|
||||||
|
result = torch.stack(images)
|
||||||
|
return result.to(samples.device, samples.dtype)
|
||||||
|
|
||||||
|
def common_upscale(samples, width, height, upscale_method, crop):
|
||||||
|
if crop == "center":
|
||||||
|
old_width = samples.shape[3]
|
||||||
|
old_height = samples.shape[2]
|
||||||
|
old_aspect = old_width / old_height
|
||||||
|
new_aspect = width / height
|
||||||
|
x = 0
|
||||||
|
y = 0
|
||||||
|
if old_aspect > new_aspect:
|
||||||
|
x = round((old_width - old_width * (new_aspect / old_aspect)) / 2)
|
||||||
|
elif old_aspect < new_aspect:
|
||||||
|
y = round((old_height - old_height * (old_aspect / new_aspect)) / 2)
|
||||||
|
s = samples[:,:,y:old_height-y,x:old_width-x]
|
||||||
|
else:
|
||||||
|
s = samples
|
||||||
|
|
||||||
|
if upscale_method == "bislerp":
|
||||||
|
return bislerp(s, width, height)
|
||||||
|
elif upscale_method == "lanczos":
|
||||||
|
return lanczos(s, width, height)
|
||||||
|
else:
|
||||||
|
return torch.nn.functional.interpolate(s, size=(height, width), mode=upscale_method)
|
||||||
|
|
||||||
|
def get_tiled_scale_steps(width, height, tile_x, tile_y, overlap):
|
||||||
|
return math.ceil((height / (tile_y - overlap))) * math.ceil((width / (tile_x - overlap)))
|
||||||
|
|
||||||
|
@torch.inference_mode()
|
||||||
|
def tiled_scale(samples, function, tile_x=64, tile_y=64, overlap = 8, upscale_amount = 4, out_channels = 3, output_device="cpu", pbar = None):
|
||||||
|
output = torch.empty((samples.shape[0], out_channels, round(samples.shape[2] * upscale_amount), round(samples.shape[3] * upscale_amount)), device=output_device)
|
||||||
|
for b in range(samples.shape[0]):
|
||||||
|
s = samples[b:b+1]
|
||||||
|
out = torch.zeros((s.shape[0], out_channels, round(s.shape[2] * upscale_amount), round(s.shape[3] * upscale_amount)), device=output_device)
|
||||||
|
out_div = torch.zeros((s.shape[0], out_channels, round(s.shape[2] * upscale_amount), round(s.shape[3] * upscale_amount)), device=output_device)
|
||||||
|
for y in range(0, s.shape[2], tile_y - overlap):
|
||||||
|
for x in range(0, s.shape[3], tile_x - overlap):
|
||||||
|
s_in = s[:,:,y:y+tile_y,x:x+tile_x]
|
||||||
|
|
||||||
|
ps = function(s_in).to(output_device)
|
||||||
|
mask = torch.ones_like(ps)
|
||||||
|
feather = round(overlap * upscale_amount)
|
||||||
|
for t in range(feather):
|
||||||
|
mask[:,:,t:1+t,:] *= ((1.0/feather) * (t + 1))
|
||||||
|
mask[:,:,mask.shape[2] -1 -t: mask.shape[2]-t,:] *= ((1.0/feather) * (t + 1))
|
||||||
|
mask[:,:,:,t:1+t] *= ((1.0/feather) * (t + 1))
|
||||||
|
mask[:,:,:,mask.shape[3]- 1 - t: mask.shape[3]- t] *= ((1.0/feather) * (t + 1))
|
||||||
|
out[:,:,round(y*upscale_amount):round((y+tile_y)*upscale_amount),round(x*upscale_amount):round((x+tile_x)*upscale_amount)] += ps * mask
|
||||||
|
out_div[:,:,round(y*upscale_amount):round((y+tile_y)*upscale_amount),round(x*upscale_amount):round((x+tile_x)*upscale_amount)] += mask
|
||||||
|
if pbar is not None:
|
||||||
|
pbar.update(1)
|
||||||
|
|
||||||
|
output[b:b+1] = out/out_div
|
||||||
|
return output
|
||||||
|
|
||||||
|
PROGRESS_BAR_ENABLED = True
|
||||||
|
def set_progress_bar_enabled(enabled):
|
||||||
|
global PROGRESS_BAR_ENABLED
|
||||||
|
PROGRESS_BAR_ENABLED = enabled
|
||||||
|
|
||||||
|
PROGRESS_BAR_HOOK = None
|
||||||
|
def set_progress_bar_global_hook(function):
|
||||||
|
global PROGRESS_BAR_HOOK
|
||||||
|
PROGRESS_BAR_HOOK = function
|
||||||
|
|
||||||
|
class ProgressBar:
|
||||||
|
def __init__(self, total):
|
||||||
|
global PROGRESS_BAR_HOOK
|
||||||
|
self.total = total
|
||||||
|
self.current = 0
|
||||||
|
self.hook = PROGRESS_BAR_HOOK
|
||||||
|
|
||||||
|
def update_absolute(self, value, total=None, preview=None):
|
||||||
|
if total is not None:
|
||||||
|
self.total = total
|
||||||
|
if value > self.total:
|
||||||
|
value = self.total
|
||||||
|
self.current = value
|
||||||
|
if self.hook is not None:
|
||||||
|
self.hook(self.current, self.total, preview)
|
||||||
|
|
||||||
|
def update(self, value):
|
||||||
|
self.update_absolute(self.current + value)
|
0
ldm_patched/pfn/__init__.py
Normal file
0
ldm_patched/pfn/__init__.py
Normal file
1182
ldm_patched/pfn/architecture/DAT.py
Normal file
1182
ldm_patched/pfn/architecture/DAT.py
Normal file
File diff suppressed because it is too large
Load Diff
1277
ldm_patched/pfn/architecture/HAT.py
Normal file
1277
ldm_patched/pfn/architecture/HAT.py
Normal file
File diff suppressed because it is too large
Load Diff
201
ldm_patched/pfn/architecture/LICENSE-DAT
Normal file
201
ldm_patched/pfn/architecture/LICENSE-DAT
Normal file
@ -0,0 +1,201 @@
|
|||||||
|
Apache License
|
||||||
|
Version 2.0, January 2004
|
||||||
|
http://www.apache.org/licenses/
|
||||||
|
|
||||||
|
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||||
|
|
||||||
|
1. Definitions.
|
||||||
|
|
||||||
|
"License" shall mean the terms and conditions for use, reproduction,
|
||||||
|
and distribution as defined by Sections 1 through 9 of this document.
|
||||||
|
|
||||||
|
"Licensor" shall mean the copyright owner or entity authorized by
|
||||||
|
the copyright owner that is granting the License.
|
||||||
|
|
||||||
|
"Legal Entity" shall mean the union of the acting entity and all
|
||||||
|
other entities that control, are controlled by, or are under common
|
||||||
|
control with that entity. For the purposes of this definition,
|
||||||
|
"control" means (i) the power, direct or indirect, to cause the
|
||||||
|
direction or management of such entity, whether by contract or
|
||||||
|
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||||
|
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||||
|
|
||||||
|
"You" (or "Your") shall mean an individual or Legal Entity
|
||||||
|
exercising permissions granted by this License.
|
||||||
|
|
||||||
|
"Source" form shall mean the preferred form for making modifications,
|
||||||
|
including but not limited to software source code, documentation
|
||||||
|
source, and configuration files.
|
||||||
|
|
||||||
|
"Object" form shall mean any form resulting from mechanical
|
||||||
|
transformation or translation of a Source form, including but
|
||||||
|
not limited to compiled object code, generated documentation,
|
||||||
|
and conversions to other media types.
|
||||||
|
|
||||||
|
"Work" shall mean the work of authorship, whether in Source or
|
||||||
|
Object form, made available under the License, as indicated by a
|
||||||
|
copyright notice that is included in or attached to the work
|
||||||
|
(an example is provided in the Appendix below).
|
||||||
|
|
||||||
|
"Derivative Works" shall mean any work, whether in Source or Object
|
||||||
|
form, that is based on (or derived from) the Work and for which the
|
||||||
|
editorial revisions, annotations, elaborations, or other modifications
|
||||||
|
represent, as a whole, an original work of authorship. For the purposes
|
||||||
|
of this License, Derivative Works shall not include works that remain
|
||||||
|
separable from, or merely link (or bind by name) to the interfaces of,
|
||||||
|
the Work and Derivative Works thereof.
|
||||||
|
|
||||||
|
"Contribution" shall mean any work of authorship, including
|
||||||
|
the original version of the Work and any modifications or additions
|
||||||
|
to that Work or Derivative Works thereof, that is intentionally
|
||||||
|
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||||
|
or by an individual or Legal Entity authorized to submit on behalf of
|
||||||
|
the copyright owner. For the purposes of this definition, "submitted"
|
||||||
|
means any form of electronic, verbal, or written communication sent
|
||||||
|
to the Licensor or its representatives, including but not limited to
|
||||||
|
communication on electronic mailing lists, source code control systems,
|
||||||
|
and issue tracking systems that are managed by, or on behalf of, the
|
||||||
|
Licensor for the purpose of discussing and improving the Work, but
|
||||||
|
excluding communication that is conspicuously marked or otherwise
|
||||||
|
designated in writing by the copyright owner as "Not a Contribution."
|
||||||
|
|
||||||
|
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||||
|
on behalf of whom a Contribution has been received by Licensor and
|
||||||
|
subsequently incorporated within the Work.
|
||||||
|
|
||||||
|
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||||
|
this License, each Contributor hereby grants to You a perpetual,
|
||||||
|
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||||
|
copyright license to reproduce, prepare Derivative Works of,
|
||||||
|
publicly display, publicly perform, sublicense, and distribute the
|
||||||
|
Work and such Derivative Works in Source or Object form.
|
||||||
|
|
||||||
|
3. Grant of Patent License. Subject to the terms and conditions of
|
||||||
|
this License, each Contributor hereby grants to You a perpetual,
|
||||||
|
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||||
|
(except as stated in this section) patent license to make, have made,
|
||||||
|
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||||
|
where such license applies only to those patent claims licensable
|
||||||
|
by such Contributor that are necessarily infringed by their
|
||||||
|
Contribution(s) alone or by combination of their Contribution(s)
|
||||||
|
with the Work to which such Contribution(s) was submitted. If You
|
||||||
|
institute patent litigation against any entity (including a
|
||||||
|
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||||
|
or a Contribution incorporated within the Work constitutes direct
|
||||||
|
or contributory patent infringement, then any patent licenses
|
||||||
|
granted to You under this License for that Work shall terminate
|
||||||
|
as of the date such litigation is filed.
|
||||||
|
|
||||||
|
4. Redistribution. You may reproduce and distribute copies of the
|
||||||
|
Work or Derivative Works thereof in any medium, with or without
|
||||||
|
modifications, and in Source or Object form, provided that You
|
||||||
|
meet the following conditions:
|
||||||
|
|
||||||
|
(a) You must give any other recipients of the Work or
|
||||||
|
Derivative Works a copy of this License; and
|
||||||
|
|
||||||
|
(b) You must cause any modified files to carry prominent notices
|
||||||
|
stating that You changed the files; and
|
||||||
|
|
||||||
|
(c) You must retain, in the Source form of any Derivative Works
|
||||||
|
that You distribute, all copyright, patent, trademark, and
|
||||||
|
attribution notices from the Source form of the Work,
|
||||||
|
excluding those notices that do not pertain to any part of
|
||||||
|
the Derivative Works; and
|
||||||
|
|
||||||
|
(d) If the Work includes a "NOTICE" text file as part of its
|
||||||
|
distribution, then any Derivative Works that You distribute must
|
||||||
|
include a readable copy of the attribution notices contained
|
||||||
|
within such NOTICE file, excluding those notices that do not
|
||||||
|
pertain to any part of the Derivative Works, in at least one
|
||||||
|
of the following places: within a NOTICE text file distributed
|
||||||
|
as part of the Derivative Works; within the Source form or
|
||||||
|
documentation, if provided along with the Derivative Works; or,
|
||||||
|
within a display generated by the Derivative Works, if and
|
||||||
|
wherever such third-party notices normally appear. The contents
|
||||||
|
of the NOTICE file are for informational purposes only and
|
||||||
|
do not modify the License. You may add Your own attribution
|
||||||
|
notices within Derivative Works that You distribute, alongside
|
||||||
|
or as an addendum to the NOTICE text from the Work, provided
|
||||||
|
that such additional attribution notices cannot be construed
|
||||||
|
as modifying the License.
|
||||||
|
|
||||||
|
You may add Your own copyright statement to Your modifications and
|
||||||
|
may provide additional or different license terms and conditions
|
||||||
|
for use, reproduction, or distribution of Your modifications, or
|
||||||
|
for any such Derivative Works as a whole, provided Your use,
|
||||||
|
reproduction, and distribution of the Work otherwise complies with
|
||||||
|
the conditions stated in this License.
|
||||||
|
|
||||||
|
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||||
|
any Contribution intentionally submitted for inclusion in the Work
|
||||||
|
by You to the Licensor shall be under the terms and conditions of
|
||||||
|
this License, without any additional terms or conditions.
|
||||||
|
Notwithstanding the above, nothing herein shall supersede or modify
|
||||||
|
the terms of any separate license agreement you may have executed
|
||||||
|
with Licensor regarding such Contributions.
|
||||||
|
|
||||||
|
6. Trademarks. This License does not grant permission to use the trade
|
||||||
|
names, trademarks, service marks, or product names of the Licensor,
|
||||||
|
except as required for reasonable and customary use in describing the
|
||||||
|
origin of the Work and reproducing the content of the NOTICE file.
|
||||||
|
|
||||||
|
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||||
|
agreed to in writing, Licensor provides the Work (and each
|
||||||
|
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||||
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||||
|
implied, including, without limitation, any warranties or conditions
|
||||||
|
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||||
|
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||||
|
appropriateness of using or redistributing the Work and assume any
|
||||||
|
risks associated with Your exercise of permissions under this License.
|
||||||
|
|
||||||
|
8. Limitation of Liability. In no event and under no legal theory,
|
||||||
|
whether in tort (including negligence), contract, or otherwise,
|
||||||
|
unless required by applicable law (such as deliberate and grossly
|
||||||
|
negligent acts) or agreed to in writing, shall any Contributor be
|
||||||
|
liable to You for damages, including any direct, indirect, special,
|
||||||
|
incidental, or consequential damages of any character arising as a
|
||||||
|
result of this License or out of the use or inability to use the
|
||||||
|
Work (including but not limited to damages for loss of goodwill,
|
||||||
|
work stoppage, computer failure or malfunction, or any and all
|
||||||
|
other commercial damages or losses), even if such Contributor
|
||||||
|
has been advised of the possibility of such damages.
|
||||||
|
|
||||||
|
9. Accepting Warranty or Additional Liability. While redistributing
|
||||||
|
the Work or Derivative Works thereof, You may choose to offer,
|
||||||
|
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||||
|
or other liability obligations and/or rights consistent with this
|
||||||
|
License. However, in accepting such obligations, You may act only
|
||||||
|
on Your own behalf and on Your sole responsibility, not on behalf
|
||||||
|
of any other Contributor, and only if You agree to indemnify,
|
||||||
|
defend, and hold each Contributor harmless for any liability
|
||||||
|
incurred by, or claims asserted against, such Contributor by reason
|
||||||
|
of your accepting any such warranty or additional liability.
|
||||||
|
|
||||||
|
END OF TERMS AND CONDITIONS
|
||||||
|
|
||||||
|
APPENDIX: How to apply the Apache License to your work.
|
||||||
|
|
||||||
|
To apply the Apache License to your work, attach the following
|
||||||
|
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||||
|
replaced with your own identifying information. (Don't include
|
||||||
|
the brackets!) The text should be enclosed in the appropriate
|
||||||
|
comment syntax for the file format. We also recommend that a
|
||||||
|
file or class name and description of purpose be included on the
|
||||||
|
same "printed page" as the copyright notice for easier
|
||||||
|
identification within third-party archives.
|
||||||
|
|
||||||
|
Copyright [yyyy] [name of copyright owner]
|
||||||
|
|
||||||
|
Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
you may not use this file except in compliance with the License.
|
||||||
|
You may obtain a copy of the License at
|
||||||
|
|
||||||
|
http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
|
||||||
|
Unless required by applicable law or agreed to in writing, software
|
||||||
|
distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
See the License for the specific language governing permissions and
|
||||||
|
limitations under the License.
|
201
ldm_patched/pfn/architecture/LICENSE-ESRGAN
Normal file
201
ldm_patched/pfn/architecture/LICENSE-ESRGAN
Normal file
@ -0,0 +1,201 @@
|
|||||||
|
Apache License
|
||||||
|
Version 2.0, January 2004
|
||||||
|
http://www.apache.org/licenses/
|
||||||
|
|
||||||
|
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||||
|
|
||||||
|
1. Definitions.
|
||||||
|
|
||||||
|
"License" shall mean the terms and conditions for use, reproduction,
|
||||||
|
and distribution as defined by Sections 1 through 9 of this document.
|
||||||
|
|
||||||
|
"Licensor" shall mean the copyright owner or entity authorized by
|
||||||
|
the copyright owner that is granting the License.
|
||||||
|
|
||||||
|
"Legal Entity" shall mean the union of the acting entity and all
|
||||||
|
other entities that control, are controlled by, or are under common
|
||||||
|
control with that entity. For the purposes of this definition,
|
||||||
|
"control" means (i) the power, direct or indirect, to cause the
|
||||||
|
direction or management of such entity, whether by contract or
|
||||||
|
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||||
|
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||||
|
|
||||||
|
"You" (or "Your") shall mean an individual or Legal Entity
|
||||||
|
exercising permissions granted by this License.
|
||||||
|
|
||||||
|
"Source" form shall mean the preferred form for making modifications,
|
||||||
|
including but not limited to software source code, documentation
|
||||||
|
source, and configuration files.
|
||||||
|
|
||||||
|
"Object" form shall mean any form resulting from mechanical
|
||||||
|
transformation or translation of a Source form, including but
|
||||||
|
not limited to compiled object code, generated documentation,
|
||||||
|
and conversions to other media types.
|
||||||
|
|
||||||
|
"Work" shall mean the work of authorship, whether in Source or
|
||||||
|
Object form, made available under the License, as indicated by a
|
||||||
|
copyright notice that is included in or attached to the work
|
||||||
|
(an example is provided in the Appendix below).
|
||||||
|
|
||||||
|
"Derivative Works" shall mean any work, whether in Source or Object
|
||||||
|
form, that is based on (or derived from) the Work and for which the
|
||||||
|
editorial revisions, annotations, elaborations, or other modifications
|
||||||
|
represent, as a whole, an original work of authorship. For the purposes
|
||||||
|
of this License, Derivative Works shall not include works that remain
|
||||||
|
separable from, or merely link (or bind by name) to the interfaces of,
|
||||||
|
the Work and Derivative Works thereof.
|
||||||
|
|
||||||
|
"Contribution" shall mean any work of authorship, including
|
||||||
|
the original version of the Work and any modifications or additions
|
||||||
|
to that Work or Derivative Works thereof, that is intentionally
|
||||||
|
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||||
|
or by an individual or Legal Entity authorized to submit on behalf of
|
||||||
|
the copyright owner. For the purposes of this definition, "submitted"
|
||||||
|
means any form of electronic, verbal, or written communication sent
|
||||||
|
to the Licensor or its representatives, including but not limited to
|
||||||
|
communication on electronic mailing lists, source code control systems,
|
||||||
|
and issue tracking systems that are managed by, or on behalf of, the
|
||||||
|
Licensor for the purpose of discussing and improving the Work, but
|
||||||
|
excluding communication that is conspicuously marked or otherwise
|
||||||
|
designated in writing by the copyright owner as "Not a Contribution."
|
||||||
|
|
||||||
|
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||||
|
on behalf of whom a Contribution has been received by Licensor and
|
||||||
|
subsequently incorporated within the Work.
|
||||||
|
|
||||||
|
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||||
|
this License, each Contributor hereby grants to You a perpetual,
|
||||||
|
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||||
|
copyright license to reproduce, prepare Derivative Works of,
|
||||||
|
publicly display, publicly perform, sublicense, and distribute the
|
||||||
|
Work and such Derivative Works in Source or Object form.
|
||||||
|
|
||||||
|
3. Grant of Patent License. Subject to the terms and conditions of
|
||||||
|
this License, each Contributor hereby grants to You a perpetual,
|
||||||
|
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||||
|
(except as stated in this section) patent license to make, have made,
|
||||||
|
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||||
|
where such license applies only to those patent claims licensable
|
||||||
|
by such Contributor that are necessarily infringed by their
|
||||||
|
Contribution(s) alone or by combination of their Contribution(s)
|
||||||
|
with the Work to which such Contribution(s) was submitted. If You
|
||||||
|
institute patent litigation against any entity (including a
|
||||||
|
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||||
|
or a Contribution incorporated within the Work constitutes direct
|
||||||
|
or contributory patent infringement, then any patent licenses
|
||||||
|
granted to You under this License for that Work shall terminate
|
||||||
|
as of the date such litigation is filed.
|
||||||
|
|
||||||
|
4. Redistribution. You may reproduce and distribute copies of the
|
||||||
|
Work or Derivative Works thereof in any medium, with or without
|
||||||
|
modifications, and in Source or Object form, provided that You
|
||||||
|
meet the following conditions:
|
||||||
|
|
||||||
|
(a) You must give any other recipients of the Work or
|
||||||
|
Derivative Works a copy of this License; and
|
||||||
|
|
||||||
|
(b) You must cause any modified files to carry prominent notices
|
||||||
|
stating that You changed the files; and
|
||||||
|
|
||||||
|
(c) You must retain, in the Source form of any Derivative Works
|
||||||
|
that You distribute, all copyright, patent, trademark, and
|
||||||
|
attribution notices from the Source form of the Work,
|
||||||
|
excluding those notices that do not pertain to any part of
|
||||||
|
the Derivative Works; and
|
||||||
|
|
||||||
|
(d) If the Work includes a "NOTICE" text file as part of its
|
||||||
|
distribution, then any Derivative Works that You distribute must
|
||||||
|
include a readable copy of the attribution notices contained
|
||||||
|
within such NOTICE file, excluding those notices that do not
|
||||||
|
pertain to any part of the Derivative Works, in at least one
|
||||||
|
of the following places: within a NOTICE text file distributed
|
||||||
|
as part of the Derivative Works; within the Source form or
|
||||||
|
documentation, if provided along with the Derivative Works; or,
|
||||||
|
within a display generated by the Derivative Works, if and
|
||||||
|
wherever such third-party notices normally appear. The contents
|
||||||
|
of the NOTICE file are for informational purposes only and
|
||||||
|
do not modify the License. You may add Your own attribution
|
||||||
|
notices within Derivative Works that You distribute, alongside
|
||||||
|
or as an addendum to the NOTICE text from the Work, provided
|
||||||
|
that such additional attribution notices cannot be construed
|
||||||
|
as modifying the License.
|
||||||
|
|
||||||
|
You may add Your own copyright statement to Your modifications and
|
||||||
|
may provide additional or different license terms and conditions
|
||||||
|
for use, reproduction, or distribution of Your modifications, or
|
||||||
|
for any such Derivative Works as a whole, provided Your use,
|
||||||
|
reproduction, and distribution of the Work otherwise complies with
|
||||||
|
the conditions stated in this License.
|
||||||
|
|
||||||
|
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||||
|
any Contribution intentionally submitted for inclusion in the Work
|
||||||
|
by You to the Licensor shall be under the terms and conditions of
|
||||||
|
this License, without any additional terms or conditions.
|
||||||
|
Notwithstanding the above, nothing herein shall supersede or modify
|
||||||
|
the terms of any separate license agreement you may have executed
|
||||||
|
with Licensor regarding such Contributions.
|
||||||
|
|
||||||
|
6. Trademarks. This License does not grant permission to use the trade
|
||||||
|
names, trademarks, service marks, or product names of the Licensor,
|
||||||
|
except as required for reasonable and customary use in describing the
|
||||||
|
origin of the Work and reproducing the content of the NOTICE file.
|
||||||
|
|
||||||
|
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||||
|
agreed to in writing, Licensor provides the Work (and each
|
||||||
|
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||||
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||||
|
implied, including, without limitation, any warranties or conditions
|
||||||
|
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||||
|
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||||
|
appropriateness of using or redistributing the Work and assume any
|
||||||
|
risks associated with Your exercise of permissions under this License.
|
||||||
|
|
||||||
|
8. Limitation of Liability. In no event and under no legal theory,
|
||||||
|
whether in tort (including negligence), contract, or otherwise,
|
||||||
|
unless required by applicable law (such as deliberate and grossly
|
||||||
|
negligent acts) or agreed to in writing, shall any Contributor be
|
||||||
|
liable to You for damages, including any direct, indirect, special,
|
||||||
|
incidental, or consequential damages of any character arising as a
|
||||||
|
result of this License or out of the use or inability to use the
|
||||||
|
Work (including but not limited to damages for loss of goodwill,
|
||||||
|
work stoppage, computer failure or malfunction, or any and all
|
||||||
|
other commercial damages or losses), even if such Contributor
|
||||||
|
has been advised of the possibility of such damages.
|
||||||
|
|
||||||
|
9. Accepting Warranty or Additional Liability. While redistributing
|
||||||
|
the Work or Derivative Works thereof, You may choose to offer,
|
||||||
|
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||||
|
or other liability obligations and/or rights consistent with this
|
||||||
|
License. However, in accepting such obligations, You may act only
|
||||||
|
on Your own behalf and on Your sole responsibility, not on behalf
|
||||||
|
of any other Contributor, and only if You agree to indemnify,
|
||||||
|
defend, and hold each Contributor harmless for any liability
|
||||||
|
incurred by, or claims asserted against, such Contributor by reason
|
||||||
|
of your accepting any such warranty or additional liability.
|
||||||
|
|
||||||
|
END OF TERMS AND CONDITIONS
|
||||||
|
|
||||||
|
APPENDIX: How to apply the Apache License to your work.
|
||||||
|
|
||||||
|
To apply the Apache License to your work, attach the following
|
||||||
|
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||||
|
replaced with your own identifying information. (Don't include
|
||||||
|
the brackets!) The text should be enclosed in the appropriate
|
||||||
|
comment syntax for the file format. We also recommend that a
|
||||||
|
file or class name and description of purpose be included on the
|
||||||
|
same "printed page" as the copyright notice for easier
|
||||||
|
identification within third-party archives.
|
||||||
|
|
||||||
|
Copyright [yyyy] [name of copyright owner]
|
||||||
|
|
||||||
|
Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
you may not use this file except in compliance with the License.
|
||||||
|
You may obtain a copy of the License at
|
||||||
|
|
||||||
|
http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
|
||||||
|
Unless required by applicable law or agreed to in writing, software
|
||||||
|
distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
See the License for the specific language governing permissions and
|
||||||
|
limitations under the License.
|
21
ldm_patched/pfn/architecture/LICENSE-HAT
Normal file
21
ldm_patched/pfn/architecture/LICENSE-HAT
Normal file
@ -0,0 +1,21 @@
|
|||||||
|
MIT License
|
||||||
|
|
||||||
|
Copyright (c) 2022 Xiangyu Chen
|
||||||
|
|
||||||
|
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||||
|
of this software and associated documentation files (the "Software"), to deal
|
||||||
|
in the Software without restriction, including without limitation the rights
|
||||||
|
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||||
|
copies of the Software, and to permit persons to whom the Software is
|
||||||
|
furnished to do so, subject to the following conditions:
|
||||||
|
|
||||||
|
The above copyright notice and this permission notice shall be included in all
|
||||||
|
copies or substantial portions of the Software.
|
||||||
|
|
||||||
|
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||||
|
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||||
|
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||||
|
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||||
|
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||||
|
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||||
|
SOFTWARE.
|
29
ldm_patched/pfn/architecture/LICENSE-RealESRGAN
Normal file
29
ldm_patched/pfn/architecture/LICENSE-RealESRGAN
Normal file
@ -0,0 +1,29 @@
|
|||||||
|
BSD 3-Clause License
|
||||||
|
|
||||||
|
Copyright (c) 2021, Xintao Wang
|
||||||
|
All rights reserved.
|
||||||
|
|
||||||
|
Redistribution and use in source and binary forms, with or without
|
||||||
|
modification, are permitted provided that the following conditions are met:
|
||||||
|
|
||||||
|
1. Redistributions of source code must retain the above copyright notice, this
|
||||||
|
list of conditions and the following disclaimer.
|
||||||
|
|
||||||
|
2. Redistributions in binary form must reproduce the above copyright notice,
|
||||||
|
this list of conditions and the following disclaimer in the documentation
|
||||||
|
and/or other materials provided with the distribution.
|
||||||
|
|
||||||
|
3. Neither the name of the copyright holder nor the names of its
|
||||||
|
contributors may be used to endorse or promote products derived from
|
||||||
|
this software without specific prior written permission.
|
||||||
|
|
||||||
|
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
|
||||||
|
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
|
||||||
|
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
||||||
|
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
|
||||||
|
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
|
||||||
|
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
|
||||||
|
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
|
||||||
|
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
|
||||||
|
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
|
||||||
|
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
201
ldm_patched/pfn/architecture/LICENSE-SCUNet
Normal file
201
ldm_patched/pfn/architecture/LICENSE-SCUNet
Normal file
@ -0,0 +1,201 @@
|
|||||||
|
Apache License
|
||||||
|
Version 2.0, January 2004
|
||||||
|
http://www.apache.org/licenses/
|
||||||
|
|
||||||
|
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||||
|
|
||||||
|
1. Definitions.
|
||||||
|
|
||||||
|
"License" shall mean the terms and conditions for use, reproduction,
|
||||||
|
and distribution as defined by Sections 1 through 9 of this document.
|
||||||
|
|
||||||
|
"Licensor" shall mean the copyright owner or entity authorized by
|
||||||
|
the copyright owner that is granting the License.
|
||||||
|
|
||||||
|
"Legal Entity" shall mean the union of the acting entity and all
|
||||||
|
other entities that control, are controlled by, or are under common
|
||||||
|
control with that entity. For the purposes of this definition,
|
||||||
|
"control" means (i) the power, direct or indirect, to cause the
|
||||||
|
direction or management of such entity, whether by contract or
|
||||||
|
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||||
|
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||||
|
|
||||||
|
"You" (or "Your") shall mean an individual or Legal Entity
|
||||||
|
exercising permissions granted by this License.
|
||||||
|
|
||||||
|
"Source" form shall mean the preferred form for making modifications,
|
||||||
|
including but not limited to software source code, documentation
|
||||||
|
source, and configuration files.
|
||||||
|
|
||||||
|
"Object" form shall mean any form resulting from mechanical
|
||||||
|
transformation or translation of a Source form, including but
|
||||||
|
not limited to compiled object code, generated documentation,
|
||||||
|
and conversions to other media types.
|
||||||
|
|
||||||
|
"Work" shall mean the work of authorship, whether in Source or
|
||||||
|
Object form, made available under the License, as indicated by a
|
||||||
|
copyright notice that is included in or attached to the work
|
||||||
|
(an example is provided in the Appendix below).
|
||||||
|
|
||||||
|
"Derivative Works" shall mean any work, whether in Source or Object
|
||||||
|
form, that is based on (or derived from) the Work and for which the
|
||||||
|
editorial revisions, annotations, elaborations, or other modifications
|
||||||
|
represent, as a whole, an original work of authorship. For the purposes
|
||||||
|
of this License, Derivative Works shall not include works that remain
|
||||||
|
separable from, or merely link (or bind by name) to the interfaces of,
|
||||||
|
the Work and Derivative Works thereof.
|
||||||
|
|
||||||
|
"Contribution" shall mean any work of authorship, including
|
||||||
|
the original version of the Work and any modifications or additions
|
||||||
|
to that Work or Derivative Works thereof, that is intentionally
|
||||||
|
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||||
|
or by an individual or Legal Entity authorized to submit on behalf of
|
||||||
|
the copyright owner. For the purposes of this definition, "submitted"
|
||||||
|
means any form of electronic, verbal, or written communication sent
|
||||||
|
to the Licensor or its representatives, including but not limited to
|
||||||
|
communication on electronic mailing lists, source code control systems,
|
||||||
|
and issue tracking systems that are managed by, or on behalf of, the
|
||||||
|
Licensor for the purpose of discussing and improving the Work, but
|
||||||
|
excluding communication that is conspicuously marked or otherwise
|
||||||
|
designated in writing by the copyright owner as "Not a Contribution."
|
||||||
|
|
||||||
|
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||||
|
on behalf of whom a Contribution has been received by Licensor and
|
||||||
|
subsequently incorporated within the Work.
|
||||||
|
|
||||||
|
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||||
|
this License, each Contributor hereby grants to You a perpetual,
|
||||||
|
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||||
|
copyright license to reproduce, prepare Derivative Works of,
|
||||||
|
publicly display, publicly perform, sublicense, and distribute the
|
||||||
|
Work and such Derivative Works in Source or Object form.
|
||||||
|
|
||||||
|
3. Grant of Patent License. Subject to the terms and conditions of
|
||||||
|
this License, each Contributor hereby grants to You a perpetual,
|
||||||
|
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||||
|
(except as stated in this section) patent license to make, have made,
|
||||||
|
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||||
|
where such license applies only to those patent claims licensable
|
||||||
|
by such Contributor that are necessarily infringed by their
|
||||||
|
Contribution(s) alone or by combination of their Contribution(s)
|
||||||
|
with the Work to which such Contribution(s) was submitted. If You
|
||||||
|
institute patent litigation against any entity (including a
|
||||||
|
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||||
|
or a Contribution incorporated within the Work constitutes direct
|
||||||
|
or contributory patent infringement, then any patent licenses
|
||||||
|
granted to You under this License for that Work shall terminate
|
||||||
|
as of the date such litigation is filed.
|
||||||
|
|
||||||
|
4. Redistribution. You may reproduce and distribute copies of the
|
||||||
|
Work or Derivative Works thereof in any medium, with or without
|
||||||
|
modifications, and in Source or Object form, provided that You
|
||||||
|
meet the following conditions:
|
||||||
|
|
||||||
|
(a) You must give any other recipients of the Work or
|
||||||
|
Derivative Works a copy of this License; and
|
||||||
|
|
||||||
|
(b) You must cause any modified files to carry prominent notices
|
||||||
|
stating that You changed the files; and
|
||||||
|
|
||||||
|
(c) You must retain, in the Source form of any Derivative Works
|
||||||
|
that You distribute, all copyright, patent, trademark, and
|
||||||
|
attribution notices from the Source form of the Work,
|
||||||
|
excluding those notices that do not pertain to any part of
|
||||||
|
the Derivative Works; and
|
||||||
|
|
||||||
|
(d) If the Work includes a "NOTICE" text file as part of its
|
||||||
|
distribution, then any Derivative Works that You distribute must
|
||||||
|
include a readable copy of the attribution notices contained
|
||||||
|
within such NOTICE file, excluding those notices that do not
|
||||||
|
pertain to any part of the Derivative Works, in at least one
|
||||||
|
of the following places: within a NOTICE text file distributed
|
||||||
|
as part of the Derivative Works; within the Source form or
|
||||||
|
documentation, if provided along with the Derivative Works; or,
|
||||||
|
within a display generated by the Derivative Works, if and
|
||||||
|
wherever such third-party notices normally appear. The contents
|
||||||
|
of the NOTICE file are for informational purposes only and
|
||||||
|
do not modify the License. You may add Your own attribution
|
||||||
|
notices within Derivative Works that You distribute, alongside
|
||||||
|
or as an addendum to the NOTICE text from the Work, provided
|
||||||
|
that such additional attribution notices cannot be construed
|
||||||
|
as modifying the License.
|
||||||
|
|
||||||
|
You may add Your own copyright statement to Your modifications and
|
||||||
|
may provide additional or different license terms and conditions
|
||||||
|
for use, reproduction, or distribution of Your modifications, or
|
||||||
|
for any such Derivative Works as a whole, provided Your use,
|
||||||
|
reproduction, and distribution of the Work otherwise complies with
|
||||||
|
the conditions stated in this License.
|
||||||
|
|
||||||
|
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||||
|
any Contribution intentionally submitted for inclusion in the Work
|
||||||
|
by You to the Licensor shall be under the terms and conditions of
|
||||||
|
this License, without any additional terms or conditions.
|
||||||
|
Notwithstanding the above, nothing herein shall supersede or modify
|
||||||
|
the terms of any separate license agreement you may have executed
|
||||||
|
with Licensor regarding such Contributions.
|
||||||
|
|
||||||
|
6. Trademarks. This License does not grant permission to use the trade
|
||||||
|
names, trademarks, service marks, or product names of the Licensor,
|
||||||
|
except as required for reasonable and customary use in describing the
|
||||||
|
origin of the Work and reproducing the content of the NOTICE file.
|
||||||
|
|
||||||
|
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||||
|
agreed to in writing, Licensor provides the Work (and each
|
||||||
|
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||||
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||||
|
implied, including, without limitation, any warranties or conditions
|
||||||
|
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||||
|
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||||
|
appropriateness of using or redistributing the Work and assume any
|
||||||
|
risks associated with Your exercise of permissions under this License.
|
||||||
|
|
||||||
|
8. Limitation of Liability. In no event and under no legal theory,
|
||||||
|
whether in tort (including negligence), contract, or otherwise,
|
||||||
|
unless required by applicable law (such as deliberate and grossly
|
||||||
|
negligent acts) or agreed to in writing, shall any Contributor be
|
||||||
|
liable to You for damages, including any direct, indirect, special,
|
||||||
|
incidental, or consequential damages of any character arising as a
|
||||||
|
result of this License or out of the use or inability to use the
|
||||||
|
Work (including but not limited to damages for loss of goodwill,
|
||||||
|
work stoppage, computer failure or malfunction, or any and all
|
||||||
|
other commercial damages or losses), even if such Contributor
|
||||||
|
has been advised of the possibility of such damages.
|
||||||
|
|
||||||
|
9. Accepting Warranty or Additional Liability. While redistributing
|
||||||
|
the Work or Derivative Works thereof, You may choose to offer,
|
||||||
|
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||||
|
or other liability obligations and/or rights consistent with this
|
||||||
|
License. However, in accepting such obligations, You may act only
|
||||||
|
on Your own behalf and on Your sole responsibility, not on behalf
|
||||||
|
of any other Contributor, and only if You agree to indemnify,
|
||||||
|
defend, and hold each Contributor harmless for any liability
|
||||||
|
incurred by, or claims asserted against, such Contributor by reason
|
||||||
|
of your accepting any such warranty or additional liability.
|
||||||
|
|
||||||
|
END OF TERMS AND CONDITIONS
|
||||||
|
|
||||||
|
APPENDIX: How to apply the Apache License to your work.
|
||||||
|
|
||||||
|
To apply the Apache License to your work, attach the following
|
||||||
|
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||||
|
replaced with your own identifying information. (Don't include
|
||||||
|
the brackets!) The text should be enclosed in the appropriate
|
||||||
|
comment syntax for the file format. We also recommend that a
|
||||||
|
file or class name and description of purpose be included on the
|
||||||
|
same "printed page" as the copyright notice for easier
|
||||||
|
identification within third-party archives.
|
||||||
|
|
||||||
|
Copyright 2022 Kai Zhang (cskaizhang@gmail.com, https://cszn.github.io/). All rights reserved.
|
||||||
|
|
||||||
|
Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
you may not use this file except in compliance with the License.
|
||||||
|
You may obtain a copy of the License at
|
||||||
|
|
||||||
|
http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
|
||||||
|
Unless required by applicable law or agreed to in writing, software
|
||||||
|
distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
See the License for the specific language governing permissions and
|
||||||
|
limitations under the License.
|
201
ldm_patched/pfn/architecture/LICENSE-SPSR
Normal file
201
ldm_patched/pfn/architecture/LICENSE-SPSR
Normal file
@ -0,0 +1,201 @@
|
|||||||
|
Apache License
|
||||||
|
Version 2.0, January 2004
|
||||||
|
http://www.apache.org/licenses/
|
||||||
|
|
||||||
|
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||||
|
|
||||||
|
1. Definitions.
|
||||||
|
|
||||||
|
"License" shall mean the terms and conditions for use, reproduction,
|
||||||
|
and distribution as defined by Sections 1 through 9 of this document.
|
||||||
|
|
||||||
|
"Licensor" shall mean the copyright owner or entity authorized by
|
||||||
|
the copyright owner that is granting the License.
|
||||||
|
|
||||||
|
"Legal Entity" shall mean the union of the acting entity and all
|
||||||
|
other entities that control, are controlled by, or are under common
|
||||||
|
control with that entity. For the purposes of this definition,
|
||||||
|
"control" means (i) the power, direct or indirect, to cause the
|
||||||
|
direction or management of such entity, whether by contract or
|
||||||
|
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||||
|
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||||
|
|
||||||
|
"You" (or "Your") shall mean an individual or Legal Entity
|
||||||
|
exercising permissions granted by this License.
|
||||||
|
|
||||||
|
"Source" form shall mean the preferred form for making modifications,
|
||||||
|
including but not limited to software source code, documentation
|
||||||
|
source, and configuration files.
|
||||||
|
|
||||||
|
"Object" form shall mean any form resulting from mechanical
|
||||||
|
transformation or translation of a Source form, including but
|
||||||
|
not limited to compiled object code, generated documentation,
|
||||||
|
and conversions to other media types.
|
||||||
|
|
||||||
|
"Work" shall mean the work of authorship, whether in Source or
|
||||||
|
Object form, made available under the License, as indicated by a
|
||||||
|
copyright notice that is included in or attached to the work
|
||||||
|
(an example is provided in the Appendix below).
|
||||||
|
|
||||||
|
"Derivative Works" shall mean any work, whether in Source or Object
|
||||||
|
form, that is based on (or derived from) the Work and for which the
|
||||||
|
editorial revisions, annotations, elaborations, or other modifications
|
||||||
|
represent, as a whole, an original work of authorship. For the purposes
|
||||||
|
of this License, Derivative Works shall not include works that remain
|
||||||
|
separable from, or merely link (or bind by name) to the interfaces of,
|
||||||
|
the Work and Derivative Works thereof.
|
||||||
|
|
||||||
|
"Contribution" shall mean any work of authorship, including
|
||||||
|
the original version of the Work and any modifications or additions
|
||||||
|
to that Work or Derivative Works thereof, that is intentionally
|
||||||
|
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||||
|
or by an individual or Legal Entity authorized to submit on behalf of
|
||||||
|
the copyright owner. For the purposes of this definition, "submitted"
|
||||||
|
means any form of electronic, verbal, or written communication sent
|
||||||
|
to the Licensor or its representatives, including but not limited to
|
||||||
|
communication on electronic mailing lists, source code control systems,
|
||||||
|
and issue tracking systems that are managed by, or on behalf of, the
|
||||||
|
Licensor for the purpose of discussing and improving the Work, but
|
||||||
|
excluding communication that is conspicuously marked or otherwise
|
||||||
|
designated in writing by the copyright owner as "Not a Contribution."
|
||||||
|
|
||||||
|
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||||
|
on behalf of whom a Contribution has been received by Licensor and
|
||||||
|
subsequently incorporated within the Work.
|
||||||
|
|
||||||
|
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||||
|
this License, each Contributor hereby grants to You a perpetual,
|
||||||
|
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||||
|
copyright license to reproduce, prepare Derivative Works of,
|
||||||
|
publicly display, publicly perform, sublicense, and distribute the
|
||||||
|
Work and such Derivative Works in Source or Object form.
|
||||||
|
|
||||||
|
3. Grant of Patent License. Subject to the terms and conditions of
|
||||||
|
this License, each Contributor hereby grants to You a perpetual,
|
||||||
|
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||||
|
(except as stated in this section) patent license to make, have made,
|
||||||
|
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||||
|
where such license applies only to those patent claims licensable
|
||||||
|
by such Contributor that are necessarily infringed by their
|
||||||
|
Contribution(s) alone or by combination of their Contribution(s)
|
||||||
|
with the Work to which such Contribution(s) was submitted. If You
|
||||||
|
institute patent litigation against any entity (including a
|
||||||
|
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||||
|
or a Contribution incorporated within the Work constitutes direct
|
||||||
|
or contributory patent infringement, then any patent licenses
|
||||||
|
granted to You under this License for that Work shall terminate
|
||||||
|
as of the date such litigation is filed.
|
||||||
|
|
||||||
|
4. Redistribution. You may reproduce and distribute copies of the
|
||||||
|
Work or Derivative Works thereof in any medium, with or without
|
||||||
|
modifications, and in Source or Object form, provided that You
|
||||||
|
meet the following conditions:
|
||||||
|
|
||||||
|
(a) You must give any other recipients of the Work or
|
||||||
|
Derivative Works a copy of this License; and
|
||||||
|
|
||||||
|
(b) You must cause any modified files to carry prominent notices
|
||||||
|
stating that You changed the files; and
|
||||||
|
|
||||||
|
(c) You must retain, in the Source form of any Derivative Works
|
||||||
|
that You distribute, all copyright, patent, trademark, and
|
||||||
|
attribution notices from the Source form of the Work,
|
||||||
|
excluding those notices that do not pertain to any part of
|
||||||
|
the Derivative Works; and
|
||||||
|
|
||||||
|
(d) If the Work includes a "NOTICE" text file as part of its
|
||||||
|
distribution, then any Derivative Works that You distribute must
|
||||||
|
include a readable copy of the attribution notices contained
|
||||||
|
within such NOTICE file, excluding those notices that do not
|
||||||
|
pertain to any part of the Derivative Works, in at least one
|
||||||
|
of the following places: within a NOTICE text file distributed
|
||||||
|
as part of the Derivative Works; within the Source form or
|
||||||
|
documentation, if provided along with the Derivative Works; or,
|
||||||
|
within a display generated by the Derivative Works, if and
|
||||||
|
wherever such third-party notices normally appear. The contents
|
||||||
|
of the NOTICE file are for informational purposes only and
|
||||||
|
do not modify the License. You may add Your own attribution
|
||||||
|
notices within Derivative Works that You distribute, alongside
|
||||||
|
or as an addendum to the NOTICE text from the Work, provided
|
||||||
|
that such additional attribution notices cannot be construed
|
||||||
|
as modifying the License.
|
||||||
|
|
||||||
|
You may add Your own copyright statement to Your modifications and
|
||||||
|
may provide additional or different license terms and conditions
|
||||||
|
for use, reproduction, or distribution of Your modifications, or
|
||||||
|
for any such Derivative Works as a whole, provided Your use,
|
||||||
|
reproduction, and distribution of the Work otherwise complies with
|
||||||
|
the conditions stated in this License.
|
||||||
|
|
||||||
|
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||||
|
any Contribution intentionally submitted for inclusion in the Work
|
||||||
|
by You to the Licensor shall be under the terms and conditions of
|
||||||
|
this License, without any additional terms or conditions.
|
||||||
|
Notwithstanding the above, nothing herein shall supersede or modify
|
||||||
|
the terms of any separate license agreement you may have executed
|
||||||
|
with Licensor regarding such Contributions.
|
||||||
|
|
||||||
|
6. Trademarks. This License does not grant permission to use the trade
|
||||||
|
names, trademarks, service marks, or product names of the Licensor,
|
||||||
|
except as required for reasonable and customary use in describing the
|
||||||
|
origin of the Work and reproducing the content of the NOTICE file.
|
||||||
|
|
||||||
|
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||||
|
agreed to in writing, Licensor provides the Work (and each
|
||||||
|
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||||
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||||
|
implied, including, without limitation, any warranties or conditions
|
||||||
|
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||||
|
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||||
|
appropriateness of using or redistributing the Work and assume any
|
||||||
|
risks associated with Your exercise of permissions under this License.
|
||||||
|
|
||||||
|
8. Limitation of Liability. In no event and under no legal theory,
|
||||||
|
whether in tort (including negligence), contract, or otherwise,
|
||||||
|
unless required by applicable law (such as deliberate and grossly
|
||||||
|
negligent acts) or agreed to in writing, shall any Contributor be
|
||||||
|
liable to You for damages, including any direct, indirect, special,
|
||||||
|
incidental, or consequential damages of any character arising as a
|
||||||
|
result of this License or out of the use or inability to use the
|
||||||
|
Work (including but not limited to damages for loss of goodwill,
|
||||||
|
work stoppage, computer failure or malfunction, or any and all
|
||||||
|
other commercial damages or losses), even if such Contributor
|
||||||
|
has been advised of the possibility of such damages.
|
||||||
|
|
||||||
|
9. Accepting Warranty or Additional Liability. While redistributing
|
||||||
|
the Work or Derivative Works thereof, You may choose to offer,
|
||||||
|
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||||
|
or other liability obligations and/or rights consistent with this
|
||||||
|
License. However, in accepting such obligations, You may act only
|
||||||
|
on Your own behalf and on Your sole responsibility, not on behalf
|
||||||
|
of any other Contributor, and only if You agree to indemnify,
|
||||||
|
defend, and hold each Contributor harmless for any liability
|
||||||
|
incurred by, or claims asserted against, such Contributor by reason
|
||||||
|
of your accepting any such warranty or additional liability.
|
||||||
|
|
||||||
|
END OF TERMS AND CONDITIONS
|
||||||
|
|
||||||
|
APPENDIX: How to apply the Apache License to your work.
|
||||||
|
|
||||||
|
To apply the Apache License to your work, attach the following
|
||||||
|
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||||
|
replaced with your own identifying information. (Don't include
|
||||||
|
the brackets!) The text should be enclosed in the appropriate
|
||||||
|
comment syntax for the file format. We also recommend that a
|
||||||
|
file or class name and description of purpose be included on the
|
||||||
|
same "printed page" as the copyright notice for easier
|
||||||
|
identification within third-party archives.
|
||||||
|
|
||||||
|
Copyright 2018-2022 BasicSR Authors
|
||||||
|
|
||||||
|
Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
you may not use this file except in compliance with the License.
|
||||||
|
You may obtain a copy of the License at
|
||||||
|
|
||||||
|
http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
|
||||||
|
Unless required by applicable law or agreed to in writing, software
|
||||||
|
distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
See the License for the specific language governing permissions and
|
||||||
|
limitations under the License.
|
121
ldm_patched/pfn/architecture/LICENSE-SwiftSRGAN
Normal file
121
ldm_patched/pfn/architecture/LICENSE-SwiftSRGAN
Normal file
@ -0,0 +1,121 @@
|
|||||||
|
Creative Commons Legal Code
|
||||||
|
|
||||||
|
CC0 1.0 Universal
|
||||||
|
|
||||||
|
CREATIVE COMMONS CORPORATION IS NOT A LAW FIRM AND DOES NOT PROVIDE
|
||||||
|
LEGAL SERVICES. DISTRIBUTION OF THIS DOCUMENT DOES NOT CREATE AN
|
||||||
|
ATTORNEY-CLIENT RELATIONSHIP. CREATIVE COMMONS PROVIDES THIS
|
||||||
|
INFORMATION ON AN "AS-IS" BASIS. CREATIVE COMMONS MAKES NO WARRANTIES
|
||||||
|
REGARDING THE USE OF THIS DOCUMENT OR THE INFORMATION OR WORKS
|
||||||
|
PROVIDED HEREUNDER, AND DISCLAIMS LIABILITY FOR DAMAGES RESULTING FROM
|
||||||
|
THE USE OF THIS DOCUMENT OR THE INFORMATION OR WORKS PROVIDED
|
||||||
|
HEREUNDER.
|
||||||
|
|
||||||
|
Statement of Purpose
|
||||||
|
|
||||||
|
The laws of most jurisdictions throughout the world automatically confer
|
||||||
|
exclusive Copyright and Related Rights (defined below) upon the creator
|
||||||
|
and subsequent owner(s) (each and all, an "owner") of an original work of
|
||||||
|
authorship and/or a database (each, a "Work").
|
||||||
|
|
||||||
|
Certain owners wish to permanently relinquish those rights to a Work for
|
||||||
|
the purpose of contributing to a commons of creative, cultural and
|
||||||
|
scientific works ("Commons") that the public can reliably and without fear
|
||||||
|
of later claims of infringement build upon, modify, incorporate in other
|
||||||
|
works, reuse and redistribute as freely as possible in any form whatsoever
|
||||||
|
and for any purposes, including without limitation commercial purposes.
|
||||||
|
These owners may contribute to the Commons to promote the ideal of a free
|
||||||
|
culture and the further production of creative, cultural and scientific
|
||||||
|
works, or to gain reputation or greater distribution for their Work in
|
||||||
|
part through the use and efforts of others.
|
||||||
|
|
||||||
|
For these and/or other purposes and motivations, and without any
|
||||||
|
expectation of additional consideration or compensation, the person
|
||||||
|
associating CC0 with a Work (the "Affirmer"), to the extent that he or she
|
||||||
|
is an owner of Copyright and Related Rights in the Work, voluntarily
|
||||||
|
elects to apply CC0 to the Work and publicly distribute the Work under its
|
||||||
|
terms, with knowledge of his or her Copyright and Related Rights in the
|
||||||
|
Work and the meaning and intended legal effect of CC0 on those rights.
|
||||||
|
|
||||||
|
1. Copyright and Related Rights. A Work made available under CC0 may be
|
||||||
|
protected by copyright and related or neighboring rights ("Copyright and
|
||||||
|
Related Rights"). Copyright and Related Rights include, but are not
|
||||||
|
limited to, the following:
|
||||||
|
|
||||||
|
i. the right to reproduce, adapt, distribute, perform, display,
|
||||||
|
communicate, and translate a Work;
|
||||||
|
ii. moral rights retained by the original author(s) and/or performer(s);
|
||||||
|
iii. publicity and privacy rights pertaining to a person's image or
|
||||||
|
likeness depicted in a Work;
|
||||||
|
iv. rights protecting against unfair competition in regards to a Work,
|
||||||
|
subject to the limitations in paragraph 4(a), below;
|
||||||
|
v. rights protecting the extraction, dissemination, use and reuse of data
|
||||||
|
in a Work;
|
||||||
|
vi. database rights (such as those arising under Directive 96/9/EC of the
|
||||||
|
European Parliament and of the Council of 11 March 1996 on the legal
|
||||||
|
protection of databases, and under any national implementation
|
||||||
|
thereof, including any amended or successor version of such
|
||||||
|
directive); and
|
||||||
|
vii. other similar, equivalent or corresponding rights throughout the
|
||||||
|
world based on applicable law or treaty, and any national
|
||||||
|
implementations thereof.
|
||||||
|
|
||||||
|
2. Waiver. To the greatest extent permitted by, but not in contravention
|
||||||
|
of, applicable law, Affirmer hereby overtly, fully, permanently,
|
||||||
|
irrevocably and unconditionally waives, abandons, and surrenders all of
|
||||||
|
Affirmer's Copyright and Related Rights and associated claims and causes
|
||||||
|
of action, whether now known or unknown (including existing as well as
|
||||||
|
future claims and causes of action), in the Work (i) in all territories
|
||||||
|
worldwide, (ii) for the maximum duration provided by applicable law or
|
||||||
|
treaty (including future time extensions), (iii) in any current or future
|
||||||
|
medium and for any number of copies, and (iv) for any purpose whatsoever,
|
||||||
|
including without limitation commercial, advertising or promotional
|
||||||
|
purposes (the "Waiver"). Affirmer makes the Waiver for the benefit of each
|
||||||
|
member of the public at large and to the detriment of Affirmer's heirs and
|
||||||
|
successors, fully intending that such Waiver shall not be subject to
|
||||||
|
revocation, rescission, cancellation, termination, or any other legal or
|
||||||
|
equitable action to disrupt the quiet enjoyment of the Work by the public
|
||||||
|
as contemplated by Affirmer's express Statement of Purpose.
|
||||||
|
|
||||||
|
3. Public License Fallback. Should any part of the Waiver for any reason
|
||||||
|
be judged legally invalid or ineffective under applicable law, then the
|
||||||
|
Waiver shall be preserved to the maximum extent permitted taking into
|
||||||
|
account Affirmer's express Statement of Purpose. In addition, to the
|
||||||
|
extent the Waiver is so judged Affirmer hereby grants to each affected
|
||||||
|
person a royalty-free, non transferable, non sublicensable, non exclusive,
|
||||||
|
irrevocable and unconditional license to exercise Affirmer's Copyright and
|
||||||
|
Related Rights in the Work (i) in all territories worldwide, (ii) for the
|
||||||
|
maximum duration provided by applicable law or treaty (including future
|
||||||
|
time extensions), (iii) in any current or future medium and for any number
|
||||||
|
of copies, and (iv) for any purpose whatsoever, including without
|
||||||
|
limitation commercial, advertising or promotional purposes (the
|
||||||
|
"License"). The License shall be deemed effective as of the date CC0 was
|
||||||
|
applied by Affirmer to the Work. Should any part of the License for any
|
||||||
|
reason be judged legally invalid or ineffective under applicable law, such
|
||||||
|
partial invalidity or ineffectiveness shall not invalidate the remainder
|
||||||
|
of the License, and in such case Affirmer hereby affirms that he or she
|
||||||
|
will not (i) exercise any of his or her remaining Copyright and Related
|
||||||
|
Rights in the Work or (ii) assert any associated claims and causes of
|
||||||
|
action with respect to the Work, in either case contrary to Affirmer's
|
||||||
|
express Statement of Purpose.
|
||||||
|
|
||||||
|
4. Limitations and Disclaimers.
|
||||||
|
|
||||||
|
a. No trademark or patent rights held by Affirmer are waived, abandoned,
|
||||||
|
surrendered, licensed or otherwise affected by this document.
|
||||||
|
b. Affirmer offers the Work as-is and makes no representations or
|
||||||
|
warranties of any kind concerning the Work, express, implied,
|
||||||
|
statutory or otherwise, including without limitation warranties of
|
||||||
|
title, merchantability, fitness for a particular purpose, non
|
||||||
|
infringement, or the absence of latent or other defects, accuracy, or
|
||||||
|
the present or absence of errors, whether or not discoverable, all to
|
||||||
|
the greatest extent permissible under applicable law.
|
||||||
|
c. Affirmer disclaims responsibility for clearing rights of other persons
|
||||||
|
that may apply to the Work or any use thereof, including without
|
||||||
|
limitation any person's Copyright and Related Rights in the Work.
|
||||||
|
Further, Affirmer disclaims responsibility for obtaining any necessary
|
||||||
|
consents, permissions or other rights required for any use of the
|
||||||
|
Work.
|
||||||
|
d. Affirmer understands and acknowledges that Creative Commons is not a
|
||||||
|
party to this document and has no duty or obligation with respect to
|
||||||
|
this CC0 or use of the Work.
|
201
ldm_patched/pfn/architecture/LICENSE-Swin2SR
Normal file
201
ldm_patched/pfn/architecture/LICENSE-Swin2SR
Normal file
@ -0,0 +1,201 @@
|
|||||||
|
Apache License
|
||||||
|
Version 2.0, January 2004
|
||||||
|
http://www.apache.org/licenses/
|
||||||
|
|
||||||
|
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||||
|
|
||||||
|
1. Definitions.
|
||||||
|
|
||||||
|
"License" shall mean the terms and conditions for use, reproduction,
|
||||||
|
and distribution as defined by Sections 1 through 9 of this document.
|
||||||
|
|
||||||
|
"Licensor" shall mean the copyright owner or entity authorized by
|
||||||
|
the copyright owner that is granting the License.
|
||||||
|
|
||||||
|
"Legal Entity" shall mean the union of the acting entity and all
|
||||||
|
other entities that control, are controlled by, or are under common
|
||||||
|
control with that entity. For the purposes of this definition,
|
||||||
|
"control" means (i) the power, direct or indirect, to cause the
|
||||||
|
direction or management of such entity, whether by contract or
|
||||||
|
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||||
|
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||||
|
|
||||||
|
"You" (or "Your") shall mean an individual or Legal Entity
|
||||||
|
exercising permissions granted by this License.
|
||||||
|
|
||||||
|
"Source" form shall mean the preferred form for making modifications,
|
||||||
|
including but not limited to software source code, documentation
|
||||||
|
source, and configuration files.
|
||||||
|
|
||||||
|
"Object" form shall mean any form resulting from mechanical
|
||||||
|
transformation or translation of a Source form, including but
|
||||||
|
not limited to compiled object code, generated documentation,
|
||||||
|
and conversions to other media types.
|
||||||
|
|
||||||
|
"Work" shall mean the work of authorship, whether in Source or
|
||||||
|
Object form, made available under the License, as indicated by a
|
||||||
|
copyright notice that is included in or attached to the work
|
||||||
|
(an example is provided in the Appendix below).
|
||||||
|
|
||||||
|
"Derivative Works" shall mean any work, whether in Source or Object
|
||||||
|
form, that is based on (or derived from) the Work and for which the
|
||||||
|
editorial revisions, annotations, elaborations, or other modifications
|
||||||
|
represent, as a whole, an original work of authorship. For the purposes
|
||||||
|
of this License, Derivative Works shall not include works that remain
|
||||||
|
separable from, or merely link (or bind by name) to the interfaces of,
|
||||||
|
the Work and Derivative Works thereof.
|
||||||
|
|
||||||
|
"Contribution" shall mean any work of authorship, including
|
||||||
|
the original version of the Work and any modifications or additions
|
||||||
|
to that Work or Derivative Works thereof, that is intentionally
|
||||||
|
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||||
|
or by an individual or Legal Entity authorized to submit on behalf of
|
||||||
|
the copyright owner. For the purposes of this definition, "submitted"
|
||||||
|
means any form of electronic, verbal, or written communication sent
|
||||||
|
to the Licensor or its representatives, including but not limited to
|
||||||
|
communication on electronic mailing lists, source code control systems,
|
||||||
|
and issue tracking systems that are managed by, or on behalf of, the
|
||||||
|
Licensor for the purpose of discussing and improving the Work, but
|
||||||
|
excluding communication that is conspicuously marked or otherwise
|
||||||
|
designated in writing by the copyright owner as "Not a Contribution."
|
||||||
|
|
||||||
|
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||||
|
on behalf of whom a Contribution has been received by Licensor and
|
||||||
|
subsequently incorporated within the Work.
|
||||||
|
|
||||||
|
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||||
|
this License, each Contributor hereby grants to You a perpetual,
|
||||||
|
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||||
|
copyright license to reproduce, prepare Derivative Works of,
|
||||||
|
publicly display, publicly perform, sublicense, and distribute the
|
||||||
|
Work and such Derivative Works in Source or Object form.
|
||||||
|
|
||||||
|
3. Grant of Patent License. Subject to the terms and conditions of
|
||||||
|
this License, each Contributor hereby grants to You a perpetual,
|
||||||
|
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||||
|
(except as stated in this section) patent license to make, have made,
|
||||||
|
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||||
|
where such license applies only to those patent claims licensable
|
||||||
|
by such Contributor that are necessarily infringed by their
|
||||||
|
Contribution(s) alone or by combination of their Contribution(s)
|
||||||
|
with the Work to which such Contribution(s) was submitted. If You
|
||||||
|
institute patent litigation against any entity (including a
|
||||||
|
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||||
|
or a Contribution incorporated within the Work constitutes direct
|
||||||
|
or contributory patent infringement, then any patent licenses
|
||||||
|
granted to You under this License for that Work shall terminate
|
||||||
|
as of the date such litigation is filed.
|
||||||
|
|
||||||
|
4. Redistribution. You may reproduce and distribute copies of the
|
||||||
|
Work or Derivative Works thereof in any medium, with or without
|
||||||
|
modifications, and in Source or Object form, provided that You
|
||||||
|
meet the following conditions:
|
||||||
|
|
||||||
|
(a) You must give any other recipients of the Work or
|
||||||
|
Derivative Works a copy of this License; and
|
||||||
|
|
||||||
|
(b) You must cause any modified files to carry prominent notices
|
||||||
|
stating that You changed the files; and
|
||||||
|
|
||||||
|
(c) You must retain, in the Source form of any Derivative Works
|
||||||
|
that You distribute, all copyright, patent, trademark, and
|
||||||
|
attribution notices from the Source form of the Work,
|
||||||
|
excluding those notices that do not pertain to any part of
|
||||||
|
the Derivative Works; and
|
||||||
|
|
||||||
|
(d) If the Work includes a "NOTICE" text file as part of its
|
||||||
|
distribution, then any Derivative Works that You distribute must
|
||||||
|
include a readable copy of the attribution notices contained
|
||||||
|
within such NOTICE file, excluding those notices that do not
|
||||||
|
pertain to any part of the Derivative Works, in at least one
|
||||||
|
of the following places: within a NOTICE text file distributed
|
||||||
|
as part of the Derivative Works; within the Source form or
|
||||||
|
documentation, if provided along with the Derivative Works; or,
|
||||||
|
within a display generated by the Derivative Works, if and
|
||||||
|
wherever such third-party notices normally appear. The contents
|
||||||
|
of the NOTICE file are for informational purposes only and
|
||||||
|
do not modify the License. You may add Your own attribution
|
||||||
|
notices within Derivative Works that You distribute, alongside
|
||||||
|
or as an addendum to the NOTICE text from the Work, provided
|
||||||
|
that such additional attribution notices cannot be construed
|
||||||
|
as modifying the License.
|
||||||
|
|
||||||
|
You may add Your own copyright statement to Your modifications and
|
||||||
|
may provide additional or different license terms and conditions
|
||||||
|
for use, reproduction, or distribution of Your modifications, or
|
||||||
|
for any such Derivative Works as a whole, provided Your use,
|
||||||
|
reproduction, and distribution of the Work otherwise complies with
|
||||||
|
the conditions stated in this License.
|
||||||
|
|
||||||
|
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||||
|
any Contribution intentionally submitted for inclusion in the Work
|
||||||
|
by You to the Licensor shall be under the terms and conditions of
|
||||||
|
this License, without any additional terms or conditions.
|
||||||
|
Notwithstanding the above, nothing herein shall supersede or modify
|
||||||
|
the terms of any separate license agreement you may have executed
|
||||||
|
with Licensor regarding such Contributions.
|
||||||
|
|
||||||
|
6. Trademarks. This License does not grant permission to use the trade
|
||||||
|
names, trademarks, service marks, or product names of the Licensor,
|
||||||
|
except as required for reasonable and customary use in describing the
|
||||||
|
origin of the Work and reproducing the content of the NOTICE file.
|
||||||
|
|
||||||
|
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||||
|
agreed to in writing, Licensor provides the Work (and each
|
||||||
|
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||||
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||||
|
implied, including, without limitation, any warranties or conditions
|
||||||
|
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||||
|
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||||
|
appropriateness of using or redistributing the Work and assume any
|
||||||
|
risks associated with Your exercise of permissions under this License.
|
||||||
|
|
||||||
|
8. Limitation of Liability. In no event and under no legal theory,
|
||||||
|
whether in tort (including negligence), contract, or otherwise,
|
||||||
|
unless required by applicable law (such as deliberate and grossly
|
||||||
|
negligent acts) or agreed to in writing, shall any Contributor be
|
||||||
|
liable to You for damages, including any direct, indirect, special,
|
||||||
|
incidental, or consequential damages of any character arising as a
|
||||||
|
result of this License or out of the use or inability to use the
|
||||||
|
Work (including but not limited to damages for loss of goodwill,
|
||||||
|
work stoppage, computer failure or malfunction, or any and all
|
||||||
|
other commercial damages or losses), even if such Contributor
|
||||||
|
has been advised of the possibility of such damages.
|
||||||
|
|
||||||
|
9. Accepting Warranty or Additional Liability. While redistributing
|
||||||
|
the Work or Derivative Works thereof, You may choose to offer,
|
||||||
|
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||||
|
or other liability obligations and/or rights consistent with this
|
||||||
|
License. However, in accepting such obligations, You may act only
|
||||||
|
on Your own behalf and on Your sole responsibility, not on behalf
|
||||||
|
of any other Contributor, and only if You agree to indemnify,
|
||||||
|
defend, and hold each Contributor harmless for any liability
|
||||||
|
incurred by, or claims asserted against, such Contributor by reason
|
||||||
|
of your accepting any such warranty or additional liability.
|
||||||
|
|
||||||
|
END OF TERMS AND CONDITIONS
|
||||||
|
|
||||||
|
APPENDIX: How to apply the Apache License to your work.
|
||||||
|
|
||||||
|
To apply the Apache License to your work, attach the following
|
||||||
|
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||||
|
replaced with your own identifying information. (Don't include
|
||||||
|
the brackets!) The text should be enclosed in the appropriate
|
||||||
|
comment syntax for the file format. We also recommend that a
|
||||||
|
file or class name and description of purpose be included on the
|
||||||
|
same "printed page" as the copyright notice for easier
|
||||||
|
identification within third-party archives.
|
||||||
|
|
||||||
|
Copyright [2021] [SwinIR Authors]
|
||||||
|
|
||||||
|
Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
you may not use this file except in compliance with the License.
|
||||||
|
You may obtain a copy of the License at
|
||||||
|
|
||||||
|
http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
|
||||||
|
Unless required by applicable law or agreed to in writing, software
|
||||||
|
distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
See the License for the specific language governing permissions and
|
||||||
|
limitations under the License.
|
201
ldm_patched/pfn/architecture/LICENSE-SwinIR
Normal file
201
ldm_patched/pfn/architecture/LICENSE-SwinIR
Normal file
@ -0,0 +1,201 @@
|
|||||||
|
Apache License
|
||||||
|
Version 2.0, January 2004
|
||||||
|
http://www.apache.org/licenses/
|
||||||
|
|
||||||
|
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||||
|
|
||||||
|
1. Definitions.
|
||||||
|
|
||||||
|
"License" shall mean the terms and conditions for use, reproduction,
|
||||||
|
and distribution as defined by Sections 1 through 9 of this document.
|
||||||
|
|
||||||
|
"Licensor" shall mean the copyright owner or entity authorized by
|
||||||
|
the copyright owner that is granting the License.
|
||||||
|
|
||||||
|
"Legal Entity" shall mean the union of the acting entity and all
|
||||||
|
other entities that control, are controlled by, or are under common
|
||||||
|
control with that entity. For the purposes of this definition,
|
||||||
|
"control" means (i) the power, direct or indirect, to cause the
|
||||||
|
direction or management of such entity, whether by contract or
|
||||||
|
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||||
|
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||||
|
|
||||||
|
"You" (or "Your") shall mean an individual or Legal Entity
|
||||||
|
exercising permissions granted by this License.
|
||||||
|
|
||||||
|
"Source" form shall mean the preferred form for making modifications,
|
||||||
|
including but not limited to software source code, documentation
|
||||||
|
source, and configuration files.
|
||||||
|
|
||||||
|
"Object" form shall mean any form resulting from mechanical
|
||||||
|
transformation or translation of a Source form, including but
|
||||||
|
not limited to compiled object code, generated documentation,
|
||||||
|
and conversions to other media types.
|
||||||
|
|
||||||
|
"Work" shall mean the work of authorship, whether in Source or
|
||||||
|
Object form, made available under the License, as indicated by a
|
||||||
|
copyright notice that is included in or attached to the work
|
||||||
|
(an example is provided in the Appendix below).
|
||||||
|
|
||||||
|
"Derivative Works" shall mean any work, whether in Source or Object
|
||||||
|
form, that is based on (or derived from) the Work and for which the
|
||||||
|
editorial revisions, annotations, elaborations, or other modifications
|
||||||
|
represent, as a whole, an original work of authorship. For the purposes
|
||||||
|
of this License, Derivative Works shall not include works that remain
|
||||||
|
separable from, or merely link (or bind by name) to the interfaces of,
|
||||||
|
the Work and Derivative Works thereof.
|
||||||
|
|
||||||
|
"Contribution" shall mean any work of authorship, including
|
||||||
|
the original version of the Work and any modifications or additions
|
||||||
|
to that Work or Derivative Works thereof, that is intentionally
|
||||||
|
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||||
|
or by an individual or Legal Entity authorized to submit on behalf of
|
||||||
|
the copyright owner. For the purposes of this definition, "submitted"
|
||||||
|
means any form of electronic, verbal, or written communication sent
|
||||||
|
to the Licensor or its representatives, including but not limited to
|
||||||
|
communication on electronic mailing lists, source code control systems,
|
||||||
|
and issue tracking systems that are managed by, or on behalf of, the
|
||||||
|
Licensor for the purpose of discussing and improving the Work, but
|
||||||
|
excluding communication that is conspicuously marked or otherwise
|
||||||
|
designated in writing by the copyright owner as "Not a Contribution."
|
||||||
|
|
||||||
|
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||||
|
on behalf of whom a Contribution has been received by Licensor and
|
||||||
|
subsequently incorporated within the Work.
|
||||||
|
|
||||||
|
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||||
|
this License, each Contributor hereby grants to You a perpetual,
|
||||||
|
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||||
|
copyright license to reproduce, prepare Derivative Works of,
|
||||||
|
publicly display, publicly perform, sublicense, and distribute the
|
||||||
|
Work and such Derivative Works in Source or Object form.
|
||||||
|
|
||||||
|
3. Grant of Patent License. Subject to the terms and conditions of
|
||||||
|
this License, each Contributor hereby grants to You a perpetual,
|
||||||
|
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||||
|
(except as stated in this section) patent license to make, have made,
|
||||||
|
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||||
|
where such license applies only to those patent claims licensable
|
||||||
|
by such Contributor that are necessarily infringed by their
|
||||||
|
Contribution(s) alone or by combination of their Contribution(s)
|
||||||
|
with the Work to which such Contribution(s) was submitted. If You
|
||||||
|
institute patent litigation against any entity (including a
|
||||||
|
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||||
|
or a Contribution incorporated within the Work constitutes direct
|
||||||
|
or contributory patent infringement, then any patent licenses
|
||||||
|
granted to You under this License for that Work shall terminate
|
||||||
|
as of the date such litigation is filed.
|
||||||
|
|
||||||
|
4. Redistribution. You may reproduce and distribute copies of the
|
||||||
|
Work or Derivative Works thereof in any medium, with or without
|
||||||
|
modifications, and in Source or Object form, provided that You
|
||||||
|
meet the following conditions:
|
||||||
|
|
||||||
|
(a) You must give any other recipients of the Work or
|
||||||
|
Derivative Works a copy of this License; and
|
||||||
|
|
||||||
|
(b) You must cause any modified files to carry prominent notices
|
||||||
|
stating that You changed the files; and
|
||||||
|
|
||||||
|
(c) You must retain, in the Source form of any Derivative Works
|
||||||
|
that You distribute, all copyright, patent, trademark, and
|
||||||
|
attribution notices from the Source form of the Work,
|
||||||
|
excluding those notices that do not pertain to any part of
|
||||||
|
the Derivative Works; and
|
||||||
|
|
||||||
|
(d) If the Work includes a "NOTICE" text file as part of its
|
||||||
|
distribution, then any Derivative Works that You distribute must
|
||||||
|
include a readable copy of the attribution notices contained
|
||||||
|
within such NOTICE file, excluding those notices that do not
|
||||||
|
pertain to any part of the Derivative Works, in at least one
|
||||||
|
of the following places: within a NOTICE text file distributed
|
||||||
|
as part of the Derivative Works; within the Source form or
|
||||||
|
documentation, if provided along with the Derivative Works; or,
|
||||||
|
within a display generated by the Derivative Works, if and
|
||||||
|
wherever such third-party notices normally appear. The contents
|
||||||
|
of the NOTICE file are for informational purposes only and
|
||||||
|
do not modify the License. You may add Your own attribution
|
||||||
|
notices within Derivative Works that You distribute, alongside
|
||||||
|
or as an addendum to the NOTICE text from the Work, provided
|
||||||
|
that such additional attribution notices cannot be construed
|
||||||
|
as modifying the License.
|
||||||
|
|
||||||
|
You may add Your own copyright statement to Your modifications and
|
||||||
|
may provide additional or different license terms and conditions
|
||||||
|
for use, reproduction, or distribution of Your modifications, or
|
||||||
|
for any such Derivative Works as a whole, provided Your use,
|
||||||
|
reproduction, and distribution of the Work otherwise complies with
|
||||||
|
the conditions stated in this License.
|
||||||
|
|
||||||
|
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||||
|
any Contribution intentionally submitted for inclusion in the Work
|
||||||
|
by You to the Licensor shall be under the terms and conditions of
|
||||||
|
this License, without any additional terms or conditions.
|
||||||
|
Notwithstanding the above, nothing herein shall supersede or modify
|
||||||
|
the terms of any separate license agreement you may have executed
|
||||||
|
with Licensor regarding such Contributions.
|
||||||
|
|
||||||
|
6. Trademarks. This License does not grant permission to use the trade
|
||||||
|
names, trademarks, service marks, or product names of the Licensor,
|
||||||
|
except as required for reasonable and customary use in describing the
|
||||||
|
origin of the Work and reproducing the content of the NOTICE file.
|
||||||
|
|
||||||
|
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||||
|
agreed to in writing, Licensor provides the Work (and each
|
||||||
|
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||||
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||||
|
implied, including, without limitation, any warranties or conditions
|
||||||
|
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||||
|
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||||
|
appropriateness of using or redistributing the Work and assume any
|
||||||
|
risks associated with Your exercise of permissions under this License.
|
||||||
|
|
||||||
|
8. Limitation of Liability. In no event and under no legal theory,
|
||||||
|
whether in tort (including negligence), contract, or otherwise,
|
||||||
|
unless required by applicable law (such as deliberate and grossly
|
||||||
|
negligent acts) or agreed to in writing, shall any Contributor be
|
||||||
|
liable to You for damages, including any direct, indirect, special,
|
||||||
|
incidental, or consequential damages of any character arising as a
|
||||||
|
result of this License or out of the use or inability to use the
|
||||||
|
Work (including but not limited to damages for loss of goodwill,
|
||||||
|
work stoppage, computer failure or malfunction, or any and all
|
||||||
|
other commercial damages or losses), even if such Contributor
|
||||||
|
has been advised of the possibility of such damages.
|
||||||
|
|
||||||
|
9. Accepting Warranty or Additional Liability. While redistributing
|
||||||
|
the Work or Derivative Works thereof, You may choose to offer,
|
||||||
|
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||||
|
or other liability obligations and/or rights consistent with this
|
||||||
|
License. However, in accepting such obligations, You may act only
|
||||||
|
on Your own behalf and on Your sole responsibility, not on behalf
|
||||||
|
of any other Contributor, and only if You agree to indemnify,
|
||||||
|
defend, and hold each Contributor harmless for any liability
|
||||||
|
incurred by, or claims asserted against, such Contributor by reason
|
||||||
|
of your accepting any such warranty or additional liability.
|
||||||
|
|
||||||
|
END OF TERMS AND CONDITIONS
|
||||||
|
|
||||||
|
APPENDIX: How to apply the Apache License to your work.
|
||||||
|
|
||||||
|
To apply the Apache License to your work, attach the following
|
||||||
|
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||||
|
replaced with your own identifying information. (Don't include
|
||||||
|
the brackets!) The text should be enclosed in the appropriate
|
||||||
|
comment syntax for the file format. We also recommend that a
|
||||||
|
file or class name and description of purpose be included on the
|
||||||
|
same "printed page" as the copyright notice for easier
|
||||||
|
identification within third-party archives.
|
||||||
|
|
||||||
|
Copyright [2021] [SwinIR Authors]
|
||||||
|
|
||||||
|
Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
you may not use this file except in compliance with the License.
|
||||||
|
You may obtain a copy of the License at
|
||||||
|
|
||||||
|
http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
|
||||||
|
Unless required by applicable law or agreed to in writing, software
|
||||||
|
distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
See the License for the specific language governing permissions and
|
||||||
|
limitations under the License.
|
201
ldm_patched/pfn/architecture/LICENSE-lama
Normal file
201
ldm_patched/pfn/architecture/LICENSE-lama
Normal file
@ -0,0 +1,201 @@
|
|||||||
|
Apache License
|
||||||
|
Version 2.0, January 2004
|
||||||
|
http://www.apache.org/licenses/
|
||||||
|
|
||||||
|
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||||
|
|
||||||
|
1. Definitions.
|
||||||
|
|
||||||
|
"License" shall mean the terms and conditions for use, reproduction,
|
||||||
|
and distribution as defined by Sections 1 through 9 of this document.
|
||||||
|
|
||||||
|
"Licensor" shall mean the copyright owner or entity authorized by
|
||||||
|
the copyright owner that is granting the License.
|
||||||
|
|
||||||
|
"Legal Entity" shall mean the union of the acting entity and all
|
||||||
|
other entities that control, are controlled by, or are under common
|
||||||
|
control with that entity. For the purposes of this definition,
|
||||||
|
"control" means (i) the power, direct or indirect, to cause the
|
||||||
|
direction or management of such entity, whether by contract or
|
||||||
|
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||||
|
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||||
|
|
||||||
|
"You" (or "Your") shall mean an individual or Legal Entity
|
||||||
|
exercising permissions granted by this License.
|
||||||
|
|
||||||
|
"Source" form shall mean the preferred form for making modifications,
|
||||||
|
including but not limited to software source code, documentation
|
||||||
|
source, and configuration files.
|
||||||
|
|
||||||
|
"Object" form shall mean any form resulting from mechanical
|
||||||
|
transformation or translation of a Source form, including but
|
||||||
|
not limited to compiled object code, generated documentation,
|
||||||
|
and conversions to other media types.
|
||||||
|
|
||||||
|
"Work" shall mean the work of authorship, whether in Source or
|
||||||
|
Object form, made available under the License, as indicated by a
|
||||||
|
copyright notice that is included in or attached to the work
|
||||||
|
(an example is provided in the Appendix below).
|
||||||
|
|
||||||
|
"Derivative Works" shall mean any work, whether in Source or Object
|
||||||
|
form, that is based on (or derived from) the Work and for which the
|
||||||
|
editorial revisions, annotations, elaborations, or other modifications
|
||||||
|
represent, as a whole, an original work of authorship. For the purposes
|
||||||
|
of this License, Derivative Works shall not include works that remain
|
||||||
|
separable from, or merely link (or bind by name) to the interfaces of,
|
||||||
|
the Work and Derivative Works thereof.
|
||||||
|
|
||||||
|
"Contribution" shall mean any work of authorship, including
|
||||||
|
the original version of the Work and any modifications or additions
|
||||||
|
to that Work or Derivative Works thereof, that is intentionally
|
||||||
|
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||||
|
or by an individual or Legal Entity authorized to submit on behalf of
|
||||||
|
the copyright owner. For the purposes of this definition, "submitted"
|
||||||
|
means any form of electronic, verbal, or written communication sent
|
||||||
|
to the Licensor or its representatives, including but not limited to
|
||||||
|
communication on electronic mailing lists, source code control systems,
|
||||||
|
and issue tracking systems that are managed by, or on behalf of, the
|
||||||
|
Licensor for the purpose of discussing and improving the Work, but
|
||||||
|
excluding communication that is conspicuously marked or otherwise
|
||||||
|
designated in writing by the copyright owner as "Not a Contribution."
|
||||||
|
|
||||||
|
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||||
|
on behalf of whom a Contribution has been received by Licensor and
|
||||||
|
subsequently incorporated within the Work.
|
||||||
|
|
||||||
|
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||||
|
this License, each Contributor hereby grants to You a perpetual,
|
||||||
|
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||||
|
copyright license to reproduce, prepare Derivative Works of,
|
||||||
|
publicly display, publicly perform, sublicense, and distribute the
|
||||||
|
Work and such Derivative Works in Source or Object form.
|
||||||
|
|
||||||
|
3. Grant of Patent License. Subject to the terms and conditions of
|
||||||
|
this License, each Contributor hereby grants to You a perpetual,
|
||||||
|
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||||
|
(except as stated in this section) patent license to make, have made,
|
||||||
|
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||||
|
where such license applies only to those patent claims licensable
|
||||||
|
by such Contributor that are necessarily infringed by their
|
||||||
|
Contribution(s) alone or by combination of their Contribution(s)
|
||||||
|
with the Work to which such Contribution(s) was submitted. If You
|
||||||
|
institute patent litigation against any entity (including a
|
||||||
|
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||||
|
or a Contribution incorporated within the Work constitutes direct
|
||||||
|
or contributory patent infringement, then any patent licenses
|
||||||
|
granted to You under this License for that Work shall terminate
|
||||||
|
as of the date such litigation is filed.
|
||||||
|
|
||||||
|
4. Redistribution. You may reproduce and distribute copies of the
|
||||||
|
Work or Derivative Works thereof in any medium, with or without
|
||||||
|
modifications, and in Source or Object form, provided that You
|
||||||
|
meet the following conditions:
|
||||||
|
|
||||||
|
(a) You must give any other recipients of the Work or
|
||||||
|
Derivative Works a copy of this License; and
|
||||||
|
|
||||||
|
(b) You must cause any modified files to carry prominent notices
|
||||||
|
stating that You changed the files; and
|
||||||
|
|
||||||
|
(c) You must retain, in the Source form of any Derivative Works
|
||||||
|
that You distribute, all copyright, patent, trademark, and
|
||||||
|
attribution notices from the Source form of the Work,
|
||||||
|
excluding those notices that do not pertain to any part of
|
||||||
|
the Derivative Works; and
|
||||||
|
|
||||||
|
(d) If the Work includes a "NOTICE" text file as part of its
|
||||||
|
distribution, then any Derivative Works that You distribute must
|
||||||
|
include a readable copy of the attribution notices contained
|
||||||
|
within such NOTICE file, excluding those notices that do not
|
||||||
|
pertain to any part of the Derivative Works, in at least one
|
||||||
|
of the following places: within a NOTICE text file distributed
|
||||||
|
as part of the Derivative Works; within the Source form or
|
||||||
|
documentation, if provided along with the Derivative Works; or,
|
||||||
|
within a display generated by the Derivative Works, if and
|
||||||
|
wherever such third-party notices normally appear. The contents
|
||||||
|
of the NOTICE file are for informational purposes only and
|
||||||
|
do not modify the License. You may add Your own attribution
|
||||||
|
notices within Derivative Works that You distribute, alongside
|
||||||
|
or as an addendum to the NOTICE text from the Work, provided
|
||||||
|
that such additional attribution notices cannot be construed
|
||||||
|
as modifying the License.
|
||||||
|
|
||||||
|
You may add Your own copyright statement to Your modifications and
|
||||||
|
may provide additional or different license terms and conditions
|
||||||
|
for use, reproduction, or distribution of Your modifications, or
|
||||||
|
for any such Derivative Works as a whole, provided Your use,
|
||||||
|
reproduction, and distribution of the Work otherwise complies with
|
||||||
|
the conditions stated in this License.
|
||||||
|
|
||||||
|
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||||
|
any Contribution intentionally submitted for inclusion in the Work
|
||||||
|
by You to the Licensor shall be under the terms and conditions of
|
||||||
|
this License, without any additional terms or conditions.
|
||||||
|
Notwithstanding the above, nothing herein shall supersede or modify
|
||||||
|
the terms of any separate license agreement you may have executed
|
||||||
|
with Licensor regarding such Contributions.
|
||||||
|
|
||||||
|
6. Trademarks. This License does not grant permission to use the trade
|
||||||
|
names, trademarks, service marks, or product names of the Licensor,
|
||||||
|
except as required for reasonable and customary use in describing the
|
||||||
|
origin of the Work and reproducing the content of the NOTICE file.
|
||||||
|
|
||||||
|
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||||
|
agreed to in writing, Licensor provides the Work (and each
|
||||||
|
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||||
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||||
|
implied, including, without limitation, any warranties or conditions
|
||||||
|
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||||
|
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||||
|
appropriateness of using or redistributing the Work and assume any
|
||||||
|
risks associated with Your exercise of permissions under this License.
|
||||||
|
|
||||||
|
8. Limitation of Liability. In no event and under no legal theory,
|
||||||
|
whether in tort (including negligence), contract, or otherwise,
|
||||||
|
unless required by applicable law (such as deliberate and grossly
|
||||||
|
negligent acts) or agreed to in writing, shall any Contributor be
|
||||||
|
liable to You for damages, including any direct, indirect, special,
|
||||||
|
incidental, or consequential damages of any character arising as a
|
||||||
|
result of this License or out of the use or inability to use the
|
||||||
|
Work (including but not limited to damages for loss of goodwill,
|
||||||
|
work stoppage, computer failure or malfunction, or any and all
|
||||||
|
other commercial damages or losses), even if such Contributor
|
||||||
|
has been advised of the possibility of such damages.
|
||||||
|
|
||||||
|
9. Accepting Warranty or Additional Liability. While redistributing
|
||||||
|
the Work or Derivative Works thereof, You may choose to offer,
|
||||||
|
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||||
|
or other liability obligations and/or rights consistent with this
|
||||||
|
License. However, in accepting such obligations, You may act only
|
||||||
|
on Your own behalf and on Your sole responsibility, not on behalf
|
||||||
|
of any other Contributor, and only if You agree to indemnify,
|
||||||
|
defend, and hold each Contributor harmless for any liability
|
||||||
|
incurred by, or claims asserted against, such Contributor by reason
|
||||||
|
of your accepting any such warranty or additional liability.
|
||||||
|
|
||||||
|
END OF TERMS AND CONDITIONS
|
||||||
|
|
||||||
|
APPENDIX: How to apply the Apache License to your work.
|
||||||
|
|
||||||
|
To apply the Apache License to your work, attach the following
|
||||||
|
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||||
|
replaced with your own identifying information. (Don't include
|
||||||
|
the brackets!) The text should be enclosed in the appropriate
|
||||||
|
comment syntax for the file format. We also recommend that a
|
||||||
|
file or class name and description of purpose be included on the
|
||||||
|
same "printed page" as the copyright notice for easier
|
||||||
|
identification within third-party archives.
|
||||||
|
|
||||||
|
Copyright [2021] Samsung Research
|
||||||
|
|
||||||
|
Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
you may not use this file except in compliance with the License.
|
||||||
|
You may obtain a copy of the License at
|
||||||
|
|
||||||
|
http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
|
||||||
|
Unless required by applicable law or agreed to in writing, software
|
||||||
|
distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
See the License for the specific language governing permissions and
|
||||||
|
limitations under the License.
|
694
ldm_patched/pfn/architecture/LaMa.py
Normal file
694
ldm_patched/pfn/architecture/LaMa.py
Normal file
@ -0,0 +1,694 @@
|
|||||||
|
# pylint: skip-file
|
||||||
|
"""
|
||||||
|
Model adapted from advimman's lama project: https://github.com/advimman/lama
|
||||||
|
"""
|
||||||
|
|
||||||
|
# Fast Fourier Convolution NeurIPS 2020
|
||||||
|
# original implementation https://github.com/pkumivision/FFC/blob/main/model_zoo/ffc.py
|
||||||
|
# paper https://proceedings.neurips.cc/paper/2020/file/2fd5d41ec6cfab47e32164d5624269b1-Paper.pdf
|
||||||
|
|
||||||
|
from typing import List
|
||||||
|
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
import torch.nn.functional as F
|
||||||
|
from torchvision.transforms.functional import InterpolationMode, rotate
|
||||||
|
|
||||||
|
|
||||||
|
class LearnableSpatialTransformWrapper(nn.Module):
|
||||||
|
def __init__(self, impl, pad_coef=0.5, angle_init_range=80, train_angle=True):
|
||||||
|
super().__init__()
|
||||||
|
self.impl = impl
|
||||||
|
self.angle = torch.rand(1) * angle_init_range
|
||||||
|
if train_angle:
|
||||||
|
self.angle = nn.Parameter(self.angle, requires_grad=True)
|
||||||
|
self.pad_coef = pad_coef
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
if torch.is_tensor(x):
|
||||||
|
return self.inverse_transform(self.impl(self.transform(x)), x)
|
||||||
|
elif isinstance(x, tuple):
|
||||||
|
x_trans = tuple(self.transform(elem) for elem in x)
|
||||||
|
y_trans = self.impl(x_trans)
|
||||||
|
return tuple(
|
||||||
|
self.inverse_transform(elem, orig_x) for elem, orig_x in zip(y_trans, x)
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
raise ValueError(f"Unexpected input type {type(x)}")
|
||||||
|
|
||||||
|
def transform(self, x):
|
||||||
|
height, width = x.shape[2:]
|
||||||
|
pad_h, pad_w = int(height * self.pad_coef), int(width * self.pad_coef)
|
||||||
|
x_padded = F.pad(x, [pad_w, pad_w, pad_h, pad_h], mode="reflect")
|
||||||
|
x_padded_rotated = rotate(
|
||||||
|
x_padded, self.angle.to(x_padded), InterpolationMode.BILINEAR, fill=0
|
||||||
|
)
|
||||||
|
|
||||||
|
return x_padded_rotated
|
||||||
|
|
||||||
|
def inverse_transform(self, y_padded_rotated, orig_x):
|
||||||
|
height, width = orig_x.shape[2:]
|
||||||
|
pad_h, pad_w = int(height * self.pad_coef), int(width * self.pad_coef)
|
||||||
|
|
||||||
|
y_padded = rotate(
|
||||||
|
y_padded_rotated,
|
||||||
|
-self.angle.to(y_padded_rotated),
|
||||||
|
InterpolationMode.BILINEAR,
|
||||||
|
fill=0,
|
||||||
|
)
|
||||||
|
y_height, y_width = y_padded.shape[2:]
|
||||||
|
y = y_padded[:, :, pad_h : y_height - pad_h, pad_w : y_width - pad_w]
|
||||||
|
return y
|
||||||
|
|
||||||
|
|
||||||
|
class SELayer(nn.Module):
|
||||||
|
def __init__(self, channel, reduction=16):
|
||||||
|
super(SELayer, self).__init__()
|
||||||
|
self.avg_pool = nn.AdaptiveAvgPool2d(1)
|
||||||
|
self.fc = nn.Sequential(
|
||||||
|
nn.Linear(channel, channel // reduction, bias=False),
|
||||||
|
nn.ReLU(inplace=True),
|
||||||
|
nn.Linear(channel // reduction, channel, bias=False),
|
||||||
|
nn.Sigmoid(),
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
b, c, _, _ = x.size()
|
||||||
|
y = self.avg_pool(x).view(b, c)
|
||||||
|
y = self.fc(y).view(b, c, 1, 1)
|
||||||
|
res = x * y.expand_as(x)
|
||||||
|
return res
|
||||||
|
|
||||||
|
|
||||||
|
class FourierUnit(nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
in_channels,
|
||||||
|
out_channels,
|
||||||
|
groups=1,
|
||||||
|
spatial_scale_factor=None,
|
||||||
|
spatial_scale_mode="bilinear",
|
||||||
|
spectral_pos_encoding=False,
|
||||||
|
use_se=False,
|
||||||
|
se_kwargs=None,
|
||||||
|
ffc3d=False,
|
||||||
|
fft_norm="ortho",
|
||||||
|
):
|
||||||
|
# bn_layer not used
|
||||||
|
super(FourierUnit, self).__init__()
|
||||||
|
self.groups = groups
|
||||||
|
|
||||||
|
self.conv_layer = torch.nn.Conv2d(
|
||||||
|
in_channels=in_channels * 2 + (2 if spectral_pos_encoding else 0),
|
||||||
|
out_channels=out_channels * 2,
|
||||||
|
kernel_size=1,
|
||||||
|
stride=1,
|
||||||
|
padding=0,
|
||||||
|
groups=self.groups,
|
||||||
|
bias=False,
|
||||||
|
)
|
||||||
|
self.bn = torch.nn.BatchNorm2d(out_channels * 2)
|
||||||
|
self.relu = torch.nn.ReLU(inplace=True)
|
||||||
|
|
||||||
|
# squeeze and excitation block
|
||||||
|
self.use_se = use_se
|
||||||
|
if use_se:
|
||||||
|
if se_kwargs is None:
|
||||||
|
se_kwargs = {}
|
||||||
|
self.se = SELayer(self.conv_layer.in_channels, **se_kwargs)
|
||||||
|
|
||||||
|
self.spatial_scale_factor = spatial_scale_factor
|
||||||
|
self.spatial_scale_mode = spatial_scale_mode
|
||||||
|
self.spectral_pos_encoding = spectral_pos_encoding
|
||||||
|
self.ffc3d = ffc3d
|
||||||
|
self.fft_norm = fft_norm
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
half_check = False
|
||||||
|
if x.type() == "torch.cuda.HalfTensor":
|
||||||
|
# half only works on gpu anyway
|
||||||
|
half_check = True
|
||||||
|
|
||||||
|
batch = x.shape[0]
|
||||||
|
|
||||||
|
if self.spatial_scale_factor is not None:
|
||||||
|
orig_size = x.shape[-2:]
|
||||||
|
x = F.interpolate(
|
||||||
|
x,
|
||||||
|
scale_factor=self.spatial_scale_factor,
|
||||||
|
mode=self.spatial_scale_mode,
|
||||||
|
align_corners=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
# (batch, c, h, w/2+1, 2)
|
||||||
|
fft_dim = (-3, -2, -1) if self.ffc3d else (-2, -1)
|
||||||
|
if half_check == True:
|
||||||
|
ffted = torch.fft.rfftn(
|
||||||
|
x.float(), dim=fft_dim, norm=self.fft_norm
|
||||||
|
) # .type(torch.cuda.HalfTensor)
|
||||||
|
else:
|
||||||
|
ffted = torch.fft.rfftn(x, dim=fft_dim, norm=self.fft_norm)
|
||||||
|
|
||||||
|
ffted = torch.stack((ffted.real, ffted.imag), dim=-1)
|
||||||
|
ffted = ffted.permute(0, 1, 4, 2, 3).contiguous() # (batch, c, 2, h, w/2+1)
|
||||||
|
ffted = ffted.view(
|
||||||
|
(
|
||||||
|
batch,
|
||||||
|
-1,
|
||||||
|
)
|
||||||
|
+ ffted.size()[3:]
|
||||||
|
)
|
||||||
|
|
||||||
|
if self.spectral_pos_encoding:
|
||||||
|
height, width = ffted.shape[-2:]
|
||||||
|
coords_vert = (
|
||||||
|
torch.linspace(0, 1, height)[None, None, :, None]
|
||||||
|
.expand(batch, 1, height, width)
|
||||||
|
.to(ffted)
|
||||||
|
)
|
||||||
|
coords_hor = (
|
||||||
|
torch.linspace(0, 1, width)[None, None, None, :]
|
||||||
|
.expand(batch, 1, height, width)
|
||||||
|
.to(ffted)
|
||||||
|
)
|
||||||
|
ffted = torch.cat((coords_vert, coords_hor, ffted), dim=1)
|
||||||
|
|
||||||
|
if self.use_se:
|
||||||
|
ffted = self.se(ffted)
|
||||||
|
|
||||||
|
if half_check == True:
|
||||||
|
ffted = self.conv_layer(ffted.half()) # (batch, c*2, h, w/2+1)
|
||||||
|
else:
|
||||||
|
ffted = self.conv_layer(
|
||||||
|
ffted
|
||||||
|
) # .type(torch.cuda.FloatTensor) # (batch, c*2, h, w/2+1)
|
||||||
|
|
||||||
|
ffted = self.relu(self.bn(ffted))
|
||||||
|
# forcing to be always float
|
||||||
|
ffted = ffted.float()
|
||||||
|
|
||||||
|
ffted = (
|
||||||
|
ffted.view(
|
||||||
|
(
|
||||||
|
batch,
|
||||||
|
-1,
|
||||||
|
2,
|
||||||
|
)
|
||||||
|
+ ffted.size()[2:]
|
||||||
|
)
|
||||||
|
.permute(0, 1, 3, 4, 2)
|
||||||
|
.contiguous()
|
||||||
|
) # (batch,c, t, h, w/2+1, 2)
|
||||||
|
|
||||||
|
ffted = torch.complex(ffted[..., 0], ffted[..., 1])
|
||||||
|
|
||||||
|
ifft_shape_slice = x.shape[-3:] if self.ffc3d else x.shape[-2:]
|
||||||
|
output = torch.fft.irfftn(
|
||||||
|
ffted, s=ifft_shape_slice, dim=fft_dim, norm=self.fft_norm
|
||||||
|
)
|
||||||
|
|
||||||
|
if half_check == True:
|
||||||
|
output = output.half()
|
||||||
|
|
||||||
|
if self.spatial_scale_factor is not None:
|
||||||
|
output = F.interpolate(
|
||||||
|
output,
|
||||||
|
size=orig_size,
|
||||||
|
mode=self.spatial_scale_mode,
|
||||||
|
align_corners=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
return output
|
||||||
|
|
||||||
|
|
||||||
|
class SpectralTransform(nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
in_channels,
|
||||||
|
out_channels,
|
||||||
|
stride=1,
|
||||||
|
groups=1,
|
||||||
|
enable_lfu=True,
|
||||||
|
separable_fu=False,
|
||||||
|
**fu_kwargs,
|
||||||
|
):
|
||||||
|
# bn_layer not used
|
||||||
|
super(SpectralTransform, self).__init__()
|
||||||
|
self.enable_lfu = enable_lfu
|
||||||
|
if stride == 2:
|
||||||
|
self.downsample = nn.AvgPool2d(kernel_size=(2, 2), stride=2)
|
||||||
|
else:
|
||||||
|
self.downsample = nn.Identity()
|
||||||
|
|
||||||
|
self.stride = stride
|
||||||
|
self.conv1 = nn.Sequential(
|
||||||
|
nn.Conv2d(
|
||||||
|
in_channels, out_channels // 2, kernel_size=1, groups=groups, bias=False
|
||||||
|
),
|
||||||
|
nn.BatchNorm2d(out_channels // 2),
|
||||||
|
nn.ReLU(inplace=True),
|
||||||
|
)
|
||||||
|
fu_class = FourierUnit
|
||||||
|
self.fu = fu_class(out_channels // 2, out_channels // 2, groups, **fu_kwargs)
|
||||||
|
if self.enable_lfu:
|
||||||
|
self.lfu = fu_class(out_channels // 2, out_channels // 2, groups)
|
||||||
|
self.conv2 = torch.nn.Conv2d(
|
||||||
|
out_channels // 2, out_channels, kernel_size=1, groups=groups, bias=False
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
x = self.downsample(x)
|
||||||
|
x = self.conv1(x)
|
||||||
|
output = self.fu(x)
|
||||||
|
|
||||||
|
if self.enable_lfu:
|
||||||
|
_, c, h, _ = x.shape
|
||||||
|
split_no = 2
|
||||||
|
split_s = h // split_no
|
||||||
|
xs = torch.cat(
|
||||||
|
torch.split(x[:, : c // 4], split_s, dim=-2), dim=1
|
||||||
|
).contiguous()
|
||||||
|
xs = torch.cat(torch.split(xs, split_s, dim=-1), dim=1).contiguous()
|
||||||
|
xs = self.lfu(xs)
|
||||||
|
xs = xs.repeat(1, 1, split_no, split_no).contiguous()
|
||||||
|
else:
|
||||||
|
xs = 0
|
||||||
|
|
||||||
|
output = self.conv2(x + output + xs)
|
||||||
|
|
||||||
|
return output
|
||||||
|
|
||||||
|
|
||||||
|
class FFC(nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
in_channels,
|
||||||
|
out_channels,
|
||||||
|
kernel_size,
|
||||||
|
ratio_gin,
|
||||||
|
ratio_gout,
|
||||||
|
stride=1,
|
||||||
|
padding=0,
|
||||||
|
dilation=1,
|
||||||
|
groups=1,
|
||||||
|
bias=False,
|
||||||
|
enable_lfu=True,
|
||||||
|
padding_type="reflect",
|
||||||
|
gated=False,
|
||||||
|
**spectral_kwargs,
|
||||||
|
):
|
||||||
|
super(FFC, self).__init__()
|
||||||
|
|
||||||
|
assert stride == 1 or stride == 2, "Stride should be 1 or 2."
|
||||||
|
self.stride = stride
|
||||||
|
|
||||||
|
in_cg = int(in_channels * ratio_gin)
|
||||||
|
in_cl = in_channels - in_cg
|
||||||
|
out_cg = int(out_channels * ratio_gout)
|
||||||
|
out_cl = out_channels - out_cg
|
||||||
|
# groups_g = 1 if groups == 1 else int(groups * ratio_gout)
|
||||||
|
# groups_l = 1 if groups == 1 else groups - groups_g
|
||||||
|
|
||||||
|
self.ratio_gin = ratio_gin
|
||||||
|
self.ratio_gout = ratio_gout
|
||||||
|
self.global_in_num = in_cg
|
||||||
|
|
||||||
|
module = nn.Identity if in_cl == 0 or out_cl == 0 else nn.Conv2d
|
||||||
|
self.convl2l = module(
|
||||||
|
in_cl,
|
||||||
|
out_cl,
|
||||||
|
kernel_size,
|
||||||
|
stride,
|
||||||
|
padding,
|
||||||
|
dilation,
|
||||||
|
groups,
|
||||||
|
bias,
|
||||||
|
padding_mode=padding_type,
|
||||||
|
)
|
||||||
|
module = nn.Identity if in_cl == 0 or out_cg == 0 else nn.Conv2d
|
||||||
|
self.convl2g = module(
|
||||||
|
in_cl,
|
||||||
|
out_cg,
|
||||||
|
kernel_size,
|
||||||
|
stride,
|
||||||
|
padding,
|
||||||
|
dilation,
|
||||||
|
groups,
|
||||||
|
bias,
|
||||||
|
padding_mode=padding_type,
|
||||||
|
)
|
||||||
|
module = nn.Identity if in_cg == 0 or out_cl == 0 else nn.Conv2d
|
||||||
|
self.convg2l = module(
|
||||||
|
in_cg,
|
||||||
|
out_cl,
|
||||||
|
kernel_size,
|
||||||
|
stride,
|
||||||
|
padding,
|
||||||
|
dilation,
|
||||||
|
groups,
|
||||||
|
bias,
|
||||||
|
padding_mode=padding_type,
|
||||||
|
)
|
||||||
|
module = nn.Identity if in_cg == 0 or out_cg == 0 else SpectralTransform
|
||||||
|
self.convg2g = module(
|
||||||
|
in_cg,
|
||||||
|
out_cg,
|
||||||
|
stride,
|
||||||
|
1 if groups == 1 else groups // 2,
|
||||||
|
enable_lfu,
|
||||||
|
**spectral_kwargs,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.gated = gated
|
||||||
|
module = (
|
||||||
|
nn.Identity if in_cg == 0 or out_cl == 0 or not self.gated else nn.Conv2d
|
||||||
|
)
|
||||||
|
self.gate = module(in_channels, 2, 1)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
x_l, x_g = x if type(x) is tuple else (x, 0)
|
||||||
|
out_xl, out_xg = 0, 0
|
||||||
|
|
||||||
|
if self.gated:
|
||||||
|
total_input_parts = [x_l]
|
||||||
|
if torch.is_tensor(x_g):
|
||||||
|
total_input_parts.append(x_g)
|
||||||
|
total_input = torch.cat(total_input_parts, dim=1)
|
||||||
|
|
||||||
|
gates = torch.sigmoid(self.gate(total_input))
|
||||||
|
g2l_gate, l2g_gate = gates.chunk(2, dim=1)
|
||||||
|
else:
|
||||||
|
g2l_gate, l2g_gate = 1, 1
|
||||||
|
|
||||||
|
if self.ratio_gout != 1:
|
||||||
|
out_xl = self.convl2l(x_l) + self.convg2l(x_g) * g2l_gate
|
||||||
|
if self.ratio_gout != 0:
|
||||||
|
out_xg = self.convl2g(x_l) * l2g_gate + self.convg2g(x_g)
|
||||||
|
|
||||||
|
return out_xl, out_xg
|
||||||
|
|
||||||
|
|
||||||
|
class FFC_BN_ACT(nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
in_channels,
|
||||||
|
out_channels,
|
||||||
|
kernel_size,
|
||||||
|
ratio_gin,
|
||||||
|
ratio_gout,
|
||||||
|
stride=1,
|
||||||
|
padding=0,
|
||||||
|
dilation=1,
|
||||||
|
groups=1,
|
||||||
|
bias=False,
|
||||||
|
norm_layer=nn.BatchNorm2d,
|
||||||
|
activation_layer=nn.Identity,
|
||||||
|
padding_type="reflect",
|
||||||
|
enable_lfu=True,
|
||||||
|
**kwargs,
|
||||||
|
):
|
||||||
|
super(FFC_BN_ACT, self).__init__()
|
||||||
|
self.ffc = FFC(
|
||||||
|
in_channels,
|
||||||
|
out_channels,
|
||||||
|
kernel_size,
|
||||||
|
ratio_gin,
|
||||||
|
ratio_gout,
|
||||||
|
stride,
|
||||||
|
padding,
|
||||||
|
dilation,
|
||||||
|
groups,
|
||||||
|
bias,
|
||||||
|
enable_lfu,
|
||||||
|
padding_type=padding_type,
|
||||||
|
**kwargs,
|
||||||
|
)
|
||||||
|
lnorm = nn.Identity if ratio_gout == 1 else norm_layer
|
||||||
|
gnorm = nn.Identity if ratio_gout == 0 else norm_layer
|
||||||
|
global_channels = int(out_channels * ratio_gout)
|
||||||
|
self.bn_l = lnorm(out_channels - global_channels)
|
||||||
|
self.bn_g = gnorm(global_channels)
|
||||||
|
|
||||||
|
lact = nn.Identity if ratio_gout == 1 else activation_layer
|
||||||
|
gact = nn.Identity if ratio_gout == 0 else activation_layer
|
||||||
|
self.act_l = lact(inplace=True)
|
||||||
|
self.act_g = gact(inplace=True)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
x_l, x_g = self.ffc(x)
|
||||||
|
x_l = self.act_l(self.bn_l(x_l))
|
||||||
|
x_g = self.act_g(self.bn_g(x_g))
|
||||||
|
return x_l, x_g
|
||||||
|
|
||||||
|
|
||||||
|
class FFCResnetBlock(nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
dim,
|
||||||
|
padding_type,
|
||||||
|
norm_layer,
|
||||||
|
activation_layer=nn.ReLU,
|
||||||
|
dilation=1,
|
||||||
|
spatial_transform_kwargs=None,
|
||||||
|
inline=False,
|
||||||
|
**conv_kwargs,
|
||||||
|
):
|
||||||
|
super().__init__()
|
||||||
|
self.conv1 = FFC_BN_ACT(
|
||||||
|
dim,
|
||||||
|
dim,
|
||||||
|
kernel_size=3,
|
||||||
|
padding=dilation,
|
||||||
|
dilation=dilation,
|
||||||
|
norm_layer=norm_layer,
|
||||||
|
activation_layer=activation_layer,
|
||||||
|
padding_type=padding_type,
|
||||||
|
**conv_kwargs,
|
||||||
|
)
|
||||||
|
self.conv2 = FFC_BN_ACT(
|
||||||
|
dim,
|
||||||
|
dim,
|
||||||
|
kernel_size=3,
|
||||||
|
padding=dilation,
|
||||||
|
dilation=dilation,
|
||||||
|
norm_layer=norm_layer,
|
||||||
|
activation_layer=activation_layer,
|
||||||
|
padding_type=padding_type,
|
||||||
|
**conv_kwargs,
|
||||||
|
)
|
||||||
|
if spatial_transform_kwargs is not None:
|
||||||
|
self.conv1 = LearnableSpatialTransformWrapper(
|
||||||
|
self.conv1, **spatial_transform_kwargs
|
||||||
|
)
|
||||||
|
self.conv2 = LearnableSpatialTransformWrapper(
|
||||||
|
self.conv2, **spatial_transform_kwargs
|
||||||
|
)
|
||||||
|
self.inline = inline
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
if self.inline:
|
||||||
|
x_l, x_g = (
|
||||||
|
x[:, : -self.conv1.ffc.global_in_num],
|
||||||
|
x[:, -self.conv1.ffc.global_in_num :],
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
x_l, x_g = x if type(x) is tuple else (x, 0)
|
||||||
|
|
||||||
|
id_l, id_g = x_l, x_g
|
||||||
|
|
||||||
|
x_l, x_g = self.conv1((x_l, x_g))
|
||||||
|
x_l, x_g = self.conv2((x_l, x_g))
|
||||||
|
|
||||||
|
x_l, x_g = id_l + x_l, id_g + x_g
|
||||||
|
out = x_l, x_g
|
||||||
|
if self.inline:
|
||||||
|
out = torch.cat(out, dim=1)
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
class ConcatTupleLayer(nn.Module):
|
||||||
|
def forward(self, x):
|
||||||
|
assert isinstance(x, tuple)
|
||||||
|
x_l, x_g = x
|
||||||
|
assert torch.is_tensor(x_l) or torch.is_tensor(x_g)
|
||||||
|
if not torch.is_tensor(x_g):
|
||||||
|
return x_l
|
||||||
|
return torch.cat(x, dim=1)
|
||||||
|
|
||||||
|
|
||||||
|
class FFCResNetGenerator(nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
input_nc,
|
||||||
|
output_nc,
|
||||||
|
ngf=64,
|
||||||
|
n_downsampling=3,
|
||||||
|
n_blocks=18,
|
||||||
|
norm_layer=nn.BatchNorm2d,
|
||||||
|
padding_type="reflect",
|
||||||
|
activation_layer=nn.ReLU,
|
||||||
|
up_norm_layer=nn.BatchNorm2d,
|
||||||
|
up_activation=nn.ReLU(True),
|
||||||
|
init_conv_kwargs={},
|
||||||
|
downsample_conv_kwargs={},
|
||||||
|
resnet_conv_kwargs={},
|
||||||
|
spatial_transform_layers=None,
|
||||||
|
spatial_transform_kwargs={},
|
||||||
|
max_features=1024,
|
||||||
|
out_ffc=False,
|
||||||
|
out_ffc_kwargs={},
|
||||||
|
):
|
||||||
|
assert n_blocks >= 0
|
||||||
|
super().__init__()
|
||||||
|
"""
|
||||||
|
init_conv_kwargs = {'ratio_gin': 0, 'ratio_gout': 0, 'enable_lfu': False}
|
||||||
|
downsample_conv_kwargs = {'ratio_gin': '${generator.init_conv_kwargs.ratio_gout}', 'ratio_gout': '${generator.downsample_conv_kwargs.ratio_gin}', 'enable_lfu': False}
|
||||||
|
resnet_conv_kwargs = {'ratio_gin': 0.75, 'ratio_gout': '${generator.resnet_conv_kwargs.ratio_gin}', 'enable_lfu': False}
|
||||||
|
spatial_transform_kwargs = {}
|
||||||
|
out_ffc_kwargs = {}
|
||||||
|
"""
|
||||||
|
"""
|
||||||
|
print(input_nc, output_nc, ngf, n_downsampling, n_blocks, norm_layer,
|
||||||
|
padding_type, activation_layer,
|
||||||
|
up_norm_layer, up_activation,
|
||||||
|
spatial_transform_layers,
|
||||||
|
add_out_act, max_features, out_ffc, file=sys.stderr)
|
||||||
|
|
||||||
|
4 3 64 3 18 <class 'torch.nn.modules.batchnorm.BatchNorm2d'>
|
||||||
|
reflect <class 'torch.nn.modules.activation.ReLU'>
|
||||||
|
<class 'torch.nn.modules.batchnorm.BatchNorm2d'>
|
||||||
|
ReLU(inplace=True)
|
||||||
|
None sigmoid 1024 False
|
||||||
|
"""
|
||||||
|
init_conv_kwargs = {"ratio_gin": 0, "ratio_gout": 0, "enable_lfu": False}
|
||||||
|
downsample_conv_kwargs = {"ratio_gin": 0, "ratio_gout": 0, "enable_lfu": False}
|
||||||
|
resnet_conv_kwargs = {
|
||||||
|
"ratio_gin": 0.75,
|
||||||
|
"ratio_gout": 0.75,
|
||||||
|
"enable_lfu": False,
|
||||||
|
}
|
||||||
|
spatial_transform_kwargs = {}
|
||||||
|
out_ffc_kwargs = {}
|
||||||
|
|
||||||
|
model = [
|
||||||
|
nn.ReflectionPad2d(3),
|
||||||
|
FFC_BN_ACT(
|
||||||
|
input_nc,
|
||||||
|
ngf,
|
||||||
|
kernel_size=7,
|
||||||
|
padding=0,
|
||||||
|
norm_layer=norm_layer,
|
||||||
|
activation_layer=activation_layer,
|
||||||
|
**init_conv_kwargs,
|
||||||
|
),
|
||||||
|
]
|
||||||
|
|
||||||
|
### downsample
|
||||||
|
for i in range(n_downsampling):
|
||||||
|
mult = 2**i
|
||||||
|
if i == n_downsampling - 1:
|
||||||
|
cur_conv_kwargs = dict(downsample_conv_kwargs)
|
||||||
|
cur_conv_kwargs["ratio_gout"] = resnet_conv_kwargs.get("ratio_gin", 0)
|
||||||
|
else:
|
||||||
|
cur_conv_kwargs = downsample_conv_kwargs
|
||||||
|
model += [
|
||||||
|
FFC_BN_ACT(
|
||||||
|
min(max_features, ngf * mult),
|
||||||
|
min(max_features, ngf * mult * 2),
|
||||||
|
kernel_size=3,
|
||||||
|
stride=2,
|
||||||
|
padding=1,
|
||||||
|
norm_layer=norm_layer,
|
||||||
|
activation_layer=activation_layer,
|
||||||
|
**cur_conv_kwargs,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
|
||||||
|
mult = 2**n_downsampling
|
||||||
|
feats_num_bottleneck = min(max_features, ngf * mult)
|
||||||
|
|
||||||
|
### resnet blocks
|
||||||
|
for i in range(n_blocks):
|
||||||
|
cur_resblock = FFCResnetBlock(
|
||||||
|
feats_num_bottleneck,
|
||||||
|
padding_type=padding_type,
|
||||||
|
activation_layer=activation_layer,
|
||||||
|
norm_layer=norm_layer,
|
||||||
|
**resnet_conv_kwargs,
|
||||||
|
)
|
||||||
|
if spatial_transform_layers is not None and i in spatial_transform_layers:
|
||||||
|
cur_resblock = LearnableSpatialTransformWrapper(
|
||||||
|
cur_resblock, **spatial_transform_kwargs
|
||||||
|
)
|
||||||
|
model += [cur_resblock]
|
||||||
|
|
||||||
|
model += [ConcatTupleLayer()]
|
||||||
|
|
||||||
|
### upsample
|
||||||
|
for i in range(n_downsampling):
|
||||||
|
mult = 2 ** (n_downsampling - i)
|
||||||
|
model += [
|
||||||
|
nn.ConvTranspose2d(
|
||||||
|
min(max_features, ngf * mult),
|
||||||
|
min(max_features, int(ngf * mult / 2)),
|
||||||
|
kernel_size=3,
|
||||||
|
stride=2,
|
||||||
|
padding=1,
|
||||||
|
output_padding=1,
|
||||||
|
),
|
||||||
|
up_norm_layer(min(max_features, int(ngf * mult / 2))),
|
||||||
|
up_activation,
|
||||||
|
]
|
||||||
|
|
||||||
|
if out_ffc:
|
||||||
|
model += [
|
||||||
|
FFCResnetBlock(
|
||||||
|
ngf,
|
||||||
|
padding_type=padding_type,
|
||||||
|
activation_layer=activation_layer,
|
||||||
|
norm_layer=norm_layer,
|
||||||
|
inline=True,
|
||||||
|
**out_ffc_kwargs,
|
||||||
|
)
|
||||||
|
]
|
||||||
|
|
||||||
|
model += [
|
||||||
|
nn.ReflectionPad2d(3),
|
||||||
|
nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0),
|
||||||
|
]
|
||||||
|
model.append(nn.Sigmoid())
|
||||||
|
self.model = nn.Sequential(*model)
|
||||||
|
|
||||||
|
def forward(self, image, mask):
|
||||||
|
return self.model(torch.cat([image, mask], dim=1))
|
||||||
|
|
||||||
|
|
||||||
|
class LaMa(nn.Module):
|
||||||
|
def __init__(self, state_dict) -> None:
|
||||||
|
super(LaMa, self).__init__()
|
||||||
|
self.model_arch = "LaMa"
|
||||||
|
self.sub_type = "Inpaint"
|
||||||
|
self.in_nc = 4
|
||||||
|
self.out_nc = 3
|
||||||
|
self.scale = 1
|
||||||
|
|
||||||
|
self.min_size = None
|
||||||
|
self.pad_mod = 8
|
||||||
|
self.pad_to_square = False
|
||||||
|
|
||||||
|
self.model = FFCResNetGenerator(self.in_nc, self.out_nc)
|
||||||
|
self.state = {
|
||||||
|
k.replace("generator.model", "model.model"): v
|
||||||
|
for k, v in state_dict.items()
|
||||||
|
}
|
||||||
|
|
||||||
|
self.supports_fp16 = False
|
||||||
|
self.support_bf16 = True
|
||||||
|
|
||||||
|
self.load_state_dict(self.state, strict=False)
|
||||||
|
|
||||||
|
def forward(self, img, mask):
|
||||||
|
masked_img = img * (1 - mask)
|
||||||
|
inpainted_mask = mask * self.model.forward(masked_img, mask)
|
||||||
|
result = inpainted_mask + (1 - mask) * img
|
||||||
|
return result
|
110
ldm_patched/pfn/architecture/OmniSR/ChannelAttention.py
Normal file
110
ldm_patched/pfn/architecture/OmniSR/ChannelAttention.py
Normal file
@ -0,0 +1,110 @@
|
|||||||
|
import math
|
||||||
|
|
||||||
|
import torch.nn as nn
|
||||||
|
|
||||||
|
|
||||||
|
class CA_layer(nn.Module):
|
||||||
|
def __init__(self, channel, reduction=16):
|
||||||
|
super(CA_layer, self).__init__()
|
||||||
|
# global average pooling
|
||||||
|
self.gap = nn.AdaptiveAvgPool2d(1)
|
||||||
|
self.fc = nn.Sequential(
|
||||||
|
nn.Conv2d(channel, channel // reduction, kernel_size=(1, 1), bias=False),
|
||||||
|
nn.GELU(),
|
||||||
|
nn.Conv2d(channel // reduction, channel, kernel_size=(1, 1), bias=False),
|
||||||
|
# nn.Sigmoid()
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
y = self.fc(self.gap(x))
|
||||||
|
return x * y.expand_as(x)
|
||||||
|
|
||||||
|
|
||||||
|
class Simple_CA_layer(nn.Module):
|
||||||
|
def __init__(self, channel):
|
||||||
|
super(Simple_CA_layer, self).__init__()
|
||||||
|
self.gap = nn.AdaptiveAvgPool2d(1)
|
||||||
|
self.fc = nn.Conv2d(
|
||||||
|
in_channels=channel,
|
||||||
|
out_channels=channel,
|
||||||
|
kernel_size=1,
|
||||||
|
padding=0,
|
||||||
|
stride=1,
|
||||||
|
groups=1,
|
||||||
|
bias=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
return x * self.fc(self.gap(x))
|
||||||
|
|
||||||
|
|
||||||
|
class ECA_layer(nn.Module):
|
||||||
|
"""Constructs a ECA module.
|
||||||
|
Args:
|
||||||
|
channel: Number of channels of the input feature map
|
||||||
|
k_size: Adaptive selection of kernel size
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, channel):
|
||||||
|
super(ECA_layer, self).__init__()
|
||||||
|
|
||||||
|
b = 1
|
||||||
|
gamma = 2
|
||||||
|
k_size = int(abs(math.log(channel, 2) + b) / gamma)
|
||||||
|
k_size = k_size if k_size % 2 else k_size + 1
|
||||||
|
self.avg_pool = nn.AdaptiveAvgPool2d(1)
|
||||||
|
self.conv = nn.Conv1d(
|
||||||
|
1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False
|
||||||
|
)
|
||||||
|
# self.sigmoid = nn.Sigmoid()
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
# x: input features with shape [b, c, h, w]
|
||||||
|
# b, c, h, w = x.size()
|
||||||
|
|
||||||
|
# feature descriptor on the global spatial information
|
||||||
|
y = self.avg_pool(x)
|
||||||
|
|
||||||
|
# Two different branches of ECA module
|
||||||
|
y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)
|
||||||
|
|
||||||
|
# Multi-scale information fusion
|
||||||
|
# y = self.sigmoid(y)
|
||||||
|
|
||||||
|
return x * y.expand_as(x)
|
||||||
|
|
||||||
|
|
||||||
|
class ECA_MaxPool_layer(nn.Module):
|
||||||
|
"""Constructs a ECA module.
|
||||||
|
Args:
|
||||||
|
channel: Number of channels of the input feature map
|
||||||
|
k_size: Adaptive selection of kernel size
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, channel):
|
||||||
|
super(ECA_MaxPool_layer, self).__init__()
|
||||||
|
|
||||||
|
b = 1
|
||||||
|
gamma = 2
|
||||||
|
k_size = int(abs(math.log(channel, 2) + b) / gamma)
|
||||||
|
k_size = k_size if k_size % 2 else k_size + 1
|
||||||
|
self.max_pool = nn.AdaptiveMaxPool2d(1)
|
||||||
|
self.conv = nn.Conv1d(
|
||||||
|
1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False
|
||||||
|
)
|
||||||
|
# self.sigmoid = nn.Sigmoid()
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
# x: input features with shape [b, c, h, w]
|
||||||
|
# b, c, h, w = x.size()
|
||||||
|
|
||||||
|
# feature descriptor on the global spatial information
|
||||||
|
y = self.max_pool(x)
|
||||||
|
|
||||||
|
# Two different branches of ECA module
|
||||||
|
y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)
|
||||||
|
|
||||||
|
# Multi-scale information fusion
|
||||||
|
# y = self.sigmoid(y)
|
||||||
|
|
||||||
|
return x * y.expand_as(x)
|
201
ldm_patched/pfn/architecture/OmniSR/LICENSE
Normal file
201
ldm_patched/pfn/architecture/OmniSR/LICENSE
Normal file
@ -0,0 +1,201 @@
|
|||||||
|
Apache License
|
||||||
|
Version 2.0, January 2004
|
||||||
|
http://www.apache.org/licenses/
|
||||||
|
|
||||||
|
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||||
|
|
||||||
|
1. Definitions.
|
||||||
|
|
||||||
|
"License" shall mean the terms and conditions for use, reproduction,
|
||||||
|
and distribution as defined by Sections 1 through 9 of this document.
|
||||||
|
|
||||||
|
"Licensor" shall mean the copyright owner or entity authorized by
|
||||||
|
the copyright owner that is granting the License.
|
||||||
|
|
||||||
|
"Legal Entity" shall mean the union of the acting entity and all
|
||||||
|
other entities that control, are controlled by, or are under common
|
||||||
|
control with that entity. For the purposes of this definition,
|
||||||
|
"control" means (i) the power, direct or indirect, to cause the
|
||||||
|
direction or management of such entity, whether by contract or
|
||||||
|
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||||
|
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||||
|
|
||||||
|
"You" (or "Your") shall mean an individual or Legal Entity
|
||||||
|
exercising permissions granted by this License.
|
||||||
|
|
||||||
|
"Source" form shall mean the preferred form for making modifications,
|
||||||
|
including but not limited to software source code, documentation
|
||||||
|
source, and configuration files.
|
||||||
|
|
||||||
|
"Object" form shall mean any form resulting from mechanical
|
||||||
|
transformation or translation of a Source form, including but
|
||||||
|
not limited to compiled object code, generated documentation,
|
||||||
|
and conversions to other media types.
|
||||||
|
|
||||||
|
"Work" shall mean the work of authorship, whether in Source or
|
||||||
|
Object form, made available under the License, as indicated by a
|
||||||
|
copyright notice that is included in or attached to the work
|
||||||
|
(an example is provided in the Appendix below).
|
||||||
|
|
||||||
|
"Derivative Works" shall mean any work, whether in Source or Object
|
||||||
|
form, that is based on (or derived from) the Work and for which the
|
||||||
|
editorial revisions, annotations, elaborations, or other modifications
|
||||||
|
represent, as a whole, an original work of authorship. For the purposes
|
||||||
|
of this License, Derivative Works shall not include works that remain
|
||||||
|
separable from, or merely link (or bind by name) to the interfaces of,
|
||||||
|
the Work and Derivative Works thereof.
|
||||||
|
|
||||||
|
"Contribution" shall mean any work of authorship, including
|
||||||
|
the original version of the Work and any modifications or additions
|
||||||
|
to that Work or Derivative Works thereof, that is intentionally
|
||||||
|
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||||
|
or by an individual or Legal Entity authorized to submit on behalf of
|
||||||
|
the copyright owner. For the purposes of this definition, "submitted"
|
||||||
|
means any form of electronic, verbal, or written communication sent
|
||||||
|
to the Licensor or its representatives, including but not limited to
|
||||||
|
communication on electronic mailing lists, source code control systems,
|
||||||
|
and issue tracking systems that are managed by, or on behalf of, the
|
||||||
|
Licensor for the purpose of discussing and improving the Work, but
|
||||||
|
excluding communication that is conspicuously marked or otherwise
|
||||||
|
designated in writing by the copyright owner as "Not a Contribution."
|
||||||
|
|
||||||
|
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||||
|
on behalf of whom a Contribution has been received by Licensor and
|
||||||
|
subsequently incorporated within the Work.
|
||||||
|
|
||||||
|
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||||
|
this License, each Contributor hereby grants to You a perpetual,
|
||||||
|
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||||
|
copyright license to reproduce, prepare Derivative Works of,
|
||||||
|
publicly display, publicly perform, sublicense, and distribute the
|
||||||
|
Work and such Derivative Works in Source or Object form.
|
||||||
|
|
||||||
|
3. Grant of Patent License. Subject to the terms and conditions of
|
||||||
|
this License, each Contributor hereby grants to You a perpetual,
|
||||||
|
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||||
|
(except as stated in this section) patent license to make, have made,
|
||||||
|
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||||
|
where such license applies only to those patent claims licensable
|
||||||
|
by such Contributor that are necessarily infringed by their
|
||||||
|
Contribution(s) alone or by combination of their Contribution(s)
|
||||||
|
with the Work to which such Contribution(s) was submitted. If You
|
||||||
|
institute patent litigation against any entity (including a
|
||||||
|
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||||
|
or a Contribution incorporated within the Work constitutes direct
|
||||||
|
or contributory patent infringement, then any patent licenses
|
||||||
|
granted to You under this License for that Work shall terminate
|
||||||
|
as of the date such litigation is filed.
|
||||||
|
|
||||||
|
4. Redistribution. You may reproduce and distribute copies of the
|
||||||
|
Work or Derivative Works thereof in any medium, with or without
|
||||||
|
modifications, and in Source or Object form, provided that You
|
||||||
|
meet the following conditions:
|
||||||
|
|
||||||
|
(a) You must give any other recipients of the Work or
|
||||||
|
Derivative Works a copy of this License; and
|
||||||
|
|
||||||
|
(b) You must cause any modified files to carry prominent notices
|
||||||
|
stating that You changed the files; and
|
||||||
|
|
||||||
|
(c) You must retain, in the Source form of any Derivative Works
|
||||||
|
that You distribute, all copyright, patent, trademark, and
|
||||||
|
attribution notices from the Source form of the Work,
|
||||||
|
excluding those notices that do not pertain to any part of
|
||||||
|
the Derivative Works; and
|
||||||
|
|
||||||
|
(d) If the Work includes a "NOTICE" text file as part of its
|
||||||
|
distribution, then any Derivative Works that You distribute must
|
||||||
|
include a readable copy of the attribution notices contained
|
||||||
|
within such NOTICE file, excluding those notices that do not
|
||||||
|
pertain to any part of the Derivative Works, in at least one
|
||||||
|
of the following places: within a NOTICE text file distributed
|
||||||
|
as part of the Derivative Works; within the Source form or
|
||||||
|
documentation, if provided along with the Derivative Works; or,
|
||||||
|
within a display generated by the Derivative Works, if and
|
||||||
|
wherever such third-party notices normally appear. The contents
|
||||||
|
of the NOTICE file are for informational purposes only and
|
||||||
|
do not modify the License. You may add Your own attribution
|
||||||
|
notices within Derivative Works that You distribute, alongside
|
||||||
|
or as an addendum to the NOTICE text from the Work, provided
|
||||||
|
that such additional attribution notices cannot be construed
|
||||||
|
as modifying the License.
|
||||||
|
|
||||||
|
You may add Your own copyright statement to Your modifications and
|
||||||
|
may provide additional or different license terms and conditions
|
||||||
|
for use, reproduction, or distribution of Your modifications, or
|
||||||
|
for any such Derivative Works as a whole, provided Your use,
|
||||||
|
reproduction, and distribution of the Work otherwise complies with
|
||||||
|
the conditions stated in this License.
|
||||||
|
|
||||||
|
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||||
|
any Contribution intentionally submitted for inclusion in the Work
|
||||||
|
by You to the Licensor shall be under the terms and conditions of
|
||||||
|
this License, without any additional terms or conditions.
|
||||||
|
Notwithstanding the above, nothing herein shall supersede or modify
|
||||||
|
the terms of any separate license agreement you may have executed
|
||||||
|
with Licensor regarding such Contributions.
|
||||||
|
|
||||||
|
6. Trademarks. This License does not grant permission to use the trade
|
||||||
|
names, trademarks, service marks, or product names of the Licensor,
|
||||||
|
except as required for reasonable and customary use in describing the
|
||||||
|
origin of the Work and reproducing the content of the NOTICE file.
|
||||||
|
|
||||||
|
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||||
|
agreed to in writing, Licensor provides the Work (and each
|
||||||
|
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||||
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||||
|
implied, including, without limitation, any warranties or conditions
|
||||||
|
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||||
|
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||||
|
appropriateness of using or redistributing the Work and assume any
|
||||||
|
risks associated with Your exercise of permissions under this License.
|
||||||
|
|
||||||
|
8. Limitation of Liability. In no event and under no legal theory,
|
||||||
|
whether in tort (including negligence), contract, or otherwise,
|
||||||
|
unless required by applicable law (such as deliberate and grossly
|
||||||
|
negligent acts) or agreed to in writing, shall any Contributor be
|
||||||
|
liable to You for damages, including any direct, indirect, special,
|
||||||
|
incidental, or consequential damages of any character arising as a
|
||||||
|
result of this License or out of the use or inability to use the
|
||||||
|
Work (including but not limited to damages for loss of goodwill,
|
||||||
|
work stoppage, computer failure or malfunction, or any and all
|
||||||
|
other commercial damages or losses), even if such Contributor
|
||||||
|
has been advised of the possibility of such damages.
|
||||||
|
|
||||||
|
9. Accepting Warranty or Additional Liability. While redistributing
|
||||||
|
the Work or Derivative Works thereof, You may choose to offer,
|
||||||
|
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||||
|
or other liability obligations and/or rights consistent with this
|
||||||
|
License. However, in accepting such obligations, You may act only
|
||||||
|
on Your own behalf and on Your sole responsibility, not on behalf
|
||||||
|
of any other Contributor, and only if You agree to indemnify,
|
||||||
|
defend, and hold each Contributor harmless for any liability
|
||||||
|
incurred by, or claims asserted against, such Contributor by reason
|
||||||
|
of your accepting any such warranty or additional liability.
|
||||||
|
|
||||||
|
END OF TERMS AND CONDITIONS
|
||||||
|
|
||||||
|
APPENDIX: How to apply the Apache License to your work.
|
||||||
|
|
||||||
|
To apply the Apache License to your work, attach the following
|
||||||
|
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||||
|
replaced with your own identifying information. (Don't include
|
||||||
|
the brackets!) The text should be enclosed in the appropriate
|
||||||
|
comment syntax for the file format. We also recommend that a
|
||||||
|
file or class name and description of purpose be included on the
|
||||||
|
same "printed page" as the copyright notice for easier
|
||||||
|
identification within third-party archives.
|
||||||
|
|
||||||
|
Copyright [yyyy] [name of copyright owner]
|
||||||
|
|
||||||
|
Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
you may not use this file except in compliance with the License.
|
||||||
|
You may obtain a copy of the License at
|
||||||
|
|
||||||
|
http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
|
||||||
|
Unless required by applicable law or agreed to in writing, software
|
||||||
|
distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
See the License for the specific language governing permissions and
|
||||||
|
limitations under the License.
|
577
ldm_patched/pfn/architecture/OmniSR/OSA.py
Normal file
577
ldm_patched/pfn/architecture/OmniSR/OSA.py
Normal file
@ -0,0 +1,577 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
# -*- coding:utf-8 -*-
|
||||||
|
#############################################################
|
||||||
|
# File: OSA.py
|
||||||
|
# Created Date: Tuesday April 28th 2022
|
||||||
|
# Author: Chen Xuanhong
|
||||||
|
# Email: chenxuanhongzju@outlook.com
|
||||||
|
# Last Modified: Sunday, 23rd April 2023 3:07:42 pm
|
||||||
|
# Modified By: Chen Xuanhong
|
||||||
|
# Copyright (c) 2020 Shanghai Jiao Tong University
|
||||||
|
#############################################################
|
||||||
|
|
||||||
|
import torch
|
||||||
|
import torch.nn.functional as F
|
||||||
|
from einops import rearrange, repeat
|
||||||
|
from einops.layers.torch import Rearrange, Reduce
|
||||||
|
from torch import einsum, nn
|
||||||
|
|
||||||
|
from .layernorm import LayerNorm2d
|
||||||
|
|
||||||
|
# helpers
|
||||||
|
|
||||||
|
|
||||||
|
def exists(val):
|
||||||
|
return val is not None
|
||||||
|
|
||||||
|
|
||||||
|
def default(val, d):
|
||||||
|
return val if exists(val) else d
|
||||||
|
|
||||||
|
|
||||||
|
def cast_tuple(val, length=1):
|
||||||
|
return val if isinstance(val, tuple) else ((val,) * length)
|
||||||
|
|
||||||
|
|
||||||
|
# helper classes
|
||||||
|
|
||||||
|
|
||||||
|
class PreNormResidual(nn.Module):
|
||||||
|
def __init__(self, dim, fn):
|
||||||
|
super().__init__()
|
||||||
|
self.norm = nn.LayerNorm(dim)
|
||||||
|
self.fn = fn
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
return self.fn(self.norm(x)) + x
|
||||||
|
|
||||||
|
|
||||||
|
class Conv_PreNormResidual(nn.Module):
|
||||||
|
def __init__(self, dim, fn):
|
||||||
|
super().__init__()
|
||||||
|
self.norm = LayerNorm2d(dim)
|
||||||
|
self.fn = fn
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
return self.fn(self.norm(x)) + x
|
||||||
|
|
||||||
|
|
||||||
|
class FeedForward(nn.Module):
|
||||||
|
def __init__(self, dim, mult=2, dropout=0.0):
|
||||||
|
super().__init__()
|
||||||
|
inner_dim = int(dim * mult)
|
||||||
|
self.net = nn.Sequential(
|
||||||
|
nn.Linear(dim, inner_dim),
|
||||||
|
nn.GELU(),
|
||||||
|
nn.Dropout(dropout),
|
||||||
|
nn.Linear(inner_dim, dim),
|
||||||
|
nn.Dropout(dropout),
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
return self.net(x)
|
||||||
|
|
||||||
|
|
||||||
|
class Conv_FeedForward(nn.Module):
|
||||||
|
def __init__(self, dim, mult=2, dropout=0.0):
|
||||||
|
super().__init__()
|
||||||
|
inner_dim = int(dim * mult)
|
||||||
|
self.net = nn.Sequential(
|
||||||
|
nn.Conv2d(dim, inner_dim, 1, 1, 0),
|
||||||
|
nn.GELU(),
|
||||||
|
nn.Dropout(dropout),
|
||||||
|
nn.Conv2d(inner_dim, dim, 1, 1, 0),
|
||||||
|
nn.Dropout(dropout),
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
return self.net(x)
|
||||||
|
|
||||||
|
|
||||||
|
class Gated_Conv_FeedForward(nn.Module):
|
||||||
|
def __init__(self, dim, mult=1, bias=False, dropout=0.0):
|
||||||
|
super().__init__()
|
||||||
|
|
||||||
|
hidden_features = int(dim * mult)
|
||||||
|
|
||||||
|
self.project_in = nn.Conv2d(dim, hidden_features * 2, kernel_size=1, bias=bias)
|
||||||
|
|
||||||
|
self.dwconv = nn.Conv2d(
|
||||||
|
hidden_features * 2,
|
||||||
|
hidden_features * 2,
|
||||||
|
kernel_size=3,
|
||||||
|
stride=1,
|
||||||
|
padding=1,
|
||||||
|
groups=hidden_features * 2,
|
||||||
|
bias=bias,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.project_out = nn.Conv2d(hidden_features, dim, kernel_size=1, bias=bias)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
x = self.project_in(x)
|
||||||
|
x1, x2 = self.dwconv(x).chunk(2, dim=1)
|
||||||
|
x = F.gelu(x1) * x2
|
||||||
|
x = self.project_out(x)
|
||||||
|
return x
|
||||||
|
|
||||||
|
|
||||||
|
# MBConv
|
||||||
|
|
||||||
|
|
||||||
|
class SqueezeExcitation(nn.Module):
|
||||||
|
def __init__(self, dim, shrinkage_rate=0.25):
|
||||||
|
super().__init__()
|
||||||
|
hidden_dim = int(dim * shrinkage_rate)
|
||||||
|
|
||||||
|
self.gate = nn.Sequential(
|
||||||
|
Reduce("b c h w -> b c", "mean"),
|
||||||
|
nn.Linear(dim, hidden_dim, bias=False),
|
||||||
|
nn.SiLU(),
|
||||||
|
nn.Linear(hidden_dim, dim, bias=False),
|
||||||
|
nn.Sigmoid(),
|
||||||
|
Rearrange("b c -> b c 1 1"),
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
return x * self.gate(x)
|
||||||
|
|
||||||
|
|
||||||
|
class MBConvResidual(nn.Module):
|
||||||
|
def __init__(self, fn, dropout=0.0):
|
||||||
|
super().__init__()
|
||||||
|
self.fn = fn
|
||||||
|
self.dropsample = Dropsample(dropout)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
out = self.fn(x)
|
||||||
|
out = self.dropsample(out)
|
||||||
|
return out + x
|
||||||
|
|
||||||
|
|
||||||
|
class Dropsample(nn.Module):
|
||||||
|
def __init__(self, prob=0):
|
||||||
|
super().__init__()
|
||||||
|
self.prob = prob
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
device = x.device
|
||||||
|
|
||||||
|
if self.prob == 0.0 or (not self.training):
|
||||||
|
return x
|
||||||
|
|
||||||
|
keep_mask = (
|
||||||
|
torch.FloatTensor((x.shape[0], 1, 1, 1), device=device).uniform_()
|
||||||
|
> self.prob
|
||||||
|
)
|
||||||
|
return x * keep_mask / (1 - self.prob)
|
||||||
|
|
||||||
|
|
||||||
|
def MBConv(
|
||||||
|
dim_in, dim_out, *, downsample, expansion_rate=4, shrinkage_rate=0.25, dropout=0.0
|
||||||
|
):
|
||||||
|
hidden_dim = int(expansion_rate * dim_out)
|
||||||
|
stride = 2 if downsample else 1
|
||||||
|
|
||||||
|
net = nn.Sequential(
|
||||||
|
nn.Conv2d(dim_in, hidden_dim, 1),
|
||||||
|
# nn.BatchNorm2d(hidden_dim),
|
||||||
|
nn.GELU(),
|
||||||
|
nn.Conv2d(
|
||||||
|
hidden_dim, hidden_dim, 3, stride=stride, padding=1, groups=hidden_dim
|
||||||
|
),
|
||||||
|
# nn.BatchNorm2d(hidden_dim),
|
||||||
|
nn.GELU(),
|
||||||
|
SqueezeExcitation(hidden_dim, shrinkage_rate=shrinkage_rate),
|
||||||
|
nn.Conv2d(hidden_dim, dim_out, 1),
|
||||||
|
# nn.BatchNorm2d(dim_out)
|
||||||
|
)
|
||||||
|
|
||||||
|
if dim_in == dim_out and not downsample:
|
||||||
|
net = MBConvResidual(net, dropout=dropout)
|
||||||
|
|
||||||
|
return net
|
||||||
|
|
||||||
|
|
||||||
|
# attention related classes
|
||||||
|
class Attention(nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
dim,
|
||||||
|
dim_head=32,
|
||||||
|
dropout=0.0,
|
||||||
|
window_size=7,
|
||||||
|
with_pe=True,
|
||||||
|
):
|
||||||
|
super().__init__()
|
||||||
|
assert (
|
||||||
|
dim % dim_head
|
||||||
|
) == 0, "dimension should be divisible by dimension per head"
|
||||||
|
|
||||||
|
self.heads = dim // dim_head
|
||||||
|
self.scale = dim_head**-0.5
|
||||||
|
self.with_pe = with_pe
|
||||||
|
|
||||||
|
self.to_qkv = nn.Linear(dim, dim * 3, bias=False)
|
||||||
|
|
||||||
|
self.attend = nn.Sequential(nn.Softmax(dim=-1), nn.Dropout(dropout))
|
||||||
|
|
||||||
|
self.to_out = nn.Sequential(
|
||||||
|
nn.Linear(dim, dim, bias=False), nn.Dropout(dropout)
|
||||||
|
)
|
||||||
|
|
||||||
|
# relative positional bias
|
||||||
|
if self.with_pe:
|
||||||
|
self.rel_pos_bias = nn.Embedding((2 * window_size - 1) ** 2, self.heads)
|
||||||
|
|
||||||
|
pos = torch.arange(window_size)
|
||||||
|
grid = torch.stack(torch.meshgrid(pos, pos))
|
||||||
|
grid = rearrange(grid, "c i j -> (i j) c")
|
||||||
|
rel_pos = rearrange(grid, "i ... -> i 1 ...") - rearrange(
|
||||||
|
grid, "j ... -> 1 j ..."
|
||||||
|
)
|
||||||
|
rel_pos += window_size - 1
|
||||||
|
rel_pos_indices = (rel_pos * torch.tensor([2 * window_size - 1, 1])).sum(
|
||||||
|
dim=-1
|
||||||
|
)
|
||||||
|
|
||||||
|
self.register_buffer("rel_pos_indices", rel_pos_indices, persistent=False)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
batch, height, width, window_height, window_width, _, device, h = (
|
||||||
|
*x.shape,
|
||||||
|
x.device,
|
||||||
|
self.heads,
|
||||||
|
)
|
||||||
|
|
||||||
|
# flatten
|
||||||
|
|
||||||
|
x = rearrange(x, "b x y w1 w2 d -> (b x y) (w1 w2) d")
|
||||||
|
|
||||||
|
# project for queries, keys, values
|
||||||
|
|
||||||
|
q, k, v = self.to_qkv(x).chunk(3, dim=-1)
|
||||||
|
|
||||||
|
# split heads
|
||||||
|
|
||||||
|
q, k, v = map(lambda t: rearrange(t, "b n (h d ) -> b h n d", h=h), (q, k, v))
|
||||||
|
|
||||||
|
# scale
|
||||||
|
|
||||||
|
q = q * self.scale
|
||||||
|
|
||||||
|
# sim
|
||||||
|
|
||||||
|
sim = einsum("b h i d, b h j d -> b h i j", q, k)
|
||||||
|
|
||||||
|
# add positional bias
|
||||||
|
if self.with_pe:
|
||||||
|
bias = self.rel_pos_bias(self.rel_pos_indices)
|
||||||
|
sim = sim + rearrange(bias, "i j h -> h i j")
|
||||||
|
|
||||||
|
# attention
|
||||||
|
|
||||||
|
attn = self.attend(sim)
|
||||||
|
|
||||||
|
# aggregate
|
||||||
|
|
||||||
|
out = einsum("b h i j, b h j d -> b h i d", attn, v)
|
||||||
|
|
||||||
|
# merge heads
|
||||||
|
|
||||||
|
out = rearrange(
|
||||||
|
out, "b h (w1 w2) d -> b w1 w2 (h d)", w1=window_height, w2=window_width
|
||||||
|
)
|
||||||
|
|
||||||
|
# combine heads out
|
||||||
|
|
||||||
|
out = self.to_out(out)
|
||||||
|
return rearrange(out, "(b x y) ... -> b x y ...", x=height, y=width)
|
||||||
|
|
||||||
|
|
||||||
|
class Block_Attention(nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
dim,
|
||||||
|
dim_head=32,
|
||||||
|
bias=False,
|
||||||
|
dropout=0.0,
|
||||||
|
window_size=7,
|
||||||
|
with_pe=True,
|
||||||
|
):
|
||||||
|
super().__init__()
|
||||||
|
assert (
|
||||||
|
dim % dim_head
|
||||||
|
) == 0, "dimension should be divisible by dimension per head"
|
||||||
|
|
||||||
|
self.heads = dim // dim_head
|
||||||
|
self.ps = window_size
|
||||||
|
self.scale = dim_head**-0.5
|
||||||
|
self.with_pe = with_pe
|
||||||
|
|
||||||
|
self.qkv = nn.Conv2d(dim, dim * 3, kernel_size=1, bias=bias)
|
||||||
|
self.qkv_dwconv = nn.Conv2d(
|
||||||
|
dim * 3,
|
||||||
|
dim * 3,
|
||||||
|
kernel_size=3,
|
||||||
|
stride=1,
|
||||||
|
padding=1,
|
||||||
|
groups=dim * 3,
|
||||||
|
bias=bias,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.attend = nn.Sequential(nn.Softmax(dim=-1), nn.Dropout(dropout))
|
||||||
|
|
||||||
|
self.to_out = nn.Conv2d(dim, dim, kernel_size=1, bias=bias)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
# project for queries, keys, values
|
||||||
|
b, c, h, w = x.shape
|
||||||
|
|
||||||
|
qkv = self.qkv_dwconv(self.qkv(x))
|
||||||
|
q, k, v = qkv.chunk(3, dim=1)
|
||||||
|
|
||||||
|
# split heads
|
||||||
|
|
||||||
|
q, k, v = map(
|
||||||
|
lambda t: rearrange(
|
||||||
|
t,
|
||||||
|
"b (h d) (x w1) (y w2) -> (b x y) h (w1 w2) d",
|
||||||
|
h=self.heads,
|
||||||
|
w1=self.ps,
|
||||||
|
w2=self.ps,
|
||||||
|
),
|
||||||
|
(q, k, v),
|
||||||
|
)
|
||||||
|
|
||||||
|
# scale
|
||||||
|
|
||||||
|
q = q * self.scale
|
||||||
|
|
||||||
|
# sim
|
||||||
|
|
||||||
|
sim = einsum("b h i d, b h j d -> b h i j", q, k)
|
||||||
|
|
||||||
|
# attention
|
||||||
|
attn = self.attend(sim)
|
||||||
|
|
||||||
|
# aggregate
|
||||||
|
|
||||||
|
out = einsum("b h i j, b h j d -> b h i d", attn, v)
|
||||||
|
|
||||||
|
# merge heads
|
||||||
|
out = rearrange(
|
||||||
|
out,
|
||||||
|
"(b x y) head (w1 w2) d -> b (head d) (x w1) (y w2)",
|
||||||
|
x=h // self.ps,
|
||||||
|
y=w // self.ps,
|
||||||
|
head=self.heads,
|
||||||
|
w1=self.ps,
|
||||||
|
w2=self.ps,
|
||||||
|
)
|
||||||
|
|
||||||
|
out = self.to_out(out)
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
class Channel_Attention(nn.Module):
|
||||||
|
def __init__(self, dim, heads, bias=False, dropout=0.0, window_size=7):
|
||||||
|
super(Channel_Attention, self).__init__()
|
||||||
|
self.heads = heads
|
||||||
|
|
||||||
|
self.temperature = nn.Parameter(torch.ones(heads, 1, 1))
|
||||||
|
|
||||||
|
self.ps = window_size
|
||||||
|
|
||||||
|
self.qkv = nn.Conv2d(dim, dim * 3, kernel_size=1, bias=bias)
|
||||||
|
self.qkv_dwconv = nn.Conv2d(
|
||||||
|
dim * 3,
|
||||||
|
dim * 3,
|
||||||
|
kernel_size=3,
|
||||||
|
stride=1,
|
||||||
|
padding=1,
|
||||||
|
groups=dim * 3,
|
||||||
|
bias=bias,
|
||||||
|
)
|
||||||
|
self.project_out = nn.Conv2d(dim, dim, kernel_size=1, bias=bias)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
b, c, h, w = x.shape
|
||||||
|
|
||||||
|
qkv = self.qkv_dwconv(self.qkv(x))
|
||||||
|
qkv = qkv.chunk(3, dim=1)
|
||||||
|
|
||||||
|
q, k, v = map(
|
||||||
|
lambda t: rearrange(
|
||||||
|
t,
|
||||||
|
"b (head d) (h ph) (w pw) -> b (h w) head d (ph pw)",
|
||||||
|
ph=self.ps,
|
||||||
|
pw=self.ps,
|
||||||
|
head=self.heads,
|
||||||
|
),
|
||||||
|
qkv,
|
||||||
|
)
|
||||||
|
|
||||||
|
q = F.normalize(q, dim=-1)
|
||||||
|
k = F.normalize(k, dim=-1)
|
||||||
|
|
||||||
|
attn = (q @ k.transpose(-2, -1)) * self.temperature
|
||||||
|
attn = attn.softmax(dim=-1)
|
||||||
|
out = attn @ v
|
||||||
|
|
||||||
|
out = rearrange(
|
||||||
|
out,
|
||||||
|
"b (h w) head d (ph pw) -> b (head d) (h ph) (w pw)",
|
||||||
|
h=h // self.ps,
|
||||||
|
w=w // self.ps,
|
||||||
|
ph=self.ps,
|
||||||
|
pw=self.ps,
|
||||||
|
head=self.heads,
|
||||||
|
)
|
||||||
|
|
||||||
|
out = self.project_out(out)
|
||||||
|
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
class Channel_Attention_grid(nn.Module):
|
||||||
|
def __init__(self, dim, heads, bias=False, dropout=0.0, window_size=7):
|
||||||
|
super(Channel_Attention_grid, self).__init__()
|
||||||
|
self.heads = heads
|
||||||
|
|
||||||
|
self.temperature = nn.Parameter(torch.ones(heads, 1, 1))
|
||||||
|
|
||||||
|
self.ps = window_size
|
||||||
|
|
||||||
|
self.qkv = nn.Conv2d(dim, dim * 3, kernel_size=1, bias=bias)
|
||||||
|
self.qkv_dwconv = nn.Conv2d(
|
||||||
|
dim * 3,
|
||||||
|
dim * 3,
|
||||||
|
kernel_size=3,
|
||||||
|
stride=1,
|
||||||
|
padding=1,
|
||||||
|
groups=dim * 3,
|
||||||
|
bias=bias,
|
||||||
|
)
|
||||||
|
self.project_out = nn.Conv2d(dim, dim, kernel_size=1, bias=bias)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
b, c, h, w = x.shape
|
||||||
|
|
||||||
|
qkv = self.qkv_dwconv(self.qkv(x))
|
||||||
|
qkv = qkv.chunk(3, dim=1)
|
||||||
|
|
||||||
|
q, k, v = map(
|
||||||
|
lambda t: rearrange(
|
||||||
|
t,
|
||||||
|
"b (head d) (h ph) (w pw) -> b (ph pw) head d (h w)",
|
||||||
|
ph=self.ps,
|
||||||
|
pw=self.ps,
|
||||||
|
head=self.heads,
|
||||||
|
),
|
||||||
|
qkv,
|
||||||
|
)
|
||||||
|
|
||||||
|
q = F.normalize(q, dim=-1)
|
||||||
|
k = F.normalize(k, dim=-1)
|
||||||
|
|
||||||
|
attn = (q @ k.transpose(-2, -1)) * self.temperature
|
||||||
|
attn = attn.softmax(dim=-1)
|
||||||
|
out = attn @ v
|
||||||
|
|
||||||
|
out = rearrange(
|
||||||
|
out,
|
||||||
|
"b (ph pw) head d (h w) -> b (head d) (h ph) (w pw)",
|
||||||
|
h=h // self.ps,
|
||||||
|
w=w // self.ps,
|
||||||
|
ph=self.ps,
|
||||||
|
pw=self.ps,
|
||||||
|
head=self.heads,
|
||||||
|
)
|
||||||
|
|
||||||
|
out = self.project_out(out)
|
||||||
|
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
class OSA_Block(nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
channel_num=64,
|
||||||
|
bias=True,
|
||||||
|
ffn_bias=True,
|
||||||
|
window_size=8,
|
||||||
|
with_pe=False,
|
||||||
|
dropout=0.0,
|
||||||
|
):
|
||||||
|
super(OSA_Block, self).__init__()
|
||||||
|
|
||||||
|
w = window_size
|
||||||
|
|
||||||
|
self.layer = nn.Sequential(
|
||||||
|
MBConv(
|
||||||
|
channel_num,
|
||||||
|
channel_num,
|
||||||
|
downsample=False,
|
||||||
|
expansion_rate=1,
|
||||||
|
shrinkage_rate=0.25,
|
||||||
|
),
|
||||||
|
Rearrange(
|
||||||
|
"b d (x w1) (y w2) -> b x y w1 w2 d", w1=w, w2=w
|
||||||
|
), # block-like attention
|
||||||
|
PreNormResidual(
|
||||||
|
channel_num,
|
||||||
|
Attention(
|
||||||
|
dim=channel_num,
|
||||||
|
dim_head=channel_num // 4,
|
||||||
|
dropout=dropout,
|
||||||
|
window_size=window_size,
|
||||||
|
with_pe=with_pe,
|
||||||
|
),
|
||||||
|
),
|
||||||
|
Rearrange("b x y w1 w2 d -> b d (x w1) (y w2)"),
|
||||||
|
Conv_PreNormResidual(
|
||||||
|
channel_num, Gated_Conv_FeedForward(dim=channel_num, dropout=dropout)
|
||||||
|
),
|
||||||
|
# channel-like attention
|
||||||
|
Conv_PreNormResidual(
|
||||||
|
channel_num,
|
||||||
|
Channel_Attention(
|
||||||
|
dim=channel_num, heads=4, dropout=dropout, window_size=window_size
|
||||||
|
),
|
||||||
|
),
|
||||||
|
Conv_PreNormResidual(
|
||||||
|
channel_num, Gated_Conv_FeedForward(dim=channel_num, dropout=dropout)
|
||||||
|
),
|
||||||
|
Rearrange(
|
||||||
|
"b d (w1 x) (w2 y) -> b x y w1 w2 d", w1=w, w2=w
|
||||||
|
), # grid-like attention
|
||||||
|
PreNormResidual(
|
||||||
|
channel_num,
|
||||||
|
Attention(
|
||||||
|
dim=channel_num,
|
||||||
|
dim_head=channel_num // 4,
|
||||||
|
dropout=dropout,
|
||||||
|
window_size=window_size,
|
||||||
|
with_pe=with_pe,
|
||||||
|
),
|
||||||
|
),
|
||||||
|
Rearrange("b x y w1 w2 d -> b d (w1 x) (w2 y)"),
|
||||||
|
Conv_PreNormResidual(
|
||||||
|
channel_num, Gated_Conv_FeedForward(dim=channel_num, dropout=dropout)
|
||||||
|
),
|
||||||
|
# channel-like attention
|
||||||
|
Conv_PreNormResidual(
|
||||||
|
channel_num,
|
||||||
|
Channel_Attention_grid(
|
||||||
|
dim=channel_num, heads=4, dropout=dropout, window_size=window_size
|
||||||
|
),
|
||||||
|
),
|
||||||
|
Conv_PreNormResidual(
|
||||||
|
channel_num, Gated_Conv_FeedForward(dim=channel_num, dropout=dropout)
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
out = self.layer(x)
|
||||||
|
return out
|
60
ldm_patched/pfn/architecture/OmniSR/OSAG.py
Normal file
60
ldm_patched/pfn/architecture/OmniSR/OSAG.py
Normal file
@ -0,0 +1,60 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
# -*- coding:utf-8 -*-
|
||||||
|
#############################################################
|
||||||
|
# File: OSAG.py
|
||||||
|
# Created Date: Tuesday April 28th 2022
|
||||||
|
# Author: Chen Xuanhong
|
||||||
|
# Email: chenxuanhongzju@outlook.com
|
||||||
|
# Last Modified: Sunday, 23rd April 2023 3:08:49 pm
|
||||||
|
# Modified By: Chen Xuanhong
|
||||||
|
# Copyright (c) 2020 Shanghai Jiao Tong University
|
||||||
|
#############################################################
|
||||||
|
|
||||||
|
|
||||||
|
import torch.nn as nn
|
||||||
|
|
||||||
|
from .esa import ESA
|
||||||
|
from .OSA import OSA_Block
|
||||||
|
|
||||||
|
|
||||||
|
class OSAG(nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
channel_num=64,
|
||||||
|
bias=True,
|
||||||
|
block_num=4,
|
||||||
|
ffn_bias=False,
|
||||||
|
window_size=0,
|
||||||
|
pe=False,
|
||||||
|
):
|
||||||
|
super(OSAG, self).__init__()
|
||||||
|
|
||||||
|
# print("window_size: %d" % (window_size))
|
||||||
|
# print("with_pe", pe)
|
||||||
|
# print("ffn_bias: %d" % (ffn_bias))
|
||||||
|
|
||||||
|
# block_script_name = kwargs.get("block_script_name", "OSA")
|
||||||
|
# block_class_name = kwargs.get("block_class_name", "OSA_Block")
|
||||||
|
|
||||||
|
# script_name = "." + block_script_name
|
||||||
|
# package = __import__(script_name, fromlist=True)
|
||||||
|
block_class = OSA_Block # getattr(package, block_class_name)
|
||||||
|
group_list = []
|
||||||
|
for _ in range(block_num):
|
||||||
|
temp_res = block_class(
|
||||||
|
channel_num,
|
||||||
|
bias,
|
||||||
|
ffn_bias=ffn_bias,
|
||||||
|
window_size=window_size,
|
||||||
|
with_pe=pe,
|
||||||
|
)
|
||||||
|
group_list.append(temp_res)
|
||||||
|
group_list.append(nn.Conv2d(channel_num, channel_num, 1, 1, 0, bias=bias))
|
||||||
|
self.residual_layer = nn.Sequential(*group_list)
|
||||||
|
esa_channel = max(channel_num // 4, 16)
|
||||||
|
self.esa = ESA(esa_channel, channel_num)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
out = self.residual_layer(x)
|
||||||
|
out = out + x
|
||||||
|
return self.esa(out)
|
143
ldm_patched/pfn/architecture/OmniSR/OmniSR.py
Normal file
143
ldm_patched/pfn/architecture/OmniSR/OmniSR.py
Normal file
@ -0,0 +1,143 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
# -*- coding:utf-8 -*-
|
||||||
|
#############################################################
|
||||||
|
# File: OmniSR.py
|
||||||
|
# Created Date: Tuesday April 28th 2022
|
||||||
|
# Author: Chen Xuanhong
|
||||||
|
# Email: chenxuanhongzju@outlook.com
|
||||||
|
# Last Modified: Sunday, 23rd April 2023 3:06:36 pm
|
||||||
|
# Modified By: Chen Xuanhong
|
||||||
|
# Copyright (c) 2020 Shanghai Jiao Tong University
|
||||||
|
#############################################################
|
||||||
|
|
||||||
|
import math
|
||||||
|
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
import torch.nn.functional as F
|
||||||
|
|
||||||
|
from .OSAG import OSAG
|
||||||
|
from .pixelshuffle import pixelshuffle_block
|
||||||
|
|
||||||
|
|
||||||
|
class OmniSR(nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
state_dict,
|
||||||
|
**kwargs,
|
||||||
|
):
|
||||||
|
super(OmniSR, self).__init__()
|
||||||
|
self.state = state_dict
|
||||||
|
|
||||||
|
bias = True # Fine to assume this for now
|
||||||
|
block_num = 1 # Fine to assume this for now
|
||||||
|
ffn_bias = True
|
||||||
|
pe = True
|
||||||
|
|
||||||
|
num_feat = state_dict["input.weight"].shape[0] or 64
|
||||||
|
num_in_ch = state_dict["input.weight"].shape[1] or 3
|
||||||
|
num_out_ch = num_in_ch # we can just assume this for now. pixelshuffle smh
|
||||||
|
|
||||||
|
pixelshuffle_shape = state_dict["up.0.weight"].shape[0]
|
||||||
|
up_scale = math.sqrt(pixelshuffle_shape / num_out_ch)
|
||||||
|
if up_scale - int(up_scale) > 0:
|
||||||
|
print(
|
||||||
|
"out_nc is probably different than in_nc, scale calculation might be wrong"
|
||||||
|
)
|
||||||
|
up_scale = int(up_scale)
|
||||||
|
res_num = 0
|
||||||
|
for key in state_dict.keys():
|
||||||
|
if "residual_layer" in key:
|
||||||
|
temp_res_num = int(key.split(".")[1])
|
||||||
|
if temp_res_num > res_num:
|
||||||
|
res_num = temp_res_num
|
||||||
|
res_num = res_num + 1 # zero-indexed
|
||||||
|
|
||||||
|
residual_layer = []
|
||||||
|
self.res_num = res_num
|
||||||
|
|
||||||
|
if (
|
||||||
|
"residual_layer.0.residual_layer.0.layer.2.fn.rel_pos_bias.weight"
|
||||||
|
in state_dict.keys()
|
||||||
|
):
|
||||||
|
rel_pos_bias_weight = state_dict[
|
||||||
|
"residual_layer.0.residual_layer.0.layer.2.fn.rel_pos_bias.weight"
|
||||||
|
].shape[0]
|
||||||
|
self.window_size = int((math.sqrt(rel_pos_bias_weight) + 1) / 2)
|
||||||
|
else:
|
||||||
|
self.window_size = 8
|
||||||
|
|
||||||
|
self.up_scale = up_scale
|
||||||
|
|
||||||
|
for _ in range(res_num):
|
||||||
|
temp_res = OSAG(
|
||||||
|
channel_num=num_feat,
|
||||||
|
bias=bias,
|
||||||
|
block_num=block_num,
|
||||||
|
ffn_bias=ffn_bias,
|
||||||
|
window_size=self.window_size,
|
||||||
|
pe=pe,
|
||||||
|
)
|
||||||
|
residual_layer.append(temp_res)
|
||||||
|
self.residual_layer = nn.Sequential(*residual_layer)
|
||||||
|
self.input = nn.Conv2d(
|
||||||
|
in_channels=num_in_ch,
|
||||||
|
out_channels=num_feat,
|
||||||
|
kernel_size=3,
|
||||||
|
stride=1,
|
||||||
|
padding=1,
|
||||||
|
bias=bias,
|
||||||
|
)
|
||||||
|
self.output = nn.Conv2d(
|
||||||
|
in_channels=num_feat,
|
||||||
|
out_channels=num_feat,
|
||||||
|
kernel_size=3,
|
||||||
|
stride=1,
|
||||||
|
padding=1,
|
||||||
|
bias=bias,
|
||||||
|
)
|
||||||
|
self.up = pixelshuffle_block(num_feat, num_out_ch, up_scale, bias=bias)
|
||||||
|
|
||||||
|
# self.tail = pixelshuffle_block(num_feat,num_out_ch,up_scale,bias=bias)
|
||||||
|
|
||||||
|
# for m in self.modules():
|
||||||
|
# if isinstance(m, nn.Conv2d):
|
||||||
|
# n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
|
||||||
|
# m.weight.data.normal_(0, sqrt(2. / n))
|
||||||
|
|
||||||
|
# chaiNNer specific stuff
|
||||||
|
self.model_arch = "OmniSR"
|
||||||
|
self.sub_type = "SR"
|
||||||
|
self.in_nc = num_in_ch
|
||||||
|
self.out_nc = num_out_ch
|
||||||
|
self.num_feat = num_feat
|
||||||
|
self.scale = up_scale
|
||||||
|
|
||||||
|
self.supports_fp16 = True # TODO: Test this
|
||||||
|
self.supports_bfp16 = True
|
||||||
|
self.min_size_restriction = 16
|
||||||
|
|
||||||
|
self.load_state_dict(state_dict, strict=False)
|
||||||
|
|
||||||
|
def check_image_size(self, x):
|
||||||
|
_, _, h, w = x.size()
|
||||||
|
# import pdb; pdb.set_trace()
|
||||||
|
mod_pad_h = (self.window_size - h % self.window_size) % self.window_size
|
||||||
|
mod_pad_w = (self.window_size - w % self.window_size) % self.window_size
|
||||||
|
# x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), 'reflect')
|
||||||
|
x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), "constant", 0)
|
||||||
|
return x
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
H, W = x.shape[2:]
|
||||||
|
x = self.check_image_size(x)
|
||||||
|
|
||||||
|
residual = self.input(x)
|
||||||
|
out = self.residual_layer(residual)
|
||||||
|
|
||||||
|
# origin
|
||||||
|
out = torch.add(self.output(out), residual)
|
||||||
|
out = self.up(out)
|
||||||
|
|
||||||
|
out = out[:, :, : H * self.up_scale, : W * self.up_scale]
|
||||||
|
return out
|
294
ldm_patched/pfn/architecture/OmniSR/esa.py
Normal file
294
ldm_patched/pfn/architecture/OmniSR/esa.py
Normal file
@ -0,0 +1,294 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
# -*- coding:utf-8 -*-
|
||||||
|
#############################################################
|
||||||
|
# File: esa.py
|
||||||
|
# Created Date: Tuesday April 28th 2022
|
||||||
|
# Author: Chen Xuanhong
|
||||||
|
# Email: chenxuanhongzju@outlook.com
|
||||||
|
# Last Modified: Thursday, 20th April 2023 9:28:06 am
|
||||||
|
# Modified By: Chen Xuanhong
|
||||||
|
# Copyright (c) 2020 Shanghai Jiao Tong University
|
||||||
|
#############################################################
|
||||||
|
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
import torch.nn.functional as F
|
||||||
|
|
||||||
|
from .layernorm import LayerNorm2d
|
||||||
|
|
||||||
|
|
||||||
|
def moment(x, dim=(2, 3), k=2):
|
||||||
|
assert len(x.size()) == 4
|
||||||
|
mean = torch.mean(x, dim=dim).unsqueeze(-1).unsqueeze(-1)
|
||||||
|
mk = (1 / (x.size(2) * x.size(3))) * torch.sum(torch.pow(x - mean, k), dim=dim)
|
||||||
|
return mk
|
||||||
|
|
||||||
|
|
||||||
|
class ESA(nn.Module):
|
||||||
|
"""
|
||||||
|
Modification of Enhanced Spatial Attention (ESA), which is proposed by
|
||||||
|
`Residual Feature Aggregation Network for Image Super-Resolution`
|
||||||
|
Note: `conv_max` and `conv3_` are NOT used here, so the corresponding codes
|
||||||
|
are deleted.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self, esa_channels, n_feats, conv=nn.Conv2d):
|
||||||
|
super(ESA, self).__init__()
|
||||||
|
f = esa_channels
|
||||||
|
self.conv1 = conv(n_feats, f, kernel_size=1)
|
||||||
|
self.conv_f = conv(f, f, kernel_size=1)
|
||||||
|
self.conv2 = conv(f, f, kernel_size=3, stride=2, padding=0)
|
||||||
|
self.conv3 = conv(f, f, kernel_size=3, padding=1)
|
||||||
|
self.conv4 = conv(f, n_feats, kernel_size=1)
|
||||||
|
self.sigmoid = nn.Sigmoid()
|
||||||
|
self.relu = nn.ReLU(inplace=True)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
c1_ = self.conv1(x)
|
||||||
|
c1 = self.conv2(c1_)
|
||||||
|
v_max = F.max_pool2d(c1, kernel_size=7, stride=3)
|
||||||
|
c3 = self.conv3(v_max)
|
||||||
|
c3 = F.interpolate(
|
||||||
|
c3, (x.size(2), x.size(3)), mode="bilinear", align_corners=False
|
||||||
|
)
|
||||||
|
cf = self.conv_f(c1_)
|
||||||
|
c4 = self.conv4(c3 + cf)
|
||||||
|
m = self.sigmoid(c4)
|
||||||
|
return x * m
|
||||||
|
|
||||||
|
|
||||||
|
class LK_ESA(nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self, esa_channels, n_feats, conv=nn.Conv2d, kernel_expand=1, bias=True
|
||||||
|
):
|
||||||
|
super(LK_ESA, self).__init__()
|
||||||
|
f = esa_channels
|
||||||
|
self.conv1 = conv(n_feats, f, kernel_size=1)
|
||||||
|
self.conv_f = conv(f, f, kernel_size=1)
|
||||||
|
|
||||||
|
kernel_size = 17
|
||||||
|
kernel_expand = kernel_expand
|
||||||
|
padding = kernel_size // 2
|
||||||
|
|
||||||
|
self.vec_conv = nn.Conv2d(
|
||||||
|
in_channels=f * kernel_expand,
|
||||||
|
out_channels=f * kernel_expand,
|
||||||
|
kernel_size=(1, kernel_size),
|
||||||
|
padding=(0, padding),
|
||||||
|
groups=2,
|
||||||
|
bias=bias,
|
||||||
|
)
|
||||||
|
self.vec_conv3x1 = nn.Conv2d(
|
||||||
|
in_channels=f * kernel_expand,
|
||||||
|
out_channels=f * kernel_expand,
|
||||||
|
kernel_size=(1, 3),
|
||||||
|
padding=(0, 1),
|
||||||
|
groups=2,
|
||||||
|
bias=bias,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.hor_conv = nn.Conv2d(
|
||||||
|
in_channels=f * kernel_expand,
|
||||||
|
out_channels=f * kernel_expand,
|
||||||
|
kernel_size=(kernel_size, 1),
|
||||||
|
padding=(padding, 0),
|
||||||
|
groups=2,
|
||||||
|
bias=bias,
|
||||||
|
)
|
||||||
|
self.hor_conv1x3 = nn.Conv2d(
|
||||||
|
in_channels=f * kernel_expand,
|
||||||
|
out_channels=f * kernel_expand,
|
||||||
|
kernel_size=(3, 1),
|
||||||
|
padding=(1, 0),
|
||||||
|
groups=2,
|
||||||
|
bias=bias,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.conv4 = conv(f, n_feats, kernel_size=1)
|
||||||
|
self.sigmoid = nn.Sigmoid()
|
||||||
|
self.relu = nn.ReLU(inplace=True)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
c1_ = self.conv1(x)
|
||||||
|
|
||||||
|
res = self.vec_conv(c1_) + self.vec_conv3x1(c1_)
|
||||||
|
res = self.hor_conv(res) + self.hor_conv1x3(res)
|
||||||
|
|
||||||
|
cf = self.conv_f(c1_)
|
||||||
|
c4 = self.conv4(res + cf)
|
||||||
|
m = self.sigmoid(c4)
|
||||||
|
return x * m
|
||||||
|
|
||||||
|
|
||||||
|
class LK_ESA_LN(nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self, esa_channels, n_feats, conv=nn.Conv2d, kernel_expand=1, bias=True
|
||||||
|
):
|
||||||
|
super(LK_ESA_LN, self).__init__()
|
||||||
|
f = esa_channels
|
||||||
|
self.conv1 = conv(n_feats, f, kernel_size=1)
|
||||||
|
self.conv_f = conv(f, f, kernel_size=1)
|
||||||
|
|
||||||
|
kernel_size = 17
|
||||||
|
kernel_expand = kernel_expand
|
||||||
|
padding = kernel_size // 2
|
||||||
|
|
||||||
|
self.norm = LayerNorm2d(n_feats)
|
||||||
|
|
||||||
|
self.vec_conv = nn.Conv2d(
|
||||||
|
in_channels=f * kernel_expand,
|
||||||
|
out_channels=f * kernel_expand,
|
||||||
|
kernel_size=(1, kernel_size),
|
||||||
|
padding=(0, padding),
|
||||||
|
groups=2,
|
||||||
|
bias=bias,
|
||||||
|
)
|
||||||
|
self.vec_conv3x1 = nn.Conv2d(
|
||||||
|
in_channels=f * kernel_expand,
|
||||||
|
out_channels=f * kernel_expand,
|
||||||
|
kernel_size=(1, 3),
|
||||||
|
padding=(0, 1),
|
||||||
|
groups=2,
|
||||||
|
bias=bias,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.hor_conv = nn.Conv2d(
|
||||||
|
in_channels=f * kernel_expand,
|
||||||
|
out_channels=f * kernel_expand,
|
||||||
|
kernel_size=(kernel_size, 1),
|
||||||
|
padding=(padding, 0),
|
||||||
|
groups=2,
|
||||||
|
bias=bias,
|
||||||
|
)
|
||||||
|
self.hor_conv1x3 = nn.Conv2d(
|
||||||
|
in_channels=f * kernel_expand,
|
||||||
|
out_channels=f * kernel_expand,
|
||||||
|
kernel_size=(3, 1),
|
||||||
|
padding=(1, 0),
|
||||||
|
groups=2,
|
||||||
|
bias=bias,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.conv4 = conv(f, n_feats, kernel_size=1)
|
||||||
|
self.sigmoid = nn.Sigmoid()
|
||||||
|
self.relu = nn.ReLU(inplace=True)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
c1_ = self.norm(x)
|
||||||
|
c1_ = self.conv1(c1_)
|
||||||
|
|
||||||
|
res = self.vec_conv(c1_) + self.vec_conv3x1(c1_)
|
||||||
|
res = self.hor_conv(res) + self.hor_conv1x3(res)
|
||||||
|
|
||||||
|
cf = self.conv_f(c1_)
|
||||||
|
c4 = self.conv4(res + cf)
|
||||||
|
m = self.sigmoid(c4)
|
||||||
|
return x * m
|
||||||
|
|
||||||
|
|
||||||
|
class AdaGuidedFilter(nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self, esa_channels, n_feats, conv=nn.Conv2d, kernel_expand=1, bias=True
|
||||||
|
):
|
||||||
|
super(AdaGuidedFilter, self).__init__()
|
||||||
|
|
||||||
|
self.gap = nn.AdaptiveAvgPool2d(1)
|
||||||
|
self.fc = nn.Conv2d(
|
||||||
|
in_channels=n_feats,
|
||||||
|
out_channels=1,
|
||||||
|
kernel_size=1,
|
||||||
|
padding=0,
|
||||||
|
stride=1,
|
||||||
|
groups=1,
|
||||||
|
bias=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.r = 5
|
||||||
|
|
||||||
|
def box_filter(self, x, r):
|
||||||
|
channel = x.shape[1]
|
||||||
|
kernel_size = 2 * r + 1
|
||||||
|
weight = 1.0 / (kernel_size**2)
|
||||||
|
box_kernel = weight * torch.ones(
|
||||||
|
(channel, 1, kernel_size, kernel_size), dtype=torch.float32, device=x.device
|
||||||
|
)
|
||||||
|
output = F.conv2d(x, weight=box_kernel, stride=1, padding=r, groups=channel)
|
||||||
|
return output
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
_, _, H, W = x.shape
|
||||||
|
N = self.box_filter(
|
||||||
|
torch.ones((1, 1, H, W), dtype=x.dtype, device=x.device), self.r
|
||||||
|
)
|
||||||
|
|
||||||
|
# epsilon = self.fc(self.gap(x))
|
||||||
|
# epsilon = torch.pow(epsilon, 2)
|
||||||
|
epsilon = 1e-2
|
||||||
|
|
||||||
|
mean_x = self.box_filter(x, self.r) / N
|
||||||
|
var_x = self.box_filter(x * x, self.r) / N - mean_x * mean_x
|
||||||
|
|
||||||
|
A = var_x / (var_x + epsilon)
|
||||||
|
b = (1 - A) * mean_x
|
||||||
|
m = A * x + b
|
||||||
|
|
||||||
|
# mean_A = self.box_filter(A, self.r) / N
|
||||||
|
# mean_b = self.box_filter(b, self.r) / N
|
||||||
|
# m = mean_A * x + mean_b
|
||||||
|
return x * m
|
||||||
|
|
||||||
|
|
||||||
|
class AdaConvGuidedFilter(nn.Module):
|
||||||
|
def __init__(
|
||||||
|
self, esa_channels, n_feats, conv=nn.Conv2d, kernel_expand=1, bias=True
|
||||||
|
):
|
||||||
|
super(AdaConvGuidedFilter, self).__init__()
|
||||||
|
f = esa_channels
|
||||||
|
|
||||||
|
self.conv_f = conv(f, f, kernel_size=1)
|
||||||
|
|
||||||
|
kernel_size = 17
|
||||||
|
kernel_expand = kernel_expand
|
||||||
|
padding = kernel_size // 2
|
||||||
|
|
||||||
|
self.vec_conv = nn.Conv2d(
|
||||||
|
in_channels=f,
|
||||||
|
out_channels=f,
|
||||||
|
kernel_size=(1, kernel_size),
|
||||||
|
padding=(0, padding),
|
||||||
|
groups=f,
|
||||||
|
bias=bias,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.hor_conv = nn.Conv2d(
|
||||||
|
in_channels=f,
|
||||||
|
out_channels=f,
|
||||||
|
kernel_size=(kernel_size, 1),
|
||||||
|
padding=(padding, 0),
|
||||||
|
groups=f,
|
||||||
|
bias=bias,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.gap = nn.AdaptiveAvgPool2d(1)
|
||||||
|
self.fc = nn.Conv2d(
|
||||||
|
in_channels=f,
|
||||||
|
out_channels=f,
|
||||||
|
kernel_size=1,
|
||||||
|
padding=0,
|
||||||
|
stride=1,
|
||||||
|
groups=1,
|
||||||
|
bias=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
y = self.vec_conv(x)
|
||||||
|
y = self.hor_conv(y)
|
||||||
|
|
||||||
|
sigma = torch.pow(y, 2)
|
||||||
|
epsilon = self.fc(self.gap(y))
|
||||||
|
|
||||||
|
weight = sigma / (sigma + epsilon)
|
||||||
|
|
||||||
|
m = weight * x + (1 - weight)
|
||||||
|
|
||||||
|
return x * m
|
70
ldm_patched/pfn/architecture/OmniSR/layernorm.py
Normal file
70
ldm_patched/pfn/architecture/OmniSR/layernorm.py
Normal file
@ -0,0 +1,70 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
# -*- coding:utf-8 -*-
|
||||||
|
#############################################################
|
||||||
|
# File: layernorm.py
|
||||||
|
# Created Date: Tuesday April 28th 2022
|
||||||
|
# Author: Chen Xuanhong
|
||||||
|
# Email: chenxuanhongzju@outlook.com
|
||||||
|
# Last Modified: Thursday, 20th April 2023 9:28:20 am
|
||||||
|
# Modified By: Chen Xuanhong
|
||||||
|
# Copyright (c) 2020 Shanghai Jiao Tong University
|
||||||
|
#############################################################
|
||||||
|
|
||||||
|
import torch
|
||||||
|
import torch.nn as nn
|
||||||
|
|
||||||
|
|
||||||
|
class LayerNormFunction(torch.autograd.Function):
|
||||||
|
@staticmethod
|
||||||
|
def forward(ctx, x, weight, bias, eps):
|
||||||
|
ctx.eps = eps
|
||||||
|
N, C, H, W = x.size()
|
||||||
|
mu = x.mean(1, keepdim=True)
|
||||||
|
var = (x - mu).pow(2).mean(1, keepdim=True)
|
||||||
|
y = (x - mu) / (var + eps).sqrt()
|
||||||
|
ctx.save_for_backward(y, var, weight)
|
||||||
|
y = weight.view(1, C, 1, 1) * y + bias.view(1, C, 1, 1)
|
||||||
|
return y
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def backward(ctx, grad_output):
|
||||||
|
eps = ctx.eps
|
||||||
|
|
||||||
|
N, C, H, W = grad_output.size()
|
||||||
|
y, var, weight = ctx.saved_variables
|
||||||
|
g = grad_output * weight.view(1, C, 1, 1)
|
||||||
|
mean_g = g.mean(dim=1, keepdim=True)
|
||||||
|
|
||||||
|
mean_gy = (g * y).mean(dim=1, keepdim=True)
|
||||||
|
gx = 1.0 / torch.sqrt(var + eps) * (g - y * mean_gy - mean_g)
|
||||||
|
return (
|
||||||
|
gx,
|
||||||
|
(grad_output * y).sum(dim=3).sum(dim=2).sum(dim=0),
|
||||||
|
grad_output.sum(dim=3).sum(dim=2).sum(dim=0),
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
class LayerNorm2d(nn.Module):
|
||||||
|
def __init__(self, channels, eps=1e-6):
|
||||||
|
super(LayerNorm2d, self).__init__()
|
||||||
|
self.register_parameter("weight", nn.Parameter(torch.ones(channels)))
|
||||||
|
self.register_parameter("bias", nn.Parameter(torch.zeros(channels)))
|
||||||
|
self.eps = eps
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
return LayerNormFunction.apply(x, self.weight, self.bias, self.eps)
|
||||||
|
|
||||||
|
|
||||||
|
class GRN(nn.Module):
|
||||||
|
"""GRN (Global Response Normalization) layer"""
|
||||||
|
|
||||||
|
def __init__(self, dim):
|
||||||
|
super().__init__()
|
||||||
|
self.gamma = nn.Parameter(torch.zeros(1, dim, 1, 1))
|
||||||
|
self.beta = nn.Parameter(torch.zeros(1, dim, 1, 1))
|
||||||
|
|
||||||
|
def forward(self, x):
|
||||||
|
Gx = torch.norm(x, p=2, dim=(2, 3), keepdim=True)
|
||||||
|
Nx = Gx / (Gx.mean(dim=1, keepdim=True) + 1e-6)
|
||||||
|
return self.gamma * (x * Nx) + self.beta + x
|
31
ldm_patched/pfn/architecture/OmniSR/pixelshuffle.py
Normal file
31
ldm_patched/pfn/architecture/OmniSR/pixelshuffle.py
Normal file
@ -0,0 +1,31 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
# -*- coding:utf-8 -*-
|
||||||
|
#############################################################
|
||||||
|
# File: pixelshuffle.py
|
||||||
|
# Created Date: Friday July 1st 2022
|
||||||
|
# Author: Chen Xuanhong
|
||||||
|
# Email: chenxuanhongzju@outlook.com
|
||||||
|
# Last Modified: Friday, 1st July 2022 10:18:39 am
|
||||||
|
# Modified By: Chen Xuanhong
|
||||||
|
# Copyright (c) 2022 Shanghai Jiao Tong University
|
||||||
|
#############################################################
|
||||||
|
|
||||||
|
import torch.nn as nn
|
||||||
|
|
||||||
|
|
||||||
|
def pixelshuffle_block(
|
||||||
|
in_channels, out_channels, upscale_factor=2, kernel_size=3, bias=False
|
||||||
|
):
|
||||||
|
"""
|
||||||
|
Upsample features according to `upscale_factor`.
|
||||||
|
"""
|
||||||
|
padding = kernel_size // 2
|
||||||
|
conv = nn.Conv2d(
|
||||||
|
in_channels,
|
||||||
|
out_channels * (upscale_factor**2),
|
||||||
|
kernel_size,
|
||||||
|
padding=1,
|
||||||
|
bias=bias,
|
||||||
|
)
|
||||||
|
pixel_shuffle = nn.PixelShuffle(upscale_factor)
|
||||||
|
return nn.Sequential(*[conv, pixel_shuffle])
|
Some files were not shown because too many files have changed in this diff Show More
Loading…
Reference in New Issue
Block a user