File size: 1,174 Bytes
dceb455 | 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 | """Put this folder and optional OmniVAE sources on the import path."""
import os
import sys
from pathlib import Path
def prepare() -> Path:
repo = Path(__file__).resolve().parent
os.chdir(repo)
paths = [str(repo)]
omni = os.environ.get("OMNIVAE_CODE_ROOT", "").strip()
if not omni:
fallback = repo.parent.parent / "model" / "OmniVAE_code"
if fallback.is_dir():
omni = str(fallback)
if omni:
paths.append(omni)
os.environ.setdefault("OMNIVAE_CODE_ROOT", omni)
if not os.environ.get("OMNIVAE_AUDIO_CKPT_DIR", "").strip():
audio = (
repo.parent.parent
/ "model"
/ "OmniVAE"
/ "models"
/ "vae"
/ "audio_only"
/ "recon_distill_avclip_ft_decoder"
)
if audio.is_dir():
os.environ.setdefault("OMNIVAE_AUDIO_CKPT_DIR", str(audio))
for name in (
"DAVINCI_DISABLE_MAGI_COMPILE",
"MAGI_DISABLE_COMPILE",
"TORCHDYNAMO_DISABLE",
):
os.environ.setdefault(name, "1")
sys.path[:0] = paths
return repo
|