| | |
| |
|
| | import math |
| | from einops import rearrange |
| | |
| | from torch import randint |
| |
|
| | def random_divisor(value: int, min_value: int, /, max_options: int = 1) -> int: |
| | min_value = min(min_value, value) |
| |
|
| | |
| | divisors = [i for i in range(min_value, value + 1) if value % i == 0] |
| |
|
| | ns = [value // i for i in divisors[:max_options]] |
| |
|
| | if len(ns) - 1 > 0: |
| | idx = randint(low=0, high=len(ns) - 1, size=(1,)).item() |
| | else: |
| | idx = 0 |
| |
|
| | return ns[idx] |
| |
|
| | class HyperTile: |
| | @classmethod |
| | def INPUT_TYPES(s): |
| | return {"required": { "model": ("MODEL",), |
| | "tile_size": ("INT", {"default": 256, "min": 1, "max": 2048}), |
| | "swap_size": ("INT", {"default": 2, "min": 1, "max": 128}), |
| | "max_depth": ("INT", {"default": 0, "min": 0, "max": 10}), |
| | "scale_depth": ("BOOLEAN", {"default": False}), |
| | }} |
| | RETURN_TYPES = ("MODEL",) |
| | FUNCTION = "patch" |
| |
|
| | CATEGORY = "model_patches/unet" |
| |
|
| | 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, |
| | } |
| |
|
