File size: 3,145 Bytes
c0aa76e | 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 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | #!/usr/bin/env python3
"""
dav_loader.py — import MiniMaxMusic3DAV without executing SimpleTuner's
package __init__ chain.
`simpletuner/__init__.py` pulls in the whole training stack, which monkeypatches
diffusers attention on import. Against a pinned diffusers commit that patching
can fail (e.g. TemplatedRingAnythingAttention not present), taking the DAV
class down with it.
vocoder.py itself imports only math / torch / diffusers, so it loads fine
standalone. This registers stub parent packages and execs just that file.
from dav_loader import get_dav_class
MiniMaxMusic3DAV = get_dav_class(Path("~/gpu_work/m3").expanduser())
"""
from __future__ import annotations
import importlib.util
import sys
import types
from pathlib import Path
_CACHE = {}
_PKG_CHAIN = [
"simpletuner",
"simpletuner.helpers",
"simpletuner.helpers.models",
"simpletuner.helpers.models.minimaxmusic",
]
_MODNAME = "simpletuner.helpers.models.minimaxmusic.vocoder"
def get_dav_class(workdir: Path | str, simpletuner_dir: Path | str | None = None):
"""Return the MiniMaxMusic3DAV class, loaded in isolation."""
key = str(workdir)
if key in _CACHE:
return _CACHE[key]
root = Path(simpletuner_dir) if simpletuner_dir else Path(workdir) / "SimpleTuner"
vpath = root / "simpletuner" / "helpers" / "models" / "minimaxmusic" / "vocoder.py"
if not vpath.is_file():
raise FileNotFoundError(
f"vocoder.py not found at {vpath}\n"
f"Re-fetch it with:\n"
f" git -C {root} fetch --depth 1 origin refs/pull/3074/head && "
f"git -C {root} checkout -q FETCH_HEAD"
)
# Stub the parent packages so the dotted module name resolves without
# any real __init__.py being executed.
for name in _PKG_CHAIN:
if name not in sys.modules:
stub = types.ModuleType(name)
stub.__path__ = [] # mark as a package
sys.modules[name] = stub
spec = importlib.util.spec_from_file_location(_MODNAME, str(vpath))
if spec is None or spec.loader is None:
raise ImportError(f"could not build a spec for {vpath}")
mod = importlib.util.module_from_spec(spec)
sys.modules[_MODNAME] = mod
spec.loader.exec_module(mod)
cls = getattr(mod, "MiniMaxMusic3DAV", None)
if cls is None:
raise ImportError(f"MiniMaxMusic3DAV not found in {vpath}")
_CACHE[key] = cls
return cls
if __name__ == "__main__":
import argparse
ap = argparse.ArgumentParser()
ap.add_argument("--workdir", type=Path, required=True)
a = ap.parse_args()
cls = get_dav_class(a.workdir.expanduser().resolve())
print(f"loaded {cls.__name__} OK")
ckpt = a.workdir.expanduser().resolve() / "minimax_music3" / "dav.pth"
if ckpt.exists():
dav = cls.from_original_dav(str(ckpt))
n = sum(p.numel() for p in dav.parameters())
print(f"weights OK | {n/1e6:.1f}M params | hop_length={dav.hop_length} "
f"| latent_channels={dav.config.latent_channels}")
else:
print(f"(checkpoint not at {ckpt}; class import verified only)")
|