mirror of
https://github.com/lllyasviel/stable-diffusion-webui-forge.git
synced 2026-07-21 21:01:24 +08:00
option - optimizations
This commit is contained in:
Haoming
parent
a62654b9be
commit
4e8dc22b09
@ -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)
|
||||
|
||||
@ -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])),
|
||||
|
||||
@ -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
|
||||
)
|
||||
|
||||
@ -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(
|
||||
"""
|
||||
<b>Token Merging</b> speeds up the diffusion process by fusing "redundant" tokens together, but also reduces quality as a result.
|
||||
[<a href="https://github.com/dbolya/tomesd">GitHub</a>] <br>
|
||||
<b>Note:</b> 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"),
|
||||
},
|
||||
|
||||
@ -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()
|
||||
|
||||
|
||||
Loading…
Reference in New Issue
Block a user