diff --git a/CHANGELOG.md b/CHANGELOG.md index 6a243073e..94116c1fc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,9 @@ - allow deployment without git clone for example, zip of the sdnext folder can be used - control: configurable output folder in settings + - hypertile: enable vae tiling + - hypertile: add autodetect optimial value + set tile size to 0 to use autodetected value - cli: sdapi.py allow manual api invoke example: `python cli/sdapi.py /sdapi/v1/sd-models` - **Fixes**: diff --git a/modules/control/units/reference.py b/modules/control/units/reference.py index ccc3b2277..e4d622235 100644 --- a/modules/control/units/reference.py +++ b/modules/control/units/reference.py @@ -1,3 +1,4 @@ +from typing import Union import time from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline from modules.control.proc.reference_sd15 import StableDiffusionReferencePipeline @@ -13,7 +14,7 @@ def list_models(): class ReferencePipeline(): - def __init__(self, pipeline: StableDiffusionXLPipeline | StableDiffusionPipeline, dtype = None): + def __init__(self, pipeline: Union[StableDiffusionXLPipeline, StableDiffusionPipeline], dtype = None): t0 = time.time() self.orig_pipeline = pipeline self.pipeline = None diff --git a/modules/sd_hijack_hypertile.py b/modules/sd_hijack_hypertile.py index dfd54dc64..97c5fed0e 100644 --- a/modules/sd_hijack_hypertile.py +++ b/modules/sd_hijack_hypertile.py @@ -11,6 +11,7 @@ import torch import torch.nn as nn from einops import rearrange from modules.shared import log +from functools import cache # global variables to keep track of changing image size in multiple passes @@ -20,6 +21,45 @@ max_h = 0 max_w = 0 error_reported = False reset_needed = False +RNG_INSTANCE = random.Random() + + +def set_seed(seed: int) -> None: + RNG_INSTANCE.seed(seed) + + +def iterative_closest_divisors(hw:int, aspect_ratio:float) -> tuple[int, int]: + """ + Finds h and w such that h*w = hw and h/w = aspect_ratio + We check all possible divisors of hw and return the closest to the aspect ratio + """ + divisors = [i for i in range(2, hw + 1) if hw % i == 0] # all divisors of hw + pairs = [(i, hw // i) for i in divisors] # all pairs of divisors of hw + ratios = [w/h for h, w in pairs] # all ratios of pairs of divisors of hw + closest_ratio = min(ratios, key=lambda x: abs(x - aspect_ratio)) # closest ratio to aspect_ratio + closest_pair = pairs[ratios.index(closest_ratio)] # closest pair of divisors to aspect_ratio + return closest_pair + +@cache +def find_hw_candidates(hw:int, aspect_ratio:float) -> tuple[int, int]: + """ + Finds h and w such that h*w = hw and h/w = aspect_ratio + """ + h, w = round(math.sqrt(hw * aspect_ratio)), round(math.sqrt(hw / aspect_ratio)) + # find h and w such that h*w = hw and h/w = aspect_ratio + if h * w != hw: + w_candidate = hw / h + # check if w is an integer + if not w_candidate.is_integer(): + h_candidate = hw / w + # check if h is an integer + if not h_candidate.is_integer(): + return iterative_closest_divisors(hw, aspect_ratio) + else: + h = int(h_candidate) + else: + w = int(w_candidate) + return h, w def possible_tile_sizes(dimension: int, tile_size: int, min_tile_size: int, tile_options: int) -> list[int]: @@ -41,7 +81,7 @@ def parse_list(x: list[int], /) -> str: @contextmanager -def split_attention(layer: nn.Module, tile_size: int=256, min_tile_size: int=256, swap_size: int=1, depth: int=0): +def split_attention(layer: nn.Module, tile_size: int=256, min_tile_size: int=128, swap_size: int=1, depth: int=0): # hijacks AttnBlock from ldm and attention from diffusers global reset_needed # pylint: disable=global-statement ar = height / width # Aspect ratio @@ -65,8 +105,10 @@ def split_attention(layer: nn.Module, tile_size: int=256, min_tile_size: int=256 global height, width, max_h, max_w, reset_needed, error_reported # pylint: disable=global-statement x = args[0] try: - nh = nhs[random.randint(0, len(nhs) - 1)] - nw = nws[random.randint(0, len(nws) - 1)] + nh = RNG_INSTANCE.randint(0, len(nhs) - 1) + nw = RNG_INSTANCE.randint(0, len(nws) - 1) + # nh = nhs[random.randint(0, len(nhs) - 1)] + # nw = nws[random.randint(0, len(nws) - 1)] except Exception as e: if not error_reported: error_reported = True @@ -83,6 +125,7 @@ def split_attention(layer: nn.Module, tile_size: int=256, min_tile_size: int=256 else: # Unet hw = x.size(1) h, w = round(math.sqrt(ar * hw)), round(math.sqrt(hw / ar)) + h, w = find_hw_candidates(hw, ar) # dynamic height/width based on fact that first two forward calls contain actual height/width # and reset if latest hw is larger since we're never downscaling in 2nd pass if reset_needed: @@ -155,9 +198,11 @@ def context_hypertile_vae(p): # shared.log.warning('Hypertile VAE is enabled but no VAE model was found') return nullcontext() else: - shared.log.info(f'Applying hypertile: vae={shared.opts.hypertile_vae_tile}') - p.extra_generation_params['Hypertile VAE'] = shared.opts.hypertile_vae_tile - return split_attention(vae, tile_size=shared.opts.hypertile_vae_tile, min_tile_size=128, swap_size=1) + tile_size = shared.opts.hypertile_vae_tile if shared.opts.hypertile_vae_tile > 0 else max(256, 64 * min(p.width // 128, p.height // 128)) + shared.log.info(f'Applying hypertile: vae={tile_size}') + p.extra_generation_params['Hypertile VAE'] = tile_size + return split_attention(vae, tile_size=tile_size, min_tile_size=128, swap_size=1) + def context_hypertile_unet(p): @@ -179,9 +224,10 @@ def context_hypertile_unet(p): # shared.log.warning('Hypertile UNet is enabled but no Unet model was found') return nullcontext() else: - shared.log.info(f'Applying hypertile: unet={shared.opts.hypertile_unet_tile}') - p.extra_generation_params['Hypertile UNet'] = shared.opts.hypertile_unet_tile - return split_attention(unet, tile_size=shared.opts.hypertile_unet_tile, min_tile_size=128, swap_size=1) + tile_size = shared.opts.hypertile_unet_tile if shared.opts.hypertile_unet_tile > 0 else max(256, 64 * min(p.width // 128, p.height // 128)) + shared.log.info(f'Applying hypertile: unet={tile_size}') + p.extra_generation_params['Hypertile UNet'] = tile_size + return split_attention(unet, tile_size=tile_size, min_tile_size=128, swap_size=1) def hypertile_set(p, hr=False): @@ -198,4 +244,5 @@ def hypertile_set(p, hr=False): else: width=p.width height=p.height + set_seed(p.all_seeds[0]) reset_needed = True diff --git a/modules/shared.py b/modules/shared.py index da2ab3714..28abf6b8b 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -357,10 +357,10 @@ options_templates.update(options_section(('advanced', "Inference Settings"), { "freeu_s2": OptionInfo(0.2, "2nd stage skip factor", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}), "hypertile_sep": OptionInfo("

HyperTile

", "", gr.HTML), - "hypertile_vae_enabled": OptionInfo(False, "HyperTile for VAE enabled", gr.Checkbox, {"visible": False}), - "hypertile_vae_tile": OptionInfo(128, "HyperTile for VAE tile size", gr.Slider, {"minimum": 128, "maximum": 512, "step": 8, "visible": False}), "hypertile_unet_enabled": OptionInfo(False, "HyperTile for UNet enabled"), - "hypertile_unet_tile": OptionInfo(256, "HyperTile for UNet tile size", gr.Slider, {"minimum": 256, "maximum": 1024, "step": 8}), + "hypertile_unet_tile": OptionInfo(256, "HyperTile for UNet tile size", gr.Slider, {"minimum": 0, "maximum": 1024, "step": 8}), + "hypertile_vae_enabled": OptionInfo(False, "HyperTile for VAE enabled", gr.Checkbox), + "hypertile_vae_tile": OptionInfo(128, "HyperTile for VAE tile size", gr.Slider, {"minimum": 0, "maximum": 1024, "step": 8}), "inference_other_sep": OptionInfo("

Other

", "", gr.HTML), "batch_frame_mode": OptionInfo(False, "Process multiple images in batch in parallel"),