diff --git a/ldm/simplet2i.py b/ldm/simplet2i.py index 50816ec403..ed43e9a598 100644 --- a/ldm/simplet2i.py +++ b/ldm/simplet2i.py @@ -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() diff --git a/scripts/dream.py b/scripts/dream.py index 5d5e8db4a5..fda92310e4 100755 --- a/scripts/dream.py +++ b/scripts/dream.py @@ -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