107 lines
4.2 KiB
Python
107 lines
4.2 KiB
Python
from __future__ import annotations
|
|
|
|
import math, torch
|
|
from typing import Any, Protocol, cast
|
|
from ...utils.logging import log
|
|
from ...utils.image import blend_processed_batch_item, pil_to_tensor, tensor_to_pil
|
|
from ..processing_args import UpscaleCoreArgs
|
|
from ..tile_processing import TileBatchArgs, extract_and_process_tile_batch
|
|
|
|
|
|
class _SingleGpuOps(Protocol):
|
|
def round_to_multiple(self, value: int) -> int: ...
|
|
def calculate_tiles(self, width: int, height: int, tile_width: int, tile_height: int, force_uniform_tiles: bool): ...
|
|
def create_tile_mask(self, width: int, height: int, tx: int, ty: int, tile_width: int, tile_height: int, mask_blur: int): ...
|
|
def blend_tile(self, base_image: Any, tile_pil: Any, x1: int, y1: int, size: tuple[int, int], tile_mask: Any, padding: int): ...
|
|
def extract_batch_tile_with_padding(self, source_batch: Any, tx: int, ty: int, tile_width: int, tile_height: int, padding: int, force_uniform_tiles: bool): ...
|
|
def process_tiles_batch(self, *args: Any, **kwargs: Any): ...
|
|
|
|
|
|
class SingleGpuModeMixin:
|
|
def process_single_gpu(
|
|
self,
|
|
upscaled_image: torch.Tensor,
|
|
core_args: UpscaleCoreArgs,
|
|
tile_width: int,
|
|
tile_height: int,
|
|
padding: int,
|
|
mask_blur: int,
|
|
force_uniform_tiles: bool,
|
|
) -> tuple[torch.Tensor]:
|
|
"""Process all tiles on a single GPU (no distribution), batching per tile like USDU."""
|
|
ops = cast(_SingleGpuOps, self)
|
|
# Round tile dimensions
|
|
tile_width = ops.round_to_multiple(tile_width)
|
|
tile_height = ops.round_to_multiple(tile_height)
|
|
|
|
# Get image dimensions and batch size
|
|
batch_size, height, width, _ = upscaled_image.shape
|
|
|
|
# Calculate all tiles
|
|
all_tiles = ops.calculate_tiles(width, height, tile_width, tile_height, force_uniform_tiles)
|
|
|
|
rows = math.ceil(height / tile_height)
|
|
cols = math.ceil(width / tile_width)
|
|
log(
|
|
f"USDU Dist: Single GPU | Canvas {width}x{height} | Tile {tile_width}x{tile_height} | Grid {rows}x{cols} ({len(all_tiles)} tiles/image) | Batch {batch_size}"
|
|
)
|
|
tile_batch_args = TileBatchArgs(
|
|
core=core_args,
|
|
tile_width=tile_width,
|
|
tile_height=tile_height,
|
|
padding=padding,
|
|
force_uniform_tiles=force_uniform_tiles,
|
|
width=width,
|
|
height=height,
|
|
)
|
|
|
|
# Prepare result images list
|
|
result_images = []
|
|
for b in range(batch_size):
|
|
image_pil = tensor_to_pil(upscaled_image[b:b+1], 0).convert('RGB')
|
|
result_images.append(image_pil.copy())
|
|
|
|
# Precompute tile masks once
|
|
tile_masks = []
|
|
for tx, ty in all_tiles:
|
|
tile_masks.append(ops.create_tile_mask(width, height, tx, ty, tile_width, tile_height, mask_blur))
|
|
|
|
# Process tiles batched across images
|
|
for tile_idx, (tx, ty) in enumerate(all_tiles):
|
|
# Progressive state parity: extract each tile from the current updated image batch.
|
|
source_batch = torch.cat([pil_to_tensor(img) for img in result_images], dim=0)
|
|
if upscaled_image.is_cuda:
|
|
source_batch = source_batch.cuda()
|
|
|
|
processed_batch, x1, y1, ew, eh = extract_and_process_tile_batch(
|
|
node=ops,
|
|
upscaled_image=source_batch,
|
|
tx=tx,
|
|
ty=ty,
|
|
args=tile_batch_args,
|
|
)
|
|
|
|
# Blend results back into each image using cached mask
|
|
tile_mask = tile_masks[tile_idx]
|
|
for b in range(batch_size):
|
|
blend_processed_batch_item(
|
|
result_images,
|
|
processed_batch,
|
|
b,
|
|
ops.blend_tile,
|
|
x1,
|
|
y1,
|
|
ew,
|
|
eh,
|
|
tile_mask,
|
|
padding,
|
|
)
|
|
|
|
# Convert back to tensor
|
|
result_tensors = [pil_to_tensor(img) for img in result_images]
|
|
result_tensor = torch.cat(result_tensors, dim=0)
|
|
if upscaled_image.is_cuda:
|
|
result_tensor = result_tensor.cuda()
|
|
|
|
return (result_tensor,)
|