Add files via upload

utils/blending.py

@@ -0,0 +1,106 @@
+import torch
+import math
+
+def generate_weight_mask(shape, blend_mode="gaussian", feather=0.25, device="cpu"):
+    """
+    生成用于分块融合的权重遮罩
+    shape: (H, W) or (C, H, W)
+    blend_mode: 'linear', 'gaussian', 'cosine', 'none'
+    feather: 0.0 - 1.0 (羽化范围比例)
+    """
+    if len(shape) == 3:
+        h, w = shape[1], shape[2]
+    else:
+        h, w = shape[0], shape[1]
+        
+    # 创建基础网格
+    # linspace 生成从 0 到 1 的线性空间
+    x = torch.linspace(-1, 1, w, device=device)
+    y = torch.linspace(-1, 1, h, device=device)
+    grid_y, grid_x = torch.meshgrid(y, x, indexing='ij')
+    
+    # 调整羽化范围，feather越大，有效中心区域越小，边缘衰减越宽
+    # 我们通过缩放网格坐标来控制边缘
+    if feather <= 0:
+        mask = torch.ones((h, w), device=device)
+    else:
+        # 映射使得边缘 feather% 的区域从 1 衰减到 0
+        # 这里的逻辑是：在 1-feather 范围内是 1，超过这个范围开始衰减
+        scale = 1.0 / (1.0 - feather + 1e-6) # 避免除零
+        
+        if blend_mode == "gaussian":
+            # 高斯分布: exp(-(x^2 + y^2) / sigma)
+            # 调整 sigma 使得在边缘处接近 0
+            d = torch.sqrt(grid_x**2 + grid_y**2)
+            sigma = 0.5  # 标准差
+            mask = torch.exp(-(d**2) / (2 * sigma**2))
+            # 归一化到 0-1 并截断边缘
+            mask = (mask - mask.min()) / (mask.max() - mask.min())
+            
+        elif blend_mode == "cosine":
+            # 余弦窗口
+            mask_x = torch.cos(grid_x * math.pi / 2)
+            mask_y = torch.cos(grid_y * math.pi / 2)
+            mask = mask_x * mask_y
+            mask = torch.clamp(mask, 0, 1)
+            
+        elif blend_mode == "linear":
+            # 线性金字塔
+            mask_x = 1.0 - torch.abs(grid_x)
+            mask_y = 1.0 - torch.abs(grid_y)
+            mask = torch.min(mask_x, mask_y)
+            mask = torch.clamp(mask, 0, 1)
+            
+        else: # none
+            mask = torch.ones((h, w), device=device)
+
+    # 扩展维度以匹配 Image (H,W,C) 或 Latent (C,H,W)
+    # 假设输入 shape 决定了输出 mask 的维度
+    if len(shape) == 3:
+        # Latent: (C, H, W), mask 需要是 (1, H, W) 用于广播
+        return mask.unsqueeze(0)
+    else:
+        # Image 通常在 Comfy 中处理时如果是单张 mask 是 (H, W)
+        return mask
+
+def get_tiled_splits(full_height, full_width, tile_size, overlap):
+    """
+    计算分块坐标。采用 "覆盖优先" 策略，最后一块可能会有较大的重叠，
+    以避免产生极小的边缘切片。
+    Returns: list of (y, x, h, w)
+    """
+    stride = tile_size - overlap
+    splits = []
+    
+    # 简单的网格生成
+    y_coords = []
+    current_y = 0
+    while current_y < full_height:
+        y_coords.append(current_y)
+        if current_y + tile_size >= full_height:
+            break
+        current_y += stride
+    
+    # 修正最后一个坐标，确保不越界且覆盖边缘
+    if y_coords[-1] + tile_size > full_height:
+        y_coords[-1] = max(0, full_height - tile_size)
+        
+    x_coords = []
+    current_x = 0
+    while current_x < full_width:
+        x_coords.append(current_x)
+        if current_x + tile_size >= full_width:
+            break
+        current_x += stride
+        
+    if x_coords[-1] + tile_size > full_width:
+        x_coords[-1] = max(0, full_width - tile_size)
+
+    for y in y_coords:
+        for x in x_coords:
+            # 实际切片大小（通常等于 tile_size，除非原图小于 tile_size）
+            h = min(tile_size, full_height - y)
+            w = min(tile_size, full_width - x)
+            splits.append((y, x, h, w))
+            
+    return splits