Expose latent upscale sampler controls

2026-09-03 13:05:32 +00:00
parent c1d937e0e2
commit 5c12cd18a3
5 changed files with 126 additions and 52 deletions
@@ -33,9 +33,9 @@ and your VRAM, chains them, and returns the finished video + audio.
 Requirements: H3 is CFG-free (cfg 1) and needs no negative prompt -- the node
 makes an empty one internally. The main pass keeps denoise fixed at 1.0: a
 partial denoise desyncs the joint audio/video schedule. An optional latent
-upscale stage can rebuild the conditioning at a larger target size, run a short
-refinement pass, and keep the output video-only before the final pixel-space
-upscale options.
+upscale stage can rebuild the conditioning at a target size, run a short
+refinement pass with its own sampler controls, and keep the output video-only
+before the final pixel-space upscale options.

 Verified against ComfyUI core (comfy_extras/nodes_minimax_h3.py, model_base.py,
 ldm/minimax/model.py, text_encoders/minimax.py, sd.py).
@@ -4068,11 +4068,22 @@ def _video_only_refined_latent(base_latent, refined_latent):
        refined_parts = refined.unbind()
        if len(base_parts) >= 2 and len(refined_parts) >= 1:
            return {"samples": comfy.nested_tensor.NestedTensor((refined_parts[0], base_parts[-1]))}
-    except Exception:
+    except Exception:
        return refined_latent
    return refined_latent


+def _latent_upscale_target_size(base_w, base_h, param):
+    width = int(param.get("width", 0) or 0)
+    height = int(param.get("height", 0) or 0)
+    if width > 0 and height > 0:
+        return width, height
+    megapixels = float(param.get("megapixels", 0.0) or 0.0)
+    if megapixels > 0:
+        return scale_to_megapixels(base_w, base_h, megapixels)
+    return int(base_w), int(base_h)
+
+
 def _nested_tensor_parts(samples):
    if samples is None:
        return ()
@@ -6107,8 +6118,9 @@ class H3LongVideos:
                               "Try 256 on a tight card at 1344x768."}),
                "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 hard-coded 2-step refinement pass before "
-                               "decode. Leave unconnected to skip latent upscaling entirely."}),
+                               "latent is upscaled and run through a short refinement pass before decode, "
+                               "using the sampler, scheduler, steps, denoise, and megapixel target 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 "
                               "Resolution (Tensor Cores -- fastest and best for video; needs the "
@@ -6422,11 +6434,11 @@ class H3LongVideos:
        # Conditioning is built, so the text encoder and VAEs are dead weight for the
        # whole sampling loop -- evict them and keep only the DiT on the card.
        _evict_all_but(model)
-        try:
-            sample_start = time.perf_counter()
-            (out,) = nodes.common_ksampler(model, seed, steps, cfg, sn, sch, positive, negative,
-                                           latent, denoise=denoise)
-            timing["sample"] += time.perf_counter() - sample_start
+        try:
+            sample_start = time.perf_counter()
+            (out,) = nodes.common_ksampler(model, seed, steps, cfg, sn, sch, positive, negative,
+                                           latent, denoise=denoise)
+            timing["sample"] += time.perf_counter() - sample_start
        except Exception as e:
            # Mark WHERE this failed. `tiled` only affects the DECODE, so the caller's
            # OOM retry cannot help an OOM raised here -- it just re-runs the whole
@@ -6440,8 +6452,6 @@ class H3LongVideos:
        if (
            isinstance(latent_upscale_param, dict)
            and str(latent_upscale_param.get("mode", "off")) != "off"
-            and int(latent_upscale_param.get("width", 0) or 0) > 0
-            and int(latent_upscale_param.get("height", 0) or 0) > 0
        ):
            try:
                latent_start = time.perf_counter()
@@ -6463,9 +6473,14 @@ class H3LongVideos:
                    ref_noise_aug=ref_noise_aug, audio_vae=audio_vae, silent=silent)
                upscale_latent["samples"] = comfy.nested_tensor.NestedTensor(
                    (upscaled_video, parts[1]))
+                refine_steps = int(latent_upscale_param.get("steps", 2) or 2)
+                refine_sampler = latent_upscale_param.get("sampler_name", sn)
+                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, 2, cfg, sn, sch, upscale_cond, negative, upscale_latent,
-                    denoise=float(latent_upscale_param.get("refine_denoise", 0.25) or 0.25))
+                    model, seed, refine_steps, cfg, refine_sampler, refine_scheduler, upscale_cond, negative, upscale_latent,
+                    denoise=refine_denoise)
                timing["latent_upscale_sample"] += time.perf_counter() - latent_start
                refined_out = _video_only_refined_latent(out, refined_out)
            except Exception as e:
@@ -6591,19 +6606,21 @@ class H3LongVideos:
            model, ms_note = apply_h3_model_sampling(model, shift_video, shift_audio)
        latent_upscale_note = ""
        if isinstance(latent_upscale_param, dict) and str(latent_upscale_param.get("mode", "off")) != "off":
-            target_w = int(latent_upscale_param.get("width", 0) or 0)
-            target_h = int(latent_upscale_param.get("height", 0) or 0)
-            if target_w > 0 and target_h > 0:
-                mode = str(latent_upscale_param.get("mode", "off"))
-                detail = f" via {mode}"
-                if mode == "model":
-                    detail += f"/{latent_upscale_param.get('model_name', 'none')}"
-                else:
-                    detail += f"/{latent_upscale_param.get('method', 'bilinear')}"
-                latent_upscale_note = (
-                    f" latent upscale: target {target_w}x{target_h}px{detail}; "
-                    f"2-step refinement denoise {float(latent_upscale_param.get('refine_denoise', 0.25) or 0.25):.2f}"
-                )
+            target_w, target_h = _latent_upscale_target_size(w, h, latent_upscale_param)
+            mode = str(latent_upscale_param.get("mode", "off"))
+            detail = f" via {mode}"
+            if mode == "model":
+                detail += f"/{latent_upscale_param.get('model_name', 'none')}"
+            else:
+                detail += f"/{latent_upscale_param.get('method', 'bilinear')}"
+            denoise_value = latent_upscale_param.get("denoise", latent_upscale_param.get("refine_denoise", 0.2))
+            denoise = 0.2 if denoise_value is None else float(denoise_value)
+            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"{latent_upscale_param.get('sampler_name', sn)}/{latent_upscale_param.get('scheduler', sch)} "
+                f"denoise {denoise:.2f}"
+            )

        paras = split_paragraphs(prompt, "##")
        if anchor_override.strip():