|
|
|
|
|
|
|
|
|
|
|
|
|
""" |
|
Functions for downloading pre-trained DiT models |
|
""" |
|
from torchvision.datasets.utils import download_url |
|
import torch |
|
import os |
|
|
|
|
|
pretrained_models = {'DiT-XL-2-512x512.pt', 'DiT-XL-2-256x256.pt'} |
|
|
|
|
|
def find_model(model_name): |
|
""" |
|
Finds a pre-trained G.pt model, downloading it if necessary. Alternatively, loads a model from a local path. |
|
""" |
|
if model_name in pretrained_models: |
|
return download_model(model_name) |
|
else: |
|
assert os.path.isfile(model_name), f'Could not find DiT checkpoint at {model_name}' |
|
return torch.load(model_name, map_location=lambda storage, loc: storage) |
|
|
|
|
|
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__": |
|
|
|
for model in pretrained_models: |
|
download_model(model) |
|
print('Done.') |
|
|