Dynamic GPU duration from the packed sequence, cached examples, torchaudio note

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/motion.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/voice.wav filter=lfs diff=lfs merge=lfs -text

README.md

@@ -151,3 +151,13 @@ up to 362, i.e. 15.083 s, and is refused.
 MiniMax-H3 is modular-only and not in a released `diffusers`, so the integration branch's `src/diffusers` tree is
 vendored here as a top-level `diffusers/` package; the working directory comes first on `sys.path`, so there is no
 install step. `requirements.txt` only carries what that tree imports.
+
+Two of those are `ref2va`-only and easy to miss. PyAV decodes a reference video or audio file as the reference is
+built, and **`torchaudio`** resamples a soundtrack that is not already at the audio VAE's 32 kHz — a 32 kHz
+reference skips the resample entirely, so the dependency only shows up once someone brings audio at another rate:
+
+```
+ImportError: Resampling a MiniMax-H3 reference soundtrack from 24000 Hz to 32000 Hz needs `torchaudio`.
+```
+
+The conditioner Space needs it as well: its `setup` step prepares the very same waveforms this one does.

app.py

@@ -36,8 +36,12 @@ PLACEMENT = os.environ.get("H3_PLACEMENT", "lazy").lower()
 # cuDNN's fused attention is 10-20% faster than the SDPA default on this pool and needs nothing installed.
 # flash-attention 3 is sm90-only and this card is sm120 (the `zero-a10g` flavour name is legacy).
 ATTENTION = os.environ.get("H3_ATTENTION", "_native_cudnn").lower()
-GPU_DURATION = int(os.environ.get("H3_GPU_DURATION", "900"))
 GPU_SIZE = os.environ.get("H3_GPU_SIZE", "xlarge")
+# Bounds on what `get_duration` may ask the pool to reserve. A request that runs out of GPU time is a total loss, so
+# the estimate is deliberately generous — but a *fixed* 900 s ceiling for every request is what makes the account hit
+# "too many ZeroGPU credits allocated to running tasks", because the pool reserves the number it is given.
+MIN_GPU_DURATION = int(os.environ.get("H3_GPU_DURATION_MIN", "120"))
+MAX_GPU_DURATION = int(os.environ.get("H3_GPU_DURATION_MAX", "1500"))
 
 # MiniMax-H3's own canvases, i.e. `resolve_canvas_size` from `diffusers.modular_pipelines.minimax_h3.packing`
 # evaluated for the six released aspect ratios. Hardcoded so the UI renders before `diffusers` is importable.
@@ -70,6 +74,28 @@ MIN_REFERENCE_VIDEO, MAX_REFERENCE_VIDEO = 2.0, 15.0
 # 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
 
+# --- What a request costs, for `get_duration` -------------------------------------------------------------------
+#
+# MiniMax-H3 attends over one packed sequence, so the cost of a step is a function of its length `S` alone. Fitted
+# on the `t2va` half (`minimax-h3`, AoTI on, same pool and same silicon) over two canvases at 124 frames:
+#
+#     544x544, S = 10693 -> 2.4 s/step
+#     960x544, S = 18870 -> 4.4 s/step
+#
+# through `s = LINEAR * S + QUADRATIC * S**2` — linear for the matmuls, quadratic for the attention. Checked against
+# two live `ref2va` requests on this Space, which is the regime the reference rows actually put it in:
+#
+#     one 1344x768 image reference,             S ~= 33232 -> 8.3 predicted, ~8.5 measured
+#     that image plus a 2.5 s video reference,  S ~= 54039 -> 14.6 predicted, ~16.1 measured
+#
+# so the fit holds to about 10% three times past the canvas it was taken from, and `SAFETY` covers the rest.
+STEP_LINEAR, STEP_QUADRATIC, SAFETY = 2.13e-4, 1.069e-9, 1.3
+# The lazy 72.16 GiB `PIPE.to("cuda")` a cold worker pays inside its first GPU call. Measured at ~45 s; every request
+# has to carry it, because nothing on this side knows whether the worker it lands on is cold.
+PLACEMENT_ALLOWANCE = int(os.environ.get("H3_PLACEMENT_ALLOWANCE", "90"))
+AUDIO_LATENTS_PER_SECOND, AUDIO_CHANNELS = 40, 2
+REFERENCE_IMAGE_SHORT_EDGE, CANVAS_MULTIPLE = 2048, 32
+
 
 def snap_frames(seconds: float) -> int:
     """The frame count MiniMax-H3's video VAE can decode: the next `17 * n + 5` at 24 fps."""
@@ -79,6 +105,86 @@ def snap_frames(seconds: float) -> int:
     return frames
 
 
+def video_latent_frames(num_frames: int) -> int:
+    """`17 * n + 5` frames become `5 * n + 2` video latents."""
+    return 5 * ((num_frames - LATENTS_PER_CHUNK) // FRAMES_PER_CHUNK) + 2
+
+
+def target_rows(height: int, width: int, num_frames: int) -> int:
+    """The generated rows of the packed sequence: video patched `(1, 2, 2)`, plus two audio rows per latent."""
+    video = video_latent_frames(num_frames) * (height // CANVAS_MULTIPLE) * (width // CANVAS_MULTIPLE)
+    return video + round(num_frames / FPS * AUDIO_LATENTS_PER_SECOND) * AUDIO_CHANNELS
+
+
+def reference_rows(references: list[tuple[str, str]], num_frames: int) -> int:
+    """The rows the reference blocks add, from metadata alone — no decode.
+
+    This mirrors what `MiniMaxH3Ref2VASetupStep` and the reference encoder will do, closely enough to size a GPU
+    reservation with. An image is resized to a 2048 pixel short edge and encoded as a single frame; a video is put on
+    the canvas *its own* aspect ratio resolves to, truncated to the generated frame count and then snapped **down**
+    to a `17 * n + 5` the VAE encodes without padding; a soundtrack contributes two rows per 1/40 s.
+    """
+    from PIL import Image
+
+    from diffusers.modular_pipelines.minimax_h3.packing import resolve_canvas_size
+
+    rows = 0
+    for kind, path in references:
+        if kind == "image":
+            width, height = Image.open(path).size
+            scale = REFERENCE_IMAGE_SHORT_EDGE / min(width, height)
+            resolved = [
+                max(CANVAS_MULTIPLE, round(edge * scale / CANVAS_MULTIPLE) * CANVAS_MULTIPLE)
+                for edge in (height, width)
+            ]
+            rows += (resolved[0] // CANVAS_MULTIPLE) * (resolved[1] // CANVAS_MULTIPLE)
+            continue
+
+        video_seconds, audio_seconds = probe(path)
+        if kind == "video" and video_seconds is not None:
+            import av
+
+            with av.open(path) as container:
+                stream = container.streams.video[0]
+                source_height, source_width = stream.height, stream.width
+            canvas_height, canvas_width = resolve_canvas_size(source_width, source_height)
+            # Resampled onto 24 fps and capped at the generated length, then snapped down to `17 * n + 5`.
+            frames = min(round(video_seconds * FPS), num_frames)
+            snapped = max(1, (frames - LATENTS_PER_CHUNK) // FRAMES_PER_CHUNK) * FRAMES_PER_CHUNK + LATENTS_PER_CHUNK
+            rows += (
+                video_latent_frames(snapped)
+                * (canvas_height // CANVAS_MULTIPLE)
+                * (canvas_width // CANVAS_MULTIPLE)
+            )
+        if audio_seconds is not None:
+            seconds = min(audio_seconds, num_frames / FPS)
+            rows += round(seconds * AUDIO_LATENTS_PER_SECOND) * AUDIO_CHANNELS
+    return rows
+
+
+def get_duration(prompt_embeds, text_token_tags, references, height, width, num_frames, steps, seed, **_):
+    """Seconds of GPU to reserve for one request, from the packed sequence it is about to build.
+
+    Takes the arguments of the `@spaces.GPU` function it decorates — and tolerates the `gr.Progress` `spaces`
+    injects — so it can price the request rather than reserve a flat ceiling for all of them.
+
+    The text rows are exact: `text_token_tags` is the conditioner's own answer, already on this side. The reference
+    and target rows come from `reference_rows` and `target_rows`.
+    """
+    sequence = int(text_token_tags.shape[0]) + reference_rows(references, num_frames) + target_rows(
+        height, width, num_frames
+    )
+    denoise = int(steps) * (STEP_LINEAR * sequence + STEP_QUADRATIC * sequence**2) * SAFETY
+    # The two reference encoders, ahead of the loop, and the two decoders plus the mux after it. Both scale with what
+    # they are handed rather than with the step count.
+    encode = 5 + reference_rows(references, num_frames) * 1e-3
+    decode = 15 + 25 * (height * width * num_frames) / (960 * 544 * 124)
+    total = PLACEMENT_ALLOWANCE + encode + denoise + decode + 10
+    duration = max(MIN_GPU_DURATION, min(MAX_GPU_DURATION, int(total)))
+    print(f"[ref2va] S={sequence} -> reserving {duration}s ({denoise:.0f}s of denoise at {steps} steps)", flush=True)
+    return duration
+
+
 PIPE = None
 MANAGER = None
 LOAD_ERROR: str | None = None
@@ -301,7 +407,7 @@ def encode_remote(prompt, references, canvas, num_frames):
         return handle.get_tensor("prompt_embeds"), handle.get_tensor("text_token_tags"), handle.metadata(), plan
 
 
-@spaces.GPU(duration=GPU_DURATION, size=GPU_SIZE)
+@spaces.GPU(duration=get_duration, size=GPU_SIZE)
 def _generate(prompt_embeds, text_token_tags, references, height, width, num_frames, steps, seed):
     """The only thing on GPU time: the two reference encoders, the packed-sequence denoise loop and the decoders.
 
@@ -335,8 +441,13 @@ def _generate(prompt_embeds, text_token_tags, references, height, width, num_fra
 
 
 def generate(
+    # The first four are the columns `gr.Examples` varies, and they lead the signature for that reason: an example
+    # row is applied to `inputs` positionally, so the exampled components have to be the leading parameters. Every
+    # parameter has a default, which is what lets a four-column row call this at all.
     prompt,
     image_1=None,
+    audio_path=None,
+    video_path=None,
     image_2=None,
     image_3=None,
     image_4=None,
@@ -345,8 +456,6 @@ def generate(
     image_7=None,
     image_8=None,
     image_9=None,
-    audio_path=None,
-    video_path=None,
     canvas=DEFAULT_CANVAS,
     match=True,
     duration=5,
@@ -493,13 +602,39 @@ with gr.Blocks(title="MiniMax-H3 Reference") as demo:
             duration_controls, [audio, video, match], [match, duration], show_progress="hidden", api_name=False
         )
 
-    run.click(
-        generate,
-        [prompt, *images, audio, video, canvas, match, duration, steps, seed],
-        result,
-        api_name="generate",
+    # Same order as `generate`'s signature: the exampled four first, then the remaining image slots.
+    request = [prompt, images[0], audio, video, *images[1:], canvas, match, duration, steps, seed]
+
+    gr.Examples(
+        examples=[
+            [
+                "The character walks through a neon-lit street in the rain, humming to themselves",
+                "examples/subject.png",
+                None,
+                None,
+            ],
+            [
+                "The character speaks to camera in a quiet room, lips matching every word",
+                "examples/subject.png",
+                "examples/voice.wav",
+                None,
+            ],
+            [
+                "The character moves with the same camera push, down a rainy alley at night",
+                "examples/subject.png",
+                None,
+                "examples/motion.mp4",
+            ],
+        ],
+        inputs=[prompt, images[0], audio, video],
+        outputs=result,
+        fn=generate,
+        cache_examples=True,
+        cache_mode="lazy",
     )
 
+    run.click(generate, request, result, api_name="generate")
+
 
 if __name__ == "__main__":
     demo.launch(show_error=True, theme=gr.themes.Citrus(), css=CSS)

examples/motion.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cb7e2f675639e86235afd8ead114b4bfbea19ee69719225b57c26fc90ff72363
+size 762637

examples/voice.wav

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d041cbaa771ee3fa0295d09e79a08cd306579850d4a698ba2593cc9fa693647a
+size 768044