From ef909cd003f7dd8a8a5379a724e7d46e23d80ffa Mon Sep 17 00:00:00 2001
From: Disty0 <semihgulec2005@gmail.com>
Date: Sat, 10 Feb 2024 22:12:47 +0300
Subject: [PATCH] HyperTile add swap size and depth options

---
 modules/intel/ipex/attention.py | 2 +-
 modules/sd_hijack_hypertile.py  | 4 ++--
 modules/shared.py               | 3 +++
 3 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/modules/intel/ipex/attention.py b/modules/intel/ipex/attention.py
index 1ca5dd20d..45c0525e6 100644
--- a/modules/intel/ipex/attention.py
+++ b/modules/intel/ipex/attention.py
@@ -122,9 +122,9 @@ def torch_bmm_32_bit(input, mat2, *, out=None):
                     mat2[start_idx:end_idx],
                     out=out
                 )
+        torch.xpu.synchronize(input.device)
     else:
         return original_torch_bmm(input, mat2, out=out)
-    torch.xpu.synchronize(input.device)
     return hidden_states
 
 original_scaled_dot_product_attention = torch.nn.functional.scaled_dot_product_attention
diff --git a/modules/sd_hijack_hypertile.py b/modules/sd_hijack_hypertile.py
index 7087b1436..c27c0305f 100644
--- a/modules/sd_hijack_hypertile.py
+++ b/modules/sd_hijack_hypertile.py
@@ -193,7 +193,7 @@ def context_hypertile_vae(p):
         tile_size = shared.opts.hypertile_vae_tile if shared.opts.hypertile_vae_tile > 0 else max(128, 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)
+        return split_attention(vae, tile_size=tile_size, min_tile_size=128, swap_size=shared.opts.hypertile_vae_swap_size)
 
 
 
@@ -219,7 +219,7 @@ def context_hypertile_unet(p):
         tile_size = shared.opts.hypertile_unet_tile if shared.opts.hypertile_unet_tile > 0 else max(128, 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)
+        return split_attention(unet, tile_size=tile_size, min_tile_size=128, swap_size=shared.opts.hypertile_unet_swap_size, depth=shared.opts.hypertile_unet_depth)
 
 
 def hypertile_set(p, hr=False):
diff --git a/modules/shared.py b/modules/shared.py
index 0d53d312c..1b8feb576 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -420,8 +420,11 @@ options_templates.update(options_section(('advanced', "Inference Settings"), {
     "hypertile_sep": OptionInfo("<h2>HyperTile</h2>", "", gr.HTML),
     "hypertile_unet_enabled": OptionInfo(False, "HyperTile UNet"),
     "hypertile_unet_tile": OptionInfo(0, "HyperTile UNet tile size", gr.Slider, {"minimum": 0, "maximum": 1024, "step": 8}),
+    "hypertile_unet_swap_size": OptionInfo(1, "HyperTile UNet swap size", gr.Slider, {"minimum": 1, "maximum": 10, "step": 1}),
+    "hypertile_unet_depth": OptionInfo(0, "HyperTile UNet depth", gr.Slider, {"minimum": 0, "maximum": 4, "step": 1}),
     "hypertile_vae_enabled": OptionInfo(False, "HyperTile VAE", gr.Checkbox),
     "hypertile_vae_tile": OptionInfo(128, "HyperTile VAE tile size", gr.Slider, {"minimum": 0, "maximum": 1024, "step": 8}),
+    "hypertile_vae_swap_size": OptionInfo(1, "HyperTile VAE swap size", gr.Slider, {"minimum": 1, "maximum": 10, "step": 1}),
 
     "inference_other_sep": OptionInfo("<h2>Other</h2>", "", gr.HTML),
     "batch_frame_mode": OptionInfo(False, "Parallel process images in batch"),