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

empty cache if cuda available

Browse Source
This commit is contained in:
Brett Kuprel 2022-07-04 07:21:54 -04:00
parent 9602055eaa
commit 256fd1aab0
2 changed files with 2 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
README.md vendored
View File

@ -81,7 +81,7 @@ display(image)
Use `image_from_text.py` to generate images from the command line. Use `image_from_text.py` to generate images from the command line.
```bash ```bash
$ python image_from_text.py --text='artificial intelligence' --seed=7 $ python image_from_text.py --text='artificial intelligence' --no-mega --seed=7
``` ```
<img src="https://github.com/kuprel/min-dalle/raw/main/examples/artificial_intelligence.jpg" alt="min-dalle" width="200"/> <img src="https://github.com/kuprel/min-dalle/raw/main/examples/artificial_intelligence.jpg" alt="min-dalle" width="200"/>

3
min_dalle/min_dalle.py
View File

@ -142,7 +142,6 @@ class MinDalle:
params = torch.load(self.detoker_params_path) params = torch.load(self.detoker_params_path)
self.detokenizer.load_state_dict(params) self.detokenizer.load_state_dict(params)
del params del params
torch.cuda.empty_cache()
if torch.cuda.is_available(): self.detokenizer = self.detokenizer.cuda() if torch.cuda.is_available(): self.detokenizer = self.detokenizer.cuda()
@ -176,7 +175,7 @@ class MinDalle:
encoder_state encoder_state
) )
if not self.is_reusable: del self.decoder if not self.is_reusable: del self.decoder
torch.cuda.empty_cache() if torch.cuda.is_available(): torch.cuda.empty_cache()
return image_tokens return image_tokens