tanner/min-dalle-test
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

save converted detokenizer params

Browse Source
This commit is contained in:
Brett Kuprel
2022-07-01 10:17:29 -04:00
parent 8b5960b687
commit e4c2be54cb
7 changed files with 35 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
min_dalle/load_params.py
+5 -1
View File
@@ -128,5 +128,9 @@ def convert_and_save_torch_params(is_mega: bool, model_path: str):
for i in encoder_params:
encoder_params[i] = encoder_params[i].to(torch.float16)
detoker_params = load_vqgan_torch_params('./pretrained/vqgan')
detoker_path = os.path.join('pretrained', 'vqgan', 'detokenizer.pt')
torch.save(encoder_params, os.path.join(model_path, 'encoder.pt'))
torch.save(decoder_params, os.path.join(model_path, 'decoder.pt'))
torch.save(decoder_params, os.path.join(model_path, 'decoder.pt'))
torch.save(detoker_params, detoker_path)