model_loader.py

"""
Model download utilities for ZeroGPU deployment.
Actual model loading happens lazily inside @spaces.GPU in app.py.
"""

import os
from pathlib import Path


def download_phoenix_weights(output_dir: str):
    """Download Phoenix weights from HuggingFace."""
    from huggingface_hub import hf_hub_download

    os.makedirs(output_dir, exist_ok=True)

    files = {
        "flow_model.pth": "weights/flow/tenx/multi/cell/20x/discrete/flow_model.pth",
        "stats_table.npz": "statistics/tenx/multi/cell/discrete/stats_table.npz",
        "xenium_human_multi.npy": "panels/xenium_human_multi.npy",
    }

    for local_name, repo_path in files.items():
        target = Path(output_dir) / local_name
        if target.exists():
            continue
        print(f"  Downloading {local_name}...")
        hf_hub_download(
            repo_id="peng-lab/phoenix",
            filename=repo_path,
            local_dir=output_dir,
            local_dir_use_symlinks=False,
        )
        # hf_hub_download preserves subdir structure, move to flat
        downloaded = Path(output_dir) / repo_path
        if downloaded.exists() and downloaded != target:
            downloaded.rename(target)

    print(f"  Phoenix weights ready in {output_dir}")