Add temporal chunking to latent upscale

This commit is contained in:
2026-09-04 06:41:03 +00:00
parent fbb5f799cd
commit a0d80bcc47
4 changed files with 107 additions and 8 deletions
+95 -5
View File
@@ -372,6 +372,43 @@ def _resolve_target_size(param, h_in, w_in):
return int(h_in), int(w_in)
def _temporal_segments(frame_count, chunk_length, overlap):
frame_count = int(frame_count)
chunk_length = int(chunk_length)
overlap = int(overlap)
if frame_count <= 0:
return [(0, 0)]
if chunk_length <= 0:
raise ValueError("chunk_length must be positive")
if overlap < 0:
raise ValueError("temporal_overlap must be non-negative")
if chunk_length <= overlap:
raise ValueError("temporal_overlap must be smaller than chunk_length")
if frame_count <= chunk_length:
return [(0, frame_count)]
hop = chunk_length - overlap
bounds = []
start = 0
while start < frame_count:
end = min(start + chunk_length, frame_count)
bounds.append((start, end))
if end >= frame_count:
break
start += hop
return bounds
def _temporal_blend_weights(length, anchor_strength):
length = int(length)
if length <= 0:
return None
anchor_strength = float(anchor_strength)
anchor_strength = min(1.0, max(0.0, anchor_strength))
start = max(0.0, 1.0 - anchor_strength)
return torch.linspace(start, 1.0, length)
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)
@@ -633,13 +670,55 @@ def upscale_video_interp(video, param):
return up, h_out, w_out
def _upscale_video_temporal_chunks(video, param, upscaler):
chunk_length = int(param.get("chunk_length", 0) or 0)
temporal_overlap = int(param.get("temporal_overlap", 0) or 0)
anchor_strength = float(param.get("anchor_strength", 0.999) or 0.999)
t = int(video.shape[2])
if chunk_length <= 0 or t <= chunk_length:
return upscaler(video, param)
bounds = _temporal_segments(t, chunk_length, temporal_overlap)
if len(bounds) <= 1:
return upscaler(video, param)
orig_dtype = video.dtype
out = None
out_h = out_w = None
for i, (t0, t1) in enumerate(bounds):
chunk = video[:, :, t0:t1].contiguous()
chunk_out, chunk_h, chunk_w = upscaler(chunk, param)
chunk_out = chunk_out.to(device="cpu", dtype=orig_dtype)
if out is None:
out_h, out_w = chunk_h, chunk_w
out = torch.zeros((video.shape[0], video.shape[1], t, out_h, out_w), device="cpu", dtype=orig_dtype)
out[:, :, t0:t1] = chunk_out
continue
ov = min(temporal_overlap, t1 - t0, t - t0)
if ov <= 0:
out[:, :, t0:t1] = chunk_out
continue
prev_region = out[:, :, t0:t0 + ov]
new_region = chunk_out[:, :, :ov]
w = _temporal_blend_weights(ov, anchor_strength).to(device=prev_region.device, dtype=prev_region.dtype)
out[:, :, t0:t0 + ov] = (
prev_region * (1.0 - w[None, None, :, None, None]) +
new_region * w[None, None, :, None, None]
)
out[:, :, t0 + ov:t1] = chunk_out[:, :, ov:]
return out, out_h, out_w
def upscale_latent_video(video, param):
mode = str(param.get("mode") or "off")
if mode == "off":
return video, video.shape[-2], video.shape[-1]
if mode == "model":
return upscale_video_model(video, param)
return upscale_video_interp(video, param)
return _upscale_video_temporal_chunks(video, param, upscale_video_model)
return _upscale_video_temporal_chunks(video, param, upscale_video_interp)
class H3LatentUpscaleParams:
@@ -711,8 +790,10 @@ class H3LatentUpscaleParams:
"tooltip": "Temporal fade schedule over each tile's sampling. Off keeps the fade fixed; narrowing shrinks it over steps; widening grows it over steps."}),
"dynamic_fade_min": ("INT", {"default": 32, "min": 0, "max": 4096, "step": 32,
"tooltip": "Minimum fade width used by dynamic_fade when it is enabled."}),
"chunk_length": ("INT", {"default": 136, "min": 17, "max": 100000, "step": 17,
"tooltip": "Reserved for upstream split compatibility. 136 matches the original chunk-length default."}),
"chunk_length": ("INT", {"default": 85, "min": 17, "max": 100000, "step": 17,
"tooltip": "Temporal chunk length for the latent upscale stage. 85 matches the practical upstream example and helps lower peak VRAM."}),
"temporal_overlap": ("INT", {"default": 17, "min": 0, "max": 100000, "step": 17,
"tooltip": "Temporal overlap between latent chunks. 17 matches the upstream split example and reduces seam risk."}),
"resize_conditioning": ("BOOLEAN", {"default": False,
"tooltip": "Reserved for upstream split compatibility. Leave OFF unless you need the original fallback behavior."}),
"anchor_strength": ("FLOAT", {"default": 0.999, "min": 0.0, "max": 1.0, "step": 0.01,
@@ -720,7 +801,7 @@ class H3LatentUpscaleParams:
}
}
def build(self, mode, model_name, method, width, height, device, precision, sampler_name, scheduler, steps, denoise, megapixels, tile_width, tile_height, overlap, fade_width, fade_height, overlap_mode, overlap_blend, tile_size_mode, grid_rows, grid_cols, spatial_w_overlap, spatial_h_overlap, min_tile_size, masked_area_noise, brightness_match, dynamic_fade, dynamic_fade_min, chunk_length, resize_conditioning, anchor_strength):
def build(self, mode, model_name, method, width, height, device, precision, sampler_name, scheduler, steps, denoise, megapixels, tile_width, tile_height, overlap, fade_width, fade_height, overlap_mode, overlap_blend, tile_size_mode, grid_rows, grid_cols, spatial_w_overlap, spatial_h_overlap, min_tile_size, masked_area_noise, brightness_match, dynamic_fade, dynamic_fade_min, chunk_length, temporal_overlap, resize_conditioning, anchor_strength):
width = int(width)
height = int(height)
steps = int(steps)
@@ -737,6 +818,7 @@ class H3LatentUpscaleParams:
masked_area_noise = float(masked_area_noise)
dynamic_fade_min = int(dynamic_fade_min)
chunk_length = int(chunk_length)
temporal_overlap = int(temporal_overlap)
anchor_strength = float(anchor_strength)
if mode == "off":
return ({
@@ -767,9 +849,16 @@ class H3LatentUpscaleParams:
"dynamic_fade": dynamic_fade,
"dynamic_fade_min": dynamic_fade_min,
"chunk_length": chunk_length,
"temporal_overlap": temporal_overlap,
"resize_conditioning": bool(resize_conditioning),
"anchor_strength": anchor_strength,
},)
if chunk_length % 17 != 0:
raise ValueError("chunk_length must be a multiple of 17 pixels")
if temporal_overlap % 17 != 0:
raise ValueError("temporal_overlap must be a multiple of 17 pixels")
if temporal_overlap >= chunk_length:
raise ValueError("temporal_overlap must be smaller than chunk_length")
if width > 0:
width = int(round(width / 32.0)) * 32
if height > 0:
@@ -806,6 +895,7 @@ class H3LatentUpscaleParams:
"dynamic_fade": dynamic_fade,
"dynamic_fade_min": dynamic_fade_min,
"chunk_length": chunk_length,
"temporal_overlap": temporal_overlap,
"resize_conditioning": bool(resize_conditioning),
"anchor_strength": anchor_strength,
},)