import modal

from config import cnf

data_volume = modal.Volume.from_name(
    "garments2look-data",
    create_if_missing=True,
    version=2,
)

output_volume = modal.Volume.from_name(
    "garments2look-checkpoints",
    create_if_missing=True,
)

hf_cache_volume = modal.Volume.from_name(
    "huggingface-cache",
    create_if_missing=True,
)

training_image = (
    modal.Image.debian_slim(python_version="3.11")
    .uv_pip_install(
        "torch==2.7.1",
        "torchvision==0.22.1",
        "sentence-transformers==5.5.1",
        "huggingface-hub>=1.0",
        "pillow==11.2.1",
        "tqdm==4.67.1",
        "pydantic==2.13.4",
    )
    .env(
        {
            "HF_HOME": cnf.train_hf_cache_mnt,
            "TOKENIZERS_PARALLELISM": "false",
            "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
        }
    )
    .add_local_python_source(
        "config",
        "embedding_train",
    )
)