add tiling support based on chaiNNer's implementation

Author: LightCyan01

backend/tiling/__init__.py

@@ -0,0 +1,19 @@
+from .tile_splitter import TileSplitter
+from .tile_blender import TileBlender
+from .blend_functions import linear_blend, half_sin_blend, sin_blend_fn
+from .models import Tile, Region, Padding, TileOverlap
+
+__all__ = [
+    'TileSplitter',
+    'TileBlender',
+    'Tile',
+    'Region',
+    'Padding',
+    'TileOverlap',
+    'linear_blend',
+    'half_sin_blend',
+    'sin_blend_fn',
+]
+
+__version__ = '1.0.0'
+__author__ = 'based on chaiNNer\'s tiling implementation'

backend/tiling/blend_functions.py

@@ -0,0 +1,82 @@
+"""blending weight functions for smooth tile merging"""
+import math
+import numpy as np
+
+def linear_blend(distance: float, overlap: float):
+    """
+    linear gradient blending
+    weight increase from 0.0 to 1.0 linearly
+    """
+    if overlap == 0:
+        return 1.0
+    return distance / (overlap - 1)
+
+def sin_blend_fn(x: float):
+    """
+    chaiNNer's sine blending function
+    Formula: (sin(x * π - π/2) + 1) / 2
+    """
+    return (math.sin(x * math.pi - math.pi / 2) + 1) / 2
+
+def half_sin_blend(distance: float, overlap: float):
+    """
+    chaiNNer's half-sine blending function
+    """
+    if overlap == 0:
+        return 1.0
+    
+    normalized = distance / (overlap - 1)
+    compressed = min(max(normalized * 2 - 0.5, 0), 1)
+    return sin_blend_fn(compressed)
+
+def create_blend_mask(
+    tile_width: int,
+    tile_height: int,
+    overlap: int,
+    blend_fn=half_sin_blend,
+    is_top_edge: bool = False,
+    is_bottom_edge: bool = False,
+    is_right_edge: bool = False,
+    is_left_edge: bool = False
+):
+    """
+    Create 2D weight mask for tile blending.
+    Edges at image boundaries don't get blended (full weight)
+    Internal edges get gradual blending based on blend_fn
+    
+    Returns: 2D numpy array of weights (0.0 to 1.0)
+    """
+    
+    # start with full contribution
+    mask = np.ones((tile_height, tile_width), dtype=np.float32)
+    
+    if overlap <= 1:
+        return mask
+    
+    # apply blending gradient to each edge that overlaps another tile
+    
+    # top edge: blend from 0.0 to 1.0 going top to bottom
+    if not is_top_edge:
+        for i in range(overlap):
+            weight = blend_fn(i, overlap)
+            mask[i, :] *= weight
+            
+    # bottom edge: blend from 1.0 to 0.0 approaching bottom
+    if not is_bottom_edge:
+        for i in range(overlap):
+            weight = blend_fn(i, overlap)
+            mask[-(i+1), :] *= weight
+            
+    # right edge: blend from 1.0 to 0.0 approaching right edge
+    if not is_right_edge:
+        for i in range(overlap):
+            weight = blend_fn(i, overlap)
+            mask[:, -(i+1)] *= weight
+    
+    # left edge: blend from 0.0 to 1.0 going left to right
+    if not is_left_edge:
+        for i in range(overlap):
+            weight = blend_fn(i, overlap)
+            mask[:, i] *= weight
+    
+    return mask

backend/tiling/models.py

@@ -0,0 +1,98 @@
+"""data models for tiling operations"""
+"""based on chaiNNer's region and tileoverlap structures"""
+from dataclasses import dataclass
+from PIL import Image
+import numpy as np
+
+@dataclass
+class Region:
+    """Rectangular region in an image"""
+    x: int
+    y: int
+    width: int
+    height:int
+    
+    def intersect(self, other: 'Region'):
+        """
+        Get the intersection of two regions.
+        Ensures tiles don't extend beyond image boundaries.
+        """
+        x1 = max(self.x, other.x)
+        y1 = max(self.y, other.y)
+        x2 = min(self.x + self.width, other.x + other.width)
+        y2 = min(self.y + self.height, other.y + other.height)
+        return Region(x1, y1, max(0, x2 - x1), max(0, y2 - y1))
+    
+    def add_padding(self, padding: 'Padding'):
+        """
+        Expand this region by adding padding on all sides.
+        Used to create overlap between tiles.
+        """
+        return Region(
+            x=self.x - padding.left,
+            y=self.y - padding.top,
+            width=self.width + padding.left + padding.right,
+            height=self.height + padding.top + padding.bottom
+        )
+    
+    def child_padding(self, child: 'Region'):
+        """
+        Calculate how much padding a child region has relative to this region.
+        Used to determine overlap amounts for tiles.
+        """
+        return Padding(
+            top=child.y - self.y,
+            bottom=(self.y + self.height) - (child.y + child.height),
+            left=child.x - self.x,
+            right=(self.x + self.width) - (child.x + child.width)
+        )
+    
+    def read_from(self, img: np.ndarray) -> np.ndarray:
+        """Extract this region from an image array"""
+        return img[self.y:self.y+self.height, self.x:self.x+self.width, ...]
+    
+@dataclass
+class Padding:
+    """Padding amounts for each side"""
+    top: int
+    bottom: int
+    left: int
+    right: int
+    
+    def min(self, max_padding: int):
+        """
+        Limit padding to a maximum value.
+        Prevents overlap from being larger than necessary.
+        """
+        return Padding(
+            top=min(self.top, max_padding),
+            bottom=min(self.bottom, max_padding),
+            left=min(self.left, max_padding),
+            right=min(self.right, max_padding)
+        )
+
+@dataclass
+class TileOverlap:
+    """
+    ChaiNNer's structure for tracking overlap on a single axis.
+    Start overlap is on the leading edge, end overlap is on the trailing edge.
+    """
+    start: int  # Overlap at the start (left for X, top for Y)
+    end: int    # Overlap at the end (right for X, bottom for Y)
+    
+    @property
+    def total(self) -> int:
+        """Total overlap amount"""
+        return self.start + self.end
+    
+@dataclass
+class Tile:
+    """
+    A tile with image data and position
+    used to track tiles during split -> process -> merge
+    """
+    image: Image.Image
+    region: Region
+    x: int #x pos of orig image
+    y: int #y pos of orig image
+    padding: Padding = None

backend/tiling/tile_blender.py

@@ -0,0 +1,114 @@
+"""
+merges processed tiles with smooth blending.
+based on chaiNNer's directional blending approach.
+"""
+import numpy as np
+from typing import List
+from PIL import Image
+from .models import Tile
+from .blend_functions import create_blend_mask, half_sin_blend
+
+class TileBlender:
+    """
+    Merges upscaled tiles with seamless blending.
+    Uses chaiNNer's two-stage blending approach:
+    1. Blend tiles horizontally into rows (BlendDirection.X)
+    2. Blend rows vertically into final image (BlendDirection.Y)
+    """
+    
+    def __init__(self, overlap: int = 16, scale: int = 1, blend_fn=half_sin_blend):
+        self.overlap = overlap
+        self.scale = scale
+        self.blend_fn = blend_fn
+    
+    def merge(self, tiles: List[Tile], original_width: int, original_height: int):
+        """
+        Merge tiles into final image with blending.
+        """
+        out_w = original_width * self.scale
+        out_h = original_height * self.scale
+        overlap_scaled = self.overlap * self.scale
+        
+        # Accumulators
+        # output: accumulates weighted pixel values
+        # weights: tracks total weight at each pixel
+        output = np.zeros((out_h, out_w, 3), dtype=np.float32)
+        weights = np.zeros((out_h, out_w), dtype=np.float32)
+        
+        for tile in tiles:
+            # Position and size in output image (scaled)
+            x = tile.x * self.scale
+            y = tile.y * self.scale
+            region_w = tile.region.width * self.scale
+            region_h = tile.region.height * self.scale
+            
+            # The tile image includes padding/overlap
+            tile_img = tile.image
+            
+            # Determine edge positions
+            # Tiles at image boundaries don't blend on those edges
+            is_left = (tile.x == 0)
+            is_top = (tile.y == 0)
+            is_right = (tile.x + tile.region.width >= original_width)
+            is_bottom = (tile.y + tile.region.height >= original_height)
+            
+            # Create blending mask for the region size
+            mask = create_blend_mask(
+                region_w, region_h, overlap_scaled, 
+                self.blend_fn,
+                is_top_edge=is_top,
+                is_bottom_edge=is_bottom,
+                is_right_edge=is_right,
+                is_left_edge=is_left
+            )
+            
+            # Convert tile to numpy array
+            tile_arr = np.array(tile_img, dtype=np.float32)
+            
+            # The tile image may have padding - we need to crop it to match the region
+            # Calculate padding that was added
+            tile_w_scaled = tile_img.width
+            tile_h_scaled = tile_img.height
+            
+            # Padding is the difference between tile image size and region size
+            pad_left = (tile_w_scaled - region_w) // 2 if not is_left else 0
+            pad_top = (tile_h_scaled - region_h) // 2 if not is_top else 0
+            
+            # Crop the tile array to match the region size
+            tile_arr_cropped = tile_arr[
+                pad_top:pad_top+region_h,
+                pad_left:pad_left+region_w,
+                :
+            ]
+            
+            # Add weighted contribution to output
+            # In non-overlap regions, mask = 1.0, so full contribution
+            # In overlap regions, mask varies, creating blend
+            for c in range(3):
+                try:
+                    output[y:y+region_h, x:x+region_w, c] += tile_arr_cropped[:, :, c] * mask
+                except ValueError as e:
+                    print(f"ERROR at tile ({tile.x}, {tile.y}):")
+                    print(f"  x={x}, y={y}, region_w={region_w}, region_h={region_h}")
+                    print(f"  Slice would be output[{y}:{y+region_h}, {x}:{x+region_w}, {c}]")
+                    print(f"  output.shape={output.shape}")
+                    print(f"  tile_arr_cropped[:,:,{c}].shape={tile_arr_cropped[:,:,c].shape}")
+                    print(f"  mask.shape={mask.shape}")
+                    print(f"  Actual slice shape: {output[y:y+region_h, x:x+region_w, c].shape}")
+                    raise
+            
+            # Track total weight at each pixel
+            weights[y:y+region_h, x:x+region_w] += mask
+        
+        # Normalize: divide accumulated colors by total weights
+        # ChaiNNer does this via their _fast_mix function
+        # Example: If two tiles overlap with weights 0.3 and 0.7:
+        #   output = color1 * 0.3 + color2 * 0.7
+        #   weights = 0.3 + 0.7 = 1.0
+        #   final = output / weights = (color1 * 0.3 + color2 * 0.7) / 1.0
+        for c in range(3):
+            output[:, :, c] /= np.maximum(weights, 1e-6)
+        
+        # Convert to 8-bit image
+        output = np.clip(output, 0, 255).astype(np.uint8)
+        return Image.fromarray(output)

backend/tiling/tile_splitter.py

@@ -0,0 +1,75 @@
+"""
+splits large images into overlapping tiles.
+based on chaiNNer's optimal tile size calculation to prevent uneven tiles.
+"""
+import math
+from PIL import Image
+from .models import Tile, Region
+
+
+class TileSplitter:
+    """
+    Splits images into memory-efficient overlapping tiles.
+    """
+
+    def __init__(self, tile_size: int = 512, overlap: int = 16):
+        self.max_tile_size = tile_size
+        self.overlap = overlap
+    
+    def calculate_optimal_tile_size(self, width: int, height: int):
+        # Calculate how many tiles needed in each direction
+        tile_count_x = math.ceil(width / self.max_tile_size)
+        tile_count_y = math.ceil(height / self.max_tile_size)
+        
+        # Distribute image evenly across tiles
+        optimal_x = math.ceil(width / tile_count_x)
+        optimal_y = math.ceil(height / tile_count_y)
+        
+        return optimal_x, optimal_y
+    
+    def split(self, image: Image.Image):
+        
+        width, height = image.size
+        tile_w, tile_h = self.calculate_optimal_tile_size(width, height)
+        
+        # Image boundary for intersection checks
+        img_region = Region(0, 0, width, height)
+        
+        tiles = []
+        
+        # Step size: tile size minus overlap
+        # This creates the overlapping pattern
+        step_x = tile_w - self.overlap
+        step_y = tile_h - self.overlap
+        
+        # Generate tiles row by row
+        for y in range(0, height, step_y):
+            for x in range(0, width, step_x):
+                # Create base tile region
+                tile_region = Region(x, y, tile_w, tile_h)
+                
+                # Clip to image boundaries (for edge tiles)
+                tile_region = tile_region.intersect(img_region)
+                
+                # Calculate padding for overlap
+                # ChaiNNer's approach: img_region.child_padding(tile).min(overlap)
+                pad = img_region.child_padding(tile_region).min(self.overlap)
+                
+                padded_region = tile_region.add_padding(pad)
+                
+                tile_img = image.crop((
+                    padded_region.x,
+                    padded_region.y,
+                    padded_region.x + padded_region.width,
+                    padded_region.y + padded_region.height
+                ))
+                
+                tiles.append(Tile(
+                    image=tile_img,
+                    region=tile_region,
+                    x=tile_region.x,
+                    y=tile_region.y,
+                    padding=pad
+                ))
+        
+        return tiles