| """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 | |