Upload app.py

app.py

@@ -109,7 +109,6 @@ PIPE = None
 LOAD_ERROR: str | None = None
 LOADED_IN: float | None = None
 
-
 def snap_frames(seconds: float) -> int:
     """The frame count MiniMax-H3's video VAE can decode: the next ``17 * n + 5`` at 24 fps."""
     frames = max(1, round(float(seconds) * FPS))
@@ -117,14 +116,12 @@ def snap_frames(seconds: float) -> int:
         frames += 1
     return frames
 
-
 def lower_duration_floor(seconds: float = MIN_UI_DURATION) -> None:
     """Let the pipeline generate below its 5 s floor."""
     from diffusers.modular_pipelines.minimax_h3.modular_pipeline import MiniMaxH3ModularPipeline
 
     MiniMaxH3ModularPipeline.min_duration = property(lambda self: float(seconds))
 
-
 def shifted_schedule(shift: float) -> list[float]:
     """``BASE_SCHEDULE`` pushed through H3's exponential sigma shift ``s*σ / (1 + (s - 1)*σ)``.
 
@@ -142,7 +139,6 @@ def shifted_schedule(shift: float) -> list[float]:
     """
     return [shift * sigma / (1 + (shift - 1) * sigma) for sigma in BASE_SCHEDULE]
 
-
 def pin_distilled_schedule(scheduler) -> None:
     """Pin the shifted ``BASE_SCHEDULE`` onto a scheduler so *every* ``set_timesteps`` call lands on it.
 
@@ -173,7 +169,6 @@ def pin_distilled_schedule(scheduler) -> None:
     scheduler.set_timesteps = set_timesteps
     scheduler._flashgen_pinned = True
 
-
 # ── LoRA loading (weight folding) ─────────────────────────────────────────────
 #
 # This mirrors ``merge_lora_ckpt.py``, the merge script the LoRA repo itself ships, rather than a
@@ -201,7 +196,6 @@ def pin_distilled_schedule(scheduler) -> None:
 #     contiguous thirds directly — the naive reading — scatters every head's q/k/v across all three
 #     projections, which turns the adapter into structured noise on all 52 attention blocks.
 
-
 def _reorder_interleaved_qkv(weight, num_heads: int, head_dim: int):
     """De-interleave per-head fused-QKV rows into ``[q_all; k_all; v_all]``.
 
@@ -215,7 +209,6 @@ def _reorder_interleaved_qkv(weight, num_heads: int, head_dim: int):
         [part.reshape(num_heads * head_dim, *weight.shape[1:]) for part in (query, key, value)], dim=0
     )
 
-
 def _lora_target_name(source_name: str) -> str:
     """Map an original-checkpoint base name (without ``.lora_A/B.*``) to the diffusers parameter path."""
     if source_name.startswith("token_refiner.blocks."):
@@ -226,7 +219,6 @@ def _lora_target_name(source_name: str) -> str:
         target = source_name
     return target.replace("final_layer.adaln_proj.linear", "norm_out.linear")
 
-
 def _lora_targets(name: str, b_weight, num_heads: int, head_dim: int):
     """Yield ``(diffusers_param_key, row_transformed_B)`` for one original LoRA base name."""
     target = _lora_target_name(name)
@@ -250,7 +242,6 @@ def _lora_targets(name: str, b_weight, num_heads: int, head_dim: int):
         # adaln_proj.linear (block-level) and norm_out.linear: identical row layout.
         yield target + ".weight", b_weight
 
-
 def load_and_apply_lora(transformer) -> str:
     """Download the FlashGen LoRA, remap keys, and fold ``scale * (B @ A)`` into the base weights."""
     import torch
@@ -317,7 +308,6 @@ def load_and_apply_lora(transformer) -> str:
         f"-> {folded} weight deltas · {LORA_FILE}"
     )
 
-
 # ── Additional LoRAs ─────────────────────────────────────────────────────────
 #
 # These are folded after FlashGen, so the original 4-step FlashGen schedule and
@@ -422,7 +412,6 @@ def _apply_extra_lora(
         f"{folded} weight deltas · scale {float(scale):g}"
     )
 
-
 def load_and_apply_extra_loras(transformer) -> list[str]:
     return [
         _apply_extra_lora(
@@ -441,10 +430,8 @@ def load_and_apply_extra_loras(transformer) -> list[str]:
         ),
     ]
 
-
 # ── Model loading ───────────────────────────────────────────────────────────
 
-
 def load_models() -> str | None:
     """Load the denoising half at startup: transformer + VAEs + schedulers, then fold the LoRA in."""
     global PIPE, LOAD_ERROR, LOADED_IN
@@ -504,17 +491,14 @@ def load_models() -> str | None:
         LOAD_ERROR = f"**Loading failed** after {time.time() - started:.0f}s: `{type(error).__name__}: {error}`"
     return LOAD_ERROR
 
-
 # ── Remote conditioning ──────────────────────────────────────────────────────
 
-
 @cache
 def conditioner():
     from gradio_client import Client
 
     return Client(CONDITIONER_SPACE)
 
-
 def encode_remote(prompt, image_path, last_image_path, canvas, num_frames, ip_token=None, rewrite_prompt=False):
     """Call the conditioner Space to get ``prompt_embeds`` + ``text_token_tags`` + resolved geometry.
 
@@ -550,7 +534,6 @@ def encode_remote(prompt, image_path, last_image_path, canvas, num_frames, ip_to
         metadata = handle.metadata()
         return handle.get_tensor("prompt_embeds"), handle.get_tensor("text_token_tags"), metadata, plan
 
-
 def conditioner_client(ip_token):
     if not ip_token:
         return conditioner()
@@ -558,10 +541,8 @@ def conditioner_client(ip_token):
 
     return Client(CONDITIONER_SPACE, headers={"x-ip-token": ip_token})
 
-
 # ── Geometry helpers ────────────────────────────────────────────────────────
 
-
 def _snap_canvas(aspect: float, current_canvas: str) -> str:
     """The supported canvas whose aspect ratio is closest to ``aspect``.
 
@@ -582,7 +563,6 @@ def _snap_canvas(aspect: float, current_canvas: str) -> str:
         return current_canvas
     return fastest[ratio][0]
 
-
 def _cover_crop(image_path: str, canvas_label: str) -> str:
     """Center cover-crop the file at ``image_path`` to ``canvas_label``'s aspect ratio, in place.
 
@@ -607,7 +587,6 @@ def _cover_crop(image_path: str, canvas_label: str) -> str:
     img.save(image_path)
     return image_path
 
-
 def _fit_keyframe(image_path, current_canvas):
     """Snap the canvas to a keyframe's aspect ratio, then cover-crop the keyframe to that canvas."""
     from PIL import Image as _Image
@@ -617,7 +596,6 @@ def _fit_keyframe(image_path, current_canvas):
     label = _snap_canvas(aspect, current_canvas)
     return _cover_crop(image_path, label), label
 
-
 def _fit_keyframe_ui(image_path, current_canvas):
     """``image.upload`` handler — the reference Space's behaviour, ported.
 
@@ -631,7 +609,6 @@ def _fit_keyframe_ui(image_path, current_canvas):
     cropped, label = _fit_keyframe(path, current_canvas)
     return gr.update(value=cropped), gr.update(value=label)
 
-
 def _fit_last_keyframe_ui(last_image_path, image_path, current_canvas):
     """``last_image.upload`` handler. The first frame owns the canvas when there is one, so a last
     frame uploaded on top of it is only cover-cropped to the canvas already chosen."""
@@ -642,20 +619,17 @@ def _fit_last_keyframe_ui(last_image_path, image_path, current_canvas):
         return gr.update(value=_cover_crop(path, current_canvas)), gr.update()
     return _fit_keyframe_ui(path, current_canvas)
 
-
 def _as_path(value):
     if isinstance(value, dict):
         value = value.get("path") or (value.get("url") or "").removeprefix("/gradio_api/file=")
     return value or None
 
-
 # ── GPU duration estimation ─────────────────────────────────────────────────
 
 _DUR_B, _DUR_C = 1.1745e-4, 3.8396e-9
 _DECODE_BASE, _DECODE_PER_DEFAULT_CANVAS, _DEFAULT_CANVAS_PIXELS = 15, 15, 960 * 544 * 124
 _PLACEMENT_ALLOWANCE, _PAD = 12, 10
 
-
 def get_duration(prompt_embeds, text_token_tags, image, last_image, height, width, num_frames, seed, *a, **k):
     """Estimate GPU duration for the ``_generate`` call. Steps is fixed at NUM_STEPS (4)."""
     height, width, num_frames = int(height), int(width), int(num_frames)
@@ -667,10 +641,8 @@ def get_duration(prompt_embeds, text_token_tags, image, last_image, height, widt
     decode = _DECODE_BASE + _DECODE_PER_DEFAULT_CANVAS * (height * width * num_frames) / _DEFAULT_CANVAS_PIXELS
     return max(60, int(denoise + decode) + _PLACEMENT_ALLOWANCE + _PAD)
 
-
 # ── Inference ───────────────────────────────────────────────────────────────
 
-
 @spaces.GPU(duration=get_duration, size=GPU_SIZE)
 def _generate(prompt_embeds, text_token_tags, image, last_image, height, width, num_frames, seed):
     """Denoise + decode on GPU time. The LoRA is already folded at startup; 4 steps with custom sigmas."""
@@ -713,14 +685,12 @@ def _generate(prompt_embeds, text_token_tags, image, last_image, height, width,
         sigmas_used,
     )
 
-
 def _caller_ip_token() -> str | None:
     from gradio.context import LocalContext
 
     request = LocalContext.request.get()
     return request.headers.get("x-ip-token") if request is not None else None
 
-
 def generate(
     prompt,
     image_path=None,
@@ -820,10 +790,8 @@ def generate(
         )
     return path, report, enhanced, gr.update(visible=bool(enhanced))
 
-
 # ── UI ──────────────────────────────────────────────────────────────────────
 
-
 INTRO = """# MiniMax-H3 4-Step FlashGen LoRA
 
 <div align="center">