Adding logging images to textual inversion. Issue 1837

examples/textual_inversion/.gitignore

@@ -0,0 +1,3 @@
+frida
+wandb
+textual_inversion_frida

examples/textual_inversion/textual_inversion.py

@@ -14,6 +14,7 @@
 # See the License for the specific language governing permissions and
 
 import argparse
+import importlib
 import logging
 import math
 import os
@@ -34,7 +35,14 @@
 from accelerate import Accelerator
 from accelerate.logging import get_logger
 from accelerate.utils import set_seed
-from diffusers import AutoencoderKL, DDPMScheduler, StableDiffusionPipeline, UNet2DConditionModel
+from diffusers import (
+    AutoencoderKL,
+    DDPMScheduler,
+    DiffusionPipeline,
+    DPMSolverMultistepScheduler,
+    StableDiffusionPipeline,
+    UNet2DConditionModel,
+)
 from diffusers.optimization import get_scheduler
 from diffusers.utils import check_min_version
 from diffusers.utils.import_utils import is_xformers_available
@@ -48,6 +56,28 @@
 from transformers import CLIPTextModel, CLIPTokenizer
 
 
+_wandb_available = importlib.util.find_spec("wandb") is not None
+
+
+def is_wandb_available():
+    return _wandb_available
+
+
+wandb = None
+if is_wandb_available():
+    import wandb
+
+
+def wandb_setup(
+    args: dict,
+    project_name: str = "glide-text2im-finetune",
+):
+    return wandb.init(
+        project=project_name,
+        config=args,
+    )
+
+
 if version.parse(version.parse(PIL.__version__).base_version) >= version.parse("9.1.0"):
     PIL_INTERPOLATION = {
         "linear": PIL.Image.Resampling.BILINEAR,
@@ -83,6 +113,42 @@ def save_progress(text_encoder, placeholder_token_id, accelerator, args, save_pa
 
 def parse_args():
     parser = argparse.ArgumentParser(description="Simple example of a training script.")
+    parser.add_argument(
+        "--num_inference_steps",
+        type=int,
+        default=25,
+        help="Number of steps for inference.",
+    )
+    parser.add_argument(
+        "--validation_prompt",
+        type=str,
+        default="A photo of <frida>",
+        help="Prompt to use for validation.",
+    )
+    parser.add_argument(
+        "--num_validation_images",
+        type=int,
+        default=1,
+        help="Number of images to generate for validation.",
+    )
+    parser.add_argument(
+        "--guidance_scale",
+        type=float,
+        default=7.5,
+        help="Guidance scale for inference.",
+    )
+    parser.add_argument(
+        "--project_name",
+        type=str,
+        default="huggingface_textual_inv",
+        help="Name of wandb run",
+    )
+    parser.add_argument(
+        "--log_frequency",
+        type=int,
+        default=100,
+        help="Frequency to log/save the model.",
+    )
     parser.add_argument(
         "--save_steps",
         type=int,
@@ -404,10 +470,7 @@ def __getitem__(self, i):
 
         if self.center_crop:
             crop = min(img.shape[0], img.shape[1])
-            (
-                h,
-                w,
-            ) = (
+            (h, w,) = (
                 img.shape[0],
                 img.shape[1],
             )
@@ -437,14 +500,12 @@ def get_full_repo_name(model_id: str, organization: Optional[str] = None, token:
 def main():
     args = parse_args()
     logging_dir = os.path.join(args.output_dir, args.logging_dir)
-
     accelerator = Accelerator(
         gradient_accumulation_steps=args.gradient_accumulation_steps,
         mixed_precision=args.mixed_precision,
         log_with=args.report_to,
         logging_dir=logging_dir,
     )
-
     # Make one log on every process with the configuration for debugging.
     logging.basicConfig(
         format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
@@ -724,6 +785,45 @@ def main():
             if accelerator.sync_gradients:
                 progress_bar.update(1)
                 global_step += 1
+                if args.validation_prompt is not None and global_step % args.log_frequency == 0:
+                    logger.info(
+                        f"Running validation... \n Generating {args.num_validation_images} images with prompt:"
+                        f" {args.validation_prompt}."
+                    )
+                    # create pipeline
+                    pipeline = DiffusionPipeline.from_pretrained(
+                        args.pretrained_model_name_or_path,
+                        unet=unet,
+                        vae=vae,
+                        text_encoder=accelerator.unwrap_model(text_encoder),
+                        tokenizer=tokenizer,
+                        revision=args.revision,
+                        torch_dtype=weight_dtype,
+                    )
+                    pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
+                    pipeline = pipeline.to(accelerator.device)
+                    pipeline.set_progress_bar_config(disable=True)
+                    # run inference
+                    generator = torch.Generator(device=accelerator.device).manual_seed(args.seed or 0)
+                    prompt = args.num_validation_images * [args.validation_prompt]
+                    images = pipeline(
+                        prompt,
+                        num_inference_steps=args.num_inference_steps,
+                        generator=generator,
+                        guidance_scale=args.guidance_scale,
+                    ).images
+                    for tracker in accelerator.trackers:
+                        if tracker.name == "wandb":
+                            tracker.log(
+                                {
+                                    "validation": [
+                                        wandb.Image(image, caption=f"{i}: {args.validation_prompt}")
+                                        for i, image in enumerate(images)
+                                    ]
+                                }
+                            )
+                    del pipeline
+                    torch.cuda.empty_cache()
                 if global_step % args.save_steps == 0:
                     save_path = os.path.join(args.output_dir, f"learned_embeds-steps-{global_step}.bin")
                     save_progress(text_encoder, placeholder_token_id, accelerator, args, save_path)
@@ -752,9 +852,9 @@ def main():
         if save_full_model:
             pipeline = StableDiffusionPipeline.from_pretrained(
                 args.pretrained_model_name_or_path,
-                text_encoder=accelerator.unwrap_model(text_encoder),
-                vae=vae,
                 unet=unet,
+                vae=vae,
+                text_encoder=accelerator.unwrap_model(text_encoder),
                 tokenizer=tokenizer,
             )
             pipeline.save_pretrained(args.output_dir)

examples/textual_inversion/textual_inversion_flax.py

@@ -306,10 +306,7 @@ def __getitem__(self, i):
 
         if self.center_crop:
             crop = min(img.shape[0], img.shape[1])
-            (
-                h,
-                w,
-            ) = (
+            (h, w,) = (
                 img.shape[0],
                 img.shape[1],
             )

src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py

@@ -482,7 +482,7 @@ def __call__(
         # 3. Encode input prompt
         text_embeddings = self._encode_prompt(
             prompt, device, num_images_per_prompt, do_classifier_free_guidance, negative_prompt
-        )
+        ).to(dtype=self.vae.dtype)
 
         # 4. Prepare timesteps
         self.scheduler.set_timesteps(num_inference_steps, device=device)