Update app.py

app.py

@@ -91,6 +91,9 @@ SCHEDULES = {
 DEFAULT_SCHEDULE = "linear_quadratic · PlagueKind"
 DEFAULT_VIDEO_SHIFT = 12.0
 DEFAULT_AUDIO_SHIFT = 3.0
+DEFAULT_SHARPEN = 0.3
+INTERPOLATION = {"off · 24 fps": 1, "2x · 48 fps (PlagueKind)": 2, "4x · 96 fps": 4}
+DEFAULT_INTERPOLATION = "2x · 48 fps (PlagueKind)"
 
 FPS, FRAMES_PER_CHUNK, LATENTS_PER_CHUNK = 24, 17, 5
 # It is the *snapped* frame count the ceiling has to hold for: 15 s is 360 frames, which rounds up to 362, i.e.
@@ -113,7 +116,12 @@ STEP_LINEAR, STEP_QUADRATIC, SAFETY = 1.1745e-4, 3.8396e-9, 1.3
 PLACEMENT_ALLOWANCE = int(os.environ.get("H3_PLACEMENT_ALLOWANCE", "8"))
 AUDIO_LATENTS_PER_SECOND, AUDIO_CHANNELS = 40, 2
 REFERENCE_IMAGE_SHORT_EDGE, CANVAS_MULTIPLE = 2048, 32
-DECODE_BASE, DECODE_PER_DEFAULT_CANVAS, DEFAULT_CANVAS_PIXELS = 15, 25, 960 * 544 * 124
+DDECODE_BASE, DECODE_PER_DEFAULT_CANVAS, DEFAULT_CANVAS_PIXELS = 15, 25, 960 * 544 * 124
+# The workflow's post chain. RCAS is a handful of elementwise passes over the clip; FILM is per *emitted*
+# intermediate frame; the h264 mux is per frame actually written. Ported from the fl2va Space's fitted constants
+# as a starting point, same caveat as `PLACEMENT_ALLOWANCE` above — worth checking booked-vs-actual here
+# specifically once this is testable.
+_POST_BASE, _FILM_PER_FRAME, _MUX_PER_FRAME = 2.0, 0.025, 0.02
 
 
 def snap_frames(seconds: float) -> int:
@@ -187,7 +195,8 @@ def reference_rows(references: list[tuple[str, str]], num_frames: int) -> int:
 
 
 def get_duration(
-    prompt_embeds, text_token_tags, references, height, width, num_frames, steps, seed, **_
+    prompt_embeds, text_token_tags, references, height, width, num_frames, steps, seed,
+    sampler, schedule, video_shift, audio_shift, sharpen, multiplier, **_
 ):
     """Seconds of GPU to reserve for one request. Takes the arguments of the `@spaces.GPU` function it decorates, and
     tolerates the `gr.Progress` `spaces` injects."""
@@ -199,7 +208,16 @@ def get_duration(
     # they are handed rather than with the step count.
     encode = 5 + reference_rows(references, num_frames) * 1e-3
     decode = DECODE_BASE + DECODE_PER_DEFAULT_CANVAS * (height * width * num_frames) / DEFAULT_CANVAS_PIXELS
-    total = PLACEMENT_ALLOWANCE + encode + denoise + decode + 10
+
+    multiplier = max(1, int(multiplier))
+    if multiplier > 1 and FILM is None:
+        multiplier = 1
+    pixel_ratio = (height * width) / (960 * 544)
+    out_frames = (num_frames - 1) * multiplier + 1 if multiplier > 1 else num_frames
+    film = (num_frames - 1) * (multiplier - 1) * _FILM_PER_FRAME * pixel_ratio
+    post = _POST_BASE + film + out_frames * _MUX_PER_FRAME * pixel_ratio
+
+    total = PLACEMENT_ALLOWANCE + encode + denoise + decode + post
     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
@@ -463,29 +481,62 @@ def encode_remote(prompt, references, canvas, num_frames, rewrite_prompt=False):
 
 
 @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.
-    References cross as paths and are decoded here; only the three generated outputs come back. A `@spaces.GPU`
-    argument crosses a process boundary by pickling, a 5 s 1344x768 reference video is 370 MB of expanded frames, and
-    the full `PipelineState` still holds the packed latents and the rotary grid on the card.
+def _generate(
+    prompt_embeds, text_token_tags, references, height, width, num_frames, steps, seed,
+    sampler, schedule, video_shift, audio_shift, sharpen, multiplier,
+):
+    """The only thing on GPU time: the two reference encoders, the packed-sequence denoise loop, the decoders, and
+    the RCAS + FILM post chain. References cross as paths and are decoded here; only the generated outputs come
+    back. A `@spaces.GPU` argument crosses a process boundary by pickling, a 5 s 1344x768 reference video is
+    370 MB of expanded frames, and the full `PipelineState` still holds the packed latents and the rotary grid on
+    the card.
     """
     import torch
+
+    global FILM
+
     if PLACEMENT == "lazy":
         PIPE.to("cuda")
     elif PLACEMENT == "pack":
         PIPE.vae.to("cuda")
         PIPE.audio_vae.to("cuda")
-    state = PIPE(
-        prompt_embeds=prompt_embeds.to("cuda"),
-        text_token_tags=text_token_tags,
-        references=build_references(references),
-        height=height,
-        width=width,
-        num_frames=num_frames,
-        num_inference_steps=int(steps),
-        generator=torch.Generator("cpu").manual_seed(int(seed)),
-    )
-    return state.get("videos")[0], state.get("audio")[0].cpu(), state.get("sampling_rate")
+
+    custom_schedule = schedule != "native"
+    # Any custom schedule hands `set_timesteps` a finished `steps + 1` sigma grid, so it runs `steps` forwards.
+    # The native grid counts its terminal zero as one of `num_inference_steps`, so it needs one more to match.
+    requested_steps = int(steps) if custom_schedule else int(steps) + 1
+
+    with pk.use_schedule(PIPE, int(steps), schedule, video_shift, audio_shift, sampler_name=sampler, seed=int(seed)):
+        state = PIPE(
+            prompt_embeds=prompt_embeds.to("cuda"),
+            text_token_tags=text_token_tags,
+            references=build_references(references),
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            num_inference_steps=requested_steps,
+            generator=torch.Generator("cpu").manual_seed(int(seed)),
+        )
+
+    video = state.get("videos")[0]  # (frames, 3, H, W), float in [0, 1], on the card
+    audio = state.get("audio")[0].cpu()
+    sampling_rate = state.get("sampling_rate")
+    del state
+    # The post chain runs on the allocator the denoise loop just left fragmented, and RCAS and FILM both want a
+    # few contiguous gigabytes.
+    torch.cuda.empty_cache()
+
+    video = pk.rcas(video, float(sharpen))
+    multiplier = max(1, int(multiplier))
+    if multiplier > 1:
+        if FILM is None:
+            multiplier = 1
+        else:
+            FILM = FILM.to("cuda")
+            video = pk.interpolate(FILM, video, multiplier)
+    fps = FPS * multiplier
+
+    return video, audio, sampling_rate, fps, multiplier
 
 
 def generate(
@@ -507,6 +558,12 @@ def generate(
     steps=28,
     seed=42,
     upsample=False,
+    sampler=DEFAULT_SAMPLER,
+    schedule=DEFAULT_SCHEDULE,
+    video_shift=DEFAULT_VIDEO_SHIFT,
+    audio_shift=DEFAULT_AUDIO_SHIFT,
+    sharpen=DEFAULT_SHARPEN,
+    interpolation=DEFAULT_INTERPOLATION,
     progress=gr.Progress(track_tqdm=True),
 ):
     """One request."""
@@ -525,6 +582,9 @@ def generate(
     # soundtrack. The conditioner resolves it either way and this Space pins whatever comes back.
     derivable = len(audio_bearing(references)) == 1
     requested = 0 if (match and derivable) else snap_frames(duration)
+    schedule_key = SCHEDULES.get(schedule, "linear_quadratic")
+    sampler_key = SAMPLERS.get(sampler, "euler")
+    multiplier = INTERPOLATION.get(interpolation, 2)
 
     progress(0.0, desc="Upsampling the prompt ..." if upsample else "Reading the prompt and references ...")
     conditioned = time.time()
@@ -548,19 +608,21 @@ def generate(
 
     progress(0.1, desc=f"Generating {num_frames / FPS:.1f} s at {width}x{height} ...")
     started = time.time()
-    frames, audio, sampling_rate = _generate(
-        prompt_embeds, text_token_tags, references, height, width, num_frames, steps, seed
+    video, audio, sampling_rate, fps, multiplier = _generate(
+        prompt_embeds, text_token_tags, references, height, width, num_frames, steps, seed,
+        sampler_key, schedule_key, float(video_shift), float(audio_shift), float(sharpen), multiplier,
     )
     generate_seconds = time.time() - started
 
     directory = os.path.join(tempfile.gettempdir(), "h3-outputs")
     os.makedirs(directory, exist_ok=True)
     path = os.path.join(directory, f"h3-ref2va-{int(time.time() * 1000)}.mp4")
-    encode_video(frames, fps=FPS, output_path=path, audio=audio, audio_sample_rate=sampling_rate)
+    encode_video(video, fps=fps, output_path=path, audio=audio, audio_sample_rate=sampling_rate)
 
     print(
         f"[ref2va] {[kind for kind, _ in references]} · `{width}x{height}`, {num_frames} frames "
-        f"({num_frames / FPS:.3f} s), {int(steps)} steps · conditioner {condition_seconds:.0f}s "
+        f"({num_frames / FPS:.3f} s), {int(steps)} steps of `{schedule_key}` · sampler `{sampler_key}` · "
+        f"shift {float(video_shift):.1f}/{float(audio_shift):.1f} · conditioner {condition_seconds:.0f}s "
         f"({plan['num_text_tokens']} tokens{', upsampled' if refined else ''}) · "
         f"denoise + decode {generate_seconds:.0f}s "
         f"({generate_seconds / int(steps):.1f} s/step) · seed {int(seed)}",
@@ -648,6 +710,20 @@ with gr.Blocks(title="MiniMax-H3 Reference Custom Lora") as demo:
                 audio_shift = gr.Slider(
                     label="Audio shift", minimum=0.5, maximum=20.0, step=0.5, value=DEFAULT_AUDIO_SHIFT
                 )
+                sharpen = gr.Slider(
+                    label="RCAS sharpening",
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.05,
+                    value=DEFAULT_SHARPEN,
+                    info="FidelityFX Robust Contrast Adaptive Sharpening. PlagueKind: 0.3 is very natural.",
+                )
+                interpolation = gr.Dropdown(
+                    label="FILM frame interpolation",
+                    choices=list(INTERPOLATION),
+                    value=DEFAULT_INTERPOLATION,
+                    info="MiniMax-H3 generates 24 fps; FILM synthesizes the frames in between.",
+                )
                 seed = gr.Number(label="Seed", value=42, precision=0)
 
         with gr.Column():
@@ -679,7 +755,8 @@ with gr.Blocks(title="MiniMax-H3 Reference Custom Lora") as demo:
     video.change(auto_canvas, video, canvas, show_progress="hidden", api_name=False)
 
     request = [
-        prompt, images[0], audio, video, canvas, *images[1:], match, duration, steps, seed, upsample
+        prompt, images[0], audio, video, canvas, *images[1:], match, duration, steps, seed, upsample,
+        sampler, schedule, video_shift, audio_shift, sharpen, interpolation,
     ]
     run.click(generate, request, [result, upsampled, upsampled_panel], api_name="generate")