diff --git a/extensions-builtin/sd_forge_radial/lib_radial/attn_mask.py b/extensions-builtin/sd_forge_radial/lib_radial/attn_mask.py new file mode 100644 index 00000000..c50b082b --- /dev/null +++ b/extensions-builtin/sd_forge_radial/lib_radial/attn_mask.py @@ -0,0 +1,206 @@ +# https://github.com/woct0rdho/ComfyUI-RadialAttn/blob/main/attn_mask.py +# reference: https://github.com/mit-han-lab/radial-attention + +RADIAL_ENABLE = True + +try: + from spas_sage_attn import block_sparse_sage2_attn_cuda +except ImportError: + RADIAL_ENABLE = False +else: + + from dataclasses import dataclass + from functools import lru_cache + + import torch + from einops import rearrange, repeat + + @dataclass + class MaskMap: + video_token_num: int + num_frame: int + + def get_cuda_arch_versions(): + cuda_archs = [] + for i in range(torch.cuda.device_count()): + major, minor = torch.cuda.get_device_capability(i) + cuda_archs.append(f"sm{major}{minor}") + return cuda_archs + + _cuda_archs: list[str] = get_cuda_arch_versions() + + def sparge_mask_convert(mask: torch.Tensor, block_size: int = 128, arch: str = "sm80") -> torch.Tensor: + assert block_size in [128, 64] + assert mask.shape[0] == mask.shape[1] + + if block_size == 128: + if arch == "sm90": + new_mask = torch.repeat_interleave(mask, 2, dim=0) + else: + new_mask = torch.repeat_interleave(mask, 2, dim=1) + + elif block_size == 64: + if arch == "sm90": + num_row, num_col = mask.shape + reshaped_mask = mask.view(num_row, num_col // 2, 2) + new_mask = torch.max(reshaped_mask, dim=2).values + else: + num_row, num_col = mask.shape + reshaped_mask = mask.view(num_row // 2, 2, num_col) + new_mask = torch.max(reshaped_mask, dim=1).values + + return new_mask + + def shrink_mask_strict(mask, block_size=128): + seqlen = mask.shape[0] + block_num = seqlen // block_size + mask = mask[: block_num * block_size, : block_num * block_size].view(block_num, block_size, block_num, block_size) + col_densities = mask.sum(dim=1) / block_size + + non_zero_densities = col_densities > 0 + high_density_cols = col_densities > 1 / 3 + frac_high_density_cols = high_density_cols.sum(dim=-1) / (non_zero_densities.sum(dim=-1) + 1e-9) + block_mask = frac_high_density_cols > 0.6 + block_mask[0, 0] = True + block_mask[-1, -1] = True + + return block_mask + + def get_diagonal_split_mask(i, j, token_per_frame, sparse_type, device): + assert sparse_type in ["radial"] + + dist = abs(i - j) + group = dist.bit_length() + threshold = 128 + decay_length = 2 ** token_per_frame.bit_length() / 2**group + if decay_length >= threshold: + return torch.ones((token_per_frame, token_per_frame), device=device, dtype=torch.bool) + + split_factor = int(threshold / decay_length) + modular = dist % split_factor + if modular == 0: + return torch.ones((token_per_frame, token_per_frame), device=device, dtype=torch.bool) + else: + return torch.zeros((token_per_frame, token_per_frame), device=device, dtype=torch.bool) + + def get_window_width(i, j, token_per_frame, sparse_type, decay_factor=1, block_size=128, model_type=None): + assert sparse_type in ["radial"] + + dist = abs(i - j) + if model_type == "wan": + if dist < 1: + return token_per_frame + if dist == 1: + return token_per_frame // 2 + elif model_type == "hunyuan": + if dist <= 1: + return token_per_frame + else: + raise ValueError(f"Unknown model type: {model_type}") + + group = dist.bit_length() + decay_length = 2 ** token_per_frame.bit_length() / 2**group * decay_factor + threshold = block_size + + if decay_length >= threshold: + return decay_length + else: + return threshold + + def gen_log_mask_shrinked(shape, device, video_token_num, num_frame, block_size=128, sparse_type="log", decay_factor=0.5, model_type=None): + s = shape[0] + final_log_mask = torch.zeros((s // block_size, s // block_size), device=device, dtype=torch.bool) + token_per_frame = video_token_num // num_frame + video_text_border = video_token_num // block_size + + col_indices = torch.arange(0, token_per_frame, device=device).view(1, -1) + row_indices = torch.arange(0, token_per_frame, device=device).view(-1, 1) + final_log_mask[video_text_border:] = True + final_log_mask[:, video_text_border:] = True + + for i in range(num_frame): + for j in range(num_frame): + local_mask = torch.zeros((token_per_frame, token_per_frame), device=device, dtype=torch.bool) + if j == 0 and model_type == "wan": + local_mask = torch.ones((token_per_frame, token_per_frame), device=device, dtype=torch.bool) + else: + window_width = get_window_width( + i, + j, + token_per_frame, + sparse_type, + decay_factor=decay_factor, + block_size=block_size, + model_type=model_type, + ) + local_mask = torch.abs(col_indices - row_indices) <= window_width + split_mask = get_diagonal_split_mask(i, j, token_per_frame, sparse_type, device) + local_mask = torch.logical_and(local_mask, split_mask) + + remainder_row = (i * token_per_frame) % block_size + remainder_col = (j * token_per_frame) % block_size + + all_length_row = remainder_row + ((token_per_frame - 1) // block_size + 1) * block_size + all_length_col = remainder_col + ((token_per_frame - 1) // block_size + 1) * block_size + padded_local_mask = torch.zeros((all_length_row, all_length_col), device=device, dtype=torch.bool) + padded_local_mask[remainder_row : remainder_row + token_per_frame, remainder_col : remainder_col + token_per_frame] = local_mask + + block_mask = shrink_mask_strict(padded_local_mask, block_size=block_size) + + block_row_start = (i * token_per_frame) // block_size + block_col_start = (j * token_per_frame) // block_size + block_row_end = block_row_start + block_mask.shape[0] + block_col_end = block_col_start + block_mask.shape[1] + final_log_mask[block_row_start:block_row_end, block_col_start:block_col_end] = torch.logical_or(final_log_mask[block_row_start:block_row_end, block_col_start:block_col_end], block_mask) + + return final_log_mask + + @lru_cache(maxsize=2, typed=False) + def query_log_mask(video_token_num, num_frame, shape, device, sparse_type, block_size=128, decay_factor=0.5, model_type=None): + return gen_log_mask_shrinked( + shape, + device, + video_token_num, + num_frame, + sparse_type=sparse_type, + decay_factor=decay_factor, + model_type=model_type, + block_size=block_size, + ) + + def SpargeSageAttnBackend(query, key, value, mask_map=None, video_mask=None, pre_defined_mask=None, block_size=128): + query_hnd = rearrange(query.unsqueeze(0), "b s h d -> b h s d") + key_hnd = rearrange(key.unsqueeze(0), "b s h d -> b h s d") + value_hnd = rearrange(value.unsqueeze(0), "b s h d -> b h s d") + arch = _cuda_archs[query.device.index] + converted_mask = repeat( + sparge_mask_convert(mask=video_mask, block_size=block_size, arch=arch), + "s t -> b h s t", + b=query_hnd.shape[0], + h=query_hnd.shape[1], + ) + + converted_mask = converted_mask.to(torch.int8) + output = block_sparse_sage2_attn_cuda( + query_hnd[:, :, : mask_map.video_token_num, :], + key_hnd[:, :, : mask_map.video_token_num, :], + value_hnd[:, :, : mask_map.video_token_num, :], + mask_id=converted_mask, + tensor_layout="HND", + ) + + output = rearrange(output, "b h s d -> s (b h) d", b=1) + return output + + def RadialAttention( + query, + key, + value, + mask_map=None, + block_size=128, + decay_factor=1, + model_type=None, + ): + assert mask_map is not None + video_mask = query_log_mask(mask_map.video_token_num, mask_map.num_frame, query.shape, query.device, "radial", block_size=block_size, decay_factor=decay_factor, model_type=model_type) + return SpargeSageAttnBackend(query, key, value, mask_map, video_mask, None, block_size=block_size) diff --git a/extensions-builtin/sd_forge_radial/scripts/forge_radial.py b/extensions-builtin/sd_forge_radial/scripts/forge_radial.py new file mode 100644 index 00000000..e385e337 --- /dev/null +++ b/extensions-builtin/sd_forge_radial/scripts/forge_radial.py @@ -0,0 +1,157 @@ +# https://github.com/woct0rdho/ComfyUI-RadialAttn/blob/main/nodes.py + +from functools import wraps +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from backend.patcher.unet import UnetPatcher + +import gradio as gr +import torch +import torch.nn.functional as F +from lib_radial.attn_mask import RADIAL_ENABLE, MaskMap, RadialAttention + +from backend.nn import wan +from modules import scripts +from modules.ui_components import InputAccordion + +ORIG_ATTENTION = wan.attention + + +def get_radial_attn_func(video_token_num, num_frame, block_size, decay_factor, allow_compile): + mask_map = MaskMap(video_token_num, num_frame) + + def radial_attn_func(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs): + assert q.shape == k.shape + assert mask is None + + if skip_reshape: + b, _, orig_seq_len, head_dim = q.shape + q, k, v = map(lambda t: t.permute(0, 2, 1, 3).reshape(-1, heads, head_dim), (q, k, v)) + else: + b, orig_seq_len, head_dim = q.shape + head_dim //= heads + q, k, v = map(lambda t: t.view(-1, heads, head_dim), (q, k, v)) + + padded_len = b * video_token_num + if q.shape[0] != padded_len: + q, k, v = map(lambda t: F.pad(t, (0, 0, 0, 0, 0, padded_len - t.shape[0])), (q, k, v)) + + out = RadialAttention(q, k, v, mask_map=mask_map, block_size=block_size, decay_factor=decay_factor, model_type="wan") + + out = out[: b * orig_seq_len, :, :] + + if skip_output_reshape: + out = out.reshape(b, orig_seq_len, heads, head_dim).permute(0, 2, 1, 3) + else: + out = out.reshape(b, -1, heads * head_dim) + + return out + + if not allow_compile: + radial_attn_func = torch.compiler.disable()(radial_attn_func) + + return radial_attn_func + + +class PatchRadialAttn: + + @staticmethod + def patch_radial_attn(model: "UnetPatcher", dense_block: int, dense_timestep: int, last_dense_timestep: int, block_size: int, decay_factor: float, allow_compile: bool): + model = model.clone() + + diffusion_model = model.get_model_object("diffusion_model") + assert diffusion_model.__class__.__name__ == "WanModel" + + model.model_options["transformer_options"]["radial_attn"] = { + "patch_size": diffusion_model.patch_size, + "dense_block": dense_block, + "dense_timestep": dense_timestep, + "last_dense_timestep": last_dense_timestep, + "block_size": block_size, + "decay_factor": decay_factor, + "allow_compile": allow_compile, + } + + def unet_wrapper_function(model_function, kwargs): + input = kwargs["input"] + timestep = kwargs["timestep"] + c = kwargs["c"] + sigmas = c["transformer_options"]["sampling_sigmas"] + + if len(matched_step_index := (sigmas == timestep).nonzero()) > 0: + current_step_index = matched_step_index.item() + else: + for i in range(len(sigmas) - 1): + if (sigmas[i] - timestep[0]) * (sigmas[i + 1] - timestep[0]) <= 0: + current_step_index = i + break + else: + current_step_index = 0 + + ra_options = c["transformer_options"]["radial_attn"] + + if not (ra_options["dense_timestep"] <= current_step_index < len(sigmas) - 1 - ra_options["last_dense_timestep"]): + wan.attention = ORIG_ATTENTION + else: + patch_size = ra_options["patch_size"] + num_frame = (input.shape[2] - 1) // patch_size[0] + 1 + frame_size = (input.shape[3] // patch_size[1]) * (input.shape[4] // patch_size[2]) + video_token_num = frame_size * num_frame + + padded_video_token_num = video_token_num + block_size = ra_options["block_size"] + if video_token_num % block_size != 0: + padded_video_token_num = (video_token_num // block_size + 1) * block_size + + dense_block = ra_options["dense_block"] + radial_attn_func = get_radial_attn_func(padded_video_token_num, num_frame, block_size, ra_options["decay_factor"], ra_options["allow_compile"]) + + @wraps(ORIG_ATTENTION) + def try_radial_attn(*args, **kwargs): + transformer_options = kwargs.get("transformer_options", {}) + block_index = transformer_options.get("block_index", -1) + + if block_index >= dense_block: + return radial_attn_func(*args, **kwargs) + else: + return ORIG_ATTENTION(*args, **kwargs) + + wan.attention = try_radial_attn + + return model_function(input, timestep, **c) + + model.set_model_unet_function_wrapper(unet_wrapper_function) + + return model + + +class RadialAttentionForForge(scripts.ScriptBuiltinUI): + sorting_priority = 18137 + + def title(self): + return "RadialAttention Integrated" + + def show(self, is_img2img): + return scripts.AlwaysVisible if RADIAL_ENABLE else None + + def ui(self, *args, **kwargs): + with InputAccordion(False, label=self.title()) as enable: + dense_block = gr.Slider(minimum=0, maximum=40, value=1, step=1, label="dense_block", info="Number of first few Blocks to bypass RadialAttention") + dense_timestep = gr.Slider(minimum=0, maximum=100, value=1, step=1, label="dense_timestep", info="Number of first few Steps to bypass RadialAttention") + last_dense_timestep = gr.Slider(minimum=0, maximum=100, value=1, step=1, label="last_dense_timestep", info="Number of last few Steps to bypass RadialAttention") + block_size = gr.Radio(choices=[64, 128], value=128, label="block_size") + decay_factor = gr.Slider(minimum=0.0, maximum=1.0, value=0.2, step=0.1, label="decay_factor", info="Lower is Faster ; Higher is more Accurate") + allow_compile = gr.Checkbox(False, label="allow_compile", info="Allow the use of torch.compile") + + return [enable, dense_block, dense_timestep, last_dense_timestep, block_size, decay_factor, allow_compile] + + def process_before_every_sampling(self, p, enable: bool, *args, **kwargs): + if not enable: + return + + unet = p.sd_model.forge_objects.unet + + unet = PatchRadialAttn.patch_radial_attn(unet, *args) + + p.sd_model.forge_objects.unet = unet