| """H3 Steady Lab (Noah Oct 28): the stabilizer from H3 Steady with exposed controls for manual experimentation. |
| |
| Same calculation as comfyui-h3-steady: for each second of new image tokens, mean m along direction w, m < b → + strength·(b−m)·w. |
| With the default settings (steps 1+2, blocks 42–48, threshold 90 %), results are bit-identical to H3 Steady. |
| Profile profil_lab.safetensors (film/forschung_zeit/steady_lab_profil.py): w [step 0–2, block 0–49, 5376], |
| b [quantile, step, block] for 50/75/90/95/100 % of the 28 clean segments. Affects image tokens only. |
| """ |
| import logging |
| import os |
|
|
| import torch |
| from safetensors.torch import load_file |
|
|
| from comfy_api.latest import ComfyExtension, io |
|
|
| _PROFIL = os.path.join(os.path.dirname(os.path.abspath(__file__)), "profil_lab.safetensors") |
| _FPS = 24 |
| _SCHWELLEN = ["50 % (mild)", "75 %", "90 % (Standard)", "95 %", "100 % (strict)"] |
|
|
|
|
| def _frame(k): |
| b, r = divmod(k, 5) |
| return b * 17 + (0 if r == 0 else 1 + 4 * (r - 1)) |
|
|
|
|
| def _sekunden(layout, offset, sek_von, sek_bis): |
| """[(second, start row, end row)] of new image tokens, only seconds in the window [sek_von, sek_bis).""" |
| _, latent_t, lat_h, lat_w, _ = layout.signature |
| rows = (lat_h // 2) * (lat_w // 2) |
| va, vb, _ = next(s for s in layout.segments if s[2] == "video") |
| spans = {} |
| for k in range(offset, latent_t): |
| spans.setdefault(_frame(k) // _FPS, []).append(k) |
| out = [(t, va + ks[0] * rows, va + (ks[-1] + 1) * rows) for t, ks in sorted(spans.items()) if sek_von <= t < sek_bis] |
| if out and out[-1][2] > vb: |
| raise ValueError("H3 Steady Lab: unexpected video layout") |
| return out |
|
|
|
|
| def _schritt(opts): |
| sig, sched = opts.get("sigmas"), opts.get("sample_sigmas") |
| if sig is None or sched is None: |
| return -1 |
| sched = sched.flatten().float().to(sig.device) |
| return int(torch.argmin((sched - sig.flatten()[0].float()).abs())) |
|
|
|
|
| def _add(tok, delta, gen): |
| if tok.dtype != torch.bfloat16: |
| tok.add_(delta.to(tok.dtype)) |
| return |
| bits = (tok.float() + delta).view(torch.int32) |
| bits.add_(torch.randint(0, 1 << 16, bits.shape, device=bits.device, dtype=torch.int32, generator=gen)) |
| bits.bitwise_and_(-65536) |
| tok.copy_(bits.view(torch.float32).to(torch.bfloat16)) |
|
|
|
|
| def _patch(w, b, blk, staerken, sek, log): |
| cache = {} |
|
|
| @torch.no_grad() |
| def block(args, extra): |
| out = extra["original_block"](args) |
| schritt = _schritt(args["transformer_options"]) |
| if staerken.get(schritt, 0.0) == 0.0: |
| return out |
| h, layout = out["img"], args["layout"] |
| key = (layout.signature, h.device) |
| if key not in cache: |
| cache.clear() |
| cache[key] = (_sekunden(layout, *sek), w[:, blk].to(h.device)) |
| spans, wb = cache[key] |
| gen = torch.Generator(device=h.device) |
| gen.manual_seed(1234 + 97 * schritt + (blk - 42) % 50) |
| st, schwelle = staerken[schritt], float(b[schritt, blk]) |
| korr = [] |
| for x in (h if h.dim() == 3 else [h]): |
| for t, a, z in spans: |
| tok = x[a:z] |
| m = float((tok.float() @ wb[schritt]).mean()) |
| if m < schwelle: |
| _add(tok, st * (schwelle - m) * wb[schritt], gen) |
| korr.append(t) |
| if log is not None: |
| log(schritt, blk, korr, len(spans) * (h.shape[0] if h.dim() == 3 else 1), [t for t, _, _ in spans]) |
| return out |
|
|
| return block |
|
|
|
|
| class H3SteadyLab(io.ComfyNode): |
| @classmethod |
| def define_schema(cls) -> io.Schema: |
| return io.Schema( |
| node_id="H3SteadyLab_Apply", |
| display_name="H3 Steady Lab", |
| category="advanced/model_patches", |
| description="The stabilizer from H3 Steady with exposed controls (steps, blocks, threshold, time window). " |
| "Default settings = H3 Steady. Calibrated for er_sde, beta, 20 steps. Affects image tokens only.", |
| inputs=[ |
| io.Model.Input("model"), |
| io.Float.Input("staerke", display_name="Strength", default=2.0, min=-4.0, max=6.0, step=0.05, |
| tooltip="1.0 = up to the threshold, 2.0 = twice the distance. Negative = counter-test: " |
| "pushes at-risk seconds further toward the failure side."), |
| io.Float.Input("schritt_0", display_name="Step 0", default=0.0, min=0.0, max=3.0, step=0.05, |
| tooltip="Strength multiplier at step 0 (almost pure noise). 0 = off."), |
| io.Float.Input("schritt_1", display_name="Step 1", default=1.0, min=0.0, max=3.0, step=0.05, |
| tooltip="Strength multiplier at step 1 (where H3 establishes the plan)."), |
| io.Float.Input("schritt_2", display_name="Step 2", default=1.0, min=0.0, max=3.0, step=0.05, |
| tooltip="Strength multiplier at step 2."), |
| io.Int.Input("block_von", display_name="Start block", default=42, min=0, max=49, |
| tooltip="First block. The direction mainly separated cases in blocks 32–48; blocks 0–25 and 49 do not separate them."), |
| io.Int.Input("block_bis", display_name="End block", default=48, min=0, max=49, |
| tooltip="Last block (inclusive)."), |
| io.Combo.Input("schwelle", display_name="Threshold", options=_SCHWELLEN, default=_SCHWELLEN[2], |
| tooltip="When a second is considered at risk: quantile of the 28 clean segments. " |
| "Higher = more seconds are affected and pushed further."), |
| io.Int.Input("sek_von", display_name="Start second", default=0, min=0, max=600, |
| tooltip="Only seconds from here onward (whole seconds of the generated clip, including context)."), |
| io.Int.Input("sek_bis", display_name="End second", default=600, min=1, max=600, |
| tooltip="Only seconds before this point (exclusive). 3 to 5 = seconds 3.0 to 5.0."), |
| io.Int.Input("offset", default=0, min=0, max=4096, |
| tooltip="Initial latent frames to leave untouched (chain from stage 2 onward: 27). 0 for regular clips."), |
| io.Boolean.Input("protokoll", display_name="Logging", default=True, |
| tooltip="Logs how many cells (second × block) were corrected at each step to the console."), |
| ], |
| outputs=[io.Model.Output()], |
| ) |
|
|
| @classmethod |
| def execute(cls, model, staerke, schritt_0, schritt_1, schritt_2, block_von, block_bis, schwelle, sek_von, sek_bis, |
| offset, protokoll) -> io.NodeOutput: |
| staerken = {s: staerke * f for s, f in enumerate((schritt_0, schritt_1, schritt_2)) if staerke * f != 0.0} |
| bloecke = list(range(min(block_von, block_bis), max(block_von, block_bis) + 1)) |
| if not staerken or sek_bis <= sek_von: |
| logging.info("H3 Steady Lab: nothing to do (strength, steps, or time window empty)") |
| return io.NodeOutput(model) |
| prof = load_file(_PROFIL) |
| w, b = prof["w"], prof["b"][_SCHWELLEN.index(schwelle)] |
| zaehler = {} |
|
|
| def log(schritt, blk, korr, gesamt, sekunden): |
| z = zaehler.setdefault(schritt, [0, 0, {}]) |
| z[0] += len(korr) |
| z[1] += gesamt |
| for t in korr: |
| z[2][t] = z[2].get(t, 0) + 1 |
| if blk == bloecke[-1]: |
| je = " ".join(f"{t}s:{z[2].get(t, 0)}" for t in sekunden) |
| logging.info(f"H3 Steady Lab: step {schritt}: {z[0]} of {z[1]} cells corrected " |
| f"(blocks {bloecke[0]}–{bloecke[-1]}); per second, out of {len(bloecke)} blocks: {je}") |
| zaehler.pop(schritt) |
|
|
| m = model.clone() |
| for blk in bloecke: |
| m.set_model_patch_replace(_patch(w, b, blk, staerken, (offset, sek_von, sek_bis), log if protokoll else None), |
| "dit", "double_block", blk) |
| logging.info(f"H3 Steady Lab: strength {staerke}, steps {staerken}, blocks {bloecke[0]}–{bloecke[-1]}, " |
| f"threshold {schwelle}, seconds {sek_von}–{sek_bis}, offset {offset}") |
| return io.NodeOutput(m) |
|
|
|
|
| class H3SteadyLabExtension(ComfyExtension): |
| async def get_node_list(self) -> list[type[io.ComfyNode]]: |
| return [H3SteadyLab] |
|
|
|
|
| async def comfy_entrypoint() -> H3SteadyLabExtension: |
| return H3SteadyLabExtension() |
|
|
