def forward(self): if exists(self.start_image): tqdm.write('Preparing with initial image...') optim = DiffGrad(self.model.parameters(), lr=self.start_image_lr) pbar = trange(self.start_image_train_iters, desc='iteration') for _ in pbar: loss = self.model.model(self.start_image) loss.backward() pbar.set_description(f'loss: {loss.item():.2f}') optim.step() optim.zero_grad() if terminate: print('interrupted by keyboard, gracefully exiting') return exit() del self.start_image del optim tqdm.write(f'Imagining "{self.text}" from the depths of my weights...') if self.open_folder: open_folder('./') self.open_folder = False for epoch in trange(self.epochs, desc='epochs'): pbar = trange(self.iterations, desc='iteration') for i in pbar: loss = self.train_step(epoch, i) pbar.set_description(f'loss: {loss.item():.2f}') if terminate: print('interrupted by keyboard, gracefully exiting') return
def forward(self): if exists(self.start_image): tqdm.write('Preparing with initial image...') optim = DiffGrad(self.model.model.parameters(), lr=self.start_image_lr) pbar = trange(self.start_image_train_iters, desc='iteration') try: for _ in pbar: loss = self.model.model(self.start_image) loss.backward() pbar.set_description(f'loss: {loss.item():.2f}') optim.step() optim.zero_grad() except KeyboardInterrupt: print('interrupted by keyboard, gracefully exiting') return exit() del self.start_image del optim tqdm.write( f'Imagining "{self.textpath}" from the depths of my weights...') with torch.no_grad(): self.model( self.clip_encoding, dry_run=True ) # do one warmup step due to potential issue with CLIP and CUDA if self.open_folder: if self.output_folder: open_folder(self.output_folder) else: open_folder('./') self.open_folder = False try: for epoch in trange(self.epochs, desc='epochs'): pbar = trange(self.iterations, desc='iteration') for i in pbar: _, loss = self.train_step(epoch, i) pbar.set_description(f'loss: {loss.item():.2f}') # Update clip_encoding per epoch if we are creating a story if self.create_story: self.clip_encoding = self.update_story_encoding(epoch, i) except KeyboardInterrupt: print('interrupted by keyboard, gracefully exiting') return self.save_image(epoch, i) # one final save at end if (self.save_gif or self.save_video) and self.save_progress: self.generate_gif()
def forward(self): if exists(self.start_image): tqdm.write('Preparing with initial image...') optim = DiffGrad(self.model.parameters(), lr = self.start_image_lr) pbar = trange(self.start_image_train_iters, desc='iteration') for _ in pbar: loss = self.model.model(self.start_image) loss.backward() pbar.set_description(f'loss: {loss.item():.2f}') optim.step() optim.zero_grad() if terminate: print('interrupted by keyboard, gracefully exiting') return sys.exit() del self.start_image del optim tqdm.write(f'Imagining "{self.textpath}" from the depths of my weights...') self.model(self.clip_encoding, dry_run = True) # do one warmup step due to potential issue with CLIP and CUDA if self.open_folder: open_folder('./') self.open_folder = False for epoch in trange(self.epochs, desc='epochs'): pbar = trange(self.iterations, desc='iteration') for i in pbar: loss = self.train_step(epoch, i) pbar.set_description(f'loss: {loss.item():.2f}') if terminate: print('interrupted by keyboard, gracefully exiting') return self.save_image(self.epochs, self.iterations) # one final save at end