File size: 696 Bytes
f4a41d8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 |
from transformers import T5Tokenizer, T5ForConditionalGeneration
import torch
import os
def download_models():
model_name = "roborovski/superprompt-v1"
tokenizer = T5Tokenizer.from_pretrained(model_name)
model = T5ForConditionalGeneration.from_pretrained(model_name, torch_dtype=torch.float16)
script_dir = os.path.dirname(os.path.abspath(__file__)) # Script directory
modelDir = os.path.join(script_dir, "./model_files/" )
os.makedirs(modelDir, exist_ok=True)
tokenizer.save_pretrained(modelDir)
model.save_pretrained(modelDir)
print("Downloaded SuperPrompt-v1 model files to", modelDir)
return modelDir
if __name__ == '__main__':
download_models() |