Save sample images while training Dreambooth & Textual Inversion #1993
Description
During Dreambooth or Textual Inversion training, I want to see periodic sample images to see how the training is going. There should be an option to save every X steps (or epochs).
I've been using the following code, which I adapted from the ShivamShrirao fork, but it does not seem to be very effective at showing any changes happening in the model. So I'm not sure if the code is wrong or if I'm just not training it properly (i.e. learning rate too high or low).
torch.cuda.empty_cache()
pipeline = StableDiffusionPipeline.from_pretrained(
args.pretrained_model_name_or_path,
text_encoder=accelerator.unwrap_model(text_encoder),
vae=vae,
unet=unet,
tokenizer=tokenizer,
safety_checker=None,
requires_safety_checker=False,
revision=args.revision,
)
pipeline = pipeline.to(accelerator.device)
g_cuda = torch.Generator(device=accelerator.device).manual_seed(args.seed) if args.seed is not None else None
pipeline.set_progress_bar_config(disable=True)
with torch.autocast('cuda'), torch.inference_mode():
sample_image = pipeline(
args.save_sample_prompt,
guidance_scale=7.5,
num_inference_steps=50,
generator=g_cuda
).images[0]
sample_image.save(os.path.join(args.output_dir, f'{epoch}.png'))
del pipeline
torch.cuda.empty_cache()