3 files changed, 113 insertions, 216 deletions

diff --git a/infer.py b/infer.py
index de3d792..40720ea 100644
--- a/infer.py
+++ b/infer.py
@@ -1,18 +1,21 @@
 import argparse
 import datetime
 import logging
+import sys
+import shlex
+import cmd
 from pathlib import Path
 from torch import autocast
 import torch
 import json
 from diffusers import StableDiffusionPipeline, AutoencoderKL, UNet2DConditionModel, PNDMScheduler, DDIMScheduler, LMSDiscreteScheduler
-from transformers import CLIPModel, CLIPTextModel, CLIPTokenizer, CLIPFeatureExtractor
+from transformers import CLIPTextModel, CLIPTokenizer, CLIPFeatureExtractor
 from slugify import slugify
 from pipelines.stable_diffusion.clip_guided_stable_diffusion import CLIPGuidedStableDiffusion
 from schedulers.scheduling_euler_a import EulerAScheduler
 
 
-def parse_args():
+def create_args_parser():
     parser = argparse.ArgumentParser(
         description="Simple example of a training script."
     )
@@ -22,6 +25,30 @@ def parse_args():
         default=None,
     )
     parser.add_argument(
+        "--scheduler",
+        type=str,
+        choices=["plms", "ddim", "klms", "euler_a"],
+        default="euler_a",
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="output/inference",
+    )
+    parser.add_argument(
+        "--config",
+        type=str,
+        default=None,
+    )
+
+    return parser
+
+
+def create_cmd_parser():
+    parser = argparse.ArgumentParser(
+        description="Simple example of a training script."
+    )
+    parser.add_argument(
         "--prompt",
         type=str,
         default=None,
@@ -49,28 +76,17 @@ def parse_args():
     parser.add_argument(
         "--batch_num",
         type=int,
-        default=50,
+        default=1,
     )
     parser.add_argument(
         "--steps",
         type=int,
-        default=120,
-    )
-    parser.add_argument(
-        "--scheduler",
-        type=str,
-        choices=["plms", "ddim", "klms", "euler_a"],
-        default="euler_a",
+        default=70,
     )
     parser.add_argument(
         "--guidance_scale",
         type=int,
-        default=7.5,
-    )
-    parser.add_argument(
-        "--clip_guidance_scale",
-        type=int,
-        default=100,
+        default=7,
     )
     parser.add_argument(
         "--seed",
@@ -78,21 +94,21 @@ def parse_args():
         default=torch.random.seed(),
     )
     parser.add_argument(
-        "--output_dir",
-        type=str,
-        default="output/inference",
-    )
-    parser.add_argument(
         "--config",
         type=str,
         default=None,
     )
 
-    args = parser.parse_args()
+    return parser
+
+
+def run_parser(parser, input=None):
+    args = parser.parse_known_args(input)[0]
+
     if args.config is not None:
         with open(args.config, 'rt') as f:
-            args = parser.parse_args(
-                namespace=argparse.Namespace(**json.load(f)["args"]))
+            args = parser.parse_known_args(
+                namespace=argparse.Namespace(**json.load(f)["args"]))[0]
 
     return args
 
@@ -104,24 +120,24 @@ def save_args(basepath, args, extra={}):
         json.dump(info, f, indent=4)
 
 
-def gen(args, output_dir):
-    tokenizer = CLIPTokenizer.from_pretrained(args.model + '/tokenizer', torch_dtype=torch.bfloat16)
-    text_encoder = CLIPTextModel.from_pretrained(args.model + '/text_encoder', torch_dtype=torch.bfloat16)
-    clip_model = CLIPModel.from_pretrained("laion/CLIP-ViT-B-32-laion2B-s34B-b79K", torch_dtype=torch.bfloat16)
-    vae = AutoencoderKL.from_pretrained(args.model + '/vae', torch_dtype=torch.bfloat16)
-    unet = UNet2DConditionModel.from_pretrained(args.model + '/unet', torch_dtype=torch.bfloat16)
-    feature_extractor = CLIPFeatureExtractor.from_pretrained(
-        "laion/CLIP-ViT-B-32-laion2B-s34B-b79K", torch_dtype=torch.bfloat16)
+def create_pipeline(model, scheduler, dtype=torch.bfloat16):
+    print("Loading Stable Diffusion pipeline...")
 
-    if args.scheduler == "plms":
+    tokenizer = CLIPTokenizer.from_pretrained(model + '/tokenizer', torch_dtype=dtype)
+    text_encoder = CLIPTextModel.from_pretrained(model + '/text_encoder', torch_dtype=dtype)
+    vae = AutoencoderKL.from_pretrained(model + '/vae', torch_dtype=dtype)
+    unet = UNet2DConditionModel.from_pretrained(model + '/unet', torch_dtype=dtype)
+    feature_extractor = CLIPFeatureExtractor.from_pretrained("openai/clip-vit-base-patch32", torch_dtype=dtype)
+
+    if scheduler == "plms":
         scheduler = PNDMScheduler(
             beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", skip_prk_steps=True
         )
-    elif args.scheduler == "klms":
+    elif scheduler == "klms":
         scheduler = LMSDiscreteScheduler(
             beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear"
         )
-    elif args.scheduler == "ddim":
+    elif scheduler == "ddim":
         scheduler = DDIMScheduler(
             beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False
         )
@@ -135,13 +151,24 @@ def gen(args, output_dir):
         vae=vae,
         unet=unet,
         tokenizer=tokenizer,
-        clip_model=clip_model,
         scheduler=scheduler,
         feature_extractor=feature_extractor
     )
     pipeline.enable_attention_slicing()
     pipeline.to("cuda")
 
+    print("Pipeline loaded.")
+
+    return pipeline
+
+
+def generate(output_dir, pipeline, args):
+    now = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S")
+    output_dir = output_dir.joinpath(f"{now}_{slugify(args.prompt)[:100]}")
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    save_args(output_dir, args)
+
     with autocast("cuda"):
         for i in range(args.batch_num):
             generator = torch.Generator(device="cuda").manual_seed(args.seed + i)
@@ -152,7 +179,6 @@ def gen(args, output_dir):
                 negative_prompt=args.negative_prompt,
                 num_inference_steps=args.steps,
                 guidance_scale=args.guidance_scale,
-                clip_guidance_scale=args.clip_guidance_scale,
                 generator=generator,
             ).images
 
@@ -160,18 +186,56 @@ def gen(args, output_dir):
                 image.save(output_dir.joinpath(f"{args.seed + i}_{j}.jpg"))
 
 
-def main():
-    args = parse_args()
+class CmdParse(cmd.Cmd):
+    prompt = 'dream> '
+    commands = []
 
-    now = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S")
-    output_dir = Path(args.output_dir).joinpath(f"{now}_{slugify(args.prompt)[:100]}")
-    output_dir.mkdir(parents=True, exist_ok=True)
+    def __init__(self, output_dir, pipeline, parser):
+        super().__init__()
 
-    save_args(output_dir, args)
+        self.output_dir = output_dir
+        self.pipeline = pipeline
+        self.parser = parser
+
+    def default(self, line):
+        line = line.replace("'", "\\'")
+
+        try:
+            elements = shlex.split(line)
+        except ValueError as e:
+            print(str(e))
+
+        if elements[0] == 'q':
+            return True
+
+        try:
+            args = run_parser(self.parser, elements)
+        except SystemExit:
+            self.parser.print_help()
+
+        if len(args.prompt) == 0:
+            print('Try again with a prompt!')
+
+        try:
+            generate(self.output_dir, self.pipeline, args)
+        except KeyboardInterrupt:
+            print('Generation cancelled.')
+
+    def do_exit(self, line):
+        return True
+
+
+def main():
+    logging.basicConfig(stream=sys.stdout, level=logging.WARN)
 
-    logging.basicConfig(filename=output_dir.joinpath("log.txt"), level=logging.DEBUG)
+    args_parser = create_args_parser()
+    args = run_parser(args_parser)
+    output_dir = Path(args.output_dir)
 
-    gen(args, output_dir)
+    pipeline = create_pipeline(args.model, args.scheduler)
+    cmd_parser = create_cmd_parser()
+    cmd_prompt = CmdParse(output_dir, pipeline, cmd_parser)
+    cmd_prompt.cmdloop()
 
 
 if __name__ == "__main__":
diff --git a/pipelines/stable_diffusion/clip_guided_stable_diffusion.py b/pipelines/stable_diffusion/clip_guided_stable_diffusion.py
index 306d9a9..ddf7ce1 100644
--- a/pipelines/stable_diffusion/clip_guided_stable_diffusion.py
+++ b/pipelines/stable_diffusion/clip_guided_stable_diffusion.py
@@ -17,53 +17,14 @@ from schedulers.scheduling_euler_a import EulerAScheduler, CFGDenoiserForward
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
 
-class MakeCutouts(nn.Module):
-    def __init__(self, cut_size, cut_power=1.0):
-        super().__init__()
-
-        self.cut_size = cut_size
-        self.cut_power = cut_power
-
-    def forward(self, pixel_values, num_cutouts):
-        sideY, sideX = pixel_values.shape[2:4]
-        max_size = min(sideX, sideY)
-        min_size = min(sideX, sideY, self.cut_size)
-        cutouts = []
-        for _ in range(num_cutouts):
-            size = int(torch.rand([]) ** self.cut_power * (max_size - min_size) + min_size)
-            offsetx = torch.randint(0, sideX - size + 1, ())
-            offsety = torch.randint(0, sideY - size + 1, ())
-            cutout = pixel_values[:, :, offsety: offsety + size, offsetx: offsetx + size]
-            cutouts.append(F.adaptive_avg_pool2d(cutout, self.cut_size))
-        return torch.cat(cutouts)
-
-
-def spherical_dist_loss(x, y):
-    x = F.normalize(x, dim=-1)
-    y = F.normalize(y, dim=-1)
-    return (x - y).norm(dim=-1).div(2).arcsin().pow(2).mul(2)
-
-
-def set_requires_grad(model, value):
-    for param in model.parameters():
-        param.requires_grad = value
-
-
 class CLIPGuidedStableDiffusion(DiffusionPipeline):
-    """CLIP guided stable diffusion based on the amazing repo by @crowsonkb and @Jack000
-    - https://github.com/Jack000/glid-3-xl
-    - https://github.dev/crowsonkb/k-diffusion
-    """
-
     def __init__(
         self,
         vae: AutoencoderKL,
         text_encoder: CLIPTextModel,
-        clip_model: CLIPModel,
         tokenizer: CLIPTokenizer,
         unet: UNet2DConditionModel,
-        scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler],
-        feature_extractor: CLIPFeatureExtractor,
+        scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler, EulerAScheduler],
         **kwargs,
     ):
         super().__init__()
@@ -85,19 +46,11 @@ class CLIPGuidedStableDiffusion(DiffusionPipeline):
         self.register_modules(
             vae=vae,
             text_encoder=text_encoder,
-            clip_model=clip_model,
             tokenizer=tokenizer,
             unet=unet,
             scheduler=scheduler,
-            feature_extractor=feature_extractor,
         )
 
-        self.normalize = transforms.Normalize(mean=feature_extractor.image_mean, std=feature_extractor.image_std)
-        self.make_cutouts = MakeCutouts(feature_extractor.size)
-
-        set_requires_grad(self.text_encoder, False)
-        set_requires_grad(self.clip_model, False)
-
     def enable_attention_slicing(self, slice_size: Optional[Union[str, int]] = "auto"):
         r"""
         Enable sliced attention computation.
@@ -125,87 +78,6 @@ class CLIPGuidedStableDiffusion(DiffusionPipeline):
         # set slice_size = `None` to disable `attention slicing`
         self.enable_attention_slicing(None)
 
-    def freeze_vae(self):
-        set_requires_grad(self.vae, False)
-
-    def unfreeze_vae(self):
-        set_requires_grad(self.vae, True)
-
-    def freeze_unet(self):
-        set_requires_grad(self.unet, False)
-
-    def unfreeze_unet(self):
-        set_requires_grad(self.unet, True)
-
-    @torch.enable_grad()
-    def cond_fn(
-        self,
-        latents,
-        timestep,
-        index,
-        text_embeddings,
-        noise_pred_original,
-        text_embeddings_clip,
-        clip_guidance_scale,
-        num_cutouts,
-        use_cutouts=True,
-    ):
-        latents = latents.detach().requires_grad_()
-
-        if isinstance(self.scheduler, LMSDiscreteScheduler):
-            sigma = self.scheduler.sigmas[index]
-            # the model input needs to be scaled to match the continuous ODE formulation in K-LMS
-            latent_model_input = latents / ((sigma**2 + 1) ** 0.5)
-        else:
-            latent_model_input = latents
-
-        # predict the noise residual
-        noise_pred = self.unet(latent_model_input, timestep, encoder_hidden_states=text_embeddings).sample
-
-        if isinstance(self.scheduler, PNDMScheduler):
-            alpha_prod_t = self.scheduler.alphas_cumprod[timestep]
-            beta_prod_t = 1 - alpha_prod_t
-            # compute predicted original sample from predicted noise also called
-            # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
-            pred_original_sample = (latents - beta_prod_t ** (0.5) * noise_pred) / alpha_prod_t ** (0.5)
-
-            fac = torch.sqrt(beta_prod_t)
-            sample = pred_original_sample * (fac) + latents * (1 - fac)
-        elif isinstance(self.scheduler, LMSDiscreteScheduler):
-            sigma = self.scheduler.sigmas[index]
-            sample = latents - sigma * noise_pred
-        else:
-            raise ValueError(f"scheduler type {type(self.scheduler)} not supported")
-
-        sample = 1 / 0.18215 * sample
-        image = self.vae.decode(sample).sample
-        image = (image / 2 + 0.5).clamp(0, 1)
-
-        if use_cutouts:
-            image = self.make_cutouts(image, num_cutouts)
-        else:
-            image = transforms.Resize(self.feature_extractor.size)(image)
-        image = self.normalize(image)
-
-        image_embeddings_clip = self.clip_model.get_image_features(image).float()
-        image_embeddings_clip = image_embeddings_clip / image_embeddings_clip.norm(p=2, dim=-1, keepdim=True)
-
-        if use_cutouts:
-            dists = spherical_dist_loss(image_embeddings_clip, text_embeddings_clip)
-            dists = dists.view([num_cutouts, sample.shape[0], -1])
-            loss = dists.sum(2).mean(0).sum() * clip_guidance_scale
-        else:
-            loss = spherical_dist_loss(image_embeddings_clip, text_embeddings_clip).mean() * clip_guidance_scale
-
-        grads = -torch.autograd.grad(loss, latents)[0]
-
-        if isinstance(self.scheduler, LMSDiscreteScheduler):
-            latents = latents.detach() + grads * (sigma**2)
-            noise_pred = noise_pred_original
-        else:
-            noise_pred = noise_pred_original - torch.sqrt(beta_prod_t) * grads
-        return noise_pred, latents
-
     @torch.no_grad()
     def __call__(
         self,
@@ -216,10 +88,6 @@ class CLIPGuidedStableDiffusion(DiffusionPipeline):
         num_inference_steps: Optional[int] = 50,
         guidance_scale: Optional[float] = 7.5,
         eta: Optional[float] = 0.0,
-        clip_guidance_scale: Optional[float] = 100,
-        clip_prompt: Optional[Union[str, List[str]]] = None,
-        num_cutouts: Optional[int] = 4,
-        use_cutouts: Optional[bool] = True,
         generator: Optional[torch.Generator] = None,
         latents: Optional[torch.FloatTensor] = None,
         output_type: Optional[str] = "pil",
@@ -305,24 +173,10 @@ class CLIPGuidedStableDiffusion(DiffusionPipeline):
                 "The following part of your input was truncated because CLIP can only handle sequences up to"
                 f" {self.tokenizer.model_max_length} tokens: {removed_text}"
             )
+            print(f"Too many tokens: {removed_text}")
             text_input_ids = text_input_ids[:, : self.tokenizer.model_max_length]
         text_embeddings = self.text_encoder(text_input_ids.to(self.device))[0]
 
-        if clip_guidance_scale > 0:
-            if clip_prompt is not None:
-                clip_text_inputs = self.tokenizer(
-                    clip_prompt,
-                    padding="max_length",
-                    max_length=self.tokenizer.model_max_length,
-                    truncation=True,
-                    return_tensors="pt",
-                )
-                clip_text_input_ids = clip_text_inputs.input_ids
-            else:
-                clip_text_input_ids = text_input_ids
-            text_embeddings_clip = self.clip_model.get_text_features(clip_text_input_ids.to(self.device))
-            text_embeddings_clip = text_embeddings_clip / text_embeddings_clip.norm(p=2, dim=-1, keepdim=True)
-
         # 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.
@@ -357,7 +211,7 @@ class CLIPGuidedStableDiffusion(DiffusionPipeline):
         else:
             if latents.shape != latents_shape:
                 raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}")
-        latents = latents.to(self.device)
+            latents = latents.to(self.device)
 
         # set timesteps
         self.scheduler.set_timesteps(num_inference_steps)
@@ -414,23 +268,6 @@ class CLIPGuidedStableDiffusion(DiffusionPipeline):
                 noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                 noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
 
-            # perform clip guidance
-            if clip_guidance_scale > 0:
-                text_embeddings_for_guidance = (
-                    text_embeddings.chunk(2)[1] if do_classifier_free_guidance else text_embeddings
-                )
-                noise_pred, latents = self.cond_fn(
-                    latents,
-                    t,
-                    i,
-                    text_embeddings_for_guidance,
-                    noise_pred,
-                    text_embeddings_clip,
-                    clip_guidance_scale,
-                    num_cutouts,
-                    use_cutouts,
-                )
-
             # compute the previous noisy sample x_t -> x_t-1
             if isinstance(self.scheduler, LMSDiscreteScheduler):
                 latents = self.scheduler.step(noise_pred, i, latents, **extra_step_kwargs).prev_sample
diff --git a/schedulers/scheduling_euler_a.py b/schedulers/scheduling_euler_a.py
index 57a56de..29ebd07 100644
--- a/schedulers/scheduling_euler_a.py
+++ b/schedulers/scheduling_euler_a.py
@@ -216,7 +216,7 @@ class EulerAScheduler(SchedulerMixin, ConfigMixin):
 
         self.num_inference_steps = num_inference_steps
         self.DSsigmas = ((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5
-        self.sigmas = get_sigmas(self.DSsigmas, self.num_inference_steps).to(device=device)
+        self.sigmas = get_sigmas(self.DSsigmas, self.num_inference_steps - 1).to(device=device)
         self.timesteps = self.sigmas
 
     def add_noise_to_input(
@@ -272,11 +272,7 @@ class EulerAScheduler(SchedulerMixin, ConfigMixin):
         """
         latents = sample
         sigma_down, sigma_up = get_ancestral_step(timestep, timestep_prev)
-
-        # if callback is not None:
-        #     callback({'x': latents, 'i': i, 'sigma': timestep, 'sigma_hat': timestep, 'denoised': model_output})
         d = to_d(latents, timestep, model_output)
-        # Euler method
         dt = sigma_down - timestep
         latents = latents + d * dt
         latents = latents + torch.randn(latents.shape, layout=latents.layout, device=latents.device,