Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

print current and max VRAM usage stats after each round of generation

Browse Source
This commit is contained in:
Lincoln Stein 2022-08-28 13:05:01 -04:00
parent 38ed6393fa
commit ba03289c14
2 changed files with 10 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
ldm/simplet2i.py
View File

@ -158,6 +158,7 @@ class T2I:
self.latent_diffusion_weights = latent_diffusion_weights
self.device = device
self.gfpgan = gfpgan
self.session_peakmem = torch.cuda.max_memory_allocated()
if seed is None:
self.seed = self._new_seed()
else:
@ -262,14 +263,15 @@ class T2I:
h = int(height / 64) * 64
if h != height or w != width:
print(
f'Height and width must be multiples of 64. Resizing to {h}x{w}'
f'Height and width must be multiples of 64. Resizing to {h}x{w}.'
)
height = h
width = w
width = w
scope = autocast if self.precision == 'autocast' else nullcontext
tic = time.time()
torch.cuda.torch.cuda.reset_peak_memory_stats()
results = list()
try:
@ -333,7 +335,11 @@ class T2I:
print('Are you sure your system has an adequate NVIDIA GPU?')
toc = time.time()
print(f'{len(results)} images generated in', '%4.2fs' % (toc - tic))
self.session_peakmem = max(self.session_peakmem,torch.cuda.max_memory_allocated() )
print('Usage stats:')
print(f' {len(results)} image(s) generated in', '%4.2fs' % (toc - tic))
print(f' Max VRAM used for this generation:', '%4.2fG' % (torch.cuda.max_memory_allocated()/1E9))
print(f' Max VRAM used since script start: ', '%4.2fG' % (self.session_peakmem/1E9))
return results
@torch.no_grad()

5
scripts/dream.py
View File

@ -10,9 +10,6 @@ import warnings
import ldm.dream.readline
from ldm.dream.pngwriter import PngWriter, PromptFormatter
debugging = False
def main():
"""Initialize command-line parsers and the diffusion model"""
arg_parser = create_argv_parser()
@ -209,10 +206,10 @@ def main_loop(t2i, outdir, parser, log, infile):
normalized_prompt = PromptFormatter(t2i, opt).normalize_prompt()
individual_images = not opt.grid
# Here is where the images are actually generated!
try:
file_writer = PngWriter(outdir, normalized_prompt, opt.batch_size)
callback = file_writer.write_image if individual_images else None
image_list = t2i.prompt2image(image_callback=callback, **vars(opt))
results = (
file_writer.files_written if individual_images else image_list