Images/audio/video tab order, expandable image slots, native attention for the VAEs

app.py

@@ -66,6 +66,9 @@ MAX_UI_DURATION = 14
 MIN_DURATION = 5
 # A reference video shorter than 2 s gives the model almost no motion to read, and 15 s is the checkpoint's ceiling.
 MIN_REFERENCE_VIDEO, MAX_REFERENCE_VIDEO = 2.0, 15.0
+# `MINIMAX_H3_MAX_REFERENCE_IMAGES`, hardcoded so the UI renders before `diffusers` is importable. The slots are all
+# built up front and revealed one at a time, because a demo asking for two subjects should not open with nine boxes.
+MAX_IMAGE_SLOTS, OPEN_IMAGE_SLOTS = 9, 2
 
 
 def snap_frames(seconds: float) -> int:
@@ -121,6 +124,23 @@ def load_models() -> str | None:
         print(f"[ref2va] loading {[c.name for c in blocks.expected_components]} from {MODEL_REPO} ...", flush=True)
         pipe = blocks.init_pipeline(MODEL_REPO, components_manager=manager, collection="h3")
         pipe.load_components(dtype=torch.bfloat16, token=token)
+
+        # Pin the two autoencoders to torch SDPA *before* the transformer takes cuDNN, and in that order.
+        #
+        # `set_attention_backend` does two things: it stamps the backend onto every attention processor of the model
+        # it is called on, and it sets the registry's **global** active backend, which every processor that was not
+        # stamped then falls through to. Both VAEs carry `AttentionModuleMixin` attention with `_attention_backend =
+        # None`, so stamping only the transformer leaves them inheriting cuDNN — and they are float32, for which
+        # cuDNN has no kernel:
+        #
+        #     RuntimeError: No available kernel. Aborting execution.    # audio_vae pre_block, is_causal=True
+        #
+        # It is `ref2va` that exposes this. The keyframe half only ever *decodes* audio, and the audio VAE's
+        # attention is on its encoder side, so nothing reached it until a reference brought a soundtrack along.
+        # Stamping the VAEs first leaves both explicitly on `native`; the transformer then stamps itself and takes
+        # the global with it, which no longer matters to anyone.
+        pipe.vae.set_attention_backend("native")
+        pipe.audio_vae.set_attention_backend("native")
         pipe.transformer_ref.set_attention_backend(ATTENTION)
 
         # Still startup, still free: an AoTI package carries no weights and opens its compiled archive lazily inside
@@ -195,18 +215,19 @@ def probe(path: str) -> tuple[float | None, float | None]:
     return video, audio
 
 
-def collect(video_path, first_image_path, second_image_path, audio_path) -> list[tuple[str, str]]:
+def collect(image_paths, audio_path, video_path) -> list[tuple[str, str]]:
     """The `(kind, path)` references of a request, **in the order the model reads them**.
 
     That order is semantic rather than cosmetic: it numbers the labels of MiniMax-H3's prompt presentation and it
     advances the shared audio/video rotary clock, so the same references in a different order are a different
-    request. Video first, then images, then a standalone audio clip — which is also the order the packed sequence
-    lays them out in.
+    request. Images first, then a standalone audio clip, then the video — the order the tabs are laid out in, so
+    what the UI shows is what the model is handed.
     """
-    ordered = [("video", video_path)] if video_path else []
-    ordered += [("image", path) for path in (first_image_path, second_image_path) if path]
+    ordered = [("image", path) for path in image_paths if path]
     if audio_path:
         ordered.append(("audio", audio_path))
+    if video_path:
+        ordered.append(("video", video_path))
     return ordered
 
 
@@ -222,10 +243,13 @@ def audio_bearing(references: list[tuple[str, str]]) -> list[tuple[str, float]]:
     return carried
 
 
-def duration_controls(video_path, first_image_path, second_image_path, audio_path, match: bool):
-    """Show the duration slider unless a single soundtrack can set it, which is when MiniMax-H3 lets it be left out."""
+def duration_controls(audio_path, video_path, match: bool):
+    """Show the duration slider unless a single soundtrack can set it, which is when MiniMax-H3 lets it be left out.
+
+    Only the audio and video slots matter here: an image reference never carries a waveform.
+    """
     try:
-        carried = audio_bearing(collect(video_path, first_image_path, second_image_path, audio_path))
+        carried = audio_bearing(collect([], audio_path, video_path))
     except Exception:
         carried = []
     # Exactly one soundtrack, and one long enough to be a duration MiniMax-H3 generates. Anything else and the
@@ -312,10 +336,17 @@ def _generate(prompt_embeds, text_token_tags, references, height, width, num_fra
 
 def generate(
     prompt,
-    video_path=None,
-    first_image_path=None,
-    second_image_path=None,
+    image_1=None,
+    image_2=None,
+    image_3=None,
+    image_4=None,
+    image_5=None,
+    image_6=None,
+    image_7=None,
+    image_8=None,
+    image_9=None,
     audio_path=None,
+    video_path=None,
     canvas=DEFAULT_CANVAS,
     match=True,
     duration=5,
@@ -330,7 +361,8 @@ def generate(
 
     from diffusers.utils import encode_video
 
-    references = collect(video_path, first_image_path, second_image_path, audio_path)
+    images = [image_1, image_2, image_3, image_4, image_5, image_6, image_7, image_8, image_9]
+    references = collect(images, audio_path, video_path)
     check(prompt, references)
 
     # `0` is "leave it to the references" over the wire, which MiniMax-H3 accepts when exactly one of them carries a
@@ -382,8 +414,8 @@ INTRO = """# MiniMax-H3
 </div>
 
 **MiniMax-H3** is a 33B parameter state of the art video generation model that produces video and a
-fully synchronized soundtrack (ambience, foley, speech). Bring your own subject, camera move or voice as a
-reference — they are read **in tab order**, video then images then audio.
+fully synchronized soundtrack (ambience, foley, speech). Bring your own subject, voice or camera move as a
+reference — they are read **in tab order**, images then audio then video.
 """
 
 CSS = """
@@ -391,7 +423,7 @@ CSS = """
 .dark .gradio-container { color: var(--body-text-color); }
 """
 
-with gr.Blocks(title="MiniMax-H3 references") as demo:
+with gr.Blocks(title="MiniMax-H3 Reference") as demo:
     gr.Markdown(INTRO)
 
     with gr.Row():
@@ -404,14 +436,24 @@ with gr.Blocks(title="MiniMax-H3 references") as demo:
             # One tab per modality, in the order the model reads them. A reference left in a tab that is not the
             # open one is still part of the request — the tabs lay the slots out, they do not choose between them.
             with gr.Tabs():
-                with gr.Tab("Video"):
-                    video = gr.Video(label="Motion & camera, 2–15 s. Its soundtrack comes along.")
                 with gr.Tab("Images"):
+                    # One `gr.Row`, so gradio splits the width evenly and wraps once the slots hit `min_width`
+                    # rather than leaving a hole where a hidden slot used to be.
                     with gr.Row():
-                        first_image = gr.Image(label="Subject, style or scene", type="filepath")
-                        second_image = gr.Image(label="Subject, style or scene", type="filepath")
+                        images = [
+                            gr.Image(
+                                label="Subject, style or scene",
+                                type="filepath",
+                                min_width=180,
+                                visible=index < OPEN_IMAGE_SLOTS,
+                            )
+                            for index in range(MAX_IMAGE_SLOTS)
+                        ]
+                    add_image = gr.Button("+ Add another image", size="sm", variant="secondary")
                 with gr.Tab("Audio"):
                     audio = gr.Audio(label="A voice or a piece of music", type="filepath")
+                with gr.Tab("Video"):
+                    video = gr.Video(label="Motion & camera, 2–15 s. Its soundtrack comes along.")
             run = gr.Button("Generate", variant="primary")
             with gr.Accordion("Advanced options", open=False):
                 canvas = gr.Dropdown(label="Canvas", choices=list(CANVASES), value=DEFAULT_CANVAS)
@@ -425,15 +467,26 @@ with gr.Blocks(title="MiniMax-H3 references") as demo:
         with gr.Column():
             result = gr.Video(label="Video + soundtrack")
 
-    slots = [video, first_image, second_image, audio]
-    for control in [*slots, match]:
+    open_slots = gr.State(OPEN_IMAGE_SLOTS)
+
+    def reveal_image_slot(open_count):
+        open_count = min(open_count + 1, MAX_IMAGE_SLOTS)
+        return [
+            open_count,
+            *[gr.update(visible=index < open_count) for index in range(MAX_IMAGE_SLOTS)],
+            gr.update(visible=open_count < MAX_IMAGE_SLOTS),
+        ]
+
+    add_image.click(reveal_image_slot, open_slots, [open_slots, *images, add_image], api_name=False)
+
+    for control in (audio, video, match):
         control.change(
-            duration_controls, [*slots, match], [match, duration], show_progress="hidden", api_name=False
+            duration_controls, [audio, video, match], [match, duration], show_progress="hidden", api_name=False
         )
 
     run.click(
         generate,
-        [prompt, *slots, canvas, match, duration, steps, seed],
+        [prompt, *images, audio, video, canvas, match, duration, steps, seed],
         result,
         api_name="generate",
     )