Tile latent upscale model inference
This commit is contained in:
+140
-1
@@ -371,7 +371,55 @@ def _resolve_target_size(param, h_in, w_in):
|
||||
return int(h_in), int(w_in)
|
||||
|
||||
|
||||
def upscale_video_model(video, param):
|
||||
def _spatial_blend_weights(t, overlap_mode, overlap_blend="linear"):
|
||||
if overlap_blend == "overwrite":
|
||||
return torch.ones_like(t) if overlap_mode == "later" else torch.zeros_like(t)
|
||||
if overlap_blend == "midpoint":
|
||||
base = (t >= 0.5).to(t.dtype)
|
||||
elif overlap_blend == "smoothstep":
|
||||
base = t * t * (3.0 - 2.0 * t)
|
||||
else:
|
||||
base = t
|
||||
return base if overlap_mode == "later" else 1.0 - base
|
||||
|
||||
|
||||
def _grid_1d(size, tile, overlap, min_tile):
|
||||
if size <= tile:
|
||||
return [0], [size], [0]
|
||||
step = tile - overlap
|
||||
n = math.ceil((size - overlap) / step)
|
||||
if (n - 1) * step + tile < size:
|
||||
n += 1
|
||||
rows = [i * step for i in range(n)]
|
||||
tiles = [min(tile, size - r) for r in rows]
|
||||
if min_tile > 0 and n >= 2:
|
||||
edge = size - rows[-1]
|
||||
if edge < min_tile:
|
||||
new_last = size - min_tile
|
||||
if rows[-2] < new_last < rows[-2] + tiles[-2]:
|
||||
rows[-1] = new_last
|
||||
tiles[-1] = size - new_last
|
||||
overlaps = [0] * n
|
||||
for i in range(1, n):
|
||||
overlaps[i] = max(0, rows[i - 1] + tiles[i - 1] - rows[i])
|
||||
return rows, tiles, overlaps
|
||||
|
||||
|
||||
def _compute_spatial_grid(h, w, th, tw, ol_h, ol_w, min_th=0, min_tw=0):
|
||||
if th <= 0 or tw <= 0:
|
||||
raise ValueError("tile dimensions must be positive")
|
||||
if ol_h >= th or ol_w >= tw:
|
||||
raise ValueError("overlap must be smaller than the tile size")
|
||||
if min_th < 0 or min_tw < 0:
|
||||
raise ValueError("minimum tile size must be non-negative")
|
||||
if min_th > th or min_tw > tw:
|
||||
raise ValueError("minimum tile size must not exceed the tile size")
|
||||
rows, trows, row_ovl = _grid_1d(h, th, ol_h, min_th)
|
||||
cols, tcols, col_ovl = _grid_1d(w, tw, ol_w, min_tw)
|
||||
return rows, cols, trows, tcols, row_ovl, col_ovl
|
||||
|
||||
|
||||
def _upscale_video_model_core(video, param):
|
||||
model_name = param["model_name"]
|
||||
device = param.get("device", "cuda")
|
||||
precision = param.get("precision", "fp16")
|
||||
@@ -407,6 +455,97 @@ def upscale_video_model(video, param):
|
||||
return out, h_out, w_out
|
||||
|
||||
|
||||
def upscale_video_model(video, param):
|
||||
_, _, _, h_in, w_in = video.shape
|
||||
h_out, w_out = _resolve_target_size(param, h_in, w_in)
|
||||
mode = str(param.get("tile_size_mode") or "specific_size")
|
||||
tile_w = int(param.get("tile_width", 0) or 0)
|
||||
tile_h = int(param.get("tile_height", 0) or 0)
|
||||
overlap = max(0, int(param.get("overlap", 0) or 0))
|
||||
fade_w = max(0, int(param.get("fade_width", 0) or 0))
|
||||
fade_h = max(0, int(param.get("fade_height", 0) or 0))
|
||||
overlap_mode = str(param.get("overlap_mode", "earlier"))
|
||||
overlap_blend = str(param.get("overlap_blend", "linear"))
|
||||
grid_rows = max(1, int(param.get("grid_rows", 2) or 2))
|
||||
grid_cols = max(1, int(param.get("grid_cols", 2) or 2))
|
||||
spatial_w_overlap = max(0, int(param.get("spatial_w_overlap", overlap) or overlap))
|
||||
spatial_h_overlap = max(0, int(param.get("spatial_h_overlap", overlap) or overlap))
|
||||
min_tile_size = max(0, int(param.get("min_tile_size", 0) or 0))
|
||||
|
||||
if mode == "rows_cols":
|
||||
tile_w = max(32, int(round(w_out / float(grid_cols) / 32.0)) * 32)
|
||||
tile_h = max(32, int(round(h_out / float(grid_rows) / 32.0)) * 32)
|
||||
rows, cols, trows, tcols, row_ovl, col_ovl = _compute_spatial_grid(
|
||||
h_out, w_out, tile_h, tile_w, spatial_h_overlap, spatial_w_overlap, min_tile_size, min_tile_size
|
||||
)
|
||||
else:
|
||||
if tile_w <= 0 or tile_h <= 0 or (tile_w >= w_out and tile_h >= h_out):
|
||||
return _upscale_video_model_core(video, param)
|
||||
for name, value in (("tile_width", tile_w), ("tile_height", tile_h), ("overlap", overlap), ("fade_width", fade_w), ("fade_height", fade_h), ("min_tile_size", min_tile_size)):
|
||||
if value % 32 != 0:
|
||||
raise ValueError(f"'{name}' must be a multiple of 32 pixels; got {value}.")
|
||||
rows, cols, trows, tcols, row_ovl, col_ovl = _compute_spatial_grid(
|
||||
h_out, w_out, tile_h, tile_w, overlap, overlap, min_tile_size, min_tile_size
|
||||
)
|
||||
|
||||
# If the requested tile is not smaller than the target on either axis,
|
||||
# the tiled path would just duplicate work.
|
||||
if len(rows) == 1 and len(cols) == 1:
|
||||
return _upscale_video_model_core(video, param)
|
||||
|
||||
scale_h = h_out / float(h_in)
|
||||
scale_w = w_out / float(w_in)
|
||||
orig_dtype = video.dtype
|
||||
out = torch.zeros((video.shape[0], video.shape[1], video.shape[2], h_out, w_out), device="cpu", dtype=orig_dtype)
|
||||
|
||||
for i, r0 in enumerate(rows):
|
||||
tr = trows[i]
|
||||
ovh = row_ovl[i]
|
||||
src_r0 = max(0, int(round(r0 / scale_h)))
|
||||
src_r1 = min(h_in, max(src_r0 + 1, int(round((r0 + tr) / scale_h))))
|
||||
for j, c0 in enumerate(cols):
|
||||
tc = tcols[j]
|
||||
ovw = col_ovl[j]
|
||||
src_c0 = max(0, int(round(c0 / scale_w)))
|
||||
src_c1 = min(w_in, max(src_c0 + 1, int(round((c0 + tc) / scale_w))))
|
||||
tile_video = video[:, :, :, src_r0:src_r1, src_c0:src_c1].contiguous()
|
||||
tile_param = dict(param)
|
||||
tile_param["width"] = int(tc)
|
||||
tile_param["height"] = int(tr)
|
||||
tile_param["megapixels"] = 0.0
|
||||
tile_out, tile_h_out, tile_w_out = _upscale_video_model_core(tile_video, tile_param)
|
||||
if tile_h_out != tr or tile_w_out != tc:
|
||||
flat = tile_out.permute(0, 2, 1, 3, 4).reshape(-1, tile_out.shape[1], tile_h_out, tile_w_out)
|
||||
flat = F.interpolate(flat, size=(tr, tc), mode="bilinear", align_corners=False)
|
||||
tile_out = flat.reshape(tile_out.shape[0], tile_out.shape[2], tile_out.shape[1], tr, tc).permute(0, 2, 1, 3, 4).contiguous()
|
||||
region = out[:, :, :, r0:r0 + tr, c0:c0 + tc]
|
||||
base_region = region.clone()
|
||||
region.copy_(tile_out.to(dtype=orig_dtype))
|
||||
if j > 0 and ovw > 0:
|
||||
t = torch.linspace(0.0, 1.0, ovw, device=region.device, dtype=region.dtype)
|
||||
w = _spatial_blend_weights(t, overlap_mode, overlap_blend)
|
||||
if fade_w > 0:
|
||||
w = w.clone()
|
||||
w[:min(fade_w, ovw)] = 0.0
|
||||
region[:, :, :, :, :ovw] = (
|
||||
base_region[:, :, :, :, :ovw] * (1.0 - w[None, None, None, None, :]) +
|
||||
tile_out[:, :, :, :, :ovw].to(dtype=orig_dtype) * w[None, None, None, None, :]
|
||||
)
|
||||
if i > 0 and ovh > 0:
|
||||
t = torch.linspace(0.0, 1.0, ovh, device=region.device, dtype=region.dtype)
|
||||
w = _spatial_blend_weights(t, overlap_mode, overlap_blend)
|
||||
if fade_h > 0:
|
||||
w = w.clone()
|
||||
w[:min(fade_h, ovh)] = 0.0
|
||||
region[:, :, :, :ovh, :] = (
|
||||
base_region[:, :, :, :ovh, :] * (1.0 - w[None, None, None, :, None]) +
|
||||
tile_out[:, :, :, :ovh, :].to(dtype=orig_dtype) * w[None, None, None, :, None]
|
||||
)
|
||||
out[:, :, :, r0:r0 + tr, c0:c0 + tc] = region
|
||||
|
||||
return out, h_out, w_out
|
||||
|
||||
|
||||
def upscale_video_interp(video, param):
|
||||
method = str(param.get("method") or "bilinear")
|
||||
_, c, t, h_in, w_in = video.shape
|
||||
|
||||
Reference in New Issue
Block a user