diff --git a/CHANGELOG.md b/CHANGELOG.md index 7be2d2fd..8ddf202f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,7 @@ What is new in each version, in plain words. Newest first. ## [Unreleased] -- **MiniMax H3-Image** (experimental): a new "Start from" pipeline, MiniMax H3-Image · Text to Image and Edit. Write a prompt for a picture, or add reference pictures and refer to them as ``, ``, ... in the prompt (the first is the picture to edit; up to 16). The result is a PNG picture, not a clip, and Length and FPS are not offered for it (both are set to 1). You need the model, its VAE and its text encoder files; the model's licence excludes the EU, the UK, South Korea and the United States. +- **MiniMax H3-Image** (experimental): a new "Start from" pipeline, MiniMax H3-Image · Text to Image and Edit. Write a prompt for a picture, or add reference pictures and refer to them as ``, ``, ... in the prompt (the first is the picture to edit; up to 16). The result is a PNG picture, not a clip, and Length and FPS are not offered for it (both are set to 1). REINS and the Attention explorer work on it (the explorer leaves out the sound and clip regions, which it does not have); Seed explorer and Decisiveness are offered too. You need the model, its VAE and its text encoder files; the model's licence excludes the EU, the UK, South Korea and the United States. - **Noise merging is now called Step merging**, because most of its modes merge the model's answers at a step, not noise. Settings are kept. - **Alternating prompts** (Settings > general): write `[dance|zoom]` in a prompt, or `DANCE|ZOOM` between two shortcuts, for two options, A and B. "Main prompt reads" sets what the prompt says by default: "A while B" (default), "A B", only A or only B. Switch on "Read A and B separately" (off by default) to also assign steps: "Steps that read only A" and "Steps that read only B" give each step you name its own option, with no extra model pass. A `|` next to anything that is not a shortcut is left as typed. Two options per bracket. - **Simultaneous actions** (Step merging, needs "Read A and B separately"): with `[A|B]` in the prompt, each merged step is answered once on A and once on B from the same picture, and the two answers are mixed at your strength. "Second pass new noise" sets how much of the noise B starts from is new (0 keeps the same noise for both). Strength 0 reads only A, and strength 1 with new noise 1 reads only B, each with a single pass. Steps named in Step merging win over the per-step lists. Costs one extra model pass per merged step. diff --git a/modules/models/minimax_h3_image/__init__.py b/modules/models/minimax_h3_image/__init__.py index bca3c8a8..c5b113bd 100644 --- a/modules/models/minimax_h3_image/__init__.py +++ b/modules/models/minimax_h3_image/__init__.py @@ -13,7 +13,7 @@ Everything the model needs at run time (the image latent, the picture labels, th shipped inside the file) is ComfyUI's own. """ -from ..._core import traits as _traits +from ..._core import dit_hooks as _dit_hooks, traits as _traits ID = "model_minimax_h3_image" TITLE = "MiniMax H3-Image" @@ -39,12 +39,33 @@ def detect(keys) -> bool: return any(all(f"{p}{k}" in keyset for k in _PRESENT) and f"{p}{_ABSENT}" not in keyset for p in _PREFIXES) +MODEL_CLASS = "comfy.ldm.minimax.model.MiniMaxH3ImageModel" + +#: What this model announces beyond core's own reading: its blocks are the H3 blocks (REINS reads the picture's rows +#: through them, the Attention explorer scales their attention). Block repeat and Attention temperature work the same way +#: but nobody has tried them on this model, so they stay hidden. +ANNOUNCED = ["block_rows", "block_scope", "hides:block_repeat", "hides:attention_temperature"] + + +def is_image_model(model) -> bool: + return _traits.has_block(model, MODEL_CLASS) + + +def traits(model): + return ANNOUNCED if is_image_model(model) else () + + +def probe_blocks(keys): + """How many transformer blocks the file has, read off its header.""" + return _dit_hooks.count_indexed(keys, tuple(f"{p}blocks." for p in _PREFIXES)) + + def probe_traits(keys) -> list: """What core reads off the loaded model (a 3-axis latent with a length-1 time axis, flow prediction), said before it loads.""" if not detect(keys): return [] from comfy.latent_formats import MiniMaxH3Video - found = ["predict_flow"] + found = ["predict_flow", *ANNOUNCED] rank = _traits.LATENT_RANK.get(MiniMaxH3Video.latent_dimensions) if rank: found.append(rank) @@ -113,4 +134,5 @@ def presets(): "slots": text_to_image(), "project": PROJECT}] -PROVIDES = {"detect": detect, "probe_traits": probe_traits, "pipeline_presets": presets} +TRAITS = traits +PROVIDES = {"detect": detect, "probe_traits": probe_traits, "probe_blocks": probe_blocks, "pipeline_presets": presets} diff --git a/modules/models/minimax_h3_image/tests/test_features.py b/modules/models/minimax_h3_image/tests/test_features.py new file mode 100644 index 00000000..417aafda --- /dev/null +++ b/modules/models/minimax_h3_image/tests/test_features.py @@ -0,0 +1,130 @@ +"""REINS and the Attention explorer on a real (tiny, random-weight) H3-Image forward: they run, and they do not touch a sound that is not there.""" + +import random + +import pytest +import torch + + +@pytest.fixture(autouse=True) +def root(comfyui, tmp_path, monkeypatch): + from modules.system.taste import store + monkeypatch.setattr(store, "ROOT", tmp_path / "taste") + monkeypatch.setattr(store, "current_prompt_id", lambda: "run-1") + + +class TinyImage: + """ComfyUI's own MiniMaxH3ImageModel, four blocks wide, on the CPU.""" + + def __init__(self, blocks=4): + import comfy.ops + from comfy.latent_formats import MiniMaxH3Video + from comfy.ldm.minimax.model import MiniMaxH3ImageModel + from comfy.model_patcher import ModelPatcher + + torch.manual_seed(0) + dm = MiniMaxH3ImageModel( + pdd_sigmas=[1.0, 0.5, 0.0], hidden_size=64, num_layers=blocks, token_refiner_num_layers=1, + num_attention_heads=2, attention_head_dim=32, ffn_hidden_size=64, text_dim=48, timestep_input_dim=32, + time_embed_hidden_size=64, time_embed_dim=32, rope_inv_freq_len=4, dtype=torch.float32, + operations=comfy.ops.manual_cast) + for p in dm.parameters(): + torch.nn.init.normal_(p, std=0.05) + p.requires_grad_(False) + dm.rope.inv_freq.copy_(torch.rand(4)) + + class Config: + latent_format = MiniMaxH3Video + + class Base(torch.nn.Module): + pass + + base = Base() + base.diffusion_model, base.model_config = dm, Config() + self.patcher = ModelPatcher(base, load_device=torch.device("cpu"), offload_device=torch.device("cpu")) + self.latent, self.context = torch.randn(1, 24, 1, 4, 4), torch.randn(1, 5, 48) + + def run(self, patcher=None, sigma=0.5): + import comfy.patcher_extension as ext + patcher = patcher or self.patcher + to = dict(patcher.model_options.get("transformer_options", {})) + to["wrappers"] = ext.copy_nested_dicts(to.get("wrappers", {})) + ext.merge_nested_dicts(to["wrappers"], patcher.wrappers, copy_dict1=False) + to.setdefault("sample_sigmas", torch.tensor([1.0, sigma, 0.0])) + to["sigmas"] = torch.tensor([sigma]) + with torch.inference_mode(): + return patcher.model.diffusion_model(self.latent, torch.tensor([sigma * 1000.0]), self.context, transformer_options=to) + + +def _load(tiny, **settings): + from modules.sampling.modifiers.nodes import FunUILoadModifiers + return FunUILoadModifiers.execute(tiny.patcher, {"taste": {"key": "fox"}, **settings}).result + + +def _through(patched, fn): + from comfy.patcher_extension import WrappersMP + call = fn + for ws in reversed([w for ws in patched.wrappers.get(WrappersMP.OUTER_SAMPLE, {}).values() for w in ws]): + call = (lambda w, inner: (lambda: w(lambda: inner())))(ws, call) + return call() + + +def test_the_loaded_model_announces_what_the_two_features_need(): + from core import traits + from modules.models import minimax_h3, minimax_h3_image as m + + tiny = TinyImage() + assert m.is_image_model(tiny.patcher) and not minimax_h3.is_h3(tiny.patcher) + found = set(traits.announced(tiny.patcher)) if hasattr(traits, "announced") else set(m.traits(tiny.patcher)) + assert {"block_rows", "block_scope"} <= found and "audio_stream" not in found + + +def test_reins_captures_the_picture_rows_and_a_learned_direction_moves_the_picture(): + from modules.system.taste import store + tiny = TinyImage() + base = tiny.run() + patched, status = _load(tiny, reins={"enabled": True, "blocks": "3", "strength": 0.5}) + assert "reins" in status.splitlines()[0] + _through(patched, lambda: tiny.run(patched, sigma=0.5)) # the last step: captured + assert store.rate("run-1", "liked")["recorded"] == ["reins"] + hidden = tiny.patcher.model.diffusion_model.blocks[0].attn.out_proj.weight.shape[0] + liked = torch.zeros(hidden) + liked[0] = 1.0 + for i, (vec, rating) in enumerate([(liked, "liked"), (-liked, "disliked")] * 2): + store.capture("fox", "reins", {3: vec}, prompt_id=f"t{i}") + store.rate(f"t{i}", rating) + moved, _ = _load(tiny, reins={"enabled": True, "blocks": "3", "strength": 0.5}) + assert not torch.allclose(_through(moved, lambda: tiny.run(moved)), base) + quiet, _ = _load(tiny, reins={"enabled": True, "blocks": "3", "strength": 0.0}) + assert torch.equal(_through(quiet, lambda: tiny.run(quiet)), base) + + +def test_the_attention_explorer_changes_the_picture_with_no_sound_and_no_clip_regions(monkeypatch): + from core import log + from modules.sampling import attention_bandit as m + tiny = TinyImage() + base = tiny.run() + patched, _ = _load(tiny, attention_bandit={"enabled": True, "mode": "manual", "range": 0.4, "blocks": "0-3", "sound": False}) + monkeypatch.setattr(m, "draw", lambda table, rng=random, **kw: [2] * len(table)) + log.new_run() + out = _through(patched, lambda: tiny.run(patched)) + assert not torch.equal(out, base), "the explorer ran and changed nothing" + said = " ".join(e["message"] for e in log.history()) + assert "clip regions" not in said and "sound could not be told apart" not in said and "changed nothing" not in said + monkeypatch.setattr(m, "draw", lambda table, rng=random, **kw: [m.STAY] * len(table)) + plain, _ = _load(tiny, attention_bandit={"enabled": True, "blocks": "0-3"}) + assert torch.equal(_through(plain, lambda: tiny.run(plain)), base) + + +def test_what_was_rated_with_a_sound_stream_does_not_teach_a_model_without_one(): + from modules.sampling import attention_bandit as m + tiny = TinyImage() + sig = [] + orig = m.make_plan + m.make_plan = lambda rows, cond, signature, *a, **k: (sig.append(list(signature)), orig(rows, cond, signature, *a, **k))[1] + try: + patched, _ = _load(tiny, attention_bandit={"enabled": True, "blocks": "0-3"}) + _through(patched, lambda: tiny.run(patched)) + finally: + m.make_plan = orig + assert sig and sig[0][-1] == 2 diff --git a/modules/sampling/attention_bandit/__init__.py b/modules/sampling/attention_bandit/__init__.py index de6f8809..50ecd57f 100644 --- a/modules/sampling/attention_bandit/__init__.py +++ b/modules/sampling/attention_bandit/__init__.py @@ -400,6 +400,10 @@ def install(patcher, values, key): video_width, sound_width = getattr(dm, "inner_dim", None), getattr(dm, "audio_inner_dim", None) if paired: signature.append(1) # ratings made on a packed-sequence model (H3) must not teach a paired one + silent = not paired and getattr(dm, "audio_latents_dim", 1) is None # a picture-only model (H3-Image): no sound to spare, no clip time to cut into regions + if silent: + regions, sound = False, True + signature.append(2) # and what was rated on a model with a sound stream must not teach one without if focus and paired: focus = False # the prompt's rows are not one run of the sequence here: that part is H3's if not sound and not paired: