diff --git a/utils/blending.py b/utils/blending.py new file mode 100644 index 0000000..5706348 --- /dev/null +++ b/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 \ No newline at end of file