src/build_pipe.py

from .pipeline_terrain import TerrainDiffusionPipeline
from huggingface_hub import hf_hub_download, snapshot_download
import os
import torch

def build_pipe():
    print('Downloading weights...')
    try:
        os.mkdir('./weights/')
    except:
        True
    snapshot_download(repo_id="NewtNewt/MESA", local_dir="./weights")
    weight_path = './weights'
    print('[DONE]')
    
    print('Instantiating Model...')
    pipe = TerrainDiffusionPipeline.from_pretrained(weight_path, torch_dtype=torch.float16)
    pipe.to("cuda")
    print('[DONE]')
    
    return pipe