Spaces:
Running on Zero
Running on Zero
File size: 3,361 Bytes
1811aab | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | """Pure inference layout construction and strict LoRA loading."""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
import re
from typing import Any
from .constants import AUDIO_LATENT_SHAPE, VIDEO_LATENT_SHAPE
@dataclass(frozen=True)
class RuntimeBatch:
layout: Any
text_hidden_states: Any
text_token_tags: Any
def build_runtime_batch(anchor: Any, reference: Any, hidden: Any, tags: Any) -> RuntimeBatch:
import torch
from musubi_tuner.minimax_h3.packing import H3ReferenceGeometry, H3VideoGeometry, build_h3_layout
if anchor.ndim != 5 or tuple(anchor.shape[:3]) != (1, 24, 1):
raise ValueError(f"anchor must be [1,24,1,H,W], got {tuple(anchor.shape)}")
if reference.ndim != 5 or tuple(reference.shape[:3]) != (1, 24, 37):
raise ValueError(f"reference must be [1,24,37,H,W], got {tuple(reference.shape)}")
if hidden.ndim != 3 or hidden.shape[0] != 1 or hidden.shape[2] != 5120:
raise ValueError(f"text hidden states must be [1,L,5120], got {tuple(hidden.shape)}")
if tags.shape != hidden.shape[:2] or tags.dtype != torch.int64:
raise ValueError("text token tags must be int64 [1,L]")
layout = build_h3_layout(
task="ref2va",
text_length=hidden.shape[1],
target_video=H3VideoGeometry(*VIDEO_LATENT_SHAPE[2:]),
target_audio_frames=AUDIO_LATENT_SHAPE[-1],
references=(
H3ReferenceGeometry("image", video=H3VideoGeometry(*anchor.shape[2:])),
H3ReferenceGeometry("video", video=H3VideoGeometry(*reference.shape[2:])),
),
)
return RuntimeBatch(layout=layout, text_hidden_states=hidden, text_token_tags=tags)
def strict_merge_lora(transformer: Any, checkpoint: Path, expected_modules: int) -> dict[str, int]:
import torch
from safetensors.torch import load_file
from musubi_tuner.networks import lora_minimax_h3
from musubi_tuner.utils.lora_utils import filter_lora_state_dict
state = filter_lora_state_dict(load_file(str(checkpoint)), None, None)
network = lora_minimax_h3.create_arch_network_from_weights(
1.0, state, unet=transformer, for_inference=True
)
if network.text_encoder_loras:
raise ValueError("LoRA unexpectedly targets the text encoder")
if len(network.unet_loras) != expected_modules:
raise ValueError(
f"LoRA module count mismatch: expected {expected_modules}, got {len(network.unet_loras)}"
)
expected_keys: set[str] = set()
names: set[str] = set()
for module in network.unet_loras:
name = str(module.lora_name)
if name in names or module.split_dims is not None or re.search(r"[^A-Za-z0-9_.]", name):
raise ValueError(f"invalid LoRA target: {name}")
names.add(name)
expected_keys.update(
{f"{name}.alpha", f"{name}.lora_down.weight", f"{name}.lora_up.weight"}
)
missing = expected_keys - set(state)
unexpected = set(state) - expected_keys
if missing or unexpected:
raise ValueError(f"LoRA key mismatch: missing={len(missing)} unexpected={len(unexpected)}")
network.merge_to(None, transformer, state, dtype=torch.bfloat16, device="cpu")
result = {"loaded_modules": len(names), "missing_keys": 0, "unexpected_keys": 0}
del network, state
return result
|