Spaces:
Running on Zero
Running on Zero
Update app.py
Browse files
app.py
CHANGED
|
@@ -75,29 +75,6 @@ MIN_REFERENCE_VIDEO, MAX_REFERENCE_VIDEO = 2.0, 15.0
|
|
| 75 |
# for two subjects should not open with nine boxes.
|
| 76 |
MAX_IMAGE_SLOTS, OPEN_IMAGE_SLOTS = 9, 2
|
| 77 |
|
| 78 |
-
# How many LoRA slots the UI offers, and the range each strength slider covers.
|
| 79 |
-
LORA_SLOTS = 3
|
| 80 |
-
LORA_MIN_SCALE, LORA_MAX_SCALE = -2.0, 2.0
|
| 81 |
-
|
| 82 |
-
# Pre-wired Turbo LoRAs from `larryvrh/MiniMax-H3-Turbo-Lora`: a few-step distillation that renders joint video +
|
| 83 |
-
# soundtrack in 4–8 steps instead of the usual ~20. Each entry is `(repo reference, recommended steps, blurb)`. The
|
| 84 |
-
# reference is the `owner/repo/filename.safetensors` form `resolve_lora` accepts, so it downloads on first use and is
|
| 85 |
-
# cached by `huggingface_hub` thereafter — nothing is bundled in this Space.
|
| 86 |
-
LORA_PRESETS = {
|
| 87 |
-
"Turbo v4 (step 600) · 6–8 steps": (
|
| 88 |
-
"larryvrh/MiniMax-H3-Turbo-Lora/minimax_h3_turbo_v4_step600_ema.safetensors",
|
| 89 |
-
8,
|
| 90 |
-
"Recommended for most work. Strong static / small-motion, good micro-detail, no over-sharpening. "
|
| 91 |
-
"Use 6–8 steps; 4 steps can smear on heavy motion.",
|
| 92 |
-
),
|
| 93 |
-
"Turbo v1 (ckpt 850) · 4 steps": (
|
| 94 |
-
"larryvrh/MiniMax-H3-Turbo-Lora/minimax_h3_turbo_4step_ema_ckpt850.safetensors",
|
| 95 |
-
4,
|
| 96 |
-
"The friendlier pick for 4-step heavy / fast motion, where v4 can trail. Over-sharpens at higher step counts, "
|
| 97 |
-
"so keep it at 4 steps.",
|
| 98 |
-
),
|
| 99 |
-
}
|
| 100 |
-
# The lowest step count the model's own schedulers accept; the Turbo LoRAs are tuned for 4.
|
| 101 |
MIN_STEPS = 4
|
| 102 |
|
| 103 |
# Seconds of GPU one request needs, from the packed sequence it is about to denoise: linear in the rows for the
|
|
@@ -109,8 +86,6 @@ PLACEMENT_ALLOWANCE = int(os.environ.get("H3_PLACEMENT_ALLOWANCE", "90"))
|
|
| 109 |
AUDIO_LATENTS_PER_SECOND, AUDIO_CHANNELS = 40, 2
|
| 110 |
REFERENCE_IMAGE_SHORT_EDGE, CANVAS_MULTIPLE = 2048, 32
|
| 111 |
DECODE_BASE, DECODE_PER_DEFAULT_CANVAS, DEFAULT_CANVAS_PIXELS = 15, 25, 960 * 544 * 124
|
| 112 |
-
# Reading one adapter off local disk and injecting it across the 33B transformer's linear layers.
|
| 113 |
-
LORA_ALLOWANCE = 12
|
| 114 |
|
| 115 |
|
| 116 |
def snap_frames(seconds: float) -> int:
|
|
@@ -184,7 +159,7 @@ def reference_rows(references: list[tuple[str, str]], num_frames: int) -> int:
|
|
| 184 |
|
| 185 |
|
| 186 |
def get_duration(
|
| 187 |
-
prompt_embeds, text_token_tags, references, height, width, num_frames, steps, seed,
|
| 188 |
):
|
| 189 |
"""Seconds of GPU to reserve for one request. Takes the arguments of the `@spaces.GPU` function it decorates, and
|
| 190 |
tolerates the `gr.Progress` `spaces` injects."""
|
|
@@ -196,7 +171,7 @@ def get_duration(
|
|
| 196 |
# they are handed rather than with the step count.
|
| 197 |
encode = 5 + reference_rows(references, num_frames) * 1e-3
|
| 198 |
decode = DECODE_BASE + DECODE_PER_DEFAULT_CANVAS * (height * width * num_frames) / DEFAULT_CANVAS_PIXELS
|
| 199 |
-
total = PLACEMENT_ALLOWANCE + encode + denoise + decode + 10
|
| 200 |
duration = max(MIN_GPU_DURATION, min(MAX_GPU_DURATION, int(total)))
|
| 201 |
print(f"[ref2va] S={sequence} -> reserving {duration}s ({denoise:.0f}s of denoise at {steps} steps)", flush=True)
|
| 202 |
return duration
|
|
@@ -288,196 +263,6 @@ def _arm_decode_hooks(pipe):
|
|
| 288 |
setattr(module, method, armed)
|
| 289 |
|
| 290 |
|
| 291 |
-
# ----------------------------------------------------------------------------------------------------------------
|
| 292 |
-
# LoRA
|
| 293 |
-
# ----------------------------------------------------------------------------------------------------------------
|
| 294 |
-
# There is no `MiniMaxH3LoraLoaderMixin` in the diffusers integration, so adapters are attached at the *model* level,
|
| 295 |
-
# through the `PeftAdapterMixin` the transformer carries. That is the whole API this needs: `load_lora_adapter` for
|
| 296 |
-
# each file and one `set_adapters` call to give them their strengths. Here the model is `transformer_ref`, so the
|
| 297 |
-
# adapters have to be trained against the `transformer_ref/` partition — a `transformer/` adapter is a different
|
| 298 |
-
# partition and will not match.
|
| 299 |
-
|
| 300 |
-
|
| 301 |
-
def _hub_url_parts(url: str) -> tuple[str, str]:
|
| 302 |
-
"""Split a huggingface.co `blob`/`resolve` URL into its repo id and the file path inside it."""
|
| 303 |
-
from urllib.parse import unquote, urlparse
|
| 304 |
-
|
| 305 |
-
parts = unquote(urlparse(url).path).strip("/").split("/")
|
| 306 |
-
if len(parts) < 5 or parts[2] not in ("resolve", "blob"):
|
| 307 |
-
raise gr.Error(f"Не разпознавам този адрес като файл в Hugging Face: `{url}`")
|
| 308 |
-
return "/".join(parts[:2]), "/".join(parts[4:])
|
| 309 |
-
|
| 310 |
-
|
| 311 |
-
def resolve_lora(reference: str) -> str:
|
| 312 |
-
"""Turn what the user typed into a local `.safetensors` path.
|
| 313 |
-
|
| 314 |
-
Accepts a local path, a huggingface.co file URL, `owner/repo/path/to/file.safetensors`, or a bare `owner/repo`
|
| 315 |
-
whose single `.safetensors` is then picked for them. Runs outside the GPU call, so the download costs no GPU time.
|
| 316 |
-
"""
|
| 317 |
-
from huggingface_hub import hf_hub_download, list_repo_files
|
| 318 |
-
|
| 319 |
-
reference = (reference or "").strip()
|
| 320 |
-
if not reference:
|
| 321 |
-
return ""
|
| 322 |
-
if os.path.exists(reference):
|
| 323 |
-
return reference
|
| 324 |
-
if reference.startswith(("http://", "https://")):
|
| 325 |
-
repo_id, filename = _hub_url_parts(reference)
|
| 326 |
-
return hf_hub_download(repo_id, filename)
|
| 327 |
-
|
| 328 |
-
parts = [part for part in reference.split("/") if part]
|
| 329 |
-
if len(parts) > 2 and parts[-1].endswith(".safetensors"):
|
| 330 |
-
return hf_hub_download("/".join(parts[:2]), "/".join(parts[2:]))
|
| 331 |
-
if len(parts) != 2:
|
| 332 |
-
raise gr.Error(
|
| 333 |
-
f"`{reference}` не е нито съществуващ файл, нито `автор/хранилище`, нито адрес към Hugging Face."
|
| 334 |
-
)
|
| 335 |
-
|
| 336 |
-
candidates = [name for name in list_repo_files(reference) if name.endswith(".safetensors")]
|
| 337 |
-
if not candidates:
|
| 338 |
-
raise gr.Error(f"В `{reference}` няма `.safetensors` файл.")
|
| 339 |
-
if len(candidates) > 1:
|
| 340 |
-
preferred = [name for name in candidates if "lora" in name.lower()]
|
| 341 |
-
if len(preferred) != 1:
|
| 342 |
-
listed = ", ".join(f"`{name}`" for name in sorted(candidates)[:8])
|
| 343 |
-
raise gr.Error(f"`{reference}` съдържа няколко файла. Напиши `{reference}/име.safetensors`. Има: {listed}")
|
| 344 |
-
candidates = preferred
|
| 345 |
-
return hf_hub_download(reference, candidates[0])
|
| 346 |
-
|
| 347 |
-
|
| 348 |
-
def _lora_prefix(state_dict) -> str | None:
|
| 349 |
-
"""The prefix `load_lora_adapter` has to strip before the keys match the transformer's own module names."""
|
| 350 |
-
key = next(iter(state_dict))
|
| 351 |
-
for prefix in ("model.diffusion_model", "diffusion_model", "transformer_ref", "transformer"):
|
| 352 |
-
if key.startswith(f"{prefix}."):
|
| 353 |
-
return prefix
|
| 354 |
-
return None
|
| 355 |
-
|
| 356 |
-
|
| 357 |
-
def _is_comfyui_lora(state_dict) -> bool:
|
| 358 |
-
"""Whether a LoRA state dict is in ComfyUI's MiniMax-H3 naming rather than diffusers'.
|
| 359 |
-
|
| 360 |
-
ComfyUI names the block stack `blocks.N.*` and the token refiner `token_refiner.blocks.N.*`; diffusers names them
|
| 361 |
-
`transformer_blocks.N.*` and `token_refiner.refiner_blocks.N.*`. A key starting with `blocks.` is the tell.
|
| 362 |
-
"""
|
| 363 |
-
for key in state_dict:
|
| 364 |
-
if key.startswith(("blocks.", "token_refiner.blocks.", "final_layer.")):
|
| 365 |
-
return True
|
| 366 |
-
return False
|
| 367 |
-
|
| 368 |
-
|
| 369 |
-
def _convert_comfyui_lora(state_dict) -> dict:
|
| 370 |
-
"""Remap a ComfyUI-format MiniMax-H3 Turbo LoRA to the diffusers `transformer_ref` module names.
|
| 371 |
-
|
| 372 |
-
The Turbo LoRA ([`larryvrh/MiniMax-H3-Turbo-Lora`](https://huggingface.co/larryvrh/MiniMax-H3-Turbo-Lora)) is trained
|
| 373 |
-
against the ComfyUI checkpoint, whose module names differ from diffusers' in four ways:
|
| 374 |
-
|
| 375 |
-
* the block stack is `blocks.N` in ComfyUI but `transformer_blocks.N` in diffusers,
|
| 376 |
-
* the token refiner is `token_refiner.blocks.N` but `token_refiner.refiner_blocks.N`,
|
| 377 |
-
* the final AdaLN is `final_layer.adaln_proj.linear` but `norm_out.linear`,
|
| 378 |
-
* attention QKV is one fused `attn.qkv_proj` in ComfyUI but three separate `attn.to_q` / `to_k` / `to_v` in
|
| 379 |
-
diffusers, and the output projection is `attn.out_proj` but `attn.to_out.0`,
|
| 380 |
-
* the feed-forward is `mlp.fc1` / `mlp.fc2` but `ff.fc1` / `ff.fc2`.
|
| 381 |
-
|
| 382 |
-
The fused QKV `lora_B` is `[3 * inner_dim, rank]`; splitting it into three along dim 0 gives the three separate
|
| 383 |
-
`lora_B` matrices, and `lora_A` (which is `[rank, hidden_size]`) is shared verbatim across the three. The metadata
|
| 384 |
-
says `W_eff = W + lora_B @ lora_A` with alpha = rank, so the scaling is 1.0 and no alpha key is added.
|
| 385 |
-
"""
|
| 386 |
-
import torch
|
| 387 |
-
|
| 388 |
-
converted = {}
|
| 389 |
-
for key, value in state_dict.items():
|
| 390 |
-
# `blocks.N.` -> `transformer_blocks.N.`
|
| 391 |
-
if key.startswith("blocks."):
|
| 392 |
-
new_key = "transformer_blocks." + key[len("blocks."):]
|
| 393 |
-
elif key.startswith("token_refiner.blocks."):
|
| 394 |
-
new_key = "token_refiner.refiner_blocks." + key[len("token_refiner.blocks."):]
|
| 395 |
-
elif key.startswith("final_layer.adaln_proj.linear."):
|
| 396 |
-
new_key = "norm_out.linear." + key[len("final_layer.adaln_proj.linear."):]
|
| 397 |
-
else:
|
| 398 |
-
converted[key] = value
|
| 399 |
-
continue
|
| 400 |
-
|
| 401 |
-
# At this point `new_key` is a diffusers block path. Remap the leaf module names.
|
| 402 |
-
if ".attn.qkv_proj." in new_key:
|
| 403 |
-
# Fused QKV: split `lora_B` along dim 0 into q/k/v, duplicate `lora_A` verbatim.
|
| 404 |
-
leaf = new_key.split(".attn.qkv_proj.")[-1] # `lora_A.weight` or `lora_B.weight`
|
| 405 |
-
stem = new_key[: new_key.index(".attn.qkv_proj.")]
|
| 406 |
-
if leaf == "lora_A.weight":
|
| 407 |
-
for proj in ("to_q", "to_k", "to_v"):
|
| 408 |
-
converted[f"{stem}.attn.{proj}.lora_A.weight"] = value
|
| 409 |
-
else: # lora_B.weight
|
| 410 |
-
q_b, k_b, v_b = value.chunk(3, dim=0)
|
| 411 |
-
converted[f"{stem}.attn.to_q.lora_B.weight"] = q_b
|
| 412 |
-
converted[f"{stem}.attn.to_k.lora_B.weight"] = k_b
|
| 413 |
-
converted[f"{stem}.attn.to_v.lora_B.weight"] = v_b
|
| 414 |
-
elif ".attn.out_proj." in new_key:
|
| 415 |
-
converted[new_key.replace(".attn.out_proj.", ".attn.to_out.0.")] = value
|
| 416 |
-
elif ".mlp." in new_key:
|
| 417 |
-
converted[new_key.replace(".mlp.", ".ff.")] = value
|
| 418 |
-
else:
|
| 419 |
-
# `adaln_proj.linear` and the token refiner's attention/ff already match diffusers' names after the
|
| 420 |
-
# block-prefix rename above.
|
| 421 |
-
converted[new_key] = value
|
| 422 |
-
return converted
|
| 423 |
-
|
| 424 |
-
|
| 425 |
-
def apply_loras(transformer, loras) -> list[str]:
|
| 426 |
-
"""Attach `loras` (local path, strength) to `transformer` and give each its strength, replacing whatever was on it.
|
| 427 |
-
|
| 428 |
-
Every adapter already on the model is removed first, so a request is never affected by the one before it — which
|
| 429 |
-
matters when a worker is reused rather than forked fresh. A LoRA in ComfyUI's MiniMax-H3 naming is remapped to
|
| 430 |
-
diffusers' module names on the fly, so the Turbo LoRA works without a separate conversion step.
|
| 431 |
-
"""
|
| 432 |
-
import torch
|
| 433 |
-
|
| 434 |
-
from safetensors.torch import load_file
|
| 435 |
-
|
| 436 |
-
for name in list(getattr(transformer, "peft_config", None) or {}):
|
| 437 |
-
transformer.delete_adapters(name)
|
| 438 |
-
|
| 439 |
-
names, scales = [], []
|
| 440 |
-
for index, (path, scale) in enumerate(loras):
|
| 441 |
-
state_dict = load_file(path)
|
| 442 |
-
if _is_comfyui_lora(state_dict):
|
| 443 |
-
state_dict = _convert_comfyui_lora(state_dict)
|
| 444 |
-
name = f"lora{index}"
|
| 445 |
-
transformer.load_lora_adapter(state_dict, adapter_name=name, prefix=_lora_prefix(state_dict))
|
| 446 |
-
names.append(name)
|
| 447 |
-
scales.append(float(scale))
|
| 448 |
-
|
| 449 |
-
if not names:
|
| 450 |
-
return []
|
| 451 |
-
|
| 452 |
-
# PEFT builds the new layers on its own default device/dtype; the base weights are the truth here, under either
|
| 453 |
-
# placement mode (`offload` keeps them on the host and moves whole modules by hook).
|
| 454 |
-
base = next(param for key, param in transformer.named_parameters() if ".lora_" not in key)
|
| 455 |
-
with torch.no_grad():
|
| 456 |
-
for key, param in transformer.named_parameters():
|
| 457 |
-
if ".lora_" in key and (param.device != base.device or param.dtype != base.dtype):
|
| 458 |
-
param.data = param.data.to(device=base.device, dtype=base.dtype)
|
| 459 |
-
|
| 460 |
-
transformer.set_adapters(names, scales)
|
| 461 |
-
return names
|
| 462 |
-
|
| 463 |
-
|
| 464 |
-
def collect_loras(lora_fields, progress) -> tuple[list[tuple[str, float]], list[str]]:
|
| 465 |
-
"""Resolve the UI's `reference, strength, reference, strength, ...` into `(local path, strength)` pairs.
|
| 466 |
-
|
| 467 |
-
Resolved before the booking: a download that happens inside `@spaces.GPU` is billed as GPU time.
|
| 468 |
-
"""
|
| 469 |
-
loras, labels = [], []
|
| 470 |
-
for reference, scale in zip(lora_fields[::2], lora_fields[1::2]):
|
| 471 |
-
reference = (reference or "").strip()
|
| 472 |
-
if not reference or abs(float(scale)) < 1e-6:
|
| 473 |
-
continue
|
| 474 |
-
progress(0.0, desc=f"Fetching LoRA {reference} ...")
|
| 475 |
-
loras.append((resolve_lora(reference), float(scale)))
|
| 476 |
-
labels.append(f"{os.path.basename(reference)} @ {float(scale):g}")
|
| 477 |
-
if loras and os.environ.get("H3_AOTI") == "1":
|
| 478 |
-
raise gr.Error("LoRA не може да се приложи върху AoTI компилиран трансформър. Изключи `H3_AOTI`.")
|
| 479 |
-
return loras, labels
|
| 480 |
-
|
| 481 |
|
| 482 |
@cache
|
| 483 |
def conditioner():
|
|
|
|
| 75 |
# for two subjects should not open with nine boxes.
|
| 76 |
MAX_IMAGE_SLOTS, OPEN_IMAGE_SLOTS = 9, 2
|
| 77 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 78 |
MIN_STEPS = 4
|
| 79 |
|
| 80 |
# Seconds of GPU one request needs, from the packed sequence it is about to denoise: linear in the rows for the
|
|
|
|
| 86 |
AUDIO_LATENTS_PER_SECOND, AUDIO_CHANNELS = 40, 2
|
| 87 |
REFERENCE_IMAGE_SHORT_EDGE, CANVAS_MULTIPLE = 2048, 32
|
| 88 |
DECODE_BASE, DECODE_PER_DEFAULT_CANVAS, DEFAULT_CANVAS_PIXELS = 15, 25, 960 * 544 * 124
|
|
|
|
|
|
|
| 89 |
|
| 90 |
|
| 91 |
def snap_frames(seconds: float) -> int:
|
|
|
|
| 159 |
|
| 160 |
|
| 161 |
def get_duration(
|
| 162 |
+
prompt_embeds, text_token_tags, references, height, width, num_frames, steps, seed, **_
|
| 163 |
):
|
| 164 |
"""Seconds of GPU to reserve for one request. Takes the arguments of the `@spaces.GPU` function it decorates, and
|
| 165 |
tolerates the `gr.Progress` `spaces` injects."""
|
|
|
|
| 171 |
# they are handed rather than with the step count.
|
| 172 |
encode = 5 + reference_rows(references, num_frames) * 1e-3
|
| 173 |
decode = DECODE_BASE + DECODE_PER_DEFAULT_CANVAS * (height * width * num_frames) / DEFAULT_CANVAS_PIXELS
|
| 174 |
+
total = PLACEMENT_ALLOWANCE + encode + denoise + decode + 10
|
| 175 |
duration = max(MIN_GPU_DURATION, min(MAX_GPU_DURATION, int(total)))
|
| 176 |
print(f"[ref2va] S={sequence} -> reserving {duration}s ({denoise:.0f}s of denoise at {steps} steps)", flush=True)
|
| 177 |
return duration
|
|
|
|
| 263 |
setattr(module, method, armed)
|
| 264 |
|
| 265 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 266 |
|
| 267 |
@cache
|
| 268 |
def conditioner():
|