From 352d96eb823bb51ce8a7ffad5761d4f7106be706 Mon Sep 17 00:00:00 2001 From: satani99 <42287151+satani99@users.noreply.github.com> Date: Fri, 31 May 2024 04:52:22 +0530 Subject: [PATCH] Modularize train_text_to_image_lora_sdxl inferencing during and after training in example (#8335) * Modularized the train_lora_sdxl file * Modularized the train_lora_sdxl file * Modularized the train_lora_sdxl file --------- Co-authored-by: Sayak Paul --- .../train_text_to_image_lora_sdxl.py | 106 ++++++++---------- 1 file changed, 46 insertions(+), 60 deletions(-) diff --git a/examples/text_to_image/train_text_to_image_lora_sdxl.py b/examples/text_to_image/train_text_to_image_lora_sdxl.py index 077bee2e2f..3d266811ff 100644 --- a/examples/text_to_image/train_text_to_image_lora_sdxl.py +++ b/examples/text_to_image/train_text_to_image_lora_sdxl.py @@ -64,6 +64,9 @@ from diffusers.utils.import_utils import is_torch_npu_available, is_xformers_ava from diffusers.utils.torch_utils import is_compiled_module +if is_wandb_available(): + import wandb + # Will error if the minimal version of diffusers is not installed. Remove at your own risks. check_min_version("0.29.0.dev0") @@ -119,6 +122,47 @@ Special VAE used for training: {vae_path}. model_card.save(os.path.join(repo_folder, "README.md")) +def log_validation( + pipeline, + args, + accelerator, + epoch, + is_final_validation=False, +): + logger.info( + f"Running validation... \n Generating {args.num_validation_images} images with prompt:" + f" {args.validation_prompt}." + ) + pipeline = pipeline.to(accelerator.device) + pipeline.set_progress_bar_config(disable=True) + + # run inference + generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) if args.seed else None + pipeline_args = {"prompt": args.validation_prompt} + if torch.backends.mps.is_available(): + autocast_ctx = nullcontext() + else: + autocast_ctx = torch.autocast(accelerator.device.type) + + with autocast_ctx: + images = [pipeline(**pipeline_args, generator=generator).images[0] for _ in range(args.num_validation_images)] + + for tracker in accelerator.trackers: + phase_name = "test" if is_final_validation else "validation" + if tracker.name == "tensorboard": + np_images = np.stack([np.asarray(img) for img in images]) + tracker.writer.add_images(phase_name, np_images, epoch, dataformats="NHWC") + if tracker.name == "wandb": + tracker.log( + { + phase_name: [ + wandb.Image(image, caption=f"{i}: {args.validation_prompt}") for i, image in enumerate(images) + ] + } + ) + return images + + def import_model_class_from_model_name_or_path( pretrained_model_name_or_path: str, revision: str, subfolder: str = "text_encoder" ): @@ -523,11 +567,6 @@ def main(args): kwargs_handlers=[kwargs], ) - if args.report_to == "wandb": - if not is_wandb_available(): - raise ImportError("Make sure to install wandb if you want to use it for logging during training.") - import wandb - # Make one log on every process with the configuration for debugging. logging.basicConfig( format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", @@ -1196,10 +1235,6 @@ def main(args): if accelerator.is_main_process: if args.validation_prompt is not None and epoch % args.validation_epochs == 0: - logger.info( - f"Running validation... \n Generating {args.num_validation_images} images with prompt:" - f" {args.validation_prompt}." - ) # create pipeline pipeline = StableDiffusionXLPipeline.from_pretrained( args.pretrained_model_name_or_path, @@ -1212,36 +1247,7 @@ def main(args): torch_dtype=weight_dtype, ) - pipeline = pipeline.to(accelerator.device) - pipeline.set_progress_bar_config(disable=True) - - # run inference - generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) if args.seed else None - pipeline_args = {"prompt": args.validation_prompt} - if torch.backends.mps.is_available(): - autocast_ctx = nullcontext() - else: - autocast_ctx = torch.autocast(accelerator.device.type) - - with autocast_ctx: - images = [ - pipeline(**pipeline_args, generator=generator).images[0] - for _ in range(args.num_validation_images) - ] - - for tracker in accelerator.trackers: - if tracker.name == "tensorboard": - np_images = np.stack([np.asarray(img) for img in images]) - tracker.writer.add_images("validation", np_images, epoch, dataformats="NHWC") - if tracker.name == "wandb": - tracker.log( - { - "validation": [ - wandb.Image(image, caption=f"{i}: {args.validation_prompt}") - for i, image in enumerate(images) - ] - } - ) + images = log_validation(pipeline, args, accelerator, epoch) del pipeline torch.cuda.empty_cache() @@ -1288,33 +1294,13 @@ def main(args): variant=args.variant, torch_dtype=weight_dtype, ) - pipeline = pipeline.to(accelerator.device) # load attention processors pipeline.load_lora_weights(args.output_dir) # run inference - images = [] if args.validation_prompt and args.num_validation_images > 0: - generator = torch.Generator(device=accelerator.device).manual_seed(args.seed) if args.seed else None - images = [ - pipeline(args.validation_prompt, num_inference_steps=25, generator=generator).images[0] - for _ in range(args.num_validation_images) - ] - - for tracker in accelerator.trackers: - if tracker.name == "tensorboard": - np_images = np.stack([np.asarray(img) for img in images]) - tracker.writer.add_images("test", np_images, epoch, dataformats="NHWC") - if tracker.name == "wandb": - tracker.log( - { - "test": [ - wandb.Image(image, caption=f"{i}: {args.validation_prompt}") - for i, image in enumerate(images) - ] - } - ) + images = log_validation(pipeline, args, accelerator, epoch, is_final_validation=True) if args.push_to_hub: save_model_card(