diff --git a/backend/misc/tomesd.py b/backend/misc/tomesd.py
index d977217a..db81c0e6 100644
--- a/backend/misc/tomesd.py
+++ b/backend/misc/tomesd.py
@@ -1,24 +1,34 @@
+# Reference: https://github.com/comfyanonymous/ComfyUI/blob/master/comfy_extras/nodes_tomesd.py
+# Credit: https://github.com/dbolya/tomesd
+
+from typing import TYPE_CHECKING, Callable
+
+if TYPE_CHECKING:
+ from modules_forge.unet_patcher import UnetPatcher
+
import math
-from typing import Callable, Tuple
import torch
+from modules.shared import opts
-def do_nothing(x: torch.Tensor, mode: str = None):
+
+def do_nothing(x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
return x
-def mps_gather_workaround(input, dim, index):
+def mps_gather_workaround(input: torch.Tensor, dim: int, index: torch.Tensor) -> torch.Tensor:
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]:
+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.
+ Partitions the tokens into src and dst, and merges r tokens from src to dst.
+ dst tokens are partitioned by choosing one randomly in each (sx, sy) region.
+
Args:
- metric [B, N, C]: metric to use for similarity
- w: image width in tokens
@@ -30,10 +40,7 @@ def bipartite_soft_matching_random2d(metric: torch.Tensor, w: int, h: int, sx: i
"""
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
+ gather: Callable[..., torch.Tensor] = mps_gather_workaround if metric.device.type == "mps" else torch.gather
with torch.no_grad():
@@ -68,7 +75,7 @@ def bipartite_soft_matching_random2d(metric: torch.Tensor, w: int, h: int, sx: i
a_idx = rand_idx[:, num_dst:, :] # src
b_idx = rand_idx[:, :num_dst, :] # dst
- def split(x):
+ def split(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
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))
@@ -90,6 +97,7 @@ def bipartite_soft_matching_random2d(metric: torch.Tensor, w: int, h: int, sx: i
src_idx = edge_idx[..., :r, :] # Merged Tokens
dst_idx = gather(node_idx[..., None], dim=-2, index=src_idx)
+ @torch.inference_mode()
def merge(x: torch.Tensor, mode="mean") -> torch.Tensor:
src, dst = split(x)
n, t1, c = src.shape
@@ -100,6 +108,7 @@ def bipartite_soft_matching_random2d(metric: torch.Tensor, w: int, h: int, sx: i
return torch.cat([unm, dst], dim=1)
+ @torch.inference_mode()
def unmerge(x: torch.Tensor) -> torch.Tensor:
unm_len = unm_idx.shape[1]
unm, dst = x[..., :unm_len, :], x[..., unm_len:, :]
@@ -118,37 +127,40 @@ def bipartite_soft_matching_random2d(metric: torch.Tensor, w: int, h: int, sx: i
return merge, unmerge
-def get_functions(x, ratio, original_shape):
- b, c, original_h, original_w = original_shape
+def get_functions(x: torch.Tensor, ratio: float, original_shape: list[int]) -> tuple[Callable, Callable]:
+ _, _, 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
+ stride_x: int = opts.token_merging_stride
+ stride_y: int = opts.token_merging_stride
+ max_downsample: int = opts.token_merging_downsample
+ no_rand: bool = opts.token_merging_no_rand
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
+ if r <= 0 or w == 1 or h == 1:
+ return do_nothing, do_nothing
+
+ return bipartite_soft_matching_random2d(x, w, h, stride_x, stride_y, r, no_rand)
+
+ return do_nothing, do_nothing
class TomePatcher:
- def __init__(self):
- self.u = None
+ @classmethod
+ def patch(cls, model: "UnetPatcher", ratio: float):
+ cls.u = None
- def patch(self, model, ratio):
- def tomesd_m(q, k, v, extra_options):
- m, self.u = get_functions(q, ratio, extra_options["original_shape"])
+ def tomesd_m(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, extra_options: dict):
+ m, cls.u = get_functions(q, ratio, extra_options["original_shape"])
return m(q), k, v
- def tomesd_u(n, extra_options):
- return self.u(n)
+ def tomesd_u(n: torch.Tensor, *args, **kwargs):
+ return cls.u(n)
m = model.clone()
m.set_model_attn1_patch(tomesd_m)
diff --git a/modules/processing.py b/modules/processing.py
index a6088e73..88bb4d1a 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -429,8 +429,6 @@ class StableDiffusionProcessing:
opts.sdxl_crop_top,
self.width,
self.height,
- opts.fp8_storage,
- opts.cache_fp16_weight,
opts.emphasis,
)
@@ -746,8 +744,6 @@ def create_infotext(p, all_prompts, all_seeds, all_subseeds, comments=None, iter
"Size": f"{p.width}x{p.height}",
"Model hash": p.sd_model_hash if opts.add_model_hash_to_info else None,
"Model": p.sd_model_name if opts.add_model_name_to_info else None,
- "FP8 weight": opts.fp8_storage if devices.fp8 else None,
- "Cache FP16 weight for LoRA": opts.cache_fp16_weight if devices.fp8 else None,
# "VAE hash": p.sd_vae_hash if opts.add_vae_hash_to_info else None,
# "VAE": p.sd_vae_name if opts.add_vae_name_to_info else None,
"Variation seed": (None if p.subseed_strength == 0 else (p.all_subseeds[0] if use_main_prompt else all_subseeds[index])),
diff --git a/modules/sd_models.py b/modules/sd_models.py
index da082b61..cc8bf30d 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -452,7 +452,7 @@ def apply_token_merging(sd_model, token_merging_ratio):
from backend.misc.tomesd import TomePatcher
- sd_model.forge_objects.unet = TomePatcher().patch(
+ sd_model.forge_objects.unet = TomePatcher.patch(
model=sd_model.forge_objects.unet,
ratio=token_merging_ratio
)
diff --git a/modules/shared_options.py b/modules/shared_options.py
index 36bdbcbf..fb92b9cb 100644
--- a/modules/shared_options.py
+++ b/modules/shared_options.py
@@ -263,18 +263,25 @@ options_templates.update(
options_section(
("optimizations", "Optimizations", "sd"),
{
- "cross_attention_optimization": OptionInfo("Automatic", "Cross attention optimization", gr.Dropdown, lambda: {"choices": shared_items.cross_attention_optimizations()}),
- "s_min_uncond": OptionInfo(0.0, "Negative Guidance minimum sigma", gr.Slider, {"minimum": 0.0, "maximum": 15.0, "step": 0.01}, infotext="NGMS").link("PR", "https://github.com/AUTOMATIC1111/stablediffusion-webui/pull/9177").info("skip negative prompt for some steps when the image is almost ready; 0=disable, higher=faster"),
- "s_min_uncond_all": OptionInfo(False, "Negative Guidance minimum sigma all steps", infotext="NGMS all steps").info("By default, NGMS above skips every other step; this makes it skip all steps"),
- "token_merging_ratio": OptionInfo(0.0, "Token merging ratio", gr.Slider, {"minimum": 0.0, "maximum": 0.9, "step": 0.1}, infotext="Token merging ratio").link("PR", "https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/9256").info("0=disable, higher=faster"),
- "token_merging_ratio_img2img": OptionInfo(0.0, "Token merging ratio for img2img", gr.Slider, {"minimum": 0.0, "maximum": 0.9, "step": 0.1}).info("only applies if non-zero and overrides above"),
- "token_merging_ratio_hr": OptionInfo(0.0, "Token merging ratio for high-res pass", gr.Slider, {"minimum": 0.0, "maximum": 0.9, "step": 0.1}, infotext="Token merging ratio hr").info("only applies if non-zero and overrides above"),
- "pad_cond_uncond": OptionInfo(False, "Pad prompt/negative prompt", infotext="Pad conds").info("improves performance when prompt and negative prompt have different lengths; changes seeds"),
- "pad_cond_uncond_v0": OptionInfo(False, "Pad prompt/negative prompt (v0)", infotext="Pad conds v0").info("alternative implementation for the above; used prior to 1.6.0 for DDIM sampler; overrides the above if set; WARNING: truncates negative prompt if it's too long; changes seeds"),
- "persistent_cond_cache": OptionInfo(True, "Persistent cond cache").info("do not recalculate conds from prompts if prompts have not changed since previous calculation"),
- "batch_cond_uncond": OptionInfo(True, "Batch cond/uncond").info("do both conditional and unconditional denoising in one batch; uses a bit more VRAM during sampling, but improves speed; previously this was controlled by --always-batch-cond-uncond commandline argument"),
- "fp8_storage": OptionInfo("Disable", "FP8 weight", gr.Radio, {"choices": ["Disable", "Enable for SDXL", "Enable"]}).info("Use FP8 to store Linear/Conv layers' weight. Require pytorch>=2.1.0."),
- "cache_fp16_weight": OptionInfo(False, "Cache FP16 weight for LoRA").info("Cache fp16 weight when enabling FP8, will increase the quality of LoRA. Use more system ram."),
+ "cross_attention_optimization": OptionInfo("Automatic", "Cross Attention Optimization", gr.Dropdown, {"choices": ("Automatic",), "interactive": False}),
+ "persistent_cond_cache": OptionInfo(True, "Persistent Cond Cache").info("do not recalculate conds if the prompts and parameters have not changed since previous generation"),
+ "skip_early_cond": OptionInfo(0.0, "Ignore Negative Prompt during Early Steps", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.05}, infotext="Skip Early CFG").info("in percentage of total steps; 0 = disable; higher = faster"),
+ "s_min_uncond": OptionInfo(0.0, "Skip Negative Prompt during Later Steps", gr.Slider, {"minimum": 0.0, "maximum": 8.0, "step": 0.05}).info('in "sigma"; 0 = disable; higher = faster'),
+ "s_min_uncond_all": OptionInfo(False, "For the above option, skip every step", infotext="NGMS all steps").info("otherwise, only skip every other step"),
+ "div_tome": OptionDiv(),
+ "token_merging_explanation": OptionHTML(
+ """
+Token Merging speeds up the diffusion process by fusing "redundant" tokens together, but also reduces quality as a result.
+[GitHub]
+Note: Has no effect on SDXL when Max Downsample is set to 1
+ """
+ ),
+ "token_merging_ratio": OptionInfo(0.0, "Token Merging Ratio", gr.Slider, {"minimum": 0.0, "maximum": 0.9, "step": 0.05}, infotext="Token merging ratio").info("0 = disable; higher = faster"),
+ "token_merging_ratio_img2img": OptionInfo(0.0, "Token Merging Ratio for img2img", gr.Slider, {"minimum": 0.0, "maximum": 0.9, "step": 0.05}).info("overrides base ratio if non-zero"),
+ "token_merging_ratio_hr": OptionInfo(0.0, "Token Merging Ratio for Hires. fix", gr.Slider, {"minimum": 0.0, "maximum": 0.9, "step": 0.05}, infotext="Token merging ratio hr").info("overrides base ratio if non-zero"),
+ "token_merging_stride": OptionInfo(2, "Token Merging - Stride", gr.Slider, {"minimum": 1, "maximum": 8, "step": 1}).info("higher = faster"),
+ "token_merging_downsample": OptionInfo(1, "Token Merging - Max Downsample", gr.Slider, {"minimum": 1, "maximum": 4, "step": 1}).info("higher = faster"),
+ "token_merging_no_rand": OptionInfo(False, "Token Merging - No Random").info("reduce randomness by always fusing the same regions"),
},
)
)
@@ -470,7 +477,6 @@ options_templates.update(
"uni_pc_order": OptionInfo(3, "UniPC order", gr.Slider, {"minimum": 1, "maximum": 50, "step": 1}, infotext="UniPC order").info("must be < sampling steps"),
"uni_pc_lower_order_final": OptionInfo(True, "UniPC lower order final", infotext="UniPC lower order final"),
"sd_noise_schedule": OptionInfo("Default", "Noise schedule for sampling", gr.Radio, {"choices": ["Default", "Zero Terminal SNR"]}, infotext="Noise Schedule").info("for use with zero terminal SNR trained models"),
- "skip_early_cond": OptionInfo(0.0, "Ignore negative prompt during early sampling", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}, infotext="Skip Early CFG").info("disables CFG on a proportion of steps at the beginning of generation; 0=skip none; 1=skip all; can both improve sample diversity/quality and speed up sampling"),
"beta_dist_alpha": OptionInfo(0.6, "Beta scheduler - alpha", gr.Slider, {"minimum": 0.01, "maximum": 1.0, "step": 0.01}, infotext="Beta scheduler alpha").info("Default = 0.6; the alpha parameter of the beta distribution used in Beta sampling"),
"beta_dist_beta": OptionInfo(0.6, "Beta scheduler - beta", gr.Slider, {"minimum": 0.01, "maximum": 1.0, "step": 0.01}, infotext="Beta scheduler beta").info("Default = 0.6; the beta parameter of the beta distribution used in Beta sampling"),
},
diff --git a/scripts/xyz_grid.py b/scripts/xyz_grid.py
index 363199f7..6d3c0c9c 100644
--- a/scripts/xyz_grid.py
+++ b/scripts/xyz_grid.py
@@ -90,7 +90,7 @@ def apply_checkpoint(p, x, xs):
def refresh_loading_params_for_xyz_grid():
"""
- Refreshes the loading parameters for the model,
+ Refreshes the loading parameters for the model,
prompts a reload in sd_models.forge_model_reload()
"""
checkpoint_info = select_checkpoint()
@@ -303,7 +303,6 @@ axis_options = [
AxisOption("Refiner checkpoint", str, apply_field('refiner_checkpoint'), format_value=format_remove_path, confirm=confirm_checkpoints_or_none, cost=1.0, choices=lambda: ['None'] + sorted(sd_models.checkpoints_list, key=str.casefold)),
AxisOption("Refiner switch at", float, apply_field('refiner_switch_at')),
AxisOption("RNG source", str, apply_override("randn_source"), choices=lambda: ["GPU", "CPU", "NV"]),
- AxisOption("FP8 mode", str, apply_override("fp8_storage"), cost=0.9, choices=lambda: ["Disable", "Enable for SDXL", "Enable"]),
AxisOption("Size", str, apply_size),
]
@@ -806,7 +805,7 @@ class Script(scripts.Script):
second_axes_processed=second_axes_processed,
margin_size=margin_size
)
-
+
# reset loading params to previous state
refresh_loading_params_for_xyz_grid()