4 files changed, 226 insertions, 115 deletions

diff --git a/dreambooth.py b/dreambooth.py
index 7b34fce..79b3d2c 100644
--- a/dreambooth.py
+++ b/dreambooth.py
@@ -489,7 +489,7 @@ class Checkpointer:
             generator=generator,
         )
 
-        with torch.inference_mode():
+        with torch.autocast("cuda"), torch.inference_mode():
             for pool, data, latents in [("stable", val_data, stable_latents), ("val", val_data, None), ("train", train_data, None)]:
                 all_samples = []
                 file_path = samples_path.joinpath(pool, f"step_{step}.jpg")
@@ -522,7 +522,7 @@ class Checkpointer:
                         negative_prompt=nprompt,
                         height=self.sample_image_size,
                         width=self.sample_image_size,
-                        latents=latents[:len(prompt)] if latents is not None else None,
+                        latents_or_image=latents[:len(prompt)] if latents is not None else None,
                         generator=generator if latents is not None else None,
                         guidance_scale=guidance_scale,
                         eta=eta,
@@ -768,7 +768,7 @@ def main():
             ).to(accelerator.device)
             pipeline.set_progress_bar_config(dynamic_ncols=True)
 
-            with torch.inference_mode():
+            with torch.autocast("cuda"), torch.inference_mode():
                 for batch in batched_data:
                     image_name = [item.class_image_path for item in batch]
                     prompt = [item.prompt.format(identifier=args.class_identifier) for item in batch]
@@ -1018,7 +1018,7 @@ def main():
             text_encoder.eval()
             val_loss = 0.0
 
-            with torch.inference_mode():
+            with torch.autocast("cuda"), torch.inference_mode():
                 for step, batch in enumerate(val_dataloader):
                     latents = vae.encode(batch["pixel_values"]).latent_dist.sample()
                     latents = latents * 0.18215
diff --git a/infer.py b/infer.py
index 9b0ec1f..2bf9cb3 100644
--- a/infer.py
+++ b/infer.py
@@ -273,14 +273,15 @@ def generate(output_dir, pipeline, args):
 
             generator = torch.Generator(device="cuda").manual_seed(args.seed + i)
             images = pipeline(
-                prompt=args.prompt * (args.batch_size // len(args.prompt)),
+                prompt=args.prompt,
+                negative_prompt=args.negative_prompt,
                 height=args.height,
                 width=args.width,
-                negative_prompt=args.negative_prompt,
+                num_images_per_prompt=args.batch_size,
                 num_inference_steps=args.steps,
                 guidance_scale=args.guidance_scale,
                 generator=generator,
-                latents=init_image,
+                latents_or_image=init_image,
                 strength=args.image_noise,
             ).images
 
diff --git a/pipelines/stable_diffusion/vlpn_stable_diffusion.py b/pipelines/stable_diffusion/vlpn_stable_diffusion.py
index ba057ba..d6b1cb1 100644
--- a/pipelines/stable_diffusion/vlpn_stable_diffusion.py
+++ b/pipelines/stable_diffusion/vlpn_stable_diffusion.py
@@ -1,6 +1,6 @@
 import inspect
 import warnings
-from typing import List, Optional, Union
+from typing import List, Optional, Union, Callable
 
 import numpy as np
 import torch
@@ -136,11 +136,165 @@ class VlpnStableDiffusion(DiffusionPipeline):
             if cpu_offloaded_model is not None:
                 cpu_offload(cpu_offloaded_model, device)
 
+    @property
+    def execution_device(self):
+        r"""
+        Returns the device on which the pipeline's models will be executed. After calling
+        `pipeline.enable_sequential_cpu_offload()` the execution device can only be inferred from Accelerate's module
+        hooks.
+        """
+        if self.device != torch.device("meta") or not hasattr(self.unet, "_hf_hook"):
+            return self.device
+        for module in self.unet.modules():
+            if (
+                hasattr(module, "_hf_hook")
+                and hasattr(module._hf_hook, "execution_device")
+                and module._hf_hook.execution_device is not None
+            ):
+                return torch.device(module._hf_hook.execution_device)
+        return self.device
+
+    def check_inputs(self, prompt, negative_prompt, width, height, strength, callback_steps):
+        if isinstance(prompt, str):
+            prompt = [prompt]
+
+        if negative_prompt is None:
+            negative_prompt = ""
+
+        if isinstance(negative_prompt, str):
+            negative_prompt = [negative_prompt] * len(prompt)
+
+        if not isinstance(prompt, list):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+        if not isinstance(negative_prompt, list):
+            raise ValueError(f"`negative_prompt` has to be of type `str` or `list` but is {type(negative_prompt)}")
+
+        if len(negative_prompt) != len(prompt):
+            raise ValueError(
+                f"`prompt` and `negative_prompt` have to be the same length, but are {len(prompt)} and {len(negative_prompt)}")
+
+        if strength < 0 or strength > 1:
+            raise ValueError(f"`strength` should in [0.0, 1.0] but is {strength}")
+
+        if height % 8 != 0 or width % 8 != 0:
+            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
+
+        if (callback_steps is None) or (
+            callback_steps is not None and (not isinstance(callback_steps, int) or callback_steps <= 0)
+        ):
+            raise ValueError(
+                f"`callback_steps` has to be a positive integer but is {callback_steps} of type"
+                f" {type(callback_steps)}."
+            )
+
+        return prompt, negative_prompt
+
+    def encode_prompt(self, prompt, negative_prompt, num_images_per_prompt, do_classifier_free_guidance):
+        text_input_ids = self.prompt_processor.get_input_ids(prompt)
+        text_input_ids *= num_images_per_prompt
+
+        if do_classifier_free_guidance:
+            unconditional_input_ids = self.prompt_processor.get_input_ids(negative_prompt)
+            unconditional_input_ids *= num_images_per_prompt
+            text_input_ids = unconditional_input_ids + text_input_ids
+
+        text_input_ids = self.prompt_processor.unify_input_ids(text_input_ids)
+        text_embeddings = self.prompt_processor.get_embeddings(text_input_ids)
+
+        return text_embeddings
+
+    def get_timesteps(self, latents_are_image, num_inference_steps, strength, device):
+        if latents_are_image:
+            # get the original timestep using init_timestep
+            offset = self.scheduler.config.get("steps_offset", 0)
+            init_timestep = int(num_inference_steps * strength) + offset
+            init_timestep = min(init_timestep, num_inference_steps)
+
+            t_start = max(num_inference_steps - init_timestep + offset, 0)
+            timesteps = self.scheduler.timesteps[t_start:]
+        else:
+            timesteps = self.scheduler.timesteps
+
+        timesteps = timesteps.to(device)
+
+        return timesteps
+
+    def prepare_latents(self, batch_size, num_images_per_prompt, num_channels_latents, height, width, dtype, device, generator, latents=None):
+        shape = (batch_size * num_images_per_prompt, num_channels_latents, height // 8, width // 8)
+
+        if latents is None:
+            if device.type == "mps":
+                # randn does not work reproducibly on mps
+                latents = torch.randn(shape, generator=generator, device="cpu", dtype=dtype).to(device)
+            else:
+                latents = torch.randn(shape, generator=generator, device=device, dtype=dtype)
+        else:
+            if latents.shape != shape:
+                raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")
+            latents = latents.to(device)
+
+        # scale the initial noise by the standard deviation required by the scheduler
+        latents = latents * self.scheduler.init_noise_sigma
+
+        return latents
+
+    def prepare_latents_from_image(self, init_image, timestep, batch_size, num_images_per_prompt, dtype, device, generator=None):
+        init_image = init_image.to(device=device, dtype=dtype)
+        init_latent_dist = self.vae.encode(init_image).latent_dist
+        init_latents = init_latent_dist.sample(generator=generator)
+        init_latents = 0.18215 * init_latents
+
+        if batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] == 0:
+            additional_image_per_prompt = batch_size // init_latents.shape[0]
+            init_latents = torch.cat([init_latents] * additional_image_per_prompt * num_images_per_prompt, dim=0)
+        elif batch_size > init_latents.shape[0] and batch_size % init_latents.shape[0] != 0:
+            raise ValueError(
+                f"Cannot duplicate `init_image` of batch size {init_latents.shape[0]} to {batch_size} text prompts."
+            )
+        else:
+            init_latents = torch.cat([init_latents] * num_images_per_prompt, dim=0)
+
+        # add noise to latents using the timesteps
+        noise = torch.randn(init_latents.shape, generator=generator, device=device, dtype=dtype)
+
+        # get latents
+        init_latents = self.scheduler.add_noise(init_latents, noise, timestep)
+        latents = init_latents
+
+        return latents
+
+    def prepare_extra_step_kwargs(self, generator, eta):
+        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
+        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
+        # and should be between [0, 1]
+
+        accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        extra_step_kwargs = {}
+        if accepts_eta:
+            extra_step_kwargs["eta"] = eta
+
+        # check if the scheduler accepts generator
+        accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
+        if accepts_generator:
+            extra_step_kwargs["generator"] = generator
+        return extra_step_kwargs
+
+    def decode_latents(self, latents):
+        latents = 1 / 0.18215 * latents
+        image = self.vae.decode(latents).sample
+        image = (image / 2 + 0.5).clamp(0, 1)
+        # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
+        image = image.cpu().permute(0, 2, 3, 1).float().numpy()
+        return image
+
     @torch.no_grad()
     def __call__(
         self,
         prompt: Union[str, List[str], List[List[str]]],
         negative_prompt: Optional[Union[str, List[str], List[List[str]]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
         strength: float = 0.8,
         height: Optional[int] = 512,
         width: Optional[int] = 512,
@@ -148,9 +302,11 @@ class VlpnStableDiffusion(DiffusionPipeline):
         guidance_scale: Optional[float] = 7.5,
         eta: Optional[float] = 0.0,
         generator: Optional[torch.Generator] = None,
-        latents: Optional[Union[torch.FloatTensor, PIL.Image.Image]] = None,
+        latents_or_image: Optional[Union[torch.FloatTensor, PIL.Image.Image]] = None,
         output_type: Optional[str] = "pil",
         return_dict: bool = True,
+        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+        callback_steps: Optional[int] = 1,
     ):
         r"""
         Function invoked when calling the pipeline for generation.
@@ -202,110 +358,60 @@ class VlpnStableDiffusion(DiffusionPipeline):
             (nsfw) content, according to the `safety_checker`.
         """
 
-        if isinstance(prompt, str):
-            prompt = [prompt]
+        # 1. Check inputs. Raise error if not correct
+        prompt, negative_prompt = self.check_inputs(prompt, negative_prompt, width, height, strength, callback_steps)
 
+        # 2. Define call parameters
         batch_size = len(prompt)
-
-        if negative_prompt is None:
-            negative_prompt = ""
-
-        if isinstance(negative_prompt, str):
-            negative_prompt = [negative_prompt] * batch_size
-
-        if len(negative_prompt) != len(prompt):
-            raise ValueError(
-                f"`prompt` and `negative_prompt` have to be the same length, but are {len(prompt)} and {len(negative_prompt)}")
-
-        if height % 8 != 0 or width % 8 != 0:
-            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
-
-        if strength < 0 or strength > 1:
-            raise ValueError(f"`strength` should in [0.0, 1.0] but is {strength}")
-
-        # set timesteps
-        self.scheduler.set_timesteps(num_inference_steps)
-
-        # get prompt text embeddings
-        text_input_ids = self.prompt_processor.get_input_ids(prompt)
-
-        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
-        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
-        # corresponds to doing no classifier free guidance.
+        device = self.execution_device
         do_classifier_free_guidance = guidance_scale > 1.0
-        # get unconditional embeddings for classifier free guidance
-        if do_classifier_free_guidance:
-            unconditional_input_ids = self.prompt_processor.get_input_ids(negative_prompt)
-            text_input_ids = unconditional_input_ids + text_input_ids
-
-        text_input_ids = self.prompt_processor.unify_input_ids(text_input_ids)
-        text_embeddings = self.prompt_processor.get_embeddings(text_input_ids)
-
-        offset = self.scheduler.config.get("steps_offset", 0)
-        init_timestep = num_inference_steps + offset
-
-        # get the initial random noise unless the user supplied it
-
-        # Unlike in other pipelines, latents need to be generated in the target device
-        # for 1-to-1 results reproducibility with the CompVis implementation.
-        # However this currently doesn't work in `mps`.
-        latents_dtype = text_embeddings.dtype
-        latents_shape = (batch_size, self.unet.in_channels, height // 8, width // 8)
+        latents_are_image = isinstance(latents_or_image, PIL.Image.Image)
 
-        if latents is None:
-            if self.device.type == "mps":
-                # randn does not exist on mps
-                latents = torch.randn(latents_shape, generator=generator, device="cpu", dtype=latents_dtype).to(
-                    self.device
-                )
-            else:
-                latents = torch.randn(latents_shape, generator=generator, device=self.device, dtype=latents_dtype)
-        elif isinstance(latents, PIL.Image.Image):
-            latents = preprocess(latents, width, height)
-            latents = latents.to(device=self.device, dtype=latents_dtype)
-            latent_dist = self.vae.encode(latents).latent_dist
-            latents = latent_dist.sample(generator=generator)
-            latents = 0.18215 * latents
-
-            # expand init_latents for batch_size
-            latents = torch.cat([latents] * batch_size, dim=0)
-
-            # get the original timestep using init_timestep
-            init_timestep = int(num_inference_steps * strength) + offset
-            init_timestep = min(init_timestep, num_inference_steps)
+        print(f">>> {device}")
 
-            timesteps = self.scheduler.timesteps[-init_timestep]
-            timesteps = torch.tensor([timesteps] * batch_size, device=self.device)
+        # 3. Encode input prompt
+        text_embeddings = self.encode_prompt(
+            prompt,
+            negative_prompt,
+            num_images_per_prompt,
+            do_classifier_free_guidance
+        )
 
-            # add noise to latents using the timesteps
-            noise = torch.randn(latents.shape, generator=generator, device=self.device, dtype=latents_dtype)
-            latents = self.scheduler.add_noise(latents, noise, timesteps)
+        # 4. Prepare timesteps
+        self.scheduler.set_timesteps(num_inference_steps, device=device)
+        timesteps = self.get_timesteps(latents_are_image, num_inference_steps, strength, device)
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.unet.in_channels
+        if latents_are_image:
+            latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)
+            latents = self.prepare_latents_from_image(
+                latents_or_image,
+                latent_timestep,
+                batch_size,
+                num_images_per_prompt,
+                text_embeddings.dtype,
+                device,
+                generator
+            )
         else:
-            if latents.shape != latents_shape:
-                raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}")
-            if latents.device != self.device:
-                raise ValueError(f"Unexpected latents device, got {latents.device}, expected {self.device}")
-
-        t_start = max(num_inference_steps - init_timestep + offset, 0)
-
-        # Some schedulers like PNDM have timesteps as arrays
-        # It's more optimzed to move all timesteps to correct device beforehand
-        timesteps_tensor = torch.tensor(self.scheduler.timesteps[t_start:], device=self.device)
+            latents = self.prepare_latents(
+                batch_size,
+                num_images_per_prompt,
+                num_channels_latents,
+                height,
+                width,
+                text_embeddings.dtype,
+                device,
+                generator,
+                latents_or_image,
+            )
 
-        # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
-        # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
-        # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
-        # and should be between [0, 1]
-        scheduler_step_args = set(inspect.signature(self.scheduler.step).parameters.keys())
-        accepts_eta = "eta" in scheduler_step_args
-        extra_step_kwargs = {}
-        if accepts_eta:
-            extra_step_kwargs["eta"] = eta
-        accepts_generator = "generator" in scheduler_step_args
-        if generator is not None and accepts_generator:
-            extra_step_kwargs["generator"] = generator
+        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
 
-        for i, t in enumerate(self.progress_bar(timesteps_tensor)):
+        # 7. Denoising loop
+        for i, t in enumerate(self.progress_bar(timesteps)):
             # expand the latents if we are doing classifier free guidance
             latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
             latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
@@ -321,17 +427,21 @@ class VlpnStableDiffusion(DiffusionPipeline):
             # compute the previous noisy sample x_t -> x_t-1
             latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs).prev_sample
 
-        # scale and decode the image latents with vae
-        latents = 1 / 0.18215 * latents
-        image = self.vae.decode(latents.to(dtype=self.vae.dtype)).sample
+            # call the callback, if provided
+            if callback is not None and i % callback_steps == 0:
+                callback(i, t, latents)
 
-        image = (image / 2 + 0.5).clamp(0, 1)
-        image = image.cpu().permute(0, 2, 3, 1).float().numpy()
+        # 8. Post-processing
+        image = self.decode_latents(latents)
+
+        # 9. Run safety checker
+        has_nsfw_concept = None
 
+        # 10. Convert to PIL
         if output_type == "pil":
             image = self.numpy_to_pil(image)
 
         if not return_dict:
-            return (image, None)
+            return (image, has_nsfw_concept)
 
-        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=None)
+        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
diff --git a/textual_inversion.py b/textual_inversion.py
index 999161b..bf591bc 100644
--- a/textual_inversion.py
+++ b/textual_inversion.py
@@ -437,7 +437,7 @@ class Checkpointer:
             generator=generator,
         )
 
-        with torch.inference_mode():
+        with torch.autocast("cuda"), torch.inference_mode():
             for pool, data, latents in [("stable", val_data, stable_latents), ("val", val_data, None), ("train", train_data, None)]:
                 all_samples = []
                 file_path = samples_path.joinpath(pool, f"step_{step}.png")
@@ -456,7 +456,7 @@ class Checkpointer:
                         negative_prompt=nprompt,
                         height=self.sample_image_size,
                         width=self.sample_image_size,
-                        latents=latents[:len(prompt)] if latents is not None else None,
+                        latents_or_image=latents[:len(prompt)] if latents is not None else None,
                         generator=generator if latents is not None else None,
                         guidance_scale=guidance_scale,
                         eta=eta,
@@ -670,7 +670,7 @@ def main():
             ).to(accelerator.device)
             pipeline.set_progress_bar_config(dynamic_ncols=True)
 
-            with torch.inference_mode():
+            with torch.autocast("cuda"), torch.inference_mode():
                 for batch in batched_data:
                     image_name = [p.class_image_path for p in batch]
                     prompt = [p.prompt.format(identifier=args.class_identifier) for p in batch]
@@ -908,7 +908,7 @@ def main():
             text_encoder.eval()
             val_loss = 0.0
 
-            with torch.inference_mode():
+            with torch.autocast("cuda"), torch.inference_mode():
                 for step, batch in enumerate(val_dataloader):
                     latents = vae.encode(batch["pixel_values"]).latent_dist.sample()
                     latents = latents * 0.18215