From 3b63d33ec7b502372a170a489c96431f9edbc619 Mon Sep 17 00:00:00 2001 From: Chris Dumas Date: Thu, 3 Sep 2026 19:01:13 +0000 Subject: [PATCH] Tile latent upscale model inference --- dumas_h3_latent_upscale.py | 141 ++++++++++++++++++++++++++++++++++++- 1 file changed, 140 insertions(+), 1 deletion(-) diff --git a/dumas_h3_latent_upscale.py b/dumas_h3_latent_upscale.py index c526fff..9833995 100644 --- a/dumas_h3_latent_upscale.py +++ b/dumas_h3_latent_upscale.py @@ -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