Clean FID eval for SDXL

diffusion/evaluation/clean_fid_eval.py

@@ -40,6 +40,7 @@ class CleanFIDEvaluator:
             Default: ``[1.0]``.
         size (int): The size of the images to generate. Default: ``256``.
         batch_size (int): The per-device batch size to use for evaluation. Default: ``16``.
+        load_strict_model_weights (bool): Whether or not to strict load model weights. Default: ``True``.
         loggers (List[LoggerDestination], optional): The loggers to use for logging results. Default: ``None``.
         seed (int): The seed to use for evaluation. Default: ``17``.
         output_dir (str): The directory to save results to. Default: ``/tmp/``.
@@ -62,6 +63,7 @@ def __init__(self,
                  batch_size: int = 16,
                  image_key: str = 'image',
                  caption_key: str = 'caption',
+                 load_strict_model_weights: bool = True,
                  loggers: Optional[List[LoggerDestination]] = None,
                  seed: int = 17,
                  output_dir: str = '/tmp/',
@@ -78,12 +80,14 @@ def __init__(self,
         self.batch_size = batch_size
         self.image_key = image_key
         self.caption_key = caption_key
+        self.load_strict_model_weights = load_strict_model_weights
         self.loggers = loggers
         self.seed = seed
         self.output_dir = output_dir
         self.num_samples = num_samples if num_samples is not None else float('inf')
         self.precision = precision
         self.prompts = prompts if prompts is not None else ['A shiba inu wearing a blue sweater']
+        self.sdxl = model.sdxl
 
         # Init loggers
         if self.loggers and dist.get_local_rank() == 0:
@@ -95,6 +99,7 @@ def __init__(self,
         Trainer(model=self.model,
                 load_path=self.load_path,
                 load_weights_only=True,
+                load_strict_model_weights=self.load_strict_model_weights,
                 eval_dataloader=self.eval_dataloader,
                 seed=self.seed)
 
@@ -134,6 +139,13 @@ def _generate_images(self, guidance_scale: float):
 
             real_images = batch[self.image_key]
             captions = batch[self.caption_key]
+            if self.sdxl:
+                crop_params = batch['cond_crops_coords_top_left']
+                input_size_params = batch['cond_original_size']
+            else:
+                crop_params = None
+                input_size_params = None
+
             # Ensure a new seed for each batch, as randomness in model.generate is fixed.
             seed = starting_seed + batch_id
             # Generate images from the captions
@@ -143,9 +155,15 @@ def _generate_images(self, guidance_scale: float):
                                                        width=self.size,
                                                        guidance_scale=guidance_scale,
                                                        seed=seed,
+                                                       crop_params=crop_params,
+                                                       input_size_params=input_size_params,
                                                        progress_bar=False)  # type: ignore
             # Get the prompts from the tokens
-            text_captions = self.tokenizer.batch_decode(captions, skip_special_tokens=True)
+            if self.sdxl:
+                # Decode with first tokenizer
+                text_captions = self.tokenizer.tokenizer.batch_decode(captions[:, 0, :], skip_special_tokens=True)
+            else:
+                text_captions = self.tokenizer.batch_decode(captions, skip_special_tokens=True)
             self.clip_metric.update((generated_images * 255).to(torch.uint8), text_captions)
             # Save the real images
             # Verify that the real images are in the proper range

diffusion/models/stable_diffusion.py

@@ -3,7 +3,7 @@
 
 """Diffusion models."""
 
-from typing import List, Optional
+from typing import List, Optional, Tuple, Union
 
 import torch
 import torch.nn.functional as F
@@ -313,7 +313,7 @@ def update_metric(self, batch, outputs, metric):
             if self.sdxl:
                 # Decode captions with first tokenizer
                 captions = [
-                    self.tokenizer.tokenizer.decode(caption[0], skip_special_tokens=True)
+                    self.tokenizer.tokenizer.decode(caption[:, 0, :][0], skip_special_tokens=True)
                     for caption in batch[self.text_key]
                 ]
             else:
@@ -342,8 +342,8 @@ def generate(
         seed: Optional[int] = None,
         progress_bar: Optional[bool] = True,
         zero_out_negative_prompt: bool = True,
-        crop_params: Optional[list] = None,
-        size_params: Optional[list] = None,
+        crop_params: Optional[Union[Tuple[int, int], torch.Tensor]] = None,
+        input_size_params: Optional[Union[Tuple[int, int], torch.Tensor]] = None,
     ):
         """Generates image from noise.
 
@@ -388,9 +388,10 @@ def generate(
                 Default: `None`.
             zero_out_negative_prompt (bool): Whether or not to zero out negative prompt if it is
                 an empty string. Default: `True`.
-            crop_params (list, optional): Crop parameters to use when generating images with SDXL.
-                Default: `None`.
-            size_params (list, optional): Size parameters to use when generating images with SDXL.
+            crop_params (tuple or torch.FloatTensor of size [Bx2], optional): Crop parameters to use
+                when generating images with SDXL. Default: `None`.
+            input_size_params (tuple or torch.FloatTensor of size [Bx2], optional): Size parameters
+                (representing original size of input image) to use when generating images with SDXL.
                 Default: `None`.
         """
         _check_prompt_given(prompt, tokenized_prompts, prompt_embeds)
@@ -433,6 +434,9 @@ def generate(
             text_embeddings = torch.cat([unconditional_embeddings, text_embeddings])
             if self.sdxl:
                 pooled_embeddings = torch.cat([pooled_unconditional_embeddings, pooled_text_embeddings])  # type: ignore
+        else:
+            if self.sdxl:
+                pooled_embeddings = pooled_text_embeddings
 
         # prepare for diffusion generation process
         latents = torch.randn(
@@ -448,14 +452,24 @@ def generate(
         added_cond_kwargs = {}
         # if using SDXL, prepare added time ids & embeddings
         if self.sdxl and pooled_embeddings is not None:
-            if not crop_params:
-                crop_params = [0., 0.]
-            if not size_params:
-                size_params = [width, height]
-            add_time_ids = torch.tensor([[width, height, *crop_params, *size_params]], dtype=torch.float, device=device)
-            add_time_ids = add_time_ids.repeat(pooled_embeddings.shape[0], 1)
-            add_text_embeds = pooled_embeddings
+            if crop_params is None:
+                crop_params = torch.tensor([0., 0.]).repeat(batch_size, 1)
+            elif isinstance(crop_params, tuple):  # convert to tensor
+                crop_params = torch.tensor([crop_params[0], crop_params[1]], dtype=torch.float32).repeat(batch_size, 1)
+            if input_size_params is None:
+                input_size_params = torch.tensor([width, height], dtype=torch.float32).repeat(batch_size, 1)
+            elif isinstance(input_size_params, tuple):  # convert to tensor
+                input_size_params = torch.tensor([input_size_params[0], input_size_params[1]],
+                                                 dtype=torch.float32).repeat(batch_size, 1)
+            output_size_params = torch.tensor([width, height], dtype=torch.float32).repeat(batch_size, 1)
 
+            if do_classifier_free_guidance:
+                crop_params = torch.cat([crop_params, crop_params])
+                input_size_params = torch.cat([input_size_params, input_size_params])
+                output_size_params = torch.cat([output_size_params, output_size_params])
+
+            add_time_ids = torch.cat([input_size_params, crop_params, output_size_params], dim=1).to(device)
+            add_text_embeds = pooled_embeddings
             added_cond_kwargs = {'text_embeds': add_text_embeds, 'time_ids': add_time_ids}
 
         # backward diffusion process