Upload app.py

app.py

@@ -70,6 +70,10 @@ MIN_REFERENCE_VIDEO, MAX_REFERENCE_VIDEO = 2.0, 15.0
 # for two subjects should not open with nine boxes.
 MAX_IMAGE_SLOTS, OPEN_IMAGE_SLOTS = 9, 2
 
+# How many LoRA slots the UI offers, and the range each strength slider covers.
+LORA_SLOTS = 3
+LORA_MIN_SCALE, LORA_MAX_SCALE = -2.0, 2.0
+
 # Seconds of GPU one request needs, from the packed sequence it is about to denoise: linear in the rows for the
 # matmuls, quadratic for the attention, against the AoTI block package this Space runs.
 STEP_LINEAR, STEP_QUADRATIC, SAFETY = 1.1745e-4, 3.8396e-9, 1.3
@@ -79,6 +83,8 @@ 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
 DECODE_BASE, DECODE_PER_DEFAULT_CANVAS, DEFAULT_CANVAS_PIXELS = 15, 25, 960 * 544 * 124
+# Reading one adapter off local disk and injecting it across the 33B transformer's linear layers.
+LORA_ALLOWANCE = 12
 
 
 def snap_frames(seconds: float) -> int:
@@ -151,7 +157,9 @@ def reference_rows(references: list[tuple[str, str]], num_frames: int) -> int:
     return rows
 
 
-def get_duration(prompt_embeds, text_token_tags, references, height, width, num_frames, steps, seed, **_):
+def get_duration(
+    prompt_embeds, text_token_tags, references, height, width, num_frames, steps, seed, loras=(), **_
+):
     """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."""
     sequence = int(text_token_tags.shape[0]) + reference_rows(references, num_frames) + target_rows(
@@ -162,7 +170,7 @@ def get_duration(prompt_embeds, text_token_tags, references, height, width, num_
     # 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
+    total = PLACEMENT_ALLOWANCE + encode + denoise + decode + 10 + LORA_ALLOWANCE * len(loras or ())
     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
@@ -254,6 +262,126 @@ def _arm_decode_hooks(pipe):
             setattr(module, method, armed)
 
 
+# ----------------------------------------------------------------------------------------------------------------
+# LoRA
+# ----------------------------------------------------------------------------------------------------------------
+# There is no `MiniMaxH3LoraLoaderMixin` in the diffusers integration, so adapters are attached at the *model* level,
+# through the `PeftAdapterMixin` the transformer carries. That is the whole API this needs: `load_lora_adapter` for
+# each file and one `set_adapters` call to give them their strengths. Here the model is `transformer_ref`, so the
+# adapters have to be trained against the `transformer_ref/` partition — a `transformer/` adapter is a different
+# partition and will not match.
+
+
+def _hub_url_parts(url: str) -> tuple[str, str]:
+    """Split a huggingface.co `blob`/`resolve` URL into its repo id and the file path inside it."""
+    from urllib.parse import unquote, urlparse
+
+    parts = unquote(urlparse(url).path).strip("/").split("/")
+    if len(parts) < 5 or parts[2] not in ("resolve", "blob"):
+        raise gr.Error(f"Не разпознавам този адрес като файл в Hugging Face: `{url}`")
+    return "/".join(parts[:2]), "/".join(parts[4:])
+
+
+def resolve_lora(reference: str) -> str:
+    """Turn what the user typed into a local `.safetensors` path.
+
+    Accepts a local path, a huggingface.co file URL, `owner/repo/path/to/file.safetensors`, or a bare `owner/repo`
+    whose single `.safetensors` is then picked for them. Runs outside the GPU call, so the download costs no GPU time.
+    """
+    from huggingface_hub import hf_hub_download, list_repo_files
+
+    reference = (reference or "").strip()
+    if not reference:
+        return ""
+    if os.path.exists(reference):
+        return reference
+    if reference.startswith(("http://", "https://")):
+        repo_id, filename = _hub_url_parts(reference)
+        return hf_hub_download(repo_id, filename)
+
+    parts = [part for part in reference.split("/") if part]
+    if len(parts) > 2 and parts[-1].endswith(".safetensors"):
+        return hf_hub_download("/".join(parts[:2]), "/".join(parts[2:]))
+    if len(parts) != 2:
+        raise gr.Error(
+            f"`{reference}` не е нито съществуващ файл, нито `автор/хранилище`, нито адрес към Hugging Face."
+        )
+
+    candidates = [name for name in list_repo_files(reference) if name.endswith(".safetensors")]
+    if not candidates:
+        raise gr.Error(f"В `{reference}` няма `.safetensors` файл.")
+    if len(candidates) > 1:
+        preferred = [name for name in candidates if "lora" in name.lower()]
+        if len(preferred) != 1:
+            listed = ", ".join(f"`{name}`" for name in sorted(candidates)[:8])
+            raise gr.Error(f"`{reference}` съдържа няколко файла. Напиши `{reference}/име.safetensors`. Има: {listed}")
+        candidates = preferred
+    return hf_hub_download(reference, candidates[0])
+
+
+def _lora_prefix(state_dict) -> str | None:
+    """The prefix `load_lora_adapter` has to strip before the keys match the transformer's own module names."""
+    key = next(iter(state_dict))
+    for prefix in ("model.diffusion_model", "diffusion_model", "transformer_ref", "transformer"):
+        if key.startswith(f"{prefix}."):
+            return prefix
+    return None
+
+
+def apply_loras(transformer, loras) -> list[str]:
+    """Attach `loras` (local path, strength) to `transformer` and give each its strength, replacing whatever was on it.
+
+    Every adapter already on the model is removed first, so a request is never affected by the one before it — which
+    matters when a worker is reused rather than forked fresh.
+    """
+    import torch
+
+    from safetensors.torch import load_file
+
+    for name in list(getattr(transformer, "peft_config", None) or {}):
+        transformer.delete_adapters(name)
+
+    names, scales = [], []
+    for index, (path, scale) in enumerate(loras):
+        state_dict = load_file(path)
+        name = f"lora{index}"
+        transformer.load_lora_adapter(state_dict, adapter_name=name, prefix=_lora_prefix(state_dict))
+        names.append(name)
+        scales.append(float(scale))
+
+    if not names:
+        return []
+
+    # PEFT builds the new layers on its own default device/dtype; the base weights are the truth here, under either
+    # placement mode (`offload` keeps them on the host and moves whole modules by hook).
+    base = next(param for key, param in transformer.named_parameters() if ".lora_" not in key)
+    with torch.no_grad():
+        for key, param in transformer.named_parameters():
+            if ".lora_" in key and (param.device != base.device or param.dtype != base.dtype):
+                param.data = param.data.to(device=base.device, dtype=base.dtype)
+
+    transformer.set_adapters(names, scales)
+    return names
+
+
+def collect_loras(lora_fields, progress) -> tuple[list[tuple[str, float]], list[str]]:
+    """Resolve the UI's `reference, strength, reference, strength, ...` into `(local path, strength)` pairs.
+
+    Resolved before the booking: a download that happens inside `@spaces.GPU` is billed as GPU time.
+    """
+    loras, labels = [], []
+    for reference, scale in zip(lora_fields[::2], lora_fields[1::2]):
+        reference = (reference or "").strip()
+        if not reference or abs(float(scale)) < 1e-6:
+            continue
+        progress(0.0, desc=f"Fetching LoRA {reference} ...")
+        loras.append((resolve_lora(reference), float(scale)))
+        labels.append(f"{os.path.basename(reference)} @ {float(scale):g}")
+    if loras and os.environ.get("H3_AOTI") == "1":
+        raise gr.Error("LoRA не може да се приложи върху AoTI компилиран трансформър. Изключи `H3_AOTI`.")
+    return loras, labels
+
+
 @cache
 def conditioner():
     """The other half, over the gradio API. `gradio_client` attaches the caller's own ZeroGPU token per call, so the
@@ -374,18 +502,23 @@ 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):
+def _generate(prompt_embeds, text_token_tags, references, height, width, num_frames, steps, seed, loras=()):
     """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.
+
+    The adapters are attached here rather than in the caller: `spaces` runs this body in its own worker, so the
+    transformer the request sees is the one that has to carry them.
     """
     import torch
 
     if PLACEMENT == "lazy":
         PIPE.to("cuda")
 
+    apply_loras(PIPE.transformer_ref, loras or ())
+
     state = PIPE(
         prompt_embeds=prompt_embeds.to("cuda"),
         text_token_tags=text_token_tags,
@@ -421,9 +554,11 @@ def generate(
     steps=28,
     seed=42,
     upsample=False,
+    *lora_fields,
     progress=gr.Progress(track_tqdm=True),
 ):
-    """One request."""
+    """One request. The LoRA fields are last and default to empty, so a positional API client that predates them is
+    unaffected. `lora_fields` arrives as `reference, strength, reference, strength, ...`."""
     if LOAD_ERROR:
         raise gr.Error(LOAD_ERROR)
     if PIPE is None:
@@ -440,6 +575,8 @@ def generate(
     derivable = len(audio_bearing(references)) == 1
     requested = 0 if (match and derivable) else snap_frames(duration)
 
+    loras, lora_labels = collect_loras(lora_fields, progress)
+
     progress(0.0, desc="Upsampling the prompt ..." if upsample else "Reading the prompt and references ...")
     conditioned = time.time()
     try:
@@ -463,7 +600,7 @@ 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
+        prompt_embeds, text_token_tags, references, height, width, num_frames, steps, seed, loras
     )
     generate_seconds = time.time() - started
 
@@ -477,12 +614,25 @@ def generate(
         f"({num_frames / FPS:.3f} s), {int(steps)} steps · 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)}",
+        f"({generate_seconds / int(steps):.1f} s/step) · seed {int(seed)}"
+        f"{' · LoRA ' + ', '.join(lora_labels) if lora_labels else ''}",
         flush=True,
     )
     return path, refined, gr.update(visible=bool(refined))
 
 
+def _fill_lora_slots(files, *current):
+    """Drop `.safetensors` files on the uploader and their paths land in the first free slots, so a local adapter
+    needs no typing at all."""
+    slots = list(current)
+    for path in files or []:
+        for index, value in enumerate(slots):
+            if not (value or "").strip():
+                slots[index] = path
+                break
+    return [gr.update(value=value) for value in slots]
+
+
 load_models()
 
 INTRO = """# MiniMax-H3 Reference
@@ -498,6 +648,11 @@ fully synchronized soundtrack (ambience, foley, speech). Bring your own subject,
 reference.
 """
 
+LORA_HELP = """Each slot takes a Hugging Face repo (`owner/repo`), a file inside one
+(`owner/repo/name.safetensors`), a file URL, or a local path — or just drop the files below. A strength of `0`
+switches a slot off without clearing it. Adapters have to be trained against the `transformer_ref/` partition.
+"""
+
 CSS = """
 .main.fillable { max-width: 1250px !important; }
 .dark .gradio-container { color: var(--body-text-color); }
@@ -538,6 +693,32 @@ with gr.Blocks(title="MiniMax-H3 Reference") as demo:
                 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("LoRA", open=False):
+                gr.Markdown(LORA_HELP)
+                lora_references, lora_scales = [], []
+                for slot in range(LORA_SLOTS):
+                    with gr.Row():
+                        lora_references.append(
+                            gr.Textbox(label=f"LoRA {slot + 1}", placeholder="owner/repo", scale=3)
+                        )
+                        lora_scales.append(
+                            gr.Slider(
+                                label="Strength",
+                                minimum=LORA_MIN_SCALE,
+                                maximum=LORA_MAX_SCALE,
+                                step=0.05,
+                                value=1.0,
+                                scale=2,
+                            )
+                        )
+                lora_upload = gr.File(
+                    label="Drop .safetensors here to fill the slots",
+                    file_count="multiple",
+                    file_types=[".safetensors"],
+                    type="filepath",
+                )
+
             with gr.Accordion("Advanced options", open=False):
                 canvas = gr.Dropdown(label="Canvas", choices=list(CANVASES), value=DEFAULT_CANVAS)
                 match = gr.Checkbox(label="Match the reference soundtrack", value=True, visible=False)
@@ -570,8 +751,15 @@ with gr.Blocks(title="MiniMax-H3 Reference") as demo:
             duration_controls, [audio, video, match], [match, duration], show_progress="hidden", api_name=False
         )
 
-    # Same order as `generate`'s signature: the exampled five first, then the remaining image slots.
-    request = [prompt, images[0], audio, video, canvas, *images[1:], match, duration, steps, seed, upsample]
+    # `reference, strength, reference, strength, ...`, which is how `generate` unpacks them.
+    lora_inputs = [field for pair in zip(lora_references, lora_scales) for field in pair]
+    lora_upload.upload(_fill_lora_slots, [lora_upload, *lora_references], lora_references, api_name=False)
+
+    # Same order as `generate`'s signature: the exampled five first, then the remaining image slots, then the LoRA
+    # fields the `*lora_fields` tail collects.
+    request = [
+        prompt, images[0], audio, video, canvas, *images[1:], match, duration, steps, seed, upsample, *lora_inputs
+    ]
 
     gr.Examples(
         examples=[