106 lines
3.8 KiB
Python
106 lines
3.8 KiB
Python
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 |