BlueWaveSemi45/DramaboxCPU

1#!/usr/bin/env python32"""3Download Dramabox models from HuggingFace.4 5Models are cached locally after first download.6Gemma text encoder is fetched separately from Google's repo.7"""8import logging9import os10from pathlib import Path11 12from huggingface_hub import hf_hub_download, snapshot_download13 14logger = logging.getLogger(__name__)15 16DRAMABOX_REPO = "ResembleAI/Dramabox"17GEMMA_REPO = "unsloth/gemma-3-12b-it-bnb-4bit"18 19# Default cache directory20DEFAULT_CACHE = os.environ.get(21    "DRAMABOX_CACHE",22    os.path.join(os.path.expanduser("~"), ".cache", "dramabox"),23)24 25# Model files in the HF repo (flat structure)26MODEL_FILES = {27    "transformer": "dramabox-dit-v1.safetensors",28    "audio_components": "dramabox-audio-components.safetensors",29    "silence_latent": "assets/silence_latent_frame.pt",30}31 32 33def get_model_path(name: str, cache_dir: str = None) -> str:34    """Download a model file from HF and return local path.35 36    Args:37        name: One of 'transformer', 'audio_components', 'silence_latent'38        cache_dir: Local cache directory (default: ~/.cache/dramabox)39 40    Returns:41        Local file path42    """43    cache_dir = cache_dir or DEFAULT_CACHE44 45    if name not in MODEL_FILES:46        raise ValueError(f"Unknown model: {name}. Choose from: {list(MODEL_FILES.keys())}")47 48    repo_path = MODEL_FILES[name]49    logger.info(f"Fetching {name} from {DRAMABOX_REPO}/{repo_path}...")50 51    local_path = hf_hub_download(52        repo_id=DRAMABOX_REPO,53        filename=repo_path,54        cache_dir=cache_dir,55        token=os.environ.get("HF_TOKEN"),56    )57    logger.info(f"  -> {local_path}")58    return local_path59 60 61def get_gemma_path(cache_dir: str = None) -> str:62    """Download Gemma 3 12B IT (pre-quantized bnb-4bit via unsloth) and return63    the snapshot directory. Using the pre-quantized variant skips runtime64    bitsandbytes quantization and ~halves the Gemma load time.65    """66    cache_dir = cache_dir or DEFAULT_CACHE67    logger.info(f"Fetching Gemma from {GEMMA_REPO}...")68 69    local_dir = snapshot_download(70        repo_id=GEMMA_REPO,71        cache_dir=cache_dir,72        token=os.environ.get("HF_TOKEN"),73    )74    logger.info(f"  -> {local_dir}")75    return local_dir76 77 78def get_all_paths(cache_dir: str = None) -> dict:79    """Download all required models and return paths dict.80 81    Returns:82        {83            'transformer': '/path/to/transformer.safetensors',84            'audio_components': '/path/to/audio-components.safetensors',85            'silence_latent': '/path/to/silence_latent_frame.pt',86            'gemma_root': '/path/to/unsloth/gemma-3-12b-it-bnb-4bit/',87        }88    """89    cache_dir = cache_dir or DEFAULT_CACHE90    paths = {}91 92    for name in MODEL_FILES:93        paths[name] = get_model_path(name, cache_dir)94 95    paths["gemma_root"] = get_gemma_path(cache_dir)96    return paths97 98 99if __name__ == "__main__":100    logging.basicConfig(level=logging.INFO)101    paths = get_all_paths()102    print("\nAll models downloaded:")103    for k, v in paths.items():104        size = os.path.getsize(v) / 1e9 if os.path.isfile(v) else "dir"105        print(f"  {k}: {v} ({size:.2f}GB)" if isinstance(size, float) else f"  {k}: {v} (directory)")106