terrain_diffusion/common/model_utils.py

import os

MODEL_PATHS = {
    "coarse": ("checkpoints/models/diffusion_coarse", "xandergos/TerrainDiffusion-Diffusion-Coarse-128A"),
    "base": ("checkpoints/models/consistency_base-192x3", "xandergos/TerrainDiffusion-Consistency-Base-192x3"),
    "decoder": ("checkpoints/models/consistency_decoder-64x3", "xandergos/TerrainDiffusion-Consistency-Decoder-64x3"),
    "diffusion_base": ("checkpoints/models/diffusion_base-192x3", "xandergos/TerrainDiffusion-Diffusion-Base-192x3"),
    "diffusion_base_guide": ("checkpoints/models/diffusion_base-128x3", "xandergos/TerrainDiffusion-Diffusion-Base-128x3"),
}

def resolve_model_path(user_path, local_default, hf_repo):
    if user_path is not None:
        return user_path
    if os.path.exists(local_default):
        return local_default
    return hf_repo

def get_default_model_path(name):
    local, hf = MODEL_PATHS[name]
    return local if os.path.exists(local) else hf