Add temporal chunking to latent upscale
This commit is contained in:
@@ -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,
|
||||
},)
|
||||
|
||||
Reference in New Issue
Block a user