diff --git a/h3_shadow_negative.py b/h3_shadow_negative.py new file mode 100644 index 00000000..3e0eef22 --- /dev/null +++ b/h3_shadow_negative.py @@ -0,0 +1,321 @@ +"""H3 shadow negative: give the negative prompt a job at CFG 1, via attention repulsion. + +MiniMax H3 runs at CFG 1.0 always (see project_h3_sampling memory), so the negative +CONDITIONING FunPack already threads through the sampler is otherwise computed and +discarded -- comfy's CFG guider skips the uncond branch entirely at cond_scale==1.0. +[[negative_erase]] already gives that text a job by projecting its direction out of the +POSITIVE conditioning once, before sampling starts. This is a different, heavier lever: +a second "shadow" copy of the negative text is run through the SAME DiT blocks, block by +block, alongside the real (positive) stream, and the real stream's attention output is +pushed away from the shadow's, NAG-style (normalized, norm-capped so it cannot run away). + +Ported from the "H3 Shadow Negative" ComfyUI-H3NegativeLab pack (MIT), Basic-node scope +only: one shared negative drives both video and audio, T2V/FL2VA/REF2VA with best-effort +reference preservation. NOT ported: the "Advanced exact REF2VA" mode (needs a +minimax_payload/layout object this codebase's comfy fork was never confirmed to expose +the same way) and its weighted-text-encode node (FunPack already has (phrase:weight) +markup for H3 in h3_token_weights.py). + +Does NOT compose with FunPack's other per-block dit patches (REINS/Q-steer/av_decouple/ +block-repeat) -- it replaces each block's forward outright, same as upstream. See +_install_h3_shadow_negative in samplers.py for the "last installed wins, say so loudly" +policy this shares with the rest of this codebase's dit-patch stack. +""" + +import torch + + +def _as_float(x, default=0.0): + try: + if isinstance(x, torch.Tensor): + return float(x.flatten()[0]) + return float(x) + except Exception: + return float(default) + + +def sigma_progress(timestep, sigmas): + """0 at the first denoise step, 1 at the last real one, or None if unresolvable.""" + if sigmas is None: + return None + try: + cur = _as_float(timestep) + vals = sigmas.detach().flatten().tolist() if isinstance(sigmas, torch.Tensor) else list(sigmas) + if len(vals) < 2: + return None + steps = len(vals) - 1 + idx = min(range(steps), key=lambda i: abs(float(vals[i]) - cur)) + return idx / float(max(1, steps - 1)) if steps > 1 else 0.0 + except Exception: + return None + + +def nag_blend(pos, neg, scale, tau, alpha, channel_dim=1, eps=1e-6): + """Normalized extrapolation of `pos` away from `neg` (NAG-form). + + scale=1.0 is a no-op (returns pos unchanged); >1 pushes away from neg. The guided + result's norm is capped at `tau` times the positive's own norm, per channel-axis + position, so a token that would blow up under extrapolation is clamped back rather + than corrupting the stream. `alpha` blends between the positive-only and guided + result (0 = ignore this call's math entirely, 1 = fully guided). + """ + s = float(scale) + a = float(alpha) + if s == 1.0 or a == 0.0: + return pos + guided = pos * s - neg * (s - 1.0) + pos_norm = torch.linalg.vector_norm(pos.float(), ord=2, dim=channel_dim, keepdim=True) + guided_norm = torch.linalg.vector_norm(guided.float(), ord=2, dim=channel_dim, keepdim=True) + ratio = guided_norm / (pos_norm + eps) + cap = torch.clamp(float(tau) / ratio.clamp_min(eps), max=1.0).to(guided.dtype) + guided = guided * cap + return guided * a + pos * (1.0 - a) + + +class ShadowState: + """Shared, mutable state for one _sample_chunk call's worth of shadow-negative blocks. + + `neg_h` is the shadow branch's own running hidden state -- it has to go through the + SAME sequence of blocks the positive stream does (each block's own attn+mlp, in + order) to mean anything as "what the model would produce for the negative text at + this depth". Reset to None at block 0 of every forward pass (block 0 always runs + first and only once per model() call -- see _install_h3_av_decouple's own use of the + same fact), rebuilt fresh from `negative_context` on first use afterwards. + """ + __slots__ = ("dm", "negative_context", "video_scale", "audio_scale", "tau", "alpha", + "start_percent", "end_percent", "neg_h", "pos_text_len", + "warned_layout", "warned_fallback", "hard_disabled") + + def __init__(self, dm, negative_context, video_scale, audio_scale, tau, alpha, + start_percent, end_percent): + self.dm = dm + self.negative_context = negative_context + self.video_scale = float(video_scale) + self.audio_scale = float(audio_scale) + self.tau = float(tau) + self.alpha = float(alpha) + self.start_percent = float(start_percent) + self.end_percent = float(end_percent) + self.neg_h = None + self.pos_text_len = None + self.warned_layout = False + self.warned_fallback = False + self.hard_disabled = False + + +def _prepare_negative(state, x, transformer_options): + neg = state.negative_context + if not isinstance(neg, torch.Tensor): + return None + if neg.ndim == 3: + neg = neg[0] + if neg.ndim != 2 or neg.shape[0] == 0: + return None + neg = neg.to(device=x.device, dtype=x.dtype) + if neg.shape[-1] != state.dm.hidden_size: + neg = state.dm.condition_proj(neg) + neg = state.dm.token_refiner(neg, transformer_options=transformer_options) + return neg + + +def _mod_one(h, shift, scale, row): + return h.mul_(1.0 + scale[row].to(h.dtype)).add_(shift[row].to(h.dtype)) + + +def _gate_slice(x, residual, gate, a, b, row): + x[a:b].addcmul_(residual[a:b], gate[row].to(x.dtype)) + + +def _presentation_plan(mod_segments, pos_text_len): + """(runs, user_prompt_run) for H3's Qwen presentation span, or (None, None). + + MiniMax H3 tags text as modality 1, visual presentation pads as 0/2. The tokenizer + appends the user's prompt after any picture/video/audio labels, so the final tag-1 + run inside the context span is the one this replaces with the negative text; earlier + labels and all visual presentation rows stay identical to the positive. + """ + runs = [] + for seg in mod_segments: + a, b, row = seg + if a >= pos_text_len: + break + if b > pos_text_len: + return None, None + runs.append(seg) + text_runs = [r for r in runs if (r[2] % 3) == 1] + if not runs or not text_runs: + return None, None + return runs, text_runs[-1] + + +def _negative_text_rope(state, rope, start_pos, neg_len, x): + """RoPE for variable-length replacement text at the original prompt's own origin.""" + stop = start_pos + neg_len + if stop <= state.pos_text_len: + return rope[:, start_pos:stop] + try: + from comfy.ldm.minimax.model import rope_rotation_table + pos = torch.zeros(neg_len, 3, dtype=torch.float64) + pos[:, 0] = torch.arange(start_pos, stop, dtype=torch.float64) + freqs = state.dm.rope_freqs(pos, x.device) + return rope_rotation_table(freqs, x.dtype) + except Exception: + if not state.warned_layout: + state.warned_layout = True + print("[FunPackSceneChain] H3 shadow negative: negative text is longer than " + "the positive prompt span and dynamic RoPE synthesis is unavailable on " + "this core -- truncating the negative to fit.") + keep = max(1, state.pos_text_len - start_pos) + return rope[:, start_pos:start_pos + keep] + + +def _build_shadow_attention_input(state, block, neg_h, h_pos, rope, presentation_runs, + user_run, suffix_start, suffix_end, + shift_msa, scale_msa, transformer_options): + text_row = user_run[2] + neg_norm = block.norm1(neg_h) + _mod_one(neg_norm, shift_msa, scale_msa, text_row) + + neg_rope = _negative_text_rope(state, rope, user_run[0], int(neg_norm.shape[0]), h_pos) + if neg_rope.shape[1] != neg_norm.shape[0]: + neg_norm = neg_norm[:neg_rope.shape[1]] + + h_parts, r_parts, cursor = [], [], 0 + neg_a = neg_b = None + for a, b, row in presentation_runs: + if (a, b, row) == user_run: + neg_a, neg_b = cursor, cursor + neg_norm.shape[0] + h_parts.append(neg_norm) + r_parts.append(neg_rope) + cursor = neg_b + else: + h_parts.append(h_pos[a:b]) + r_parts.append(rope[:, a:b]) + cursor += b - a + + suffix_alt_start = cursor + h_parts.append(h_pos[suffix_start:suffix_end]) + r_parts.append(rope[:, suffix_start:suffix_end]) + h_alt = torch.cat(h_parts, dim=0) + rope_alt = torch.cat(r_parts, dim=1) + return h_alt, rope_alt, neg_a, neg_b, suffix_alt_start, text_row + + +def make_block_hook(block, state, index): + """One H3 double_block replacement for the shadow-negative branch. `index` is this + block's position in dm.blocks -- 0 resets the shared shadow state at the start of + every forward pass (block 0 always runs first, exactly once per model() call).""" + + def _hook(args, extra_options): + if state.hard_disabled: + return extra_options["original_block"](args) + if index == 0: + state.neg_h = None + try: + return _shadow_forward(state, block, args, extra_options) + except Exception as e: + if not state.warned_fallback: + state.warned_fallback = True + print(f"[FunPackSceneChain] H3 shadow negative: block {index} failed " + f"({type(e).__name__}: {str(e)[:120]}); falling back to stock H3 " + f"blocks for the rest of this run. Core H3 internals may have " + f"changed since this was ported.") + state.hard_disabled = True + return extra_options["original_block"](args) + + return _hook + + +def _shadow_forward(state, block, args, extra_options): + x = args["img"] + t_emb = args["t_emb"] + mod_segments = args["mod_segments"] + rope = args["rope_freqs"] + to = args["transformer_options"] + + sigmas = to.get("sample_sigmas", to.get("sigmas")) + progress = sigma_progress(args.get("timestep", to.get("timestep")), sigmas) + active = state.alpha > 0 and (state.video_scale != 1.0 or state.audio_scale != 1.0) + if progress is not None: + active = active and state.start_percent <= progress <= state.end_percent + if not active or len(mod_segments) < 3: + return extra_options["original_block"](args) + + state.pos_text_len = int(x.shape[0]) # placeholder, corrected below once c_crossattn known + cross_attn = to.get("c_crossattn") + pos_text_len = (int(cross_attn.shape[1]) + if isinstance(cross_attn, torch.Tensor) and cross_attn.ndim >= 3 + else None) + if pos_text_len is None: + return extra_options["original_block"](args) + state.pos_text_len = pos_text_len + + audio_a, audio_b, audio_row = mod_segments[-2] + video_a, video_b, video_row = mod_segments[-1] + if (int(audio_row) % 3) != 2 or (int(video_row) % 3) != 0 or audio_b != video_a \ + or video_b != x.shape[0]: + if not state.warned_layout: + state.warned_layout = True + print("[FunPackSceneChain] H3 shadow negative: could not prove the final " + "audio/video target segments on this core's H3 block layout -- " + "bypassing shadow guidance.") + return extra_options["original_block"](args) + + presentation_runs, user_run = _presentation_plan(mod_segments, pos_text_len) + if presentation_runs is None: + if not state.warned_layout: + state.warned_layout = True + print("[FunPackSceneChain] H3 shadow negative: could not identify the H3 " + "Qwen presentation span -- bypassing shadow guidance.") + return extra_options["original_block"](args) + + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = block.adaln_proj(t_emb) + + h_pos = block.norm1(x) + for a, b, row in mod_segments: + _mod_one(h_pos[a:b], shift_msa, scale_msa, row) + attn_pos = block.attn(h_pos, rope_freqs=rope, transformer_options=to) + + neg_h = state.neg_h + if neg_h is None: + neg_h = _prepare_negative(state, x, to) + if neg_h is None: + return extra_options["original_block"](args) + + h_alt, rope_alt, neg_a, neg_b, alt_suffix_start, text_row = _build_shadow_attention_input( + state, block, neg_h, h_pos, rope, presentation_runs, user_run, + pos_text_len, video_b, shift_msa, scale_msa, to) + attn_neg = block.attn(h_alt, rope_freqs=rope_alt, transformer_options=to) + a_span = (alt_suffix_start + (audio_a - pos_text_len), alt_suffix_start + (audio_b - pos_text_len)) + v_span = (alt_suffix_start + (video_a - pos_text_len), alt_suffix_start + (video_b - pos_text_len)) + + if state.video_scale != 1.0: + na, nb = v_span + attn_pos[video_a:video_b] = nag_blend(attn_pos[video_a:video_b], attn_neg[na:nb], + state.video_scale, state.tau, state.alpha) + if state.audio_scale != 1.0: + na, nb = a_span + attn_pos[audio_a:audio_b] = nag_blend(attn_pos[audio_a:audio_b], attn_neg[na:nb], + state.audio_scale, state.tau, state.alpha) + + x_out = x.clone() + for a, b, row in mod_segments: + _gate_slice(x_out, attn_pos, gate_msa, a, b, row) + h2 = block.norm2(x_out) + for a, b, row in mod_segments: + _mod_one(h2[a:b], shift_mlp, scale_mlp, row) + mlp_pos = block.mlp(h2) + for a, b, row in mod_segments: + _gate_slice(x_out, mlp_pos, gate_mlp, a, b, row) + + # Advance the shadow branch through this block too, same as the positive stream, so + # the NEXT block's repulsion target reflects the negative text at the same depth. + neg_x = neg_h[:neg_b - neg_a].clone() + attn_shadow = attn_neg[neg_a:neg_b] + neg_x.addcmul_(attn_shadow, gate_msa[text_row].to(neg_x.dtype)) + neg2 = block.norm2(neg_x) + _mod_one(neg2, shift_mlp, scale_mlp, text_row) + neg_x.addcmul_(block.mlp(neg2), gate_mlp[text_row].to(neg_x.dtype)) + state.neg_h = neg_x + + return {"img": x_out} diff --git a/movie_editor/backend/settings_card.py b/movie_editor/backend/settings_card.py index ec04bfd8..23b16507 100644 --- a/movie_editor/backend/settings_card.py +++ b/movie_editor/backend/settings_card.py @@ -203,6 +203,13 @@ _OWNED_BY = { "h3_repr_steering_strength": ("h3_repr_steering", None), "h3_q_steer_strength": ("h3_q_steer_block", None), "h3_explore_temperature_block": ("h3_explore_temperature", None), + "explore_first_step_candidates": ("explore_first_step", None), + "h3_shadow_negative_video_scale": ("h3_shadow_negative", None), + "h3_shadow_negative_audio_scale": ("h3_shadow_negative", None), + "h3_shadow_negative_tau": ("h3_shadow_negative", None), + "h3_shadow_negative_alpha": ("h3_shadow_negative", None), + "h3_shadow_negative_start_percent": ("h3_shadow_negative", None), + "h3_shadow_negative_end_percent": ("h3_shadow_negative", None), "embed_guidance_source": ("embed_guidance", None), "embed_guidance_strength": ("embed_guidance", None), "id_strength": ("identity_transfer_enabled", None), @@ -269,10 +276,19 @@ _EDITOR_DEFAULTS = { "embed_guidance": False, "embed_guidance_source": "relative", "embed_guidance_strength": 0.02, + "explore_first_step": False, + "explore_first_step_candidates": 3, "frame_overlap": 16, "h3_av_decouple": 0.0, "h3_explore_temperature": 0.0, "h3_explore_temperature_block": "40-49", + "h3_shadow_negative": False, + "h3_shadow_negative_video_scale": 3.0, + "h3_shadow_negative_audio_scale": 1.0, + "h3_shadow_negative_tau": 2.5, + "h3_shadow_negative_alpha": 0.35, + "h3_shadow_negative_start_percent": 0.0, + "h3_shadow_negative_end_percent": 0.6, "h3_block_repeat": "", "h3_block_repeat_times": 1, "h3_block_repeat_span_loop": False, diff --git a/movie_editor/frontend/engine_settings.js b/movie_editor/frontend/engine_settings.js index e1c40e73..5de56daf 100644 --- a/movie_editor/frontend/engine_settings.js +++ b/movie_editor/frontend/engine_settings.js @@ -273,6 +273,11 @@ detail: "Every other rating-driven setting only acts over the last half of a generation \u2014 after the motion and the layout have already been settled. This learns a separate memory for each quarter of the run and applies each one in its own window, so a rating about movement finally has somewhere to land. Measured across 51 rated runs, the early part carried about 88% of the signal the late part does. Costs the same as Output guidance (nothing you would notice). A quarter without enough ratings stays inert and says so." }, { name: "trajectory_guidance_strength", label: "Whole-generation strength", kind: "float", default: 0.02, min: 0.005, max: 0.1, step: 0.005, dependsOn: "trajectory_guidance", hint: "How hard each step is corrected. Same scale as Output guidance strength, but not eased off near the end \u2014 so it bites harder at the same number." }, + { name: "explore_first_step", label: "Explore at the first step", kind: "bool", default: false, + hint: "Tries a few different starting seeds for one throwaway step, scores them, and commits to whichever one the model liked best \u2014 before running the real generation from it. Needs 10+ rated generations before it does anything.", + detail: "Branches into a few candidate seeds and runs each ONE throwaway denoising step (the point where different noise still means different overall motion), scores each with the same trained memory Output guidance uses, then runs the real, full generation from the single best-scoring seed. Cost is one extra step per candidate, not a second full generation. A candidate that fails is skipped, not fatal; if nothing is ready yet the original seed runs exactly as if this were off." }, + { name: "explore_first_step_candidates", label: "Candidate seeds", kind: "int", default: 3, min: 2, max: 8, step: 1, dependsOn: "explore_first_step", + hint: "How many starting seeds to try before committing. More costs one extra throwaway step each." }, { name: "decode_noise_scale", label: "Decode noise scale", kind: "float", default: 0.0, min: 0, max: 1, step: 0.01, hint: "Adds fine detail and grain back while decoding. 0 is a clean decode, ~0.025 is a gentle restore. Free, and affects the video only — not the latent." }, { name: "decode_timestep", label: "Decode timestep", kind: "float", default: 0.05, min: 0, max: 1, step: 0.01, @@ -348,6 +353,21 @@ { name: "h3_av_decouple", label: "Video/audio attention decoupling (H3, experimental)", kind: "float", default: 0.0, min: 0.0, max: 50.0, step: 0.5, hint: "EXPERIMENTAL, unvalidated, H3 only. Weakens how much video and audio attend to each other, for the failure pattern where one comes out clean and the other doesn't. 0 = off. Not rating-driven — a manual dial, same as Repeat block(s). This is a raw attention-logit penalty, not a 0-1 fraction — try 1-5 first, then 10-20.", detail: "H3 has no separate cross-attention to mask — video, text and audio share one joint self-attention. Instead of building a mask over that whole sequence (unaffordable at H3's token count), this runs the same attention as separate video-query and audio-query passes, each carrying a penalty on the other modality's attention weight. Same total cost as the ordinary single pass. There's no way to read H3's actual trained attention scale from here (the checkpoint isn't on this machine), so this number can't be pre-calibrated — v1 tried normalizing it to 0-1 against a guessed constant and live testing came back flat across the whole range, which is why it's now the raw penalty for you to search directly. Large enough to fully exclude the other modality likely also kills legitimate audio/video sync — search up from a small value, not down from a large one." }, + { name: "h3_shadow_negative", label: "Shadow negative (H3, experimental)", kind: "bool", default: false, + hint: "Gives the negative prompt a job at CFG 1: runs it through the same blocks as a second stream and pushes the real output away from it. Heavier than 'Use the negative prompt' below (a per-step attention pass instead of a one-time edit) — needs a negative prompt connected to do anything.", + detail: "MiniMax H3 runs at CFG 1.0 always, so the negative branch is never evaluated and a plain negative prompt does nothing. This runs the negative text through the SAME DiT blocks as a second 'shadow' stream, block by block, and repels the real stream's attention output from it (NAG-style, norm-capped so it cannot run away). A different, pricier lever than negative_erase's one-time embedding projection — roughly 15-30% slower while it's on. Does NOT compose with REINS / Query steering / Video-audio decoupling / Repeat block(s) / Randomizer: whichever of those is enabled loses to this one, and the console says so. Unvalidated — ported from a community node, never judged by eye on this project's own output yet." }, + { name: "h3_shadow_negative_video_scale", label: "Video push strength", kind: "float", default: 3.0, min: 1.0, max: 10.0, step: 0.1, dependsOn: "h3_shadow_negative", + hint: "1 = off for video. Higher pushes the picture harder away from the negative." }, + { name: "h3_shadow_negative_audio_scale", label: "Audio push strength", kind: "float", default: 1.0, min: 1.0, max: 10.0, step: 0.1, dependsOn: "h3_shadow_negative", + hint: "1 = off for audio (the default — most negative prompts describe unwanted VISUALS). Higher pushes the soundtrack away from the same negative text." }, + { name: "h3_shadow_negative_tau", label: "Guidance cap (tau)", kind: "float", default: 2.5, min: 1.0, max: 8.0, step: 0.05, dependsOn: "h3_shadow_negative", + hint: "Caps how far the guided result's norm can grow relative to the un-guided one. Lower prevents overcooking; raise it only if the push strengths above feel capped too early." }, + { name: "h3_shadow_negative_alpha", label: "Blend (alpha)", kind: "float", default: 0.35, min: 0.0, max: 1.0, step: 0.05, dependsOn: "h3_shadow_negative", + hint: "0 = ignore this entirely, 1 = fully guided. If the result looks overcooked, lower this before lowering tau." }, + { name: "h3_shadow_negative_start_percent", label: "Window start", kind: "float", default: 0.0, min: 0.0, max: 1.0, step: 0.01, dependsOn: "h3_shadow_negative", + hint: "Where in the denoise schedule guidance starts (0 = the very first step)." }, + { name: "h3_shadow_negative_end_percent", label: "Window end", kind: "float", default: 0.60, min: 0.0, max: 1.0, step: 0.01, dependsOn: "h3_shadow_negative", + hint: "Where guidance stops (0.60 = the first 60% of the schedule — the community default). Guidance late in a schedule tends to fight fine detail rather than composition." }, { name: "identity_transfer_enabled", label: "Best-FaceID compatibility", kind: "bool", default: false, hint: "Feeds the identity pin image the way Best-FaceID identity LoRAs expect it. Needs an Identity pin set.", detail: "Replaces Continuity's identity-pin guide with separate, non-rendered reference tokens plus an optional ArcFace projector below. Load the LoRA itself in Models. No effect without an identity pin." }, @@ -534,9 +554,9 @@ const CHAIN_VIEW_KNOBS = { chain_continuity: ["carry_i2v_guides", "carry_overlap_through_anchor"], chain_timing: ["frame_overlap", "transition_duration", "use_same_seed", "cut_opening_frames"], - chain_guidance: ["cfg", "embed_guidance", "embed_guidance_source", "embed_guidance_strength", "score_slider", "score_slider_strength", "taste_nearest_prompt", "output_guidance", "output_guidance_strength", "trajectory_guidance", "trajectory_guidance_strength", "dynashift", "dynashift_strength", "dynashift_threshold", "h3_repr_steering", "h3_repr_steering_strength", "h3_repr_steering_block", "h3_repr_steering_passive_capture", "h3_q_steer_block", "h3_q_steer_strength", "h3_av_decouple"], + chain_guidance: ["cfg", "embed_guidance", "embed_guidance_source", "embed_guidance_strength", "score_slider", "score_slider_strength", "taste_nearest_prompt", "output_guidance", "output_guidance_strength", "trajectory_guidance", "trajectory_guidance_strength", "explore_first_step", "explore_first_step_candidates", "dynashift", "dynashift_strength", "dynashift_threshold", "h3_repr_steering", "h3_repr_steering_strength", "h3_repr_steering_block", "h3_repr_steering_passive_capture", "h3_q_steer_block", "h3_q_steer_strength", "h3_av_decouple"], chain_decode: ["decode_noise_scale", "decode_timestep", "decode_tile_size"], - chain_experimental: ["context_windows", "context_window_length", "context_window_overlap", "context_window_schedule", "context_window_fuse", "context_window_freenoise", "context_window_retain_first", "joyai_memory", "joyai_memory_size", "joyai_fix_frames", "joyai_frame_select", "joyai_memory_strength", "joyai_audio_memory", "v2a_grad_scale", "alg_blur_guides", "alg_guide_blur_strength", "alg_guide_blur_sigma_threshold", "bounded_attention_enabled", "h3_block_repeat", "h3_block_repeat_span_loop", "h3_block_repeat_video_only", "h3_block_repeat_times", "h3_block_repeat_last_steps", "h3_explore_temperature", "h3_explore_temperature_block", "identity_transfer_enabled", "source_id", "phase_scale", "id_strength", "arcface_mode", "debug_log"], + chain_experimental: ["context_windows", "context_window_length", "context_window_overlap", "context_window_schedule", "context_window_fuse", "context_window_freenoise", "context_window_retain_first", "joyai_memory", "joyai_memory_size", "joyai_fix_frames", "joyai_frame_select", "joyai_memory_strength", "joyai_audio_memory", "v2a_grad_scale", "alg_blur_guides", "alg_guide_blur_strength", "alg_guide_blur_sigma_threshold", "bounded_attention_enabled", "h3_block_repeat", "h3_block_repeat_span_loop", "h3_block_repeat_video_only", "h3_block_repeat_times", "h3_block_repeat_last_steps", "h3_explore_temperature", "h3_explore_temperature_block", "h3_shadow_negative", "h3_shadow_negative_video_scale", "h3_shadow_negative_audio_scale", "h3_shadow_negative_tau", "h3_shadow_negative_alpha", "h3_shadow_negative_start_percent", "h3_shadow_negative_end_percent", "identity_transfer_enabled", "source_id", "phase_scale", "id_strength", "arcface_mode", "debug_log"], }; function countChainView(p, id, st) { diff --git a/movie_editor/frontend/pipeline_caps.js b/movie_editor/frontend/pipeline_caps.js index 747d760a..2af5d6bc 100644 --- a/movie_editor/frontend/pipeline_caps.js +++ b/movie_editor/frontend/pipeline_caps.js @@ -170,9 +170,21 @@ // The mirror image of H3_DEAD_SAMPLER_INPUTS: settings that only mean something ON H3. // Left on against an LTX pipeline they are just as silently inert, so they get a chip // by the same rule — the user should not have to spend a generation to find out. - // key -> why it cannot run off H3. Empty for now (Picture detail, the only H3-only - // sampler input, was removed) — add an entry here if another H3-only knob needs one. - const H3_ONLY_SAMPLER_INPUTS = {}; + // key -> why it cannot run off H3. + const H3_ONLY_SAMPLER_INPUTS = { + h3_shadow_negative: + "Repels the model's attention output away from a second, negative-text pass through " + + "H3's own DiT blocks. LTXAV already evaluates a real negative prompt via CFG, so " + + "there is nothing dead here to give a job to — this only exists because H3's CFG " + + "is fixed at 1.0.", + }; + // Sub-knobs of h3_shadow_negative: meaningless without the toggle itself, so they + // travel with it rather than each needing its own entry above. + const H3_ONLY_SUB_INPUTS = [ + "h3_shadow_negative_video_scale", "h3_shadow_negative_audio_scale", + "h3_shadow_negative_tau", "h3_shadow_negative_alpha", + "h3_shadow_negative_start_percent", "h3_shadow_negative_end_percent", + ]; // Sub-settings of identity transfer: meaningless wherever the feature itself cannot run, // so they travel with it rather than each needing its own entry above. @@ -194,7 +206,7 @@ // working knob unreachable. function familyInertInputs(st) { if (!usesChainSampler(st)) return new Set(); - if (!isH3(st)) return new Set(Object.keys(H3_ONLY_SAMPLER_INPUTS)); + if (!isH3(st)) return new Set([...Object.keys(H3_ONLY_SAMPLER_INPUTS), ...H3_ONLY_SUB_INPUTS]); return new Set([...Object.keys(H3_DEAD_SAMPLER_INPUTS), ...IDENTITY_SUB_INPUTS, ...JOYAI_SUB_INPUTS, ...CONTEXT_SUB_INPUTS, ...ALG_GUIDE_SUB_INPUTS, ...H3_DEAD_VALUE_INPUTS]); diff --git a/samplers.py b/samplers.py index 53c5ef79..71526ebc 100644 --- a/samplers.py +++ b/samplers.py @@ -3097,6 +3097,34 @@ class FunPackLTXAVSceneChainSampler: "default": 3, "min": 2, "max": 8, "step": 1, "tooltip": "How many seeds to try at the first step before committing. Cost is roughly (N-1) extra single-step forward passes added to the whole schedule — 3 on a 35-step run is about 6% overhead, 8 is about 20%. No effect when explore_first_step is off.", }), + "h3_shadow_negative": ("BOOLEAN", { + "default": False, + "tooltip": "EXPERIMENTAL, unvalidated (H3 only). Ported from the community \"H3 Shadow Negative\" pack. H3's CFG is fixed at 1.0, so the negative prompt Studio already encodes is otherwise dead weight -- this runs it through the same DiT blocks as a second shadow stream and pushes the real stream's attention output away from it. Does NOT compose with REINS/Q-steer/av_decouple/block-repeat/attn_temperature -- it replaces each block's forward outright and prints which of those it superseded this run, if any. Off does not clone the model.", + }), + "h3_shadow_negative_video_scale": ("FLOAT", { + "default": 3.0, "min": 1.0, "max": 10.0, "step": 0.1, + "tooltip": "How hard to push the video rows away from the shadow negative. 1.0 = no push (identical to off). No effect when h3_shadow_negative is off.", + }), + "h3_shadow_negative_audio_scale": ("FLOAT", { + "default": 1.0, "min": 1.0, "max": 10.0, "step": 0.1, + "tooltip": "Same as h3_shadow_negative_video_scale but for the audio rows. 1.0 = no push. No effect when h3_shadow_negative is off.", + }), + "h3_shadow_negative_tau": ("FLOAT", { + "default": 2.5, "min": 1.0, "max": 8.0, "step": 0.05, + "tooltip": "NAG norm cap -- how far the pushed-away attention output is allowed to grow relative to the original before it gets clamped back down. Higher allows a stronger push before clamping kicks in. No effect when h3_shadow_negative is off.", + }), + "h3_shadow_negative_alpha": ("FLOAT", { + "default": 0.35, "min": 0.0, "max": 1.0, "step": 0.05, + "tooltip": "How much of the clamped, pushed-away result replaces the original attention output, 0-1. 0 = off (nothing changes even with the toggle on). No effect when h3_shadow_negative is off.", + }), + "h3_shadow_negative_start_percent": ("FLOAT", { + "default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, + "tooltip": "Where in the denoise schedule the shadow negative starts having any effect, as a fraction of the way through (0 = the very first step). No effect when h3_shadow_negative is off.", + }), + "h3_shadow_negative_end_percent": ("FLOAT", { + "default": 0.60, "min": 0.0, "max": 1.0, "step": 0.01, + "tooltip": "Where it stops, same units as h3_shadow_negative_start_percent. Structure is usually settled well before the end of the schedule, so this typically does not need to reach 1.0. No effect when h3_shadow_negative is off.", + }), # A connection socket, never a widget — safe at the end, and it must stay after # every widget above (see the widgets_values note at the top of this block). "second_pass_sigmas": ("SIGMAS", { @@ -3859,6 +3887,10 @@ class FunPackLTXAVSceneChainSampler: h3_block_repeat_video_only=False, h3_block_repeat_span_loop=False, h3_block_repeat_last_steps=0, h3_q_steer_strength=0.0, h3_q_steer_block="", + h3_shadow_negative=False, h3_shadow_negative_video_scale=3.0, + h3_shadow_negative_audio_scale=1.0, h3_shadow_negative_tau=2.5, + h3_shadow_negative_alpha=0.35, h3_shadow_negative_start_percent=0.0, + h3_shadow_negative_end_percent=0.60, refinement_key="",): if sampler is None: raise ValueError("sampler input is required.") @@ -4001,6 +4033,13 @@ class FunPackLTXAVSceneChainSampler: model, refinement_key, h3_q_steer_strength, _q_steer_capture, steer_block=h3_q_steer_block) + if self._is_h3 and h3_shadow_negative: + model = self._install_h3_shadow_negative( + model, negative, h3_shadow_negative_video_scale, + h3_shadow_negative_audio_scale, h3_shadow_negative_tau, + h3_shadow_negative_alpha, h3_shadow_negative_start_percent, + h3_shadow_negative_end_percent) + _phrase_probe = self._install_phrase_probe(model, positive, latent) if _phrase_probe is not None: model = _phrase_probe[0] @@ -7728,6 +7767,93 @@ class FunPackLTXAVSceneChainSampler: "not applied and not capturing this run") return model + def _install_h3_shadow_negative(self, model, negative, video_scale, audio_scale, + tau, alpha, start_percent, end_percent): + """EXPERIMENTAL, unvalidated (H3 only). Ported from the community "H3 Shadow + Negative" pack (see h3_shadow_negative.py) -- runs the negative CONDITIONING + FunPack already threads through the sampler (otherwise dead weight: H3's CFG is + fixed at 1.0, so comfy never evaluates the uncond branch) through the SAME DiT + blocks as a second "shadow" stream, and pushes the real stream's attention output + away from it, NAG-style. Heavier and more dynamic than [[negative_erase]]'s + one-time embedding projection -- see samplers.py's own H3 shadow negative + docstring in h3_shadow_negative.py for the full comparison. + + Does NOT compose with REINS/Q-steer/av_decouple/block-repeat/attn_temperature -- + it replaces each double_block's forward outright, same as upstream's own stated + limitation with block-cache nodes. Installed LAST in _sample_chunk's dit-patch + chain so an explicit request for this always wins over the others, matching what + the user just turned on rather than a mechanism that happened to install earlier. + """ + try: + video_scale = float(video_scale) + audio_scale = float(audio_scale) + tau = float(tau) + alpha = float(alpha) + start_percent = float(start_percent) + end_percent = float(end_percent) + except (TypeError, ValueError): + return model + # NaN/inf fails every ordered comparison, so `video_scale == 1.0` alone does not + # catch it (see the same trap called out for h3_av_decouple's strength guard) -- + # without this, a corrupted knob value installs successfully, prints a bogus + # "video_scale=nan" line, and NAG-blends every video/audio row to NaN with no + # warning at all, silently overwriting the block-hook's own try/except recovery. + if not all(math.isfinite(v) for v in (video_scale, audio_scale, tau, alpha, + start_percent, end_percent)): + print("[FunPackSceneChain] H3 shadow negative: a strength/tau/alpha/percent " + "value is NaN or infinite -- skipped rather than risk corrupting every " + "block downstream.") + return model + if video_scale == 1.0 and audio_scale == 1.0: + return model + if (not isinstance(negative, list) or not negative + or not isinstance(negative[0][0], torch.Tensor) + or negative[0][0].numel() == 0): + print("[FunPackSceneChain] H3 shadow negative: no negative conditioning " + "connected -- nothing to guide away from, skipped.") + return model + try: + try: + from . import h3_shadow_negative as _sn + except ImportError: + import h3_shadow_negative as _sn + patched = model.clone() + dm = patched.get_model_object("diffusion_model") + if not hasattr(dm, "blocks") or not hasattr(dm, "hidden_size"): + print("[FunPackSceneChain] H3 shadow negative: this model does not expose " + "H3 double_blocks -- skipped.") + return model + to = patched.model_options.get("transformer_options", {}).copy() + patches_replace = dict(to.get("patches_replace", {})) + dit_patches = dict(patches_replace.get("dit", {})) + _superseded = sorted({b for (kind, b) in dit_patches if kind == "double_block"}) + if _superseded: + print(f"[FunPackSceneChain] H3 shadow negative: supersedes {len(_superseded)} " + f"existing block patch(es) (REINS/Q-steer/av_decouple/block-repeat/" + f"attn_temperature) for this run -- shadow negative does not compose " + f"with those, it replaces the block outright.") + _lo, _hi = (end_percent, start_percent) if end_percent < start_percent \ + else (start_percent, end_percent) + state = _sn.ShadowState( + dm=dm, negative_context=negative[0][0].detach(), + video_scale=video_scale, audio_scale=audio_scale, + tau=tau, alpha=alpha, start_percent=_lo, end_percent=_hi) + for i, block in enumerate(dm.blocks): + dit_patches[("double_block", i)] = _tag_dit_hook( + _sn.make_block_hook(block, state, i), None) + patches_replace["dit"] = dit_patches + to["patches_replace"] = patches_replace + patched.model_options["transformer_options"] = to + print(f"[FunPackSceneChain] H3 shadow negative: installed on {len(dm.blocks)} " + f"blocks (video_scale={video_scale:g}, audio_scale={audio_scale:g}, " + f"tau={float(tau):g}, alpha={float(alpha):g}, " + f"window={_lo:.2f}-{_hi:.2f}).") + return patched + except Exception as _e: # noqa: BLE001 + _log.failed("FunPackSceneChain", "H3 shadow negative", _e, + "not applied this run") + return model + @staticmethod def _parse_block_spec(spec, n_blocks=50): """"40" | "38-42" | "10,40,44" | "" -> a set of block indices. Anything unparseable @@ -8487,6 +8613,10 @@ class FunPackLTXAVSceneChainSampler: h3_video_detail=1.0, audio_vae=None, h3_keyframes=None, explore_first_step=False, explore_first_step_candidates=3, + h3_shadow_negative=False, h3_shadow_negative_video_scale=3.0, + h3_shadow_negative_audio_scale=1.0, h3_shadow_negative_tau=2.5, + h3_shadow_negative_alpha=0.35, h3_shadow_negative_start_percent=0.0, + h3_shadow_negative_end_percent=0.60, unique_id=None, prompt=None): if not isinstance(positive, list) or not positive: raise ValueError("positive conditioning must contain at least one scene entry.") @@ -9406,6 +9536,13 @@ class FunPackLTXAVSceneChainSampler: h3_block_repeat_last_steps=h3_block_repeat_last_steps, h3_q_steer_strength=h3_q_steer_strength, h3_q_steer_block=h3_q_steer_block, + h3_shadow_negative=h3_shadow_negative, + h3_shadow_negative_video_scale=h3_shadow_negative_video_scale, + h3_shadow_negative_audio_scale=h3_shadow_negative_audio_scale, + h3_shadow_negative_tau=h3_shadow_negative_tau, + h3_shadow_negative_alpha=h3_shadow_negative_alpha, + h3_shadow_negative_start_percent=h3_shadow_negative_start_percent, + h3_shadow_negative_end_percent=h3_shadow_negative_end_percent, refinement_key=refinement_key_input, ) # cut_opening_frames: let the real, untouched i2v anchor condition the scene diff --git a/tests/test_h3_shadow_negative.py b/tests/test_h3_shadow_negative.py new file mode 100644 index 00000000..018552a3 --- /dev/null +++ b/tests/test_h3_shadow_negative.py @@ -0,0 +1,143 @@ +"""H3 shadow negative: pure-math pieces that don't need a real DiT block. + +MiniMax H3 runs at CFG 1.0 always, so a plain negative prompt is otherwise dead weight +(see negative_erase's own tests for that half of the story). This mechanism gives it a +job by pushing the positive attention output away from a shadow copy of the negative, +NAG-style. These tests pin the norm-cap math and the presentation-span detection that the +rest of h3_shadow_negative.py builds on, without needing a loaded H3 model. +""" +import sys +from pathlib import Path + +import pytest +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import h3_shadow_negative as sn # noqa: E402 + + +# --- nag_blend --------------------------------------------------------------------------- + + +def test_scale_one_is_a_no_op(): + pos = torch.randn(4, 8) + neg = torch.randn(4, 8) + out = sn.nag_blend(pos, neg, scale=1.0, tau=2.5, alpha=1.0) + assert out is pos # short-circuited, not just numerically equal + + +def test_alpha_zero_is_a_no_op(): + pos = torch.randn(4, 8) + neg = torch.randn(4, 8) + out = sn.nag_blend(pos, neg, scale=3.0, tau=2.5, alpha=0.0) + assert out is pos + + +def test_pushes_away_from_negative(): + """scale>1 should move the result further from neg than pos itself was.""" + pos = torch.ones(1, 4) + neg = torch.zeros(1, 4) + out = sn.nag_blend(pos, neg, scale=2.0, tau=8.0, alpha=1.0) + assert torch.linalg.vector_norm(out - neg) > torch.linalg.vector_norm(pos - neg) + + +def test_tau_caps_the_output_norm(): + """However hard the push, the guided result's norm cannot exceed tau times pos's own.""" + pos = torch.ones(1, 4) + neg = -100.0 * torch.ones(1, 4) # would blow the extrapolation up without a cap + out = sn.nag_blend(pos, neg, scale=5.0, tau=1.5, alpha=1.0) + pos_norm = torch.linalg.vector_norm(pos, dim=1) + out_norm = torch.linalg.vector_norm(out, dim=1) + assert torch.all(out_norm <= pos_norm * 1.5 + 1e-4) + + +def test_alpha_blends_toward_positive_only(): + pos = torch.ones(1, 4) + neg = torch.zeros(1, 4) + full = sn.nag_blend(pos, neg, scale=2.0, tau=8.0, alpha=1.0) + half = sn.nag_blend(pos, neg, scale=2.0, tau=8.0, alpha=0.5) + assert torch.allclose(half, pos * 0.5 + full * 0.5) + + +# --- sigma_progress ------------------------------------------------------------------------ + + +def test_no_sigmas_is_unresolvable(): + assert sn.sigma_progress(0.5, None) is None + + +def test_single_sigma_is_unresolvable(): + assert sn.sigma_progress(0.5, torch.tensor([1.0])) is None + + +def test_first_step_is_zero_progress(): + sigmas = torch.tensor([1.0, 0.75, 0.5, 0.25, 0.0]) + assert sn.sigma_progress(1.0, sigmas) == pytest.approx(0.0) + + +def test_last_real_step_is_full_progress(): + sigmas = torch.tensor([1.0, 0.75, 0.5, 0.25, 0.0]) + assert sn.sigma_progress(0.25, sigmas) == pytest.approx(1.0) + + +def test_nearest_sigma_wins_off_grid(): + sigmas = torch.tensor([1.0, 0.75, 0.5, 0.25, 0.0]) + assert sn.sigma_progress(0.7, sigmas) == pytest.approx(1.0 / 3.0) + + +# --- _presentation_plan -------------------------------------------------------------------- + + +def test_no_segments_inside_the_text_span_is_a_bypass(): + """Every segment starts at/after pos_text_len — there is no text run to replace.""" + runs, user_run = sn._presentation_plan([(10, 20, 0)], pos_text_len=5) + assert runs is None and user_run is None + + +def test_a_segment_straddling_the_text_boundary_is_a_bypass(): + """A run that starts inside the text span but ends past it can't be a pure label or a + pure user-prompt run — the layout assumption doesn't hold, so refuse rather than guess.""" + runs, user_run = sn._presentation_plan([(0, 8, 1)], pos_text_len=5) + assert runs is None and user_run is None + + +def test_the_last_text_tagged_run_is_the_user_prompt(): + """Row % 3 == 1 is text. An image-label run (tag 1) precedes the user's own prompt run + (also tag 1) — the LAST one is the one the tokenizer appended the user's text as.""" + segs = [(0, 3, 1), (3, 6, 0), (6, 10, 1)] + runs, user_run = sn._presentation_plan(segs, pos_text_len=10) + assert runs == segs + assert user_run == (6, 10, 1) + + +def test_no_text_tagged_run_at_all_is_a_bypass(): + runs, user_run = sn._presentation_plan([(0, 5, 0)], pos_text_len=5) + assert runs is None and user_run is None + + +# --- ShadowState ----------------------------------------------------------------------- + + +def test_shadow_state_starts_with_no_negative_prepared(): + state = sn.ShadowState(dm=None, negative_context=torch.randn(1, 3, 16), + video_scale=3.0, audio_scale=1.0, tau=2.5, alpha=0.35, + start_percent=0.0, end_percent=0.6) + assert state.neg_h is None + assert state.hard_disabled is False + + +# --- _prepare_negative ------------------------------------------------------------------- + + +def test_non_tensor_negative_context_yields_nothing(): + state = sn.ShadowState(dm=None, negative_context=None, video_scale=3.0, audio_scale=1.0, + tau=2.5, alpha=0.35, start_percent=0.0, end_percent=0.6) + assert sn._prepare_negative(state, torch.zeros(1, 4), {}) is None + + +def test_empty_negative_context_yields_nothing(): + state = sn.ShadowState(dm=None, negative_context=torch.zeros(1, 0, 4), video_scale=3.0, + audio_scale=1.0, tau=2.5, alpha=0.35, start_percent=0.0, + end_percent=0.6) + assert sn._prepare_negative(state, torch.zeros(1, 4), {}) is None diff --git a/tests/test_h3_shadow_negative_install.py b/tests/test_h3_shadow_negative_install.py new file mode 100644 index 00000000..38e89cfa --- /dev/null +++ b/tests/test_h3_shadow_negative_install.py @@ -0,0 +1,170 @@ +"""_install_h3_shadow_negative: install/no-op behaviour at the model-patching level. + +Exercised with fake model/diffusion-model stand-ins -- no real H3 weights or GPU needed. +The per-block math itself (nag_blend, presentation-span detection, ...) is covered in +tests/test_h3_shadow_negative.py; this file is about whether the installer clones the +right model, hooks the right blocks, and is honest about not composing with the other +dit-patch mechanisms. +""" +import sys +from pathlib import Path + +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) +sys.path.insert(0, ".") +import _comfy_stubs # noqa: E402, F401 +from samplers import FunPackLTXAVSceneChainSampler as S # noqa: E402 + + +class _FakeDiffusionModel: + def __init__(self, n_blocks=4): + self.blocks = [object() for _ in range(n_blocks)] + self.hidden_size = 16 + + +class _FakeModel: + def __init__(self, dm=None): + self.model_options = {} + self._dm = dm if dm is not None else _FakeDiffusionModel() + + def clone(self): + m = _FakeModel(self._dm) + m.model_options = dict(self.model_options) + return m + + def get_model_object(self, name): + if name == "diffusion_model": + return self._dm + raise KeyError(name) + + +def _negative(tensor=None): + return [(tensor if tensor is not None else torch.randn(1, 3, 16), {})] + + +def test_both_scales_at_one_is_a_true_noop(): + model = _FakeModel() + node = S() + out = node._install_h3_shadow_negative(model, _negative(), 1.0, 1.0, 2.5, 0.35, 0.0, 0.6) + assert out is model + + +def test_unparseable_scale_is_a_noop(): + model = _FakeModel() + node = S() + out = node._install_h3_shadow_negative(model, _negative(), "nope", 1.0, 2.5, 0.35, 0.0, 0.6) + assert out is model + + +def test_nan_scale_is_a_noop_not_a_silent_corruption(capsys): + """NaN fails `video_scale == 1.0`, so that alone would let it through and NAG-blend + every row to NaN with no warning -- must be caught explicitly, same as the codebase's + other dit-patch strength guards (av_decouple, q_steering).""" + model = _FakeModel() + node = S() + out = node._install_h3_shadow_negative(model, _negative(), float("nan"), 1.0, + 2.5, 0.35, 0.0, 0.6) + assert out is model + assert "NaN or infinite" in capsys.readouterr().out + + +def test_inf_tau_is_a_noop(): + model = _FakeModel() + node = S() + out = node._install_h3_shadow_negative(model, _negative(), 3.0, 1.0, + float("inf"), 0.35, 0.0, 0.6) + assert out is model + + +def test_nan_percent_is_a_noop(): + model = _FakeModel() + node = S() + out = node._install_h3_shadow_negative(model, _negative(), 3.0, 1.0, + 2.5, 0.35, float("nan"), 0.6) + assert out is model + + +def test_missing_negative_is_a_noop(): + model = _FakeModel() + node = S() + out = node._install_h3_shadow_negative(model, [], 3.0, 1.0, 2.5, 0.35, 0.0, 0.6) + assert out is model + + +def test_negative_with_empty_tensor_is_a_noop(): + model = _FakeModel() + node = S() + out = node._install_h3_shadow_negative(model, _negative(torch.zeros(1, 0, 16)), + 3.0, 1.0, 2.5, 0.35, 0.0, 0.6) + assert out is model + + +def test_a_model_with_no_double_blocks_is_a_noop(): + class _NoBlocks: + pass + model = _FakeModel(dm=_NoBlocks()) + node = S() + out = node._install_h3_shadow_negative(model, _negative(), 3.0, 1.0, 2.5, 0.35, 0.0, 0.6) + assert out is model + + +def test_enabled_clones_and_hooks_every_block(): + model = _FakeModel(dm=_FakeDiffusionModel(n_blocks=5)) + node = S() + patched = node._install_h3_shadow_negative(model, _negative(), 3.0, 1.0, 2.5, 0.35, 0.0, 0.6) + assert patched is not model + dit = patched.model_options["transformer_options"]["patches_replace"]["dit"] + assert {k for k in dit if k[0] == "double_block"} == {("double_block", i) for i in range(5)} + + +def test_swapped_start_and_end_percent_are_normalized(): + """A user typing 0.6/0.0 instead of 0.0/0.6 should not silently produce an empty + (start > end) active window -- the installer should sort them.""" + import h3_shadow_negative as sn + model = _FakeModel() + node = S() + patched = node._install_h3_shadow_negative(model, _negative(), 3.0, 1.0, 2.5, 0.35, 0.6, 0.0) + hook = patched.model_options["transformer_options"]["patches_replace"]["dit"][("double_block", 0)] + state = next(c.cell_contents for c, name in + zip(hook.__closure__, hook.__code__.co_freevars) + if name == "state" and isinstance(c.cell_contents, sn.ShadowState)) + assert state.start_percent == 0.0 and state.end_percent == 0.6 + + +def test_supersedes_an_existing_block_patch_rather_than_composing(capsys): + """Non-composing by explicit design: it must fully replace whatever dit-patch mechanism + (REINS/Q-steer/av_decouple/block-repeat/attn_temperature) already hooked a block, and + say so loudly rather than silently overwrite it.""" + model = _FakeModel(dm=_FakeDiffusionModel(n_blocks=2)) + model.model_options["transformer_options"] = {"patches_replace": {"dit": { + ("double_block", 0): lambda args, extra: {"img": args["img"] * 0.0}, + }}} + node = S() + patched = node._install_h3_shadow_negative(model, _negative(), 3.0, 1.0, 2.5, 0.35, 0.0, 0.6) + out = capsys.readouterr().out + assert "supersedes" in out + dit = patched.model_options["transformer_options"]["patches_replace"]["dit"] + # The old hook must be GONE, not chained through -- calling the new one must not + # reduce to zero the way the old one did. + img = torch.ones(3, 4) + result = dit[("double_block", 0)]({"img": img, "mod_segments": []}, + {"original_block": lambda a: {"img": a["img"]}}) + assert torch.allclose(result["img"], img), "should fall through to original_block, not the old hook" + + +def test_a_hooked_block_with_no_mod_segments_falls_through_untouched(): + model = _FakeModel(dm=_FakeDiffusionModel(n_blocks=1)) + node = S() + patched = node._install_h3_shadow_negative(model, _negative(), 3.0, 1.0, 2.5, 0.35, 0.0, 0.6) + dit = patched.model_options["transformer_options"]["patches_replace"]["dit"] + img = torch.ones(3, 4) + out = dit[("double_block", 0)]({"img": img, "mod_segments": []}, + {"original_block": lambda a: {"img": a["img"] * 2.0}}) + assert torch.allclose(out["img"], img * 2.0) + + +if __name__ == "__main__": + test_both_scales_at_one_is_a_true_noop() + test_enabled_clones_and_hooks_every_block() + print("ok (run via pytest for the rest)")