Force temporal chunks for CUDA latent upscale
This commit is contained in:
+32
-12
@@ -42,6 +42,24 @@ LATENTS_STD = [
|
|||||||
_LATENT_UPSCALE_FOLDER = "latent_upscale_models"
|
_LATENT_UPSCALE_FOLDER = "latent_upscale_models"
|
||||||
MP_UNIT = 1024 * 1024
|
MP_UNIT = 1024 * 1024
|
||||||
RES_MULTIPLE = 32
|
RES_MULTIPLE = 32
|
||||||
|
CUDA_MODEL_CHUNK_LENGTH = 17
|
||||||
|
|
||||||
|
|
||||||
|
def _uses_cuda_model_upscale(param):
|
||||||
|
mode = str(param.get("mode") or "")
|
||||||
|
model_name = str(param.get("model_name") or "")
|
||||||
|
device = str(param.get("device", "cuda") or "cuda")
|
||||||
|
has_model_name = bool(model_name and model_name != "none" and not model_name.startswith("(no upscale models"))
|
||||||
|
return device == "cuda" and (mode == "model" or has_model_name)
|
||||||
|
|
||||||
|
|
||||||
|
def _effective_temporal_params(param):
|
||||||
|
chunk_length = int(param.get("chunk_length", 0) or 0)
|
||||||
|
temporal_overlap = int(param.get("temporal_overlap", 0) or 0)
|
||||||
|
if _uses_cuda_model_upscale(param):
|
||||||
|
chunk_length = CUDA_MODEL_CHUNK_LENGTH if chunk_length <= 0 else min(chunk_length, CUDA_MODEL_CHUNK_LENGTH)
|
||||||
|
temporal_overlap = min(max(0, temporal_overlap), max(0, chunk_length - 17))
|
||||||
|
return chunk_length, temporal_overlap
|
||||||
|
|
||||||
|
|
||||||
def _models_dir():
|
def _models_dir():
|
||||||
@@ -735,17 +753,19 @@ def upscale_video_interp(video, param):
|
|||||||
def _upscale_video_temporal_chunks(video, param, upscaler):
|
def _upscale_video_temporal_chunks(video, param, upscaler):
|
||||||
if video.device.type != "cpu":
|
if video.device.type != "cpu":
|
||||||
video = video.to(device="cpu", copy=True)
|
video = video.to(device="cpu", copy=True)
|
||||||
chunk_length = int(param.get("chunk_length", 0) or 0)
|
chunk_length, temporal_overlap = _effective_temporal_params(param)
|
||||||
temporal_overlap = int(param.get("temporal_overlap", 0) or 0)
|
chunk_param = dict(param)
|
||||||
|
chunk_param["chunk_length"] = chunk_length
|
||||||
|
chunk_param["temporal_overlap"] = temporal_overlap
|
||||||
anchor_strength = float(param.get("anchor_strength", 0.999) or 0.999)
|
anchor_strength = float(param.get("anchor_strength", 0.999) or 0.999)
|
||||||
t = int(video.shape[2])
|
t = int(video.shape[2])
|
||||||
frame_count = _frames_for_tokens(t)
|
frame_count = _frames_for_tokens(t)
|
||||||
if chunk_length <= 0 or frame_count <= chunk_length:
|
if chunk_length <= 0 or frame_count <= chunk_length:
|
||||||
return upscaler(video, param)
|
return upscaler(video, chunk_param)
|
||||||
|
|
||||||
bounds = _temporal_segments(t, chunk_length, temporal_overlap)
|
bounds = _temporal_segments(t, chunk_length, temporal_overlap)
|
||||||
if len(bounds) <= 1:
|
if len(bounds) <= 1:
|
||||||
return upscaler(video, param)
|
return upscaler(video, chunk_param)
|
||||||
|
|
||||||
orig_dtype = video.dtype
|
orig_dtype = video.dtype
|
||||||
out = None
|
out = None
|
||||||
@@ -753,11 +773,11 @@ def _upscale_video_temporal_chunks(video, param, upscaler):
|
|||||||
for i, (k0, f0, k1, f1) in enumerate(bounds):
|
for i, (k0, f0, k1, f1) in enumerate(bounds):
|
||||||
chunk = video[:, :, k0:k1].contiguous()
|
chunk = video[:, :, k0:k1].contiguous()
|
||||||
try:
|
try:
|
||||||
chunk_out, chunk_h, chunk_w = upscaler(chunk, param)
|
chunk_out, chunk_h, chunk_w = upscaler(chunk, chunk_param)
|
||||||
except RuntimeError as exc:
|
except RuntimeError as exc:
|
||||||
if "out of memory" not in str(exc).lower():
|
if "out of memory" not in str(exc).lower():
|
||||||
raise
|
raise
|
||||||
smaller = _shrink_temporal_param(param)
|
smaller = _shrink_temporal_param(chunk_param)
|
||||||
if smaller is None:
|
if smaller is None:
|
||||||
raise
|
raise
|
||||||
try:
|
try:
|
||||||
@@ -872,10 +892,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."}),
|
"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,
|
"dynamic_fade_min": ("INT", {"default": 32, "min": 0, "max": 4096, "step": 32,
|
||||||
"tooltip": "Minimum fade width used by dynamic_fade when it is enabled."}),
|
"tooltip": "Minimum fade width used by dynamic_fade when it is enabled."}),
|
||||||
"chunk_length": ("INT", {"default": 85, "min": 17, "max": 100000, "step": 17,
|
"chunk_length": ("INT", {"default": 17, "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."}),
|
"tooltip": "Temporal chunk length for latent upscale. CUDA model upscale is capped to 17 internally so short long-video shots do not bypass splitting and OOM."}),
|
||||||
"temporal_overlap": ("INT", {"default": 17, "min": 0, "max": 100000, "step": 17,
|
"temporal_overlap": ("INT", {"default": 0, "min": 0, "max": 100000, "step": 17,
|
||||||
"tooltip": "Temporal overlap between latent chunks. 17 matches the upstream split example and reduces seam risk."}),
|
"tooltip": "Temporal overlap between latent chunks. CUDA model upscale uses 0 when capped to one H3 block to minimize peak VRAM."}),
|
||||||
"resize_conditioning": ("BOOLEAN", {"default": False,
|
"resize_conditioning": ("BOOLEAN", {"default": False,
|
||||||
"tooltip": "Reserved for upstream split compatibility. Leave OFF unless you need the original fallback behavior."}),
|
"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,
|
"anchor_strength": ("FLOAT", {"default": 0.999, "min": 0.0, "max": 1.0, "step": 0.01,
|
||||||
@@ -936,9 +956,9 @@ class H3LatentUpscaleParams:
|
|||||||
"anchor_strength": anchor_strength,
|
"anchor_strength": anchor_strength,
|
||||||
},)
|
},)
|
||||||
if chunk_length % 17 != 0:
|
if chunk_length % 17 != 0:
|
||||||
raise ValueError("chunk_length must be a multiple of 17 pixels")
|
raise ValueError("chunk_length must be a multiple of 17 frames")
|
||||||
if temporal_overlap % 17 != 0:
|
if temporal_overlap % 17 != 0:
|
||||||
raise ValueError("temporal_overlap must be a multiple of 17 pixels")
|
raise ValueError("temporal_overlap must be a multiple of 17 frames")
|
||||||
if temporal_overlap >= chunk_length:
|
if temporal_overlap >= chunk_length:
|
||||||
raise ValueError("temporal_overlap must be smaller than chunk_length")
|
raise ValueError("temporal_overlap must be smaller than chunk_length")
|
||||||
if width > 0:
|
if width > 0:
|
||||||
|
|||||||
@@ -1165,8 +1165,8 @@ class DumasH3LongVideosHelperTests(unittest.TestCase):
|
|||||||
self.assertFalse(required["brightness_match"][1]["default"])
|
self.assertFalse(required["brightness_match"][1]["default"])
|
||||||
self.assertEqual(required["dynamic_fade"][1]["default"], "off")
|
self.assertEqual(required["dynamic_fade"][1]["default"], "off")
|
||||||
self.assertEqual(required["dynamic_fade_min"][1]["default"], 32)
|
self.assertEqual(required["dynamic_fade_min"][1]["default"], 32)
|
||||||
self.assertEqual(required["chunk_length"][1]["default"], 85)
|
self.assertEqual(required["chunk_length"][1]["default"], 17)
|
||||||
self.assertEqual(required["temporal_overlap"][1]["default"], 17)
|
self.assertEqual(required["temporal_overlap"][1]["default"], 0)
|
||||||
self.assertFalse(required["resize_conditioning"][1]["default"])
|
self.assertFalse(required["resize_conditioning"][1]["default"])
|
||||||
self.assertEqual(required["anchor_strength"][1]["default"], 0.999)
|
self.assertEqual(required["anchor_strength"][1]["default"], 0.999)
|
||||||
|
|
||||||
@@ -1256,6 +1256,28 @@ class DumasH3LongVideosHelperTests(unittest.TestCase):
|
|||||||
self.assertEqual(smaller["chunk_length"], 17)
|
self.assertEqual(smaller["chunk_length"], 17)
|
||||||
self.assertEqual(smaller["temporal_overlap"], 0)
|
self.assertEqual(smaller["temporal_overlap"], 0)
|
||||||
|
|
||||||
|
def test_cuda_model_temporal_params_cap_saved_workflows(self):
|
||||||
|
latent = importlib.import_module("dumas_h3_latent_upscale")
|
||||||
|
chunk_length, temporal_overlap = latent._effective_temporal_params({
|
||||||
|
"mode": "model",
|
||||||
|
"device": "cuda",
|
||||||
|
"chunk_length": 85,
|
||||||
|
"temporal_overlap": 17,
|
||||||
|
})
|
||||||
|
self.assertEqual(chunk_length, 17)
|
||||||
|
self.assertEqual(temporal_overlap, 0)
|
||||||
|
|
||||||
|
def test_interp_temporal_params_preserve_upstream_defaults(self):
|
||||||
|
latent = importlib.import_module("dumas_h3_latent_upscale")
|
||||||
|
chunk_length, temporal_overlap = latent._effective_temporal_params({
|
||||||
|
"mode": "interp",
|
||||||
|
"device": "cuda",
|
||||||
|
"chunk_length": 85,
|
||||||
|
"temporal_overlap": 17,
|
||||||
|
})
|
||||||
|
self.assertEqual(chunk_length, 85)
|
||||||
|
self.assertEqual(temporal_overlap, 17)
|
||||||
|
|
||||||
def test_upscale_video_model_raises_when_gpu_cannot_shrink(self):
|
def test_upscale_video_model_raises_when_gpu_cannot_shrink(self):
|
||||||
latent = importlib.import_module("dumas_h3_latent_upscale")
|
latent = importlib.import_module("dumas_h3_latent_upscale")
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user