API tools faq
paste
Advertisement
zelenooki87

tile_blending.py

zelenooki87
May 25th, 2025
175
0
Never
Add comment
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 8.11 KB | None | 0 0
raw download clone embed print report
  1. #chaiNNer-windows-latest-nightly\resources\src\nodes\impl\upscale\tile_blending.py
  2. from __future__ import annotations
  3.  
  4. import math
  5. from dataclasses import dataclass
  6. from enum import Enum
  7. from typing import Callable
  8.  
  9. import numpy as np
  10.  
  11. from nodes.utils.utils import get_h_w_c
  12.  
  13.  
  14. def sin_blend_fn(x: np.ndarray) -> np.ndarray:
  15.     """Standardna sinusoidna blend funkcija"""
  16.     return (np.sin(x * math.pi - math.pi / 2) + 1) / 2
  17.  
  18.  
  19. def smooth_blend_fn(x: np.ndarray) -> np.ndarray:
  20.     """Hermite interpolacija za glatkije prelaze"""
  21.     return x * x * (3.0 - 2.0 * x)
  22.  
  23.  
  24. def smoother_blend_fn(x: np.ndarray) -> np.ndarray:
  25.     """Još glatkija blend funkcija (Perlin's smootherstep)"""
  26.     return x * x * x * (x * (x * 6 - 15) + 10)
  27.  
  28.  
  29. def cosine_blend_fn(x: np.ndarray) -> np.ndarray:
  30.     """Kosinusna interpolacija"""
  31.     return (1 - np.cos(x * math.pi)) / 2
  32.  
  33.  
  34. def linear_blend_fn(x: np.ndarray) -> np.ndarray:
  35.     """Linearna blend funkcija - najjednostavnija"""
  36.     return x
  37.  
  38.  
  39. def half_sin_blend_fn(i: np.ndarray) -> np.ndarray:
  40.     """Modificirana half-sin funkcija koja koristi veći deo overlap-a"""
  41.     # Umesto originalnog i * 2 - 0.5, koristimo 1.5 za bolji prelaz
  42.     i = np.clip(i * 1.5 - 0.25, 0, 1)
  43.     return sin_blend_fn(i)
  44.  
  45.  
  46. def enhanced_blend_fn(x: np.ndarray) -> np.ndarray:
  47.     """Napredna blend funkcija sa prilagodljivim prelazom"""
  48.     # Kombinacija smooth i sin funkcija za optimalan rezultat
  49.     smooth_part = x * x * (3.0 - 2.0 * x)
  50.     sin_part = (np.sin(x * math.pi - math.pi / 2) + 1) / 2
  51.     # Weighted average - 70% smooth, 30% sin
  52.     return 0.7 * smooth_part + 0.3 * sin_part
  53.  
  54.  
  55. class BlendDirection(Enum):
  56.     X = 0
  57.     Y = 1
  58.  
  59.  
  60. @dataclass(frozen=True)
  61. class TileOverlap:
  62.     start: int
  63.     end: int
  64.  
  65.     @property
  66.     def total(self) -> int:
  67.         return self.start + self.end
  68.  
  69.  
  70. def _fast_mix(a: np.ndarray, b: np.ndarray, blend: np.ndarray) -> np.ndarray:
  71.     """
  72.    Returns `a * (1 - blend) + b * blend`
  73.    """
  74.     # a * (1 - blend) + b * blend
  75.     # a - a * blend + b * blend
  76.     r = b * blend
  77.     r += a
  78.     r -= a * blend  # type: ignore
  79.     return r
  80.  
  81.  
  82. class TileBlender:
  83.     def __init__(
  84.         self,
  85.         width: int,
  86.         height: int,
  87.         channels: int,
  88.         direction: BlendDirection,
  89.         blend_fn: Callable[[np.ndarray], np.ndarray] = smooth_blend_fn,  # Promenjen default
  90.         _prev: TileBlender | None = None,
  91.     ) -> None:
  92.         self.direction: BlendDirection = direction
  93.         self.blend_fn: Callable[[np.ndarray], np.ndarray] = blend_fn
  94.         self.offset: int = 0
  95.         self.last_end_overlap: int = 0
  96.         self._last_blend: np.ndarray | None = None
  97.  
  98.         if (
  99.             _prev is not None
  100.             and _prev.direction == direction
  101.             and _prev.width == width
  102.             and _prev.height == height
  103.             and _prev.channels == channels
  104.         ):
  105.             if _prev.blend_fn == blend_fn:
  106.                 # reuse blend
  107.                 self._last_blend = _prev._last_blend  # noqa: SLF001
  108.             result = _prev.result
  109.         else:
  110.             result = np.zeros((height, width, channels), dtype=np.float32)
  111.         self.result: np.ndarray = result
  112.  
  113.     @property
  114.     def width(self) -> int:
  115.         return self.result.shape[1]
  116.  
  117.     @property
  118.     def height(self) -> int:
  119.         return self.result.shape[0]
  120.  
  121.     @property
  122.     def channels(self) -> int:
  123.         return self.result.shape[2]
  124.  
  125.     def _get_blend(self, blend_size: int) -> np.ndarray:
  126.         if self.direction == BlendDirection.X:
  127.             if self._last_blend is not None and self._last_blend.shape[1] == blend_size:
  128.                 return self._last_blend
  129.  
  130.             blend = self.blend_fn(
  131.                 np.arange(blend_size, dtype=np.float32) / (blend_size - 1)
  132.             )
  133.             blend = blend.reshape((1, blend_size, 1))
  134.             blend = np.repeat(blend, repeats=self.height, axis=0)
  135.             blend = np.repeat(blend, repeats=self.channels, axis=2)
  136.         else:
  137.             if self._last_blend is not None and self._last_blend.shape[0] == blend_size:
  138.                 return self._last_blend
  139.  
  140.             blend = self.blend_fn(
  141.                 np.arange(blend_size, dtype=np.float32) / (blend_size - 1)
  142.             )
  143.             blend = blend.reshape((blend_size, 1, 1))
  144.             blend = np.repeat(blend, repeats=self.width, axis=1)
  145.             blend = np.repeat(blend, repeats=self.channels, axis=2)
  146.  
  147.         self._last_blend = blend
  148.         return blend
  149.  
  150.     def add_tile(self, tile: np.ndarray, overlap: TileOverlap) -> None:
  151.         h, w, c = get_h_w_c(tile)
  152.         assert c == self.channels
  153.         o = overlap
  154.  
  155.         if self.direction == BlendDirection.X:
  156.             assert h == self.height
  157.             assert w > o.total
  158.  
  159.             if self.offset == 0:
  160.                 # the first tile is copied in as is
  161.                 self.result[:, :w, ...] = tile
  162.  
  163.                 assert o.start == 0
  164.                 self.offset += w - o.end
  165.                 self.last_end_overlap = o.end
  166.  
  167.             else:
  168.                 assert self.offset < self.width, "All tiles were filled in already"
  169.  
  170.                 if self.last_end_overlap < o.start:
  171.                     # we can't use all the overlap of the current tile, so we have to cut it off
  172.                     diff = o.start - self.last_end_overlap
  173.                     tile = tile[:, diff:, ...]
  174.                     h, w, c = get_h_w_c(tile)
  175.                     o = TileOverlap(self.last_end_overlap, o.end)
  176.  
  177.                 # copy over the part that doesn't need blending (yet)
  178.                 self.result[
  179.                     :, self.offset + o.start : self.offset + w - o.start, ...
  180.                 ] = tile[:, o.start * 2 :, ...]
  181.  
  182.                 # blend the overlapping part
  183.                 blend_size = o.start * 2
  184.                 blend = self._get_blend(blend_size)
  185.  
  186.                 left = self.result[
  187.                     :, self.offset - o.start : self.offset + o.start, ...
  188.                 ]
  189.                 right = tile[:, :blend_size, ...]
  190.  
  191.                 self.result[:, self.offset - o.start : self.offset + o.start, ...] = (
  192.                     _fast_mix(left, right, blend)
  193.                 )
  194.  
  195.                 self.offset += w - o.total
  196.                 self.last_end_overlap = o.end
  197.         else:
  198.             assert w == self.width
  199.             assert h > o.total
  200.  
  201.             if self.offset == 0:
  202.                 # the first tile is copied in as is
  203.                 self.result[:h, :, ...] = tile
  204.  
  205.                 assert o.start == 0
  206.                 self.offset += h - o.end
  207.                 self.last_end_overlap = o.end
  208.  
  209.             else:
  210.                 assert self.offset < self.height, "All tiles were filled in already"
  211.  
  212.                 if self.last_end_overlap < o.start:
  213.                     # we can't use all the overlap of the current tile, so we have to cut it off
  214.                     diff = o.start - self.last_end_overlap
  215.                     tile = tile[diff:, :, ...]
  216.                     h, w, c = get_h_w_c(tile)
  217.                     o = TileOverlap(self.last_end_overlap, o.end)
  218.  
  219.                 # copy over the part that doesn't need blending
  220.                 self.result[
  221.                     self.offset + o.start : self.offset + h - o.start, :, ...
  222.                 ] = tile[o.start * 2 :, :, ...]
  223.  
  224.                 # blend the overlapping part
  225.                 blend_size = o.start * 2
  226.                 blend = self._get_blend(blend_size)
  227.  
  228.                 left = self.result[
  229.                     self.offset - o.start : self.offset + o.start, :, ...
  230.                 ]
  231.                 right = tile[: o.start * 2, :, ...]
  232.  
  233.                 self.result[self.offset - o.start : self.offset + o.start, :, ...] = (
  234.                     _fast_mix(left, right, blend)
  235.                 )
  236.  
  237.                 self.offset += h - o.total
  238.                 self.last_end_overlap = o.end
  239.  
  240.     def get_result(self) -> np.ndarray:
  241.         if self.direction == BlendDirection.X:
  242.             assert self.offset == self.width
  243.         else:
  244.             assert self.offset == self.height
  245.  
  246.         return self.result
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement
Public Pastes
Advertisement
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy.  OK, I Understand
Not a member of Pastebin yet?
Sign Up, it unlocks many cool features!