diff --git a/conditioning.py b/conditioning.py index dec7ba4a..62579925 100644 --- a/conditioning.py +++ b/conditioning.py @@ -3089,7 +3089,8 @@ class FunPackVideoRefiner: # Unlike anything in conditioning space these transfer between prompts: 0.85 means the # same thing on every prompt, so a value learned on one run is worth something on the # next. That is the property that makes them worth learning at all. - H3_GAIN_KEYS = ("video", "prompt", "audio", "prompt_scale", "refiner_bias") + H3_GAIN_KEYS = ("video", "prompt", "audio", "prompt_scale", "refiner_bias", + "video_detail") # Which rated AXIS each gain answers to. A rating carries three per-axis signals in # [-1, +1] alongside its overall reward, and they say far more than the scalar does: # "Missing concept" is reward +0.10 (quality was fine) but concept_signal -1.00. Read as @@ -3106,6 +3107,12 @@ class FunPackVideoRefiner: "prompt_scale": ("concept_signal",), "audio": (), "refiner_bias": (), + # The final-layer read strength: the same axes as the block gain it is the + # audio-safe twin of. `prompt_time` is deliberately absent from these keys — its + # off value (0.0) sits at the far end of its own range rather than in the middle, + # so exploring around it would step straight into "the prompt is pure noise". + # It stays a manual dial until a rental run says what a good value looks like. + "video_detail": ("detail_signal", "quality_signal"), } H3_GAIN_MIN = 0.60 H3_GAIN_MAX = 1.40 diff --git a/minimax_h3.py b/minimax_h3.py index c0e0c0b8..53391528 100644 --- a/minimax_h3.py +++ b/minimax_h3.py @@ -1360,8 +1360,12 @@ def patch_module_forward(patcher, key, wrapper): patcher.add_object_patch(f"{key}.forward", wrapper) -class AdalnModalityGain: - """Wraps one block's ``adaln_proj.forward`` and scales its gates per modality. +class AdalnEdit: + """Wraps one block's ``adaln_proj.forward``: per-modality gate gain, prompt timestep. + + Both edits act on the same projection output through the same row indexing, so they + share one wrapper — two wrappers on one key would nest, and ``unwrapped_forward`` + deliberately strips a FunPack wrapper rather than stacking on it. Deliberately NOT an ``nn.Module``. Replacing a module with a wrapper module renames every weight beneath it in the state dict (``adaln_proj.weight`` becomes @@ -1374,16 +1378,20 @@ class AdalnModalityGain: #: marks an already-installed wrapper, so a re-install unwraps instead of nesting _funpack_wrapper = True - def __init__(self, inner, gains): + def __init__(self, inner, gains, prompt_row=None): self.inner = inner # The ORIGINAL callable, captured before the patch is applied. Calling # ``inner.forward`` at call time would re-enter this wrapper once it is installed. self.original = unwrapped_forward(inner) self.gains = {int(t): float(g) for t, g in (gains or {}).items() if float(g) != 1.0} + self.prompt_row = prompt_row def __call__(self, t_emb): - out = self.original(t_emb) - if not self.gains or getattr(self.inner, "modalities", 1) != _ADALN_MODALITIES: + if getattr(self.inner, "modalities", 1) != _ADALN_MODALITIES: + return self.original(t_emb) # FinalLayer shares the class, not the row meaning + out = self._with_prompt_timestep(t_emb) if self.prompt_row is not None \ + else self.original(t_emb) + if not self.gains: return out out = list(out) for idx in _ADALN_GATE_CHUNKS: @@ -1395,26 +1403,112 @@ class AdalnModalityGain: out[idx] = scaled return tuple(out) + def _with_prompt_timestep(self, t_emb): + """Run the projection with ONE extra timestep row, and give it to the prompt. -def apply_adaln_gains(model, gains): - """Attach per-modality gate gains to every DiT block. Returns (model, note). + The model tells every row what noise level it is at. Text is told the video's. + Appending a row for the requested level and copying its tag-1 modulation over every + text row hands the prompt its own noise level without touching the model's own row + bookkeeping — the extra row is dropped before the result is returned, so every index + the caller computed still means what it meant. + """ + rows = int(t_emb.shape[0]) + try: + extra = self.prompt_row(t_emb) + except Exception: # embedding failed — run the prompt as trained + return self.original(t_emb) + out = self.original(torch.cat([t_emb, extra], dim=0)) + source = rows * _ADALN_MODALITIES + MODALITY_TAGS["text"] + trimmed = [] + for chunk in out: + kept = chunk[:rows * _ADALN_MODALITIES].clone() + kept[MODALITY_TAGS["text"]::_ADALN_MODALITIES] = chunk[source] + trimmed.append(kept) + return tuple(trimmed) + + +class PromptTimestepRow: + """One absolute timestep, embedded in the DiT's own time-embedding space. + + Computed on first use rather than at install time: the checkpoint may still be on the + CPU when the patch is attached, and the row has to land on whatever device and dtype + the live ``t_emb`` is using. Cached per (device, dtype) and shared by all 50 blocks. + """ + + def __init__(self, timestep, table=None, embedder=None): + self.timestep = float(timestep) + self.table = table # curve-form checkpoints: the time-embedding basis + self.embedder = embedder # full-form checkpoints: the time embedder module + self._cache = None + + def __call__(self, like): + key = (like.device, like.dtype) + if self._cache is not None and self._cache[0] == key: + return self._cache[1] + row = self._embed(like).reshape(1, -1).to(dtype=like.dtype) + self._cache = (key, row) + return row + + def _embed(self, like): + t = min(1.0, max(0.0, self.timestep)) + if self.table is not None: + # the same fractional lookup the model does for its own timesteps + table = self.table.to(device=like.device, dtype=torch.float32) + pos = t * (table.shape[0] - 1) + low = min(int(pos), table.shape[0] - 2) + return torch.lerp(table[low], table[low + 1], + torch.tensor(pos - low, dtype=torch.float32, device=like.device)) + value = torch.tensor([t], dtype=torch.float32, device=like.device) + return self.embedder(value) + + +def prompt_timestep_row(model, timestep): + """Build the row for `timestep`, reading whichever time basis this checkpoint carries.""" + try: + table = model.get_model_object("diffusion_model.adaln_t_table") + except Exception: + table = None + if table is not None: + return PromptTimestepRow(timestep, table=table), None + try: + embedder = model.get_model_object("diffusion_model.time_embedder") + except Exception as error: + return None, f"no time embedding to read ({error})" + return PromptTimestepRow(timestep, embedder=embedder), None + + +def apply_adaln_edits(model, gains=None, prompt_timestep=0.0): + """Attach the per-block AdaLN edits. Returns (model, note). Uses ``add_object_patch`` rather than replacing methods by name: ComfyUI restores object patches on unpatch, so nothing survives the run. Degradation that outlives a reset has been a real bug here before, and this is the mechanism that prevents it. - `gains` maps a name in MODALITY_TAGS to a float. 1.0 (or absent) means untouched, and - an all-1.0 request returns the model unchanged rather than cloning it. + `gains` maps a name in MODALITY_TAGS to a float; 1.0 (or absent) means untouched. + `prompt_timestep` of 0 means the prompt keeps the video's. A request that changes + nothing returns the model unchanged rather than cloning it. """ live = {MODALITY_TAGS[name]: float(value) for name, value in (gains or {}).items() if name in MODALITY_TAGS and float(value) != 1.0} - if not live: + want_time = float(prompt_timestep or 0.0) > 0.0 + if not live and not want_time: return model, None try: blocks = model.get_model_object("diffusion_model.blocks") except Exception as error: # not H3, or a layout we do not know - return model, f"AdaLN gain skipped — no DiT blocks to patch ({error})." + return model, f"AdaLN edits skipped — no DiT blocks to patch ({error})." + prompt_row = None + if want_time: + prompt_row, why = prompt_timestep_row(model, prompt_timestep) + if prompt_row is None: + if not live: + return model, f"Prompt timestep skipped — {why}." + note_time = f" Prompt timestep skipped — {why}." + else: + note_time = f" Prompt read at timestep {float(prompt_timestep):g}." + else: + note_time = "" patched = model.clone() count = 0 for index in range(len(blocks)): @@ -1423,13 +1517,71 @@ def apply_adaln_gains(model, gains): inner = patched.get_model_object(key) except Exception: continue - patch_module_forward(patched, key, AdalnModalityGain(inner, live)) + patch_module_forward(patched, key, AdalnEdit(inner, live, prompt_row)) count += 1 if not count: - return model, "AdaLN gain skipped — no adaln_proj found on the blocks." + return model, "AdaLN edits skipped — no adaln_proj found on the blocks." named = ", ".join(f"{n}={v:g}" for n, v in sorted((gains or {}).items()) if float(v) != 1.0) - return patched, (f"AdaLN modality gain on {count} blocks ({named}). Keyframe pins and " - f"reference images ride the VIDEO tag, so the video gain moves them too.") + gain_note = (f"AdaLN modality gain on {count} blocks ({named}). Keyframe pins and " + f"reference images ride the VIDEO tag, so the video gain moves them too." + if live else f"AdaLN edits on {count} blocks.") + return patched, (gain_note + note_time).strip() + + +# ── final layer: the one video-only edit ───────────────────────────────────── +# +# Every block shares one attention pass, so anything written into the video rows is read by +# the audio rows in the next block — which is what made post-block injection corrupt the +# soundtrack. The final layer is past the last attention: video and audio are modulated on +# separate rows and leave through separate heads, so an edit to the video branch here has +# no path to the audio at all. +# +# The arithmetic below MIRRORS ``FinalLayer.forward`` using the module's own submodules +# rather than wrapping its result, because the value being changed (the video row's scale) +# is consumed inside that method. It is checked against the module's attributes before it +# runs, so a changed upstream shape declines instead of computing something else. +_FINAL_LAYER_PARTS = ("norm", "adaln_proj", "video_out", "audio_out") + + +class FinalLayerVideoScale: + """Scales how strongly the final layer reads the accumulated video rows.""" + + _funpack_wrapper = True + + def __init__(self, inner, scale): + self.inner = inner + self.original = unwrapped_forward(inner) + self.scale = float(scale) + + def __call__(self, x, t_emb, video_seg, audio_seg): + module = self.inner + if self.scale == 1.0 or not all(hasattr(module, part) for part in _FINAL_LAYER_PARTS): + return self.original(x, t_emb, video_seg, audio_seg) + shift, scale = module.adaln_proj(t_emb) + va, vb, vrow = video_seg + aa, ab, arow = audio_seg + hv = (module.norm(x[va:vb]) * (1.0 + scale[vrow] * self.scale) + + shift[vrow]).to(torch.float32) + ha = (module.norm(x[aa:ab]) * (1.0 + scale[arow]) + shift[arow]).to(torch.float32) + return module.video_out(hv), module.audio_out(ha) + + +def apply_final_video_scale(model, scale): + """Attach the video-only final-layer scale. Returns (model, note).""" + if float(scale) == 1.0: + return model, None + key = "diffusion_model.final_layer" + try: + inner = model.get_model_object(key) + except Exception as error: + return model, f"Final-layer video scale skipped — no final layer to patch ({error})." + if not all(hasattr(inner, part) for part in _FINAL_LAYER_PARTS): + return model, ("Final-layer video scale skipped — this final layer is not the shape " + "the edit was written against.") + patched = model.clone() + patch_module_forward(patched, key, FinalLayerVideoScale(inner, scale)) + return patched, (f"Final-layer video scale {float(scale):g} — past the last attention, " + f"so the audio stream cannot see it.") # ── token-refiner output ───────────────────────────────────────────────────── diff --git a/movie_editor/backend/timeline.py b/movie_editor/backend/timeline.py index dda731d6..ac211577 100644 --- a/movie_editor/backend/timeline.py +++ b/movie_editor/backend/timeline.py @@ -556,7 +556,7 @@ class Scene: #: drives them by hand, which is what `h3_gain_mode: manual` is for. KEY_SCOPED_SAMPLER_INPUTS = frozenset({ "h3_gain_mode", "h3_gain_video", "h3_gain_prompt", "h3_gain_audio", - "h3_prompt_scale", "h3_taste_bias", + "h3_prompt_scale", "h3_taste_bias", "h3_video_detail", "h3_prompt_time", }) diff --git a/movie_editor/frontend/engine_settings.js b/movie_editor/frontend/engine_settings.js index 8323a694..fbbcf5ba 100644 --- a/movie_editor/frontend/engine_settings.js +++ b/movie_editor/frontend/engine_settings.js @@ -134,7 +134,7 @@ dependsOn: "prompt_enhance", dependsVals: [true], hint: "Let the model reason before answering, if it supports it. Slower, and the reasoning is stripped from the result." }, { name: "h3_render_gains", label: "Rating-learned render gains (H3)", kind: "bool", default: true, - hint: "Learns five render strengths from your ratings alone: how hard each block writes into the picture, the prompt and the soundtrack, how loudly the prompt is read, and how far to push the prompt toward what you have rated well. Five numbers is a smaller search than the sigma schedule already learns, so there is nothing to tune by hand — each run is rendered slightly off the learned values so the next rating has something to credit. The taste push needs at least 3 liked runs on the key before it does anything. Free. The learned values live in the refinement key, not in this project — clearing your keys resets them. Off leaves the model at its trained strengths; the Chain Sampler's h3_gain_* and h3_taste_bias widgets then do nothing unless its h3_gain_mode is set to manual." } + hint: "Learns six render strengths from your ratings alone: how hard each block writes into the picture, the prompt and the soundtrack, how loudly the prompt is read, how far to push the prompt toward what you have rated well, and how much detail the picture gets. Six numbers is a smaller search than the sigma schedule already learns, so there is nothing to tune by hand — each run is rendered slightly off the learned values so the next rating has something to credit. The taste push needs at least 3 liked runs on the key before it does anything. Free. The learned values live in the refinement key, not in this project — clearing your keys resets them. Off leaves the model at its trained strengths; the Chain Sampler's h3_gain_*, h3_taste_bias, h3_video_detail and h3_prompt_time widgets then do nothing unless its h3_gain_mode is set to manual." } ]; function parseStudioSettings(p) { @@ -351,6 +351,14 @@ dependsOn: "h3_gain_mode", dependsValue: "manual", hint: "Pushes the prompt toward what you have rated well. Negative pushes away from it. Needs at least 3 liked runs on the refinement key; without them it does nothing and says so. 0.0 = off. Free.", detail: "The Refiner already steers along your liked direction, but on H3 that edit is mixed and renormalized by the token refiner before the model reads it. This adds the same direction after the refiner, where it lands as sent. Measured as a fraction of a typical prompt row, so one value means the same thing on every prompt." }, + { name: "h3_video_detail", label: "Picture detail (audio-safe)", kind: "float", default: 1.0, min: 0.0, max: 2.0, step: 0.01, + dependsOn: "h3_gain_mode", dependsValue: "manual", + hint: "Makes the picture crisper or softer without changing the sound at all. Above 1.0 = more detail and contrast; past about 1.3 it overcooks. 1.0 = untouched. Free. Unvalidated.", + detail: "Every block of the model shares one attention pass, so anything that moves the picture also reaches the soundtrack. This runs after the last one, where there is no path left for it to travel — it is the audio-safe twin of Write gain \u00b7 picture." }, + { name: "h3_prompt_time", label: "Prompt is settled", kind: "float", default: 0.0, min: 0.0, max: 1.0, step: 0.01, + dependsOn: "h3_gain_mode", dependsValue: "manual", + hint: "Makes the picture stick closer to the prompt, and drift away from it less as a scene goes on. 0.0 = off. Try 0.9-1.0; below about 0.5 it makes things worse, not weaker. Free. Unvalidated.", + detail: "The model tells every part of its sequence how finished it is, and it tells the prompt whatever it tells the picture — so early on, while the picture is still noise, the prompt is treated as unreliable too. This tells the prompt it is settled no matter where the picture is. Not learned from ratings yet: its off value sits at the end of its range rather than the middle, so the learner cannot explore around it safely." }, { name: "segmented_detailing", label: "Segmented detailing (region refine)", kind: "bool", default: false, hint: "Re-renders small regions like hands at higher detail after each scene. Costs about 15%.", detail: "After each scene renders, CLIPSeg finds the named regions, refines them at 2x through Lightricks' latent upsampler and pastes them back through a feathered silhouette. Final resolution never changes. Regions over 35% of the frame are refused; every skip is in the scene report." }, @@ -551,7 +559,7 @@ 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", "dynashift", "dynashift_strength", "dynashift_threshold"], 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", "plateau_cache", "plateau_cache_threshold", "h3_audio_clock", "h3_gain_mode", "h3_gain_video", "h3_gain_prompt", "h3_gain_audio", "h3_prompt_scale", "h3_taste_bias", "segmented_detailing", "detail_targets", "detail_strength", "detail_threshold", "detail_max_area", "detail_mode", "detail_denoise", "mid_scene_guide", "mid_scene_guide_strength", "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", "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", "plateau_cache", "plateau_cache_threshold", "h3_audio_clock", "h3_gain_mode", "h3_gain_video", "h3_gain_prompt", "h3_gain_audio", "h3_prompt_scale", "h3_taste_bias", "h3_video_detail", "h3_prompt_time", "segmented_detailing", "detail_targets", "detail_strength", "detail_threshold", "detail_max_area", "detail_mode", "detail_denoise", "mid_scene_guide", "mid_scene_guide_strength", "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", "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 4d83d3b8..ee330444 100644 --- a/movie_editor/frontend/pipeline_caps.js +++ b/movie_editor/frontend/pipeline_caps.js @@ -189,6 +189,8 @@ h3_gain_audio: H3_RENDER_GAIN_REASON, h3_prompt_scale: H3_RENDER_GAIN_REASON, h3_taste_bias: H3_RENDER_GAIN_REASON, + h3_video_detail: H3_RENDER_GAIN_REASON, + h3_prompt_time: H3_RENDER_GAIN_REASON, }; // Sub-settings of identity transfer: meaningless wherever the feature itself cannot run, diff --git a/samplers.py b/samplers.py index ae34b626..25523959 100644 --- a/samplers.py +++ b/samplers.py @@ -3173,12 +3173,20 @@ class FunPackLTXAVSceneChainSampler: }), "h3_gain_mode": (["learned", "manual"], { "default": "learned", - "tooltip": "Where the five H3 render strengths come from. 'learned' (default) takes them from the refinement key: Studio learns them from your ratings alone and tags them onto the conditioning, and the five widgets below are IGNORED — nothing to tune by hand, which is the point. Five scalars is a smaller search than the sigma schedule already learns from ratings, so it converges in tens of rated runs. With no key wired, or before the first rating, 'learned' renders at the model's trained strengths (untouched, and the taste push at 0). 'manual' ignores the learned values and uses the widgets below exactly as set — for deliberately probing one strength, or for a graph with no Refiner in it.", + "tooltip": "Where the seven H3 render strengths come from. 'learned' (default) takes them from the refinement key: Studio learns them from your ratings alone and tags them onto the conditioning, and the seven widgets below are IGNORED — nothing to tune by hand, which is the point. Six of the seven are learned (h3_prompt_time is not — its off value sits at the end of its own range, not the middle, so the loop cannot explore around it safely; it stays a manual dial). Six scalars is a smaller search than the sigma schedule already learns from ratings, so it converges in tens of rated runs. With no key wired, or before the first rating, 'learned' renders at the model's trained strengths (untouched, and the taste push at 0). 'manual' ignores the learned values and uses the widgets below exactly as set — for deliberately probing one strength, or for a graph with no Refiner in it.", }), "h3_prompt_scale": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05, "tooltip": "EXPERIMENTAL, MiniMax H3 only. How LOUD the prompt is when the picture reads it. H3 puts the Qwen text through condition_proj + a 2-block token refiner (ending in an RMSNorm) before the DiT sees it, and the blocks then read the text through attention — so the magnitude of those refined rows is how strongly the prompt competes for attention against the picture and the reference. This multiplies them: above 1.0 the prompt is harder to ignore, below 1.0 it recedes and the reference or the anchor gets more say. Applied AFTER the refiner's final norm, so it lands as set rather than being renormalized away. 1.0 = untouched, and the model is not cloned. Free — one multiply on the text rows, no extra model calls. Confined to the PROMPT rows: the label and vision block sit ahead of them and are left alone (read from minimax_token_tags; if those cannot be read the whole text span is scaled and the console says so). DIFFERENT from h3_gain_prompt, which controls how much each block WRITES BACK into the text rows — this controls how loudly they are READ.", }), + "h3_prompt_time": ("FLOAT", { + "default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01, + "tooltip": "EXPERIMENTAL, MiniMax H3 only. Makes the picture stick closer to the prompt, and drift away from it less as a scene goes on. 0.0 = off (default). Try 0.9-1.0; the effect gets stronger the closer to 1.0. HOW: H3 tells every part of its sequence how finished it is, and it tells the prompt whatever it tells the picture — so early on, while the picture is still noise, the prompt is treated as unreliable too. This tells the prompt it is finished no matter where the picture is, so the model leans on it from the first step. Free: one extra row through a small projection per block, no extra model calls. Does not touch the reference image or a keyframe pin. UNVALIDATED.", + }), + "h3_video_detail": ("FLOAT", { + "default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05, + "tooltip": "EXPERIMENTAL, MiniMax H3 only. Makes the picture crisper or softer WITHOUT changing the sound at all. Above 1.0 = more detail and contrast (past ~1.3 it overcooks), below 1.0 = softer and calmer. 1.0 = untouched, and the model is not cloned. This is the audio-safe twin of h3_gain_video: everything earlier in the model shares one attention pass, so a change to the picture reaches the soundtrack, but this runs after the last one — there is no path left for it to travel. Free. UNVALIDATED.", + }), "alg_anchor": ("BOOLEAN", { "default": False, "tooltip": "EXPERIMENTAL: run ALG's i2v anchor blur (arXiv:2506.08456) on WHATEVER sampler is wired — a stock KSampler with any sampler_name, Hybrid Euler 2S, a two-evals-per-step sampler like heun, anything. The blur de-statics an anchored scene by hiding the anchor's high-frequency detail during the near-pure-noise steps, so the model cannot shortcut to a video that just matches the still. It is the same guidance as the FunPack Distilled Flow sampler's own alg_enabled, and this switch drives that one too when Distilled Flow is the wired sampler, so there is one control wherever you are. The swap is decided by the step's sigma alone, which is an argument of every model call, so it does not need to run inside a sampler's loop. No effect on a scene with no i2v anchor.", @@ -3992,6 +4000,7 @@ class FunPackLTXAVSceneChainSampler: # key, no rating and no Studio state, so it works on a graph with the Refiner absent # entirely. 1.0 on all three does not even clone the model. model = self._install_h3_adaln_gains(model, positive) + model = self._install_h3_final_layer(model, positive) model = self._install_h3_token_refiner(model, positive) try: @@ -6723,7 +6732,8 @@ class FunPackLTXAVSceneChainSampler: #: the value of each render gain that means "untouched". Not all of them are 1.0: #: refiner_bias is a signed push along a learned direction, so its neutral is 0.0. H3_GAIN_NEUTRAL = {"video": 1.0, "prompt": 1.0, "audio": 1.0, - "prompt_scale": 1.0, "refiner_bias": 0.0} + "prompt_scale": 1.0, "refiner_bias": 0.0, + "prompt_time": 0.0, "video_detail": 1.0} def _h3_render_gains(self, positive): """The four render strengths for this run: learned from ratings, or the widgets. @@ -6741,7 +6751,9 @@ class FunPackLTXAVSceneChainSampler: "prompt": float(getattr(self, "_h3_gain_prompt", 1.0)), "audio": float(getattr(self, "_h3_gain_audio", 1.0)), "prompt_scale": float(getattr(self, "_h3_prompt_scale", 1.0)), - "refiner_bias": float(getattr(self, "_h3_taste_bias", 0.0))} + "refiner_bias": float(getattr(self, "_h3_taste_bias", 0.0)), + "prompt_time": float(getattr(self, "_h3_prompt_time", 0.0)), + "video_detail": float(getattr(self, "_h3_video_detail", 1.0))} if str(getattr(self, "_h3_gain_mode", "learned")).lower() == "manual": return manual learned = None @@ -6889,7 +6901,9 @@ class FunPackLTXAVSceneChainSampler: Returns `model` untouched when all three gains are 1.0, which is every default run. """ gains = self._h3_render_gains(positive) - if all(gains.get(k, 1.0) == 1.0 for k in ("video", "prompt", "audio")): + prompt_time = float(gains.get("prompt_time", 0.0)) + if all(gains.get(k, 1.0) == 1.0 for k in ("video", "prompt", "audio")) \ + and prompt_time <= 0.0: return model try: from . import minimax_h3 as h3mod @@ -6904,7 +6918,7 @@ class FunPackLTXAVSceneChainSampler: # is what it is to the person setting it. tagged = {"video": gains["video"], "text": gains["prompt"], "audio": gains["audio"]} try: - patched, note = h3mod.apply_adaln_gains(model, tagged) + patched, note = h3mod.apply_adaln_edits(model, tagged, prompt_timestep=prompt_time) except Exception as error: # noqa: BLE001 _log.failed("FunPackSceneChain", "AdaLN modality gain", error, "the blocks write at their trained strength for every modality") @@ -6913,6 +6927,36 @@ class FunPackLTXAVSceneChainSampler: print(f"[FunPackSceneChain] {note}") return patched + def _install_h3_final_layer(self, model, positive): + """The video-only detail scale, applied past the model's last attention pass. + + Everything else that moves the picture also reaches the soundtrack, because all 50 + blocks share one attention sequence. This one cannot: there is no attention after + the final layer, so the audio branch is untouched by construction. + """ + gains = self._h3_render_gains(positive) + detail = float(gains.get("video_detail", 1.0)) + if detail == 1.0: + return model + try: + from . import minimax_h3 as h3mod + except ImportError: + import minimax_h3 as h3mod + if not h3mod.is_h3_model(model): + _log.feature( + "FunPackSceneChain", "Video detail", False, + "not a MiniMax H3 model. The final-layer edit is an H3-only lane.") + return model + try: + patched, note = h3mod.apply_final_video_scale(model, detail) + except Exception as error: # noqa: BLE001 + _log.failed("FunPackSceneChain", "Video detail", error, + "the picture is rendered at the model's trained strength") + return model + if note: + print(f"[FunPackSceneChain] {note}") + return patched + def _install_h3_token_weights(self, model, positive): """Apply the rating-derived phrase emphasis Studio tagged onto the conditioning. @@ -7242,6 +7286,7 @@ class FunPackLTXAVSceneChainSampler: h3_audio_clock=False, h3_gain_video=1.0, h3_gain_prompt=1.0, h3_gain_audio=1.0, h3_prompt_scale=1.0, h3_taste_bias=0.0, h3_gain_mode="learned", + h3_prompt_time=0.0, h3_video_detail=1.0, audio_vae=None, h3_keyframes=None, unique_id=None, prompt=None): if not isinstance(positive, list) or not positive: @@ -7262,6 +7307,8 @@ class FunPackLTXAVSceneChainSampler: self._h3_gain_audio = max(0.0, min(2.0, float(h3_gain_audio))) self._h3_prompt_scale = max(0.0, min(2.0, float(h3_prompt_scale))) self._h3_taste_bias = max(-0.30, min(0.30, float(h3_taste_bias))) + self._h3_prompt_time = max(0.0, min(1.0, float(h3_prompt_time))) + self._h3_video_detail = max(0.0, min(2.0, float(h3_video_detail))) self._h3_gain_mode = str(h3_gain_mode or "learned").strip().lower() # The gate every rating-driven wrapper shares. On H3 it is read off the schedule's # own base grid; on LTX it stays the absolute-sigma gate it was validated with. diff --git a/tests/test_h3_adaln_gain.py b/tests/test_h3_adaln_gain.py index ec30a2c3..ec0de68c 100644 --- a/tests/test_h3_adaln_gain.py +++ b/tests/test_h3_adaln_gain.py @@ -44,7 +44,7 @@ def gates(out): def test_a_gain_scales_only_its_own_modality(): inner = FakeProj() - wrapped = h3.AdalnModalityGain(inner, {h3.MODALITY_TAGS["audio"]: 0.5}) + wrapped = h3.AdalnEdit(inner, {h3.MODALITY_TAGS["audio"]: 0.5}) base_msa, base_mlp = gates(inner(None)) msa, mlp = gates(wrapped(None)) for tag in (0, 1): @@ -57,7 +57,7 @@ def test_a_gain_scales_only_its_own_modality(): def test_only_the_gate_chunks_move(): """shift and scale modulate what a block READS; the gate is what it WRITES.""" inner = FakeProj() - wrapped = h3.AdalnModalityGain(inner, {0: 0.25}) + wrapped = h3.AdalnEdit(inner, {0: 0.25}) base, out = inner(None), wrapped(None) for idx in (0, 1, 3, 4): assert torch.equal(out[idx], base[idx]) @@ -66,7 +66,7 @@ def test_only_the_gate_chunks_move(): def test_every_timestep_row_of_that_modality_is_scaled(): """Rows are t_row*3+tag, so a modality appears once per unique timestep.""" inner = FakeProj(timesteps=3) - wrapped = h3.AdalnModalityGain(inner, {1: 2.0}) + wrapped = h3.AdalnEdit(inner, {1: 2.0}) base_msa, _ = gates(inner(None)) msa, _ = gates(wrapped(None)) assert msa[1::3].shape[0] == 3 @@ -77,20 +77,20 @@ def test_the_source_buffer_is_not_mutated(): """chunk() returns views onto one tensor — scaling in place would corrupt the module.""" inner = FakeProj() before = inner.buf.clone() - h3.AdalnModalityGain(inner, {0: 0.1})(None) - h3.AdalnModalityGain(inner, {0: 0.1})(None) + h3.AdalnEdit(inner, {0: 0.1})(None) + h3.AdalnEdit(inner, {0: 0.1})(None) assert torch.equal(inner.buf, before) def test_gain_of_one_is_dropped_entirely(): - wrapped = h3.AdalnModalityGain(FakeProj(), {0: 1.0, 1: 1.0}) + wrapped = h3.AdalnEdit(FakeProj(), {0: 1.0, 1: 1.0}) assert wrapped.gains == {} def test_a_single_modality_projection_is_left_alone(): """FinalLayer uses AdalnProj too, with modalities=1 — different row meaning.""" inner = FakeProj(modalities=1, expand=2) - wrapped = h3.AdalnModalityGain(inner, {0: 0.5}) + wrapped = h3.AdalnEdit(inner, {0: 0.5}) assert all(torch.equal(a, b) for a, b in zip(wrapped(None), inner(None))) @@ -127,30 +127,30 @@ class FakePatcher: def test_all_ones_returns_the_same_model_unpatched(): model = FakePatcher() - out, note = h3.apply_adaln_gains(model, {"video": 1.0, "text": 1.0, "audio": 1.0}) + out, note = h3.apply_adaln_edits(model, {"video": 1.0, "text": 1.0, "audio": 1.0}) assert out is model and note is None assert model.patched == {} def test_every_block_is_wrapped(): model = FakePatcher(n=4) - out, note = h3.apply_adaln_gains(model, {"video": 0.8}) + out, note = h3.apply_adaln_edits(model, {"video": 0.8}) assert out is not model assert len(out.patched) == 4 - assert all(isinstance(v, h3.AdalnModalityGain) for v in out.patched.values()) + assert all(isinstance(v, h3.AdalnEdit) for v in out.patched.values()) assert "4 blocks" in note def test_it_says_the_anchor_moves_with_the_video(): """Pins and reference images carry the VIDEO tag — a real limit, stated not hidden.""" - _out, note = h3.apply_adaln_gains(FakePatcher(), {"video": 0.8}) + _out, note = h3.apply_adaln_edits(FakePatcher(), {"video": 0.8}) assert "reference images ride the VIDEO tag" in note def test_object_patches_are_used_not_method_replacement(): """ComfyUI restores object patches on unpatch. Degradation outliving a reset has been a real bug here; this is the mechanism that stops it.""" - src = inspect.getsource(h3.apply_adaln_gains) + src = inspect.getsource(h3.apply_adaln_edits) assert "add_object_patch" in src assert "model.clone()" in src @@ -159,14 +159,14 @@ def test_an_unknown_model_declines_with_a_reason(): class NoBlocks: def get_model_object(self, name): raise KeyError(name) - out, note = h3.apply_adaln_gains(NoBlocks(), {"video": 0.5}) + out, note = h3.apply_adaln_edits(NoBlocks(), {"video": 0.5}) assert isinstance(out, NoBlocks) assert "no DiT blocks" in note def test_unknown_modality_names_are_ignored(): model = FakePatcher() - out, note = h3.apply_adaln_gains(model, {"nonsense": 0.5}) + out, note = h3.apply_adaln_edits(model, {"nonsense": 0.5}) assert out is model and note is None diff --git a/tests/test_h3_final_layer.py b/tests/test_h3_final_layer.py new file mode 100644 index 00000000..62d550f6 --- /dev/null +++ b/tests/test_h3_final_layer.py @@ -0,0 +1,166 @@ +"""The final layer: the one place a video-only edit stays video-only. + +All 50 DiT blocks share a single attention pass, so anything written into the video rows is +read by the audio rows in the next block — that is what made post-block injection corrupt +the soundtrack. The final layer runs after the last attention: video and audio are +modulated on separate rows and leave through separate heads, so an edit to the video branch +here has no path to the audio at all. These tests hold that guarantee, not the arithmetic. +""" +import sys +from pathlib import Path + +import pytest +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import minimax_h3 as h3 + +HIDDEN = 4 +ROWS = 2 + + +class FakeAdaln(torch.nn.Module): + """modalities=1: two chunks (shift, scale), one row per timestep.""" + + def __init__(self): + super().__init__() + self.modalities = 1 + self.shift = torch.tensor([[0.1] * HIDDEN, [0.2] * HIDDEN]) + self.scale = torch.tensor([[0.5] * HIDDEN, [0.3] * HIDDEN]) + + def forward(self, _t_emb): + return self.shift, self.scale + + +class FakeFinalLayer(torch.nn.Module): + """Mirrors comfy.ldm.minimax.model.FinalLayer.forward.""" + + def __init__(self): + super().__init__() + self.norm = torch.nn.Identity() + self.adaln_proj = FakeAdaln() + self.video_out = torch.nn.Linear(HIDDEN, 3, dtype=torch.float32) + self.audio_out = torch.nn.Linear(HIDDEN, 2, dtype=torch.float32) + + def forward(self, x, t_emb, video_seg, audio_seg): + shift, scale = self.adaln_proj(t_emb) + va, vb, vrow = video_seg + aa, ab, arow = audio_seg + hv = (self.norm(x[va:vb]) * (1.0 + scale[vrow]) + shift[vrow]).to(torch.float32) + ha = (self.norm(x[aa:ab]) * (1.0 + scale[arow]) + shift[arow]).to(torch.float32) + return self.video_out(hv), self.audio_out(ha) + + +def run(module, scale): + x = torch.arange(6 * HIDDEN, dtype=torch.float32).view(6, HIDDEN) / 10.0 + args = (x, None, (0, 3, 0), (3, 6, 1)) + base_v, base_a = module(*args) + got_v, got_a = h3.FinalLayerVideoScale(module, scale)(*args) + return (base_v, base_a), (got_v, got_a) + + +def test_the_audio_is_bit_for_bit_what_it_would_have_been(): + """The whole reason this handle exists. Not 'close' — identical.""" + (_, base_a), (_, got_a) = run(FakeFinalLayer(), 1.6) + assert torch.equal(base_a, got_a) + + +def test_the_video_actually_moves(): + (base_v, _), (got_v, _) = run(FakeFinalLayer(), 1.6) + assert not torch.allclose(base_v, got_v) + + +def test_a_scale_of_one_is_the_untouched_forward(): + (base_v, base_a), (got_v, got_a) = run(FakeFinalLayer(), 1.0) + assert torch.equal(base_v, got_v) and torch.equal(base_a, got_a) + + +def test_the_scale_multiplies_the_modulation_not_the_output(): + """It changes how strongly the final layer READS the accumulated rows. Scaling the + output instead would be a step-size change, which is a different thing entirely.""" + module = FakeFinalLayer() + x = torch.arange(6 * HIDDEN, dtype=torch.float32).view(6, HIDDEN) / 10.0 + got_v, _ = h3.FinalLayerVideoScale(module, 2.0)(x, None, (0, 3, 0), (3, 6, 1)) + shift, scale = module.adaln_proj(None) + want = module.video_out(x[0:3] * (1.0 + scale[0] * 2.0) + shift[0]) + assert torch.allclose(got_v, want) + + +def test_the_bias_is_left_alone(): + """Scale is a contrast dial; shift would push every channel off its trained centre.""" + module = FakeFinalLayer() + x = torch.zeros(6, HIDDEN) + got_v, _ = h3.FinalLayerVideoScale(module, 3.0)(x, None, (0, 3, 0), (3, 6, 1)) + base_v, _ = module(x, None, (0, 3, 0), (3, 6, 1)) + assert torch.allclose(got_v, base_v) # x is 0, so only the shift survives — unchanged + + +def test_a_final_layer_of_a_shape_we_do_not_know_declines(): + """Mirroring upstream arithmetic means upstream can move underneath it. Declining is the + only safe answer — computing something else would be a silent behaviour change.""" + class Different(torch.nn.Module): + def forward(self, x, t_emb, video_seg, audio_seg): + return x[:1], x[1:2] + module = Different() + x = torch.ones(4, HIDDEN) + got = h3.FinalLayerVideoScale(module, 1.5)(x, None, (0, 3, 0), (3, 4, 1)) + assert torch.equal(got[0], module(x, None, (0, 3, 0), (3, 4, 1))[0]) + + +# ── installing it ─────────────────────────────────────────────────────────── + +class FakePatcher: + def __init__(self, final=None): + self.objects = {"diffusion_model.final_layer": final if final is not None + else FakeFinalLayer()} + self.patched = {} + + def get_model_object(self, name): + if name in self.patched: + return self.patched[name] + return self.objects[name] + + def clone(self): + other = FakePatcher.__new__(FakePatcher) + other.objects = self.objects + other.patched = dict(self.patched) + return other + + def add_object_patch(self, name, obj): + self.patched[name] = obj + + +def test_one_is_a_no_op_and_does_not_clone(): + model = FakePatcher() + out, note = h3.apply_final_video_scale(model, 1.0) + assert out is model and note is None and model.patched == {} + + +def test_it_patches_the_forward_not_the_module(): + """Replacing a module renames every weight under it in the state dict; ComfyUI records + weight keys off the live tree, so an unwrap by another patcher would walk a path that is + no longer there.""" + out, note = h3.apply_final_video_scale(FakePatcher(), 1.4) + assert list(out.patched) == ["diffusion_model.final_layer.forward"] + assert isinstance(out.patched["diffusion_model.final_layer.forward"], + h3.FinalLayerVideoScale) + assert "audio stream cannot see it" in note + + +def test_a_model_with_no_final_layer_declines_with_a_reason(): + class Nothing: + def get_model_object(self, name): + raise KeyError(name) + out, note = h3.apply_final_video_scale(Nothing(), 1.4) + assert isinstance(out, Nothing) + assert "no final layer" in note + + +def test_an_unfamiliar_final_layer_declines_before_it_is_installed(): + class Different(torch.nn.Module): + pass + model = FakePatcher(final=Different()) + out, note = h3.apply_final_video_scale(model, 1.4) + assert out is model and model.patched == {} + assert "not the shape" in note diff --git a/tests/test_h3_prompt_timestep.py b/tests/test_h3_prompt_timestep.py new file mode 100644 index 00000000..34d9ec2b --- /dev/null +++ b/tests/test_h3_prompt_timestep.py @@ -0,0 +1,240 @@ +"""Prompt timestep: the prompt is told how finished it is, independently of the picture. + +H3 tells every row of its packed sequence what noise level it is at, and it gives the text +rows the video's. Early in a run that means the prompt is treated as unfinished too. This +gives the text rows a level of their own by putting ONE extra row through each block's +AdaLN projection and copying its tag-1 modulation over the text rows — the model's own row +bookkeeping never sees it, because the extra row is dropped before the result comes back. +""" +import sys +import types +from pathlib import Path + +import pytest +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import minimax_h3 as h3 + +MOD = 3 +EXPAND = 6 +HIDDEN = 4 + + +class RowProj(torch.nn.Module): + """AdalnProj stand-in whose every output row reports (timestep, modality). + + Row value is `t * 100 + tag`, so which timestep a modality was modulated at is readable + straight off the tensor. + """ + + def __init__(self, modalities=MOD, expand=EXPAND, hidden=HIDDEN): + super().__init__() + self.modalities = modalities + self.expand = expand + self.hidden = hidden + + def forward(self, t_emb): + rows = int(t_emb.shape[0]) + base = t_emb[:, 0].repeat_interleave(self.modalities) * 100.0 + tags = torch.arange(self.modalities, dtype=torch.float32).repeat(rows) + out = (base + tags).unsqueeze(1).expand( + rows * self.modalities, self.expand * self.hidden).contiguous() + return out.chunk(self.expand, dim=-1) + + +def t_emb(*values): + return torch.tensor(values, dtype=torch.float32).unsqueeze(1) + + +def row_values(chunk): + return chunk[:, 0].tolist() + + +def constant_row(value): + return lambda like: torch.tensor([[value]], dtype=like.dtype, device=like.device) + + +# ── the edit itself ───────────────────────────────────────────────────────── + +def test_the_prompt_is_modulated_at_its_own_timestep(): + inner = RowProj() + wrapped = h3.AdalnEdit(inner, None, prompt_row=constant_row(0.95)) + out = wrapped(t_emb(0.2, 0.7)) + # rows are t_row*3 + tag: video 0/3, text 1/4, audio 2/5. The fake reports t*100+tag, + # so 96 proves the copy came from the extra timestep's TEXT row, not its video row. + assert row_values(out[0]) == [20.0, 96.0, 22.0, 70.0, 96.0, 72.0] + + +def test_every_other_modality_keeps_the_timestep_it_had(): + inner = RowProj() + plain = row_values(inner(t_emb(0.2, 0.7))[0]) + out = row_values(h3.AdalnEdit(inner, None, prompt_row=constant_row(0.95))(t_emb(0.2, 0.7))[0]) + for row in (0, 2, 3, 5): + assert out[row] == plain[row] + + +def test_the_extra_row_never_reaches_the_caller(): + """The model indexes modulation as t_row*3+tag off its OWN row count. An extra row left + on the end would not break those indices, but a changed row COUNT would be a lie about + how many timesteps are in play.""" + out = h3.AdalnEdit(RowProj(), None, prompt_row=constant_row(0.9))(t_emb(0.2, 0.7)) + assert all(chunk.shape[0] == 2 * MOD for chunk in out) + + +def test_all_six_modulation_vectors_move_together(): + """shift/scale/gate for attention and for the MLP all belong to one timestep — moving + some and not others would describe a state the model was never trained on.""" + out = h3.AdalnEdit(RowProj(), None, prompt_row=constant_row(0.95))(t_emb(0.2)) + assert all(row_values(chunk)[1] == 96.0 for chunk in out) + + +def test_off_by_default_leaves_the_projection_alone(): + inner = RowProj() + out = h3.AdalnEdit(inner, None)(t_emb(0.2, 0.7)) + assert all(torch.equal(a, b) for a, b in zip(out, inner(t_emb(0.2, 0.7)))) + + +def test_a_single_modality_projection_is_left_alone(): + """FinalLayer shares AdalnProj with modalities=1, where tag 1 means nothing.""" + inner = RowProj(modalities=1, expand=2) + wrapped = h3.AdalnEdit(inner, None, prompt_row=constant_row(0.9)) + assert all(torch.equal(a, b) for a, b in zip(wrapped(t_emb(0.2)), inner(t_emb(0.2)))) + + +def test_a_failed_embedding_renders_the_prompt_as_trained(): + """Never a hard failure: the run continues with the model's own behaviour.""" + def boom(_like): + raise RuntimeError("no time basis") + inner = RowProj() + out = h3.AdalnEdit(inner, None, prompt_row=boom)(t_emb(0.2)) + assert all(torch.equal(a, b) for a, b in zip(out, inner(t_emb(0.2)))) + + +def test_it_composes_with_the_modality_gains(): + """Both edits ride one wrapper — installing two on the same key would nest, and + unwrapped_forward deliberately strips a FunPack wrapper rather than stacking on it.""" + wrapped = h3.AdalnEdit(RowProj(), {h3.MODALITY_TAGS["text"]: 2.0}, + prompt_row=constant_row(0.95)) + gate = row_values(wrapped(t_emb(0.2))[2]) + assert gate[1] == 96.0 * 2.0 # re-timed AND gained + assert gate[0] == 20.0 # video untouched by both + + +# ── the timestep row ──────────────────────────────────────────────────────── + +def test_the_curve_table_is_read_the_way_the_model_reads_it(): + table = torch.tensor([[0.0, 0.0], [10.0, 20.0], [20.0, 40.0]], dtype=torch.float32) + row = h3.PromptTimestepRow(0.5, table=table)(torch.zeros(1, 2)) + assert torch.allclose(row, torch.tensor([[10.0, 20.0]])) # t=0.5 -> grid row 1 + + +def test_a_fractional_position_interpolates(): + table = torch.tensor([[0.0, 0.0], [10.0, 20.0]], dtype=torch.float32) + row = h3.PromptTimestepRow(0.25, table=table)(torch.zeros(1, 2)) + assert torch.allclose(row, torch.tensor([[2.5, 5.0]])) + + +def test_the_top_of_the_range_stays_on_the_last_interval(): + """t=1.0 must read the final grid row, not one past the end of the table.""" + table = torch.tensor([[0.0], [1.0], [2.0]], dtype=torch.float32) + assert torch.allclose(h3.PromptTimestepRow(1.0, table=table)(torch.zeros(1, 1)), + torch.tensor([[2.0]])) + + +def test_a_checkpoint_without_curves_uses_its_time_embedder(): + class Embedder: + def __call__(self, value): + return torch.stack([value, value * 2.0], dim=-1) + row = h3.PromptTimestepRow(0.4, embedder=Embedder())(torch.zeros(1, 2)) + assert torch.allclose(row, torch.tensor([[0.4, 0.8]])) + + +def test_the_row_is_cast_to_whatever_the_live_embedding_is_using(): + table = torch.tensor([[1.0], [2.0]], dtype=torch.float32) + row = h3.PromptTimestepRow(0.0, table=table)(torch.zeros(1, 1, dtype=torch.bfloat16)) + assert row.dtype == torch.bfloat16 + + +def test_the_row_is_computed_once_and_reused(): + """All 50 blocks ask for the same row every step.""" + calls = [] + + class Counting: + def __call__(self, value): + calls.append(1) + return value.unsqueeze(-1) + provider = h3.PromptTimestepRow(0.5, embedder=Counting()) + like = torch.zeros(1, 1) + for _ in range(5): + provider(like) + assert len(calls) == 1 + + +# ── installing it ─────────────────────────────────────────────────────────── + +class FakeBlock(torch.nn.Module): + def __init__(self): + super().__init__() + self.adaln_proj = RowProj() + + +class FakePatcher: + def __init__(self, n=3, table=torch.tensor([[0.0], [1.0]])): + self.objects = {"diffusion_model.blocks": [FakeBlock() for _ in range(n)]} + for i in range(n): + self.objects[f"diffusion_model.blocks.{i}.adaln_proj"] = \ + self.objects["diffusion_model.blocks"][i].adaln_proj + if table is not None: + self.objects["diffusion_model.adaln_t_table"] = table + self.patched = {} + + def get_model_object(self, name): + if name in self.patched: + return self.patched[name] + return self.objects[name] + + def clone(self): + other = FakePatcher.__new__(FakePatcher) + other.objects = self.objects + other.patched = dict(self.patched) + return other + + def add_object_patch(self, name, obj): + self.patched[name] = obj + + +def test_zero_changes_nothing_and_does_not_clone(): + model = FakePatcher() + out, note = h3.apply_adaln_edits(model, None, prompt_timestep=0.0) + assert out is model and note is None and model.patched == {} + + +def test_a_timestep_alone_is_enough_to_install(): + model = FakePatcher(n=3) + out, note = h3.apply_adaln_edits(model, None, prompt_timestep=0.9) + assert len(out.patched) == 3 + assert all(v.prompt_row is not None for v in out.patched.values()) + assert "0.9" in note + + +def test_every_block_shares_one_row_provider(): + out, _ = h3.apply_adaln_edits(FakePatcher(n=3), None, prompt_timestep=0.9) + providers = {id(v.prompt_row) for v in out.patched.values()} + assert len(providers) == 1 + + +def test_a_checkpoint_with_no_time_basis_declines_with_a_reason(): + model = FakePatcher(table=None) + out, note = h3.apply_adaln_edits(model, None, prompt_timestep=0.9) + assert out is model + assert "no time embedding" in note + + +def test_the_gains_still_apply_when_the_timestep_cannot_be_read(): + """A missing time basis must not silently cancel the edit the user also asked for.""" + model = FakePatcher(table=None) + out, note = h3.apply_adaln_edits(model, {"video": 0.8}, prompt_timestep=0.9) + assert out is not model and len(out.patched) == 3 + assert "Prompt timestep skipped" in note diff --git a/tests/test_h3_token_refiner_edit.py b/tests/test_h3_token_refiner_edit.py index 71be1f74..1c85b362 100644 --- a/tests/test_h3_token_refiner_edit.py +++ b/tests/test_h3_token_refiner_edit.py @@ -136,7 +136,7 @@ def test_it_patches_the_forward_not_the_module(): def test_the_edit_is_not_a_module_so_the_state_dict_is_untouched(): assert not isinstance(h3.TokenRefinerEdit(FakeRefiner()), torch.nn.Module) - assert not isinstance(h3.AdalnModalityGain(FakeRefiner(), {}), torch.nn.Module) + assert not isinstance(h3.AdalnEdit(FakeRefiner(), {}), torch.nn.Module) def test_reinstalling_unwraps_instead_of_nesting(): @@ -327,5 +327,5 @@ def test_the_module_itself_stays_in_place(): def test_the_adaln_gain_patches_a_forward_too(): """Same hazard, same shape: adaln_proj.weight must not become adaln_proj.inner.weight.""" - src = inspect.getsource(h3.apply_adaln_gains) + src = inspect.getsource(h3.apply_adaln_edits) assert "patch_module_forward" in src diff --git a/tests/test_key_scoped_knobs.py b/tests/test_key_scoped_knobs.py index e098c585..263131c0 100644 --- a/tests/test_key_scoped_knobs.py +++ b/tests/test_key_scoped_knobs.py @@ -25,7 +25,7 @@ FRONTEND = Path(__file__).resolve().parents[1] / "movie_editor" / "frontend" def test_every_learned_render_knob_is_key_scoped(): assert KEY_SCOPED_SAMPLER_INPUTS == { "h3_gain_mode", "h3_gain_video", "h3_gain_prompt", "h3_gain_audio", - "h3_prompt_scale", "h3_taste_bias", + "h3_prompt_scale", "h3_taste_bias", "h3_video_detail", "h3_prompt_time", }