support Radial Attention

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)

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