Add spatial batching to latent upscale
This commit is contained in:
+97
-5
@@ -4084,6 +4084,32 @@ def _latent_upscale_target_size(base_w, base_h, param):
|
||||
return int(base_w), int(base_h)
|
||||
|
||||
|
||||
def _latent_spatial_grid(h, w, th, tw, ol_h, ol_w):
|
||||
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")
|
||||
sh = th - ol_h
|
||||
sw = tw - ol_w
|
||||
nrows = 1 if h <= th else math.ceil((h - ol_h) / sh)
|
||||
if (nrows - 1) * sh + th < h:
|
||||
nrows += 1
|
||||
ncols = 1 if w <= tw else math.ceil((w - ol_w) / sw)
|
||||
if (ncols - 1) * sw + tw < w:
|
||||
ncols += 1
|
||||
rows = [i * sh for i in range(nrows)]
|
||||
cols = [j * sw for j in range(ncols)]
|
||||
trows = [min(th, h - r) for r in rows]
|
||||
tcols = [min(tw, w - c) for c in cols]
|
||||
return rows, cols, trows, tcols
|
||||
|
||||
|
||||
def _latent_spatial_blend_weights(t, overlap_mode):
|
||||
if overlap_mode == "later":
|
||||
return 1.0 - t
|
||||
return t
|
||||
|
||||
|
||||
def _nested_tensor_parts(samples):
|
||||
if samples is None:
|
||||
return ()
|
||||
@@ -6119,7 +6145,8 @@ class H3LongVideos:
|
||||
"latent_upscale_param": ("DUMAS_H3_LATENT_UPSCALE_PARAM", {
|
||||
"tooltip": "Output of 'Dumas H3 Latent Upscale Params'. When connected, the first-pass "
|
||||
"latent is upscaled and run through a short refinement pass before decode, "
|
||||
"using the sampler, scheduler, steps, denoise, and megapixel target from that node. "
|
||||
"using the sampler, scheduler, steps, denoise, megapixel target, and optional "
|
||||
"spatial batching from that node. "
|
||||
"Leave unconnected to skip latent upscaling entirely."}),
|
||||
"upscale": (["off", "rtx", "model", "lanczos"], {"default": "off",
|
||||
"tooltip": "Optional post-pass on the finished frames. 'rtx' = NVIDIA RTX Video Super "
|
||||
@@ -6478,9 +6505,68 @@ class H3LongVideos:
|
||||
refine_scheduler = latent_upscale_param.get("scheduler", sch)
|
||||
refine_denoise_value = latent_upscale_param.get("denoise", latent_upscale_param.get("refine_denoise", 0.2))
|
||||
refine_denoise = 0.2 if refine_denoise_value is None else float(refine_denoise_value)
|
||||
(refined_out,) = nodes.common_ksampler(
|
||||
model, seed, refine_steps, cfg, refine_sampler, refine_scheduler, upscale_cond, negative, upscale_latent,
|
||||
denoise=refine_denoise)
|
||||
tile_w_px = int(latent_upscale_param.get("tile_width", 512) or 512)
|
||||
tile_h_px = int(latent_upscale_param.get("tile_height", 512) or 512)
|
||||
overlap_px = max(0, int(latent_upscale_param.get("overlap", 64) or 64))
|
||||
fade_px = max(0, int(latent_upscale_param.get("fade_width", 0) or 0))
|
||||
overlap_mode = str(latent_upscale_param.get("overlap_mode", "earlier"))
|
||||
tile_tw = max(1, min(int(up_w), max(1, tile_w_px // 16)))
|
||||
tile_th = max(1, min(int(up_h), max(1, tile_h_px // 16)))
|
||||
ol_tw = max(0, min(tile_tw - 1, overlap_px // 16))
|
||||
ol_th = max(0, min(tile_th - 1, overlap_px // 16))
|
||||
fw_tw = max(0, min(ol_tw, fade_px // 16))
|
||||
fw_th = max(0, min(ol_th, fade_px // 16))
|
||||
rows, cols, trows, tcols = _latent_spatial_grid(int(up_h), int(up_w), tile_th, tile_tw, ol_th, ol_tw)
|
||||
if len(rows) == 1 and len(cols) == 1:
|
||||
(refined_out,) = nodes.common_ksampler(
|
||||
model, seed, refine_steps, cfg, refine_sampler, refine_scheduler, upscale_cond, negative, upscale_latent,
|
||||
denoise=refine_denoise)
|
||||
else:
|
||||
refined_video = upscaled_video.clone()
|
||||
full_audio = parts[1]
|
||||
for row_index, r0 in enumerate(rows):
|
||||
tr = trows[row_index]
|
||||
for col_index, c0 in enumerate(cols):
|
||||
tc = tcols[col_index]
|
||||
tile_target_w = int(tc) * 16
|
||||
tile_target_h = int(tr) * 16
|
||||
tile_cond, tile_latent = _build_shot_conditioning(
|
||||
clip, vae, prompt, tile_target_w, tile_target_h, ln, fps, handoff,
|
||||
ref_images=refs, ref_image_size=ref_image_size,
|
||||
ref_noise_aug=ref_noise_aug, audio_vae=audio_vae, silent=silent)
|
||||
tile_video = upscaled_video[:, :, :, r0:r0 + tr, c0:c0 + tc].contiguous()
|
||||
tile_latent["samples"] = comfy.nested_tensor.NestedTensor((tile_video, full_audio))
|
||||
tile_out, = nodes.common_ksampler(
|
||||
model, seed, refine_steps, cfg, refine_sampler, refine_scheduler, tile_cond, negative, tile_latent,
|
||||
denoise=refine_denoise)
|
||||
tile_out = _video_only_refined_latent(
|
||||
{"samples": comfy.nested_tensor.NestedTensor((tile_video, full_audio))},
|
||||
tile_out)
|
||||
tile_video_out = tile_out["samples"].tensors[0]
|
||||
region = refined_video[:, :, :, r0:r0 + tr, c0:c0 + tc]
|
||||
base_region = region.clone()
|
||||
region.copy_(tile_video_out)
|
||||
if col_index > 0 and ol_tw > 0:
|
||||
t = torch.linspace(0.0, 1.0, ol_tw, device=region.device, dtype=region.dtype)
|
||||
w = _latent_spatial_blend_weights(t, overlap_mode)
|
||||
if fw_tw > 0:
|
||||
w = w.clone()
|
||||
w[:fw_tw] = 0.0
|
||||
region[:, :, :, :, :ol_tw] = (
|
||||
base_region[:, :, :, :, :ol_tw] * (1.0 - w[None, None, None, None, :]) +
|
||||
tile_video_out[:, :, :, :, :ol_tw] * w[None, None, None, None, :]
|
||||
)
|
||||
if row_index > 0 and ol_th > 0:
|
||||
t = torch.linspace(0.0, 1.0, ol_th, device=region.device, dtype=region.dtype)
|
||||
w = _latent_spatial_blend_weights(t, overlap_mode)
|
||||
if fw_th > 0:
|
||||
w = w.clone()
|
||||
w[:fw_th] = 0.0
|
||||
region[:, :, :, :ol_th, :] = (
|
||||
base_region[:, :, :, :ol_th, :] * (1.0 - w[None, None, None, :, None]) +
|
||||
tile_video_out[:, :, :, :ol_th, :] * w[None, None, None, :, None]
|
||||
)
|
||||
refined_out = {"samples": comfy.nested_tensor.NestedTensor((refined_video, full_audio))}
|
||||
timing["latent_upscale_sample"] += time.perf_counter() - latent_start
|
||||
refined_out = _video_only_refined_latent(out, refined_out)
|
||||
except Exception as e:
|
||||
@@ -6617,10 +6703,16 @@ class H3LongVideos:
|
||||
denoise = 0.2 if denoise_value is None else float(denoise_value)
|
||||
refine_sampler = latent_upscale_param.get("sampler_name", "euler_ancestral")
|
||||
refine_scheduler = latent_upscale_param.get("scheduler", "simple")
|
||||
batch_note = ""
|
||||
tile_w_px = int(latent_upscale_param.get("tile_width", 512) or 512)
|
||||
tile_h_px = int(latent_upscale_param.get("tile_height", 512) or 512)
|
||||
overlap_px = max(0, int(latent_upscale_param.get("overlap", 64) or 64))
|
||||
if tile_w_px > 0 and tile_h_px > 0 and (tile_w_px < target_w or tile_h_px < target_h):
|
||||
batch_note = f"; spatial batches {tile_w_px}x{tile_h_px}px overlap {overlap_px}px"
|
||||
latent_upscale_note = (
|
||||
f" latent upscale: target {target_w}x{target_h}px{detail}; "
|
||||
f"{int(latent_upscale_param.get('steps', 2) or 2)}-step refinement "
|
||||
f"{refine_sampler}/{refine_scheduler} denoise {denoise:.2f}"
|
||||
f"{refine_sampler}/{refine_scheduler} denoise {denoise:.2f}{batch_note}"
|
||||
)
|
||||
|
||||
paras = split_paragraphs(prompt, "##")
|
||||
|
||||
Reference in New Issue
Block a user