153 lines
6.9 KiB
Python
153 lines
6.9 KiB
Python
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
|
|
def resize_mask_to_latent(outpaint_mask, latent_h, latent_w, device=None):
|
|
"""Resize outpaint_mask [1, H, W, 1] to latent space [1, latent_H, latent_W]."""
|
|
mask = outpaint_mask[:, :, :, 0]
|
|
result = F.interpolate(mask.unsqueeze(1), size=(latent_h, latent_w), mode='nearest').squeeze(1)
|
|
return result.to(device) if device is not None else result
|
|
|
|
|
|
def blend_tile_into_canvas(final_tensor, tile, pos_x, pos_y, overlap_x, overlap_y, has_left, has_top):
|
|
"""
|
|
Blend a tile into the canvas with linear feathering over overlap zones.
|
|
|
|
Uses the neighbor's single edge pixel (not the full overlap strip) as the blend
|
|
start, so neighbor content is never echoed into the new tile's zone — eliminating
|
|
the double-seam artifact where the edge appears on both sides of the boundary.
|
|
"""
|
|
tile_h, tile_w = tile.shape[0], tile.shape[1]
|
|
tile_cpu = tile.cpu()
|
|
|
|
# Hard copy the full tile
|
|
final_tensor[0, pos_y:pos_y + tile_h, pos_x:pos_x + tile_w, :] = tile_cpu
|
|
|
|
# Left seam: fade from left neighbor's rightmost column into this tile's left edge
|
|
if has_left and overlap_x > 0:
|
|
left_edge = final_tensor[0, pos_y:pos_y + tile_h,
|
|
pos_x - 1:pos_x, :].clone() # single column [H,1,3]
|
|
ramp = torch.linspace(0.0, 1.0, overlap_x).view(1, overlap_x, 1).expand(tile_h, -1, 1)
|
|
final_tensor[0, pos_y:pos_y + tile_h, pos_x:pos_x + overlap_x, :] = (
|
|
left_edge * (1.0 - ramp) + tile_cpu[:, :overlap_x, :] * ramp
|
|
)
|
|
|
|
# Top seam: fade from top neighbor's bottom row into this tile's top edge
|
|
if has_top and overlap_y > 0:
|
|
top_edge = final_tensor[0, pos_y - 1:pos_y,
|
|
pos_x:pos_x + tile_w, :].clone() # single row [1,W,3]
|
|
ramp = torch.linspace(0.0, 1.0, overlap_y).view(overlap_y, 1, 1).expand(-1, tile_w, 1)
|
|
final_tensor[0, pos_y:pos_y + overlap_y, pos_x:pos_x + tile_w, :] = (
|
|
top_edge * (1.0 - ramp) + tile_cpu[:overlap_y, :, :] * ramp
|
|
)
|
|
|
|
# Corner: fade from the single corner pixel diagonally into this tile's top-left
|
|
if has_left and has_top and overlap_x > 0 and overlap_y > 0:
|
|
corner_px = final_tensor[0, pos_y - 1:pos_y,
|
|
pos_x - 1:pos_x, :].clone() # single pixel [1,1,3]
|
|
ramp_x = torch.linspace(0.0, 1.0, overlap_x).view(1, overlap_x, 1).expand(overlap_y, -1, 1)
|
|
ramp_y = torch.linspace(0.0, 1.0, overlap_y).view(overlap_y, 1, 1).expand(-1, overlap_x, 1)
|
|
alpha_corner = torch.min(ramp_x, ramp_y)
|
|
final_tensor[0, pos_y:pos_y + overlap_y, pos_x:pos_x + overlap_x, :] = (
|
|
corner_px * (1.0 - alpha_corner) + tile_cpu[:overlap_y, :overlap_x, :] * alpha_corner
|
|
)
|
|
|
|
|
|
def gaussian_blend_tiles(tiles, positions, tile_width, tile_height, overlap_x, overlap_y, final_width, final_height,
|
|
sigma=0.4):
|
|
"""
|
|
Blend tiles using Gaussian weights after all tiles have been generated.
|
|
|
|
Args:
|
|
tiles: List of image tensors in any shape
|
|
positions: List of (x, y) positions for each tile
|
|
tile_width/height: Dimensions of each tile
|
|
overlap_x/y: Overlap amounts in pixels
|
|
final_width/height: Final image dimensions
|
|
sigma: Controls the spread of the Gaussian (smaller = sharper transition)
|
|
|
|
Returns:
|
|
Blended image tensor [1, H, W, C]
|
|
"""
|
|
# Determine device from the first tile
|
|
device = tiles[0].device if len(tiles) > 0 else torch.device('cpu')
|
|
|
|
# Create empty canvas and weight accumulator
|
|
final_tensor = torch.zeros((1, final_height, final_width, 3), dtype=torch.float32, device=device)
|
|
weight_accumulator = torch.zeros((1, final_height, final_width, 1), dtype=torch.float32, device=device)
|
|
|
|
# Process each tile
|
|
for i, (original_tile, (pos_x, pos_y)) in enumerate(zip(tiles, positions)):
|
|
# Print debug info
|
|
#print(f"Processing tile {i}, shape: {original_tile.shape}, position: ({pos_x}, {pos_y})")
|
|
|
|
# Calculate valid region (handle edge cases)
|
|
h = min(tile_height, final_height - pos_y)
|
|
w = min(tile_width, final_width - pos_x)
|
|
|
|
# Create weight map for this tile
|
|
weights = np.ones((h, w), dtype=np.float32)
|
|
|
|
# Apply Gaussian falloff in overlap regions
|
|
# Left edge (if not leftmost tile)
|
|
if pos_x > 0 and overlap_x > 0:
|
|
for x in range(min(overlap_x, w)):
|
|
# Distance from edge (0 at edge, 1 at overlap boundary)
|
|
dist = x / overlap_x
|
|
weights[:, x] *= np.exp(-((1 - dist) ** 2) / (2 * sigma ** 2))
|
|
|
|
# Right edge (if not rightmost tile)
|
|
if pos_x + w < final_width and overlap_x > 0:
|
|
for x in range(max(0, w - overlap_x), w):
|
|
dist = (w - 1 - x) / overlap_x
|
|
weights[:, x] *= np.exp(-((1 - dist) ** 2) / (2 * sigma ** 2))
|
|
|
|
# Top edge (if not topmost tile)
|
|
if pos_y > 0 and overlap_y > 0:
|
|
for y in range(min(overlap_y, h)):
|
|
dist = y / overlap_y
|
|
weights[y, :] *= np.exp(-((1 - dist) ** 2) / (2 * sigma ** 2))
|
|
|
|
# Bottom edge (if not bottommost tile)
|
|
if pos_y + h < final_height and overlap_y > 0:
|
|
for y in range(max(0, h - overlap_y), h):
|
|
dist = (h - 1 - y) / overlap_y
|
|
weights[y, :] *= np.exp(-((1 - dist) ** 2) / (2 * sigma ** 2))
|
|
|
|
# Convert weights to tensor with same shape as tile region
|
|
weights_tensor = torch.from_numpy(weights).to(device).reshape(h, w, 1)
|
|
|
|
try:
|
|
# Handle different tensor shapes
|
|
if len(original_tile.shape) == 4: # [B, H, W, C]
|
|
# Extract the first element from batch
|
|
tile_section = original_tile[0, :h, :w, :]
|
|
elif len(original_tile.shape) == 3: # [H, W, C]
|
|
# No batch dimension
|
|
tile_section = original_tile[:h, :w, :]
|
|
else:
|
|
print(f"Unexpected tensor shape: {original_tile.shape}")
|
|
continue
|
|
|
|
# Apply weighted accumulation directly with matched shapes
|
|
weighted_tile = tile_section * weights_tensor
|
|
|
|
# Add to the final tensor
|
|
final_tensor[0, pos_y:pos_y + h, pos_x:pos_x + w, :] += weighted_tile
|
|
weight_accumulator[0, pos_y:pos_y + h, pos_x:pos_x + w, :] += weights_tensor
|
|
|
|
except Exception as e:
|
|
print(f"Error processing tile {i}: {e}")
|
|
print(f"Tile shape: {original_tile.shape}")
|
|
print(f"Weights shape: {weights_tensor.shape}")
|
|
print(f"Position: ({pos_x}, {pos_y}), size: ({w}, {h})")
|
|
raise
|
|
|
|
# Normalize by accumulated weights (with small epsilon to avoid division by zero)
|
|
epsilon = 1e-8
|
|
mask = weight_accumulator > epsilon
|
|
final_tensor = torch.where(mask, final_tensor / weight_accumulator, final_tensor)
|
|
|
|
return final_tensor
|