70 lines
2.3 KiB
Python
70 lines
2.3 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
import math
|
|
|
|
class SplitLatent_Tiled:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"samples": ("LATENT",),
|
|
"rows": ("INT", {"default": 2, "min": 1, "max": 100}),
|
|
"columns": ("INT", {"default": 2, "min": 1, "max": 100}),
|
|
"overlap": ("INT", {"default": 8, "min": 0, "max": 128}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT", "INT", "INT", "INT", "*")
|
|
RETURN_NAMES = ("tiles", "rows", "columns", "overlap", "tiling_spec")
|
|
OUTPUT_IS_LIST = (True, False, False, False, False)
|
|
FUNCTION = "split"
|
|
CATEGORY = "SplitImage"
|
|
|
|
def split(self, samples, rows, columns, overlap):
|
|
latent = samples["samples"]
|
|
|
|
shape = list(latent.shape)
|
|
if len(shape) < 3:
|
|
raise ValueError(f"SplitLatent_Tiled: Expected at least 3D latent tensor, got {shape}")
|
|
|
|
H = shape[-2]
|
|
W = shape[-1]
|
|
|
|
tile_w = math.ceil((W + (columns - 1) * overlap) / columns)
|
|
tile_h = math.ceil((H + (rows - 1) * overlap) / rows)
|
|
|
|
pad_w = (columns * tile_w - (columns - 1) * overlap) - W
|
|
pad_h = (rows * tile_h - (rows - 1) * overlap) - H
|
|
|
|
if pad_w > 0 or pad_h > 0:
|
|
latent = F.pad(latent, (0, pad_w, 0, pad_h), mode='constant', value=0)
|
|
H = latent.shape[-2]
|
|
W = latent.shape[-1]
|
|
|
|
tiles = []
|
|
for r in range(rows):
|
|
for c in range(columns):
|
|
x1 = c * (tile_w - overlap)
|
|
y1 = r * (tile_h - overlap)
|
|
x2 = x1 + tile_w
|
|
y2 = y1 + tile_h
|
|
|
|
tile = latent[..., y1:y2, x1:x2]
|
|
# Store each tile as a single-latent sample in the list
|
|
for i in range(tile.shape[0]):
|
|
tiles.append({"samples": tile[i:i+1]})
|
|
|
|
tiling_spec = {
|
|
"rows": rows,
|
|
"columns": columns,
|
|
"overlap": overlap * 8, # Store as pixel overlap
|
|
"tile_w": tile_w * 8,
|
|
"tile_h": tile_h * 8,
|
|
"orig_w": (W - pad_w) * 8,
|
|
"orig_h": (H - pad_h) * 8,
|
|
"padded_w": W * 8,
|
|
"padded_h": H * 8
|
|
}
|
|
|
|
return (tiles, rows, columns, overlap, tiling_spec)
|