From 9832f6ea21c645f068d9cefcb9efcb8ff0bce6dd Mon Sep 17 00:00:00 2001 From: smthemex <138738845+smthemex@users.noreply.github.com> Date: Mon, 20 Oct 2025 15:42:27 +0800 Subject: [PATCH] Update utils.py --- FlashVSR/examples/WanVSR/utils/utils.py | 112 +++++++++++++++++++++++- 1 file changed, 111 insertions(+), 1 deletion(-) diff --git a/FlashVSR/examples/WanVSR/utils/utils.py b/FlashVSR/examples/WanVSR/utils/utils.py index 00496c9..6b9d6a3 100644 --- a/FlashVSR/examples/WanVSR/utils/utils.py +++ b/FlashVSR/examples/WanVSR/utils/utils.py @@ -5,10 +5,120 @@ import torch.nn as nn import torch.nn.functional as F from tqdm import tqdm import time - +import math CACHE_T = 2 +def block_sparse_attn_func(q, k, v, cu_seqlens_q, cu_seqlens_k, head_mask_type, + streaming_info, base_blockmask, max_seqlen_q_, max_seqlen_k_, + p_dropout, deterministic=False, softmax_scale=None, + is_causal=False, exact_streaming=False, return_attn_probs=False): + """ + 使用标准注意力机制替代块稀疏注意力实现 + """ + batch_size = 1 + seq_len_q = q.shape[0] // batch_size + seq_len_k = k.shape[0] // batch_size + num_heads = q.shape[1] + head_dim = q.shape[2] + + # 重塑张量为标准注意力格式 + q = q.view(batch_size, seq_len_q, num_heads, head_dim).transpose(1, 2) # (B, H, S_q, D) + k = k.view(batch_size, seq_len_k, num_heads, head_dim).transpose(1, 2) # (B, H, S_k, D) + v = v.view(batch_size, seq_len_k, num_heads, head_dim).transpose(1, 2) # (B, H, S_k, D) + + # 使用 F.scaled_dot_product_attention，它能更好地处理各种情况 + if base_blockmask is not None: + if base_blockmask.dim() == 2: + attn_mask = base_blockmask.unsqueeze(0).unsqueeze(0) + attn_mask = attn_mask.expand(batch_size, num_heads, -1, -1) + elif base_blockmask.dim() == 4: + attn_mask = base_blockmask + else: + attn_mask = None + if attn_mask is not None: + _, _, mask_seq_q, mask_seq_k = attn_mask.shape + + target_seq_q = seq_len_q + target_seq_k = seq_len_k + #print(f"Attention mask shape: {attn_mask.shape},{target_seq_q},{target_seq_k}") + #Attention mask shape: torch.Size([1, 12, 24, 96]),3072,12288 + # 使用插值方法进行扩展，模式为 nearest 以保持稀疏结构 + if mask_seq_q != target_seq_q or mask_seq_k != target_seq_k: + try: + attn_mask = F.interpolate( + attn_mask.float(), + size=(target_seq_q, target_seq_k), + mode='nearest' + ).to(attn_mask.dtype) + except torch.OutOfMemoryError: + # 如果直接插值OOM，则分块处理 + print("Direct interpolation OOM, using chunked processing") + attn_mask = chunked_interpolate(attn_mask, (target_seq_q, target_seq_k)) + else: + attn_mask = None + training=False + try: + output = F.scaled_dot_product_attention( + q, k, v, + attn_mask=attn_mask if attn_mask is not None else None, + dropout_p=p_dropout if training else 0.0, + is_causal=is_causal + ) + except Exception: + attention_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(head_dim) + if attn_mask is not None: + if attn_mask.shape != attention_scores.shape: + min_shape = [min(a, b) for a, b in zip(attn_mask.shape, attention_scores.shape)] + attn_mask = attn_mask[:min_shape[0], :min_shape[1], :min_shape[2], :min_shape[3]] + attention_scores = attention_scores[:min_shape[0], :min_shape[1], :min_shape[2], :min_shape[3]] + if attn_mask.dtype == torch.bool: + attention_scores = attention_scores.masked_fill(~attn_mask, float('-inf')) + else: + attention_scores = attention_scores + attn_mask + attention_weights = F.softmax(attention_scores, dim=-1) + output = torch.matmul(attention_weights, v) + output = output.transpose(1, 2).contiguous().view(batch_size * seq_len_q, num_heads, head_dim) + return output.squeeze(0) + +def chunked_interpolate(tensor, target_size, chunk_size=1024): + """ + 分块插值以减少内存消耗 + """ + batch, heads, h, w = tensor.shape + target_h, target_w = target_size + + output = torch.zeros(batch, heads, target_h, target_w, dtype=tensor.dtype, device=tensor.device) + + chunk_h = min(chunk_size, h) + chunk_w = min(chunk_size, w) + + for i in range(0, h, chunk_h): + for j in range(0, w, chunk_w): + end_i = min(i + chunk_h, h) + end_j = min(j + chunk_w, w) + + chunk = tensor[:, :, i:end_i, j:end_j] + + + target_end_i = int(target_h * end_i / h) + target_start_i = int(target_h * i / h) + target_end_j = int(target_w * end_j / w) + target_start_j = int(target_w * j / w) + + + if chunk.numel() > 0: + interpolated = F.interpolate( + chunk.float(), + size=(target_end_i - target_start_i, target_end_j - target_start_j), + mode='nearest' + ).to(tensor.dtype) + + + output[:, :, target_start_i:target_end_i, target_start_j:target_end_j] = interpolated + + return output + class RMS_norm(nn.Module):