379 lines
16 KiB
Python
379 lines
16 KiB
Python
from enum import Enum
|
|
from typing import Final, Optional, Union
|
|
|
|
import numpy as np
|
|
import torch
|
|
from numpy import exp, pi, sqrt
|
|
from torch import Tensor
|
|
|
|
import comfy.model_management as model_management
|
|
from comfy.utils import common_upscale
|
|
|
|
_dev: Final[torch.device] = model_management.get_torch_device()
|
|
|
|
LATENT_SCALE_FACTOR = 8
|
|
|
|
|
|
class BlendMode(Enum):
|
|
FOREGROUND = "Foreground"
|
|
BACKGROUND = "Background"
|
|
|
|
|
|
class BBox:
|
|
def __init__(self, x: int, y: int, w: int, h: int):
|
|
self.x = x
|
|
self.y = y
|
|
self.w = w
|
|
self.h = h
|
|
self.box = [x, y, x + w, y + h]
|
|
self.slicer = slice(None), slice(None), slice(y, y + h), slice(x, x + w)
|
|
|
|
def __getitem__(self, idx: int) -> int:
|
|
return self.box[idx]
|
|
|
|
|
|
def ceildiv(big: int, small: int) -> int:
|
|
return -(big // -small)
|
|
|
|
|
|
def split_bboxes(
|
|
w: int,
|
|
h: int,
|
|
tile_w: int,
|
|
tile_h: int,
|
|
overlap: int = 16,
|
|
init_weight: Union[Tensor, float] = 1.0,
|
|
) -> tuple[list[BBox], Tensor]:
|
|
cols = ceildiv((w - overlap), (tile_w - overlap))
|
|
rows = ceildiv((h - overlap), (tile_h - overlap))
|
|
dx = (w - tile_w) / (cols - 1) if cols > 1 else 0
|
|
dy = (h - tile_h) / (rows - 1) if rows > 1 else 0
|
|
|
|
bbox_list: list[BBox] = []
|
|
weight = torch.zeros((1, 1, h, w), device=_dev, dtype=torch.float32)
|
|
for row in range(rows):
|
|
y = min(int(row * dy), h - tile_h)
|
|
for col in range(cols):
|
|
x = min(int(col * dx), w - tile_w)
|
|
bbox = BBox(x, y, tile_w, tile_h)
|
|
bbox_list.append(bbox)
|
|
weight[bbox.slicer] += init_weight
|
|
|
|
return bbox_list, weight
|
|
|
|
|
|
def _resize_hint(tensor: Tensor, width: int, height: int, like: Tensor) -> Tensor:
|
|
resized = common_upscale(tensor, width, height, "nearest-exact", "center")
|
|
return resized.to(device=like.device, dtype=like.dtype)
|
|
|
|
|
|
class AbstractDiffusion:
|
|
def __init__(self):
|
|
self.method = self.__class__.__name__
|
|
self.w = 0
|
|
self.h = 0
|
|
self.tile_width: Optional[int] = None
|
|
self.tile_height: Optional[int] = None
|
|
self.tile_overlap: Optional[int] = None
|
|
self.tile_batch_size: Optional[int] = None
|
|
self.x_buffer: Optional[Tensor] = None
|
|
self.enable_grid_bbox = False
|
|
self.tile_w: Optional[int] = None
|
|
self.tile_h: Optional[int] = None
|
|
self.tile_bs: Optional[int] = None
|
|
self.num_tiles: Optional[int] = None
|
|
self.num_batches: Optional[int] = None
|
|
self.batched_bboxes: list[list[BBox]] = []
|
|
self.enable_custom_bbox = False
|
|
self.custom_bboxes: list[BBox] = []
|
|
self.draw_background = True
|
|
self.weights: Optional[Tensor] = None
|
|
self.control_params: dict[tuple, list[list[Optional[Tensor]]]] = {}
|
|
|
|
def repeat_tensor(self, x: Tensor, n: int, concat: bool = False, concat_to: int = 0) -> Tensor:
|
|
if n == 1:
|
|
return x
|
|
bsz = x.shape[0]
|
|
r_dims = len(x.shape) - 1
|
|
if bsz == 1:
|
|
shape = [n] + [-1] * r_dims
|
|
return x.expand(shape)
|
|
if concat:
|
|
return torch.cat([x for _ in range(n)], dim=0)[:concat_to]
|
|
shape = [n] + [1] * r_dims
|
|
return x.repeat(shape)
|
|
|
|
def reset_buffer(self, x_in: Tensor):
|
|
if self.x_buffer is None or self.x_buffer.shape != x_in.shape:
|
|
self.x_buffer = torch.zeros_like(x_in, device=x_in.device, dtype=x_in.dtype)
|
|
else:
|
|
self.x_buffer.zero_()
|
|
|
|
def init_grid_bbox(self, tile_w: int, tile_h: int, overlap: int, tile_bs: int):
|
|
self.weights = torch.zeros((1, 1, self.h, self.w), device=_dev, dtype=torch.float32)
|
|
self.enable_grid_bbox = True
|
|
self.tile_w = min(tile_w, self.w)
|
|
self.tile_h = min(tile_h, self.h)
|
|
overlap = max(0, min(overlap, min(tile_w, tile_h) - 4))
|
|
bboxes, weights = split_bboxes(self.w, self.h, self.tile_w, self.tile_h, overlap, self.get_tile_weights())
|
|
self.weights += weights
|
|
self.num_tiles = len(bboxes)
|
|
self.num_batches = ceildiv(self.num_tiles, tile_bs)
|
|
self.tile_bs = ceildiv(len(bboxes), self.num_batches)
|
|
self.batched_bboxes = [bboxes[i * self.tile_bs : (i + 1) * self.tile_bs] for i in range(self.num_batches)]
|
|
|
|
def get_tile_weights(self) -> Union[Tensor, float]:
|
|
return 1.0
|
|
|
|
def init_done(self):
|
|
total_bboxes = 0
|
|
if self.enable_grid_bbox:
|
|
total_bboxes += self.num_batches or 0
|
|
if self.enable_custom_bbox:
|
|
total_bboxes += len(self.custom_bboxes)
|
|
if total_bboxes <= 0:
|
|
raise AssertionError("Nothing to paint! No background to draw and no custom bboxes were provided.")
|
|
|
|
def process_controlnet(self, x_like: Tensor, c_in: dict, cond_or_uncond: list, bboxes, batch_size: int, batch_id: int):
|
|
control = c_in.get("control")
|
|
param_id = -1
|
|
tuple_key = tuple(cond_or_uncond) + tuple(x_like.shape)
|
|
while control is not None:
|
|
param_id += 1
|
|
ph, pw = self.h * LATENT_SCALE_FACTOR, self.w * LATENT_SCALE_FACTOR
|
|
|
|
if self.control_params.get(tuple_key) is None:
|
|
self.control_params[tuple_key] = []
|
|
val = self.control_params[tuple_key]
|
|
while len(val) <= param_id:
|
|
val.append([])
|
|
while len(val[param_id]) <= batch_id:
|
|
val[param_id].append(None)
|
|
|
|
cached_hint = val[param_id][batch_id]
|
|
hint_original = getattr(control, "cond_hint_original", None)
|
|
if cached_hint is None and hint_original is not None:
|
|
dtype = getattr(control, "manual_cast_dtype", None) or getattr(getattr(control, "control_model", None), "dtype", None) or x_like.dtype
|
|
cond_hint = hint_original
|
|
if cond_hint.shape[-2:] != (ph, pw):
|
|
cond_hint = _resize_hint(cond_hint, pw, ph, x_like)
|
|
cond_hint = cond_hint.to(device=getattr(control, "device", x_like.device), dtype=dtype)
|
|
if cond_hint.shape[0] < batch_size:
|
|
cond_hint = self.repeat_tensor(cond_hint, ceildiv(batch_size, cond_hint.shape[0]))[:batch_size]
|
|
tiles = [cond_hint[:, :, bbox[1] * LATENT_SCALE_FACTOR : bbox[3] * LATENT_SCALE_FACTOR, bbox[0] * LATENT_SCALE_FACTOR : bbox[2] * LATENT_SCALE_FACTOR] for bbox in bboxes]
|
|
cached_hint = torch.cat(tiles, dim=0)
|
|
val[param_id][batch_id] = cached_hint
|
|
|
|
if cached_hint is not None:
|
|
control.cond_hint = cached_hint
|
|
control = getattr(control, "previous_controlnet", None)
|
|
|
|
|
|
def gaussian_weights(tile_w: int, tile_h: int) -> Tensor:
|
|
f = lambda x, midpoint, var=0.01: exp(-(x - midpoint) * (x - midpoint) / (tile_w * tile_w) / (2 * var)) / sqrt(2 * pi * var)
|
|
x_probs = [f(x, (tile_w - 1) / 2) for x in range(tile_w)]
|
|
y_probs = [f(y, tile_h / 2) for y in range(tile_h)]
|
|
w = np.outer(y_probs, x_probs)
|
|
return torch.from_numpy(w).to(_dev, dtype=torch.float32)
|
|
|
|
|
|
class MultiDiffusion(AbstractDiffusion):
|
|
@torch.no_grad()
|
|
def __call__(self, model_function, args: dict):
|
|
x_in: Tensor = args["input"]
|
|
t_in: Tensor = args["timestep"]
|
|
c_in: dict = args["c"]
|
|
cond_or_uncond: list = args["cond_or_uncond"]
|
|
c_crossattn: Tensor = c_in["c_crossattn"]
|
|
|
|
is_5d = x_in.ndim == 5
|
|
if is_5d:
|
|
if x_in.shape[2] != 1:
|
|
raise AssertionError("Expected a singleton frame dimension for 5D latents.")
|
|
x_in = x_in.squeeze(2)
|
|
|
|
n, _, h, w = x_in.shape
|
|
if self.weights is None or self.h != h or self.w != w:
|
|
self.h, self.w = h, w
|
|
self.init_grid_bbox(self.tile_width, self.tile_height, self.tile_overlap, self.tile_batch_size)
|
|
self.init_done()
|
|
|
|
self.h, self.w = h, w
|
|
self.reset_buffer(x_in)
|
|
if self.draw_background:
|
|
for batch_id, bboxes in enumerate(self.batched_bboxes):
|
|
x_tile = torch.cat([x_in[bbox.slicer] for bbox in bboxes], dim=0)
|
|
n_rep = len(bboxes)
|
|
ts_tile = self.repeat_tensor(t_in, n_rep)
|
|
cond_tile = self.repeat_tensor(c_crossattn, n_rep)
|
|
c_tile = c_in.copy()
|
|
c_tile["c_crossattn"] = cond_tile
|
|
if "time_context" in c_in:
|
|
c_tile["time_context"] = self.repeat_tensor(c_in["time_context"], n_rep)
|
|
for key in c_tile:
|
|
if key in ["y", "c_concat"]:
|
|
icond = c_tile[key]
|
|
if icond.ndim == 5:
|
|
if icond.shape[2] != 1:
|
|
raise AssertionError("Expected a singleton frame dimension for 5D conditioning.")
|
|
icond = icond.squeeze(2)
|
|
|
|
if icond.shape[2:] == (self.h, self.w):
|
|
c_tile[key] = torch.cat([icond[bbox.slicer] for bbox in bboxes])
|
|
else:
|
|
c_tile[key] = self.repeat_tensor(icond, n_rep)
|
|
if "control" in c_in:
|
|
self.process_controlnet(x_tile, c_in, cond_or_uncond, bboxes, n, batch_id)
|
|
c_tile["control"] = c_in["control"].get_control(x_tile, ts_tile, c_tile, len(cond_or_uncond))
|
|
|
|
if is_5d:
|
|
x_tile = x_tile.unsqueeze(2)
|
|
for key in ["y", "c_concat"]:
|
|
if key in c_tile and c_tile[key].ndim == 4:
|
|
c_tile[key] = c_tile[key].unsqueeze(2)
|
|
|
|
control = c_tile.get("control")
|
|
while control is not None:
|
|
if hasattr(control, "cond_hint") and control.cond_hint is not None and control.cond_hint.ndim == 4:
|
|
control.cond_hint = control.cond_hint.unsqueeze(2)
|
|
control = getattr(control, "previous_controlnet", None)
|
|
|
|
x_tile_out = model_function(x_tile, ts_tile, **c_tile)
|
|
|
|
if is_5d:
|
|
x_tile_out = x_tile_out.squeeze(2)
|
|
|
|
for i, bbox in enumerate(bboxes):
|
|
self.x_buffer[bbox.slicer] += x_tile_out[i * n : (i + 1) * n, :, :, :]
|
|
|
|
x_out = torch.where(self.weights > 1, self.x_buffer / self.weights, self.x_buffer)
|
|
if is_5d:
|
|
x_out = x_out.unsqueeze(2)
|
|
return x_out
|
|
|
|
|
|
class MixtureOfDiffusers(AbstractDiffusion):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.custom_weights: list[Tensor] = []
|
|
self.get_weight = gaussian_weights
|
|
|
|
def init_done(self):
|
|
super().init_done()
|
|
self.rescale_factor = 1 / self.weights
|
|
for bbox_id, bbox in enumerate(self.custom_bboxes):
|
|
if bbox.blend_mode == BlendMode.BACKGROUND:
|
|
self.custom_weights[bbox_id] *= self.rescale_factor[bbox.slicer]
|
|
|
|
def get_tile_weights(self) -> Tensor:
|
|
self.tile_weights = self.get_weight(self.tile_w, self.tile_h)
|
|
return self.tile_weights
|
|
|
|
@torch.no_grad()
|
|
def __call__(self, model_function, args: dict):
|
|
x_in: Tensor = args["input"]
|
|
t_in: Tensor = args["timestep"]
|
|
c_in: dict = args["c"]
|
|
cond_or_uncond: list = args["cond_or_uncond"]
|
|
c_crossattn: Tensor = c_in["c_crossattn"]
|
|
|
|
is_5d = x_in.ndim == 5
|
|
if is_5d:
|
|
if x_in.shape[2] != 1:
|
|
raise AssertionError("Expected a singleton frame dimension for 5D latents.")
|
|
x_in = x_in.squeeze(2)
|
|
|
|
n, _, h, w = x_in.shape
|
|
if self.weights is None or self.h != h or self.w != w:
|
|
self.h, self.w = h, w
|
|
self.init_grid_bbox(self.tile_width, self.tile_height, self.tile_overlap, self.tile_batch_size)
|
|
self.init_done()
|
|
|
|
self.h, self.w = h, w
|
|
self.reset_buffer(x_in)
|
|
if self.draw_background:
|
|
for batch_id, bboxes in enumerate(self.batched_bboxes):
|
|
x_tile_list = []
|
|
icond_map = {}
|
|
for bbox in bboxes:
|
|
x_tile_list.append(x_in[bbox.slicer])
|
|
for key in ["y", "c_concat"]:
|
|
if key in c_in:
|
|
icond = c_in[key]
|
|
if icond.ndim == 5:
|
|
if icond.shape[2] != 1:
|
|
raise AssertionError("Expected a singleton frame dimension for 5D conditioning.")
|
|
icond = icond.squeeze(2)
|
|
|
|
if icond.shape[2:] == (self.h, self.w):
|
|
icond = icond[bbox.slicer]
|
|
icond_map.setdefault(key, []).append(icond)
|
|
|
|
n_rep = len(bboxes)
|
|
x_tile = torch.cat(x_tile_list, dim=0)
|
|
t_tile = self.repeat_tensor(t_in, n_rep)
|
|
tcond_tile = self.repeat_tensor(c_crossattn, n_rep)
|
|
c_tile = c_in.copy()
|
|
c_tile["c_crossattn"] = tcond_tile
|
|
if "time_context" in c_in:
|
|
c_tile["time_context"] = self.repeat_tensor(c_in["time_context"], n_rep)
|
|
for key in ["y", "c_concat"]:
|
|
if key in icond_map:
|
|
c_tile[key] = torch.cat(icond_map[key], dim=0)
|
|
if "control" in c_in:
|
|
self.process_controlnet(x_tile, c_in, cond_or_uncond, bboxes, n, batch_id)
|
|
c_tile["control"] = c_in["control"].get_control(x_tile, t_tile, c_tile, len(cond_or_uncond))
|
|
|
|
if is_5d:
|
|
x_tile = x_tile.unsqueeze(2)
|
|
for key in ["y", "c_concat"]:
|
|
if key in c_tile and c_tile[key].ndim == 4:
|
|
c_tile[key] = c_tile[key].unsqueeze(2)
|
|
|
|
control = c_tile.get("control")
|
|
while control is not None:
|
|
if hasattr(control, "cond_hint") and control.cond_hint is not None and control.cond_hint.ndim == 4:
|
|
control.cond_hint = control.cond_hint.unsqueeze(2)
|
|
control = getattr(control, "previous_controlnet", None)
|
|
|
|
x_tile_out = model_function(x_tile, t_tile, **c_tile)
|
|
|
|
if is_5d:
|
|
x_tile_out = x_tile_out.squeeze(2)
|
|
|
|
for i, bbox in enumerate(bboxes):
|
|
w_tile = self.tile_weights * self.rescale_factor[bbox.slicer]
|
|
self.x_buffer[bbox.slicer] += x_tile_out[i * n : (i + 1) * n, :, :, :] * w_tile
|
|
|
|
x_out = self.x_buffer
|
|
if is_5d:
|
|
x_out = x_out.unsqueeze(2)
|
|
return x_out
|
|
|
|
|
|
class TiledDiffusion:
|
|
@staticmethod
|
|
def apply(model, method: str, tile_width: int, tile_height: int, tile_overlap: int, tile_batch_size: int):
|
|
if method == "MultiDiffusion":
|
|
implement = MultiDiffusion()
|
|
elif method == "Mixture of Diffusers":
|
|
implement = MixtureOfDiffusers()
|
|
else:
|
|
raise ValueError(f"Unsupported tiled diffusion method: {method}")
|
|
|
|
implement.tile_width = max(1, tile_width // LATENT_SCALE_FACTOR)
|
|
implement.tile_height = max(1, tile_height // LATENT_SCALE_FACTOR)
|
|
implement.tile_overlap = max(0, tile_overlap // LATENT_SCALE_FACTOR)
|
|
implement.tile_batch_size = max(1, tile_batch_size)
|
|
|
|
model_clone = model.clone()
|
|
model_clone.set_model_unet_function_wrapper(implement)
|
|
model_clone.model_options["tiled_diffusion"] = {
|
|
"method": method,
|
|
"tile_width": tile_width,
|
|
"tile_height": tile_height,
|
|
"tile_overlap": tile_overlap,
|
|
"tile_batch_size": tile_batch_size,
|
|
}
|
|
return model_clone
|