Spaces:
Runtime error
Runtime error
File size: 1,288 Bytes
0035a82 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
"""
Functions for downloading pre-trained DiT models
"""
from torchvision.datasets.utils import download_url
import torch
import os
def find_model(model_name):
checkpoint = torch.load(model_name, map_location=lambda storage, loc: storage)
if "ema" in checkpoint: # supports checkpoints from train.py
print('Ema existing!')
checkpoint = checkpoint["ema"]
return checkpoint
def download_model(model_name):
"""
Downloads a pre-trained DiT model from the web.
"""
assert model_name in pretrained_models
local_path = f'pretrained_models/{model_name}'
if not os.path.isfile(local_path):
os.makedirs('pretrained_models', exist_ok=True)
web_path = f'https://dl.fbaipublicfiles.com/DiT/models/{model_name}'
download_url(web_path, 'pretrained_models')
model = torch.load(local_path, map_location=lambda storage, loc: storage)
return model
if __name__ == "__main__":
# Download all DiT checkpoints
for model in pretrained_models:
download_model(model)
print('Done.')
|