Upload pipeline.py

pipeline.py

@@ -1,5 +1,5 @@
 import random
-from typing import Callable, Dict, List, Optional
+from typing import Callable, Dict
 
 import torch
 from diffusers import DiffusionPipeline
@@ -11,11 +11,10 @@ from tqdm import tqdm
 
 
 def get_scaled_coeffs():
-    """get_scaled_coeffs.
-    """
+    """get_scaled_coeffs."""
     beta_min = 0.85
     beta_max = 12.0
-    return beta_min**0.5, beta_max**0.5-beta_min**0.5
+    return beta_min**0.5, beta_max**0.5 - beta_min**0.5
 
 
 def beta(t):
@@ -27,7 +26,7 @@ def beta(t):
         t
     """
     a, b = get_scaled_coeffs()
-    return (a+t*b)**2
+    return (a + t * b) ** 2
 
 
 def int_beta(t):
@@ -39,7 +38,7 @@ def int_beta(t):
         t
     """
     a, b = get_scaled_coeffs()
-    return ((a+b*t)**3-a**3)/(3*b)
+    return ((a + b * t) ** 3 - a**3) / (3 * b)
 
 
 def sigma(t):
@@ -50,7 +49,7 @@ def sigma(t):
     t :
         t
     """
-    return torch.expm1(int_beta(t))**0.5
+    return torch.expm1(int_beta(t)) ** 0.5
 
 
 def sigma_orig(t):
@@ -61,13 +60,21 @@ def sigma_orig(t):
     t :
         t
     """
-    return (-torch.expm1(-int_beta(t)))**0.5
+    return (-torch.expm1(-int_beta(t))) ** 0.5
 
 
 class SuperDiffSDXLPipeline(DiffusionPipeline, ConfigMixin):
     """SuperDiffSDXLPipeline."""
 
-    def __init__(self, unet: Callable, vae: Callable, text_encoder: Callable, text_encoder_2: Callable, tokenizer: Callable, tokenizer_2: Callable) -> None:
+    def __init__(
+        self,
+        unet: Callable,
+        vae: Callable,
+        text_encoder: Callable,
+        text_encoder_2: Callable,
+        tokenizer: Callable,
+        tokenizer_2: Callable,
+    ) -> None:
         """__init__.
 
         Parameters
@@ -99,13 +106,14 @@ class SuperDiffSDXLPipeline(DiffusionPipeline, ConfigMixin):
         text_encoder.to(device)
         text_encoder_2.to(device)
 
-        self.register_modules(unet=unet,
-                              vae=vae,
-                              text_encoder=text_encoder,
-                              text_encoder_2=text_encoder_2,
-                              tokenizer=tokenizer,
-                              tokenizer_2=tokenizer_2,
-                              )
+        self.register_modules(
+            unet=unet,
+            vae=vae,
+            text_encoder=text_encoder,
+            text_encoder_2=text_encoder_2,
+            tokenizer=tokenizer,
+            tokenizer_2=tokenizer_2,
+        )
 
     def prepare_prompt_input(self, prompt_o, prompt_b, batch_size, height, width):
         """prepare_prompt_input.
@@ -123,44 +131,82 @@ class SuperDiffSDXLPipeline(DiffusionPipeline, ConfigMixin):
         width :
             width
         """
-        text_input = self.tokenizer(prompt_o * batch_size, padding="max_length",
-                                    max_length=self.tokenizer.model_max_length, truncation=True, return_tensors="pt")
-        text_input_2 = self.tokenizer_2(prompt_o * batch_size, padding="max_length",
-                                        max_length=self.tokenizer_2.model_max_length, truncation=True, return_tensors="pt")
+        text_input = self.tokenizer(
+            prompt_o * batch_size,
+            padding="max_length",
+            max_length=self.tokenizer.model_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+        text_input_2 = self.tokenizer_2(
+            prompt_o * batch_size,
+            padding="max_length",
+            max_length=self.tokenizer_2.model_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
         with torch.no_grad():
             text_embeddings = self.text_encoder(
-                text_input.input_ids.to(self.device), output_hidden_states=True)
+                text_input.input_ids.to(self.device), output_hidden_states=True
+            )
             text_embeddings_2 = self.text_encoder_2(
-                text_input_2.input_ids.to(self.device), output_hidden_states=True)
+                text_input_2.input_ids.to(self.device), output_hidden_states=True
+            )
         prompt_embeds_o = torch.concat(
-            (text_embeddings.hidden_states[-2], text_embeddings_2.hidden_states[-2]), dim=-1)
+            (text_embeddings.hidden_states[-2],
+             text_embeddings_2.hidden_states[-2]),
+            dim=-1,
+        )
         pooled_prompt_embeds_o = text_embeddings_2[0]
         negative_prompt_embeds = torch.zeros_like(prompt_embeds_o)
         negative_pooled_prompt_embeds = torch.zeros_like(
             pooled_prompt_embeds_o)
 
-        text_input = self.tokenizer(prompt_b * batch_size, padding="max_length",
-                                    max_length=self.tokenizer.model_max_length, truncation=True, return_tensors="pt")
-        text_input_2 = self.tokenizer_2(prompt_b * batch_size, padding="max_length",
-                                        max_length=self.tokenizer_2.model_max_length, truncation=True, return_tensors="pt")
+        text_input = self.tokenizer(
+            prompt_b * batch_size,
+            padding="max_length",
+            max_length=self.tokenizer.model_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
+        text_input_2 = self.tokenizer_2(
+            prompt_b * batch_size,
+            padding="max_length",
+            max_length=self.tokenizer_2.model_max_length,
+            truncation=True,
+            return_tensors="pt",
+        )
         with torch.no_grad():
             text_embeddings = self.text_encoder(
-                text_input.input_ids.to(self.device), output_hidden_states=True)
+                text_input.input_ids.to(self.device), output_hidden_states=True
+            )
             text_embeddings_2 = self.text_encoder_2(
-                text_input_2.input_ids.to(self.device), output_hidden_states=True)
+                text_input_2.input_ids.to(self.device), output_hidden_states=True
+            )
         prompt_embeds_b = torch.concat(
-            (text_embeddings.hidden_states[-2], text_embeddings_2.hidden_states[-2]), dim=-1)
+            (text_embeddings.hidden_states[-2],
+             text_embeddings_2.hidden_states[-2]),
+            dim=-1,
+        )
         pooled_prompt_embeds_b = text_embeddings_2[0]
         add_time_ids_o = torch.tensor([(height, width, 0, 0, height, width)])
         add_time_ids_b = torch.tensor([(height, width, 0, 0, height, width)])
         negative_add_time_ids = torch.tensor(
             [(height, width, 0, 0, height, width)])
         prompt_embeds = torch.cat(
-            [negative_prompt_embeds, prompt_embeds_o, prompt_embeds_b], dim=0)
+            [negative_prompt_embeds, prompt_embeds_o, prompt_embeds_b], dim=0
+        )
         add_text_embeds = torch.cat(
-            [negative_pooled_prompt_embeds, pooled_prompt_embeds_o, pooled_prompt_embeds_b], dim=0)
+            [
+                negative_pooled_prompt_embeds,
+                pooled_prompt_embeds_o,
+                pooled_prompt_embeds_b,
+            ],
+            dim=0,
+        )
         add_time_ids = torch.cat(
-            [negative_add_time_ids, add_time_ids_o, add_time_ids_b], dim=0)
+            [negative_add_time_ids, add_time_ids_o, add_time_ids_b], dim=0
+        )
 
         prompt_embeds = prompt_embeds.to(self.device)
         add_text_embeds = add_text_embeds.to(self.device)
@@ -234,16 +280,8 @@ class SuperDiffSDXLPipeline(DiffusionPipeline, ConfigMixin):
         embeddings : Callable
             embeddings
         """
-        def v(_x, _e): return self.model(
-            """v.
 
-            Parameters
-            ----------
-            _x :
-                _x
-            _e :
-                _e
-            """
+        def v(_x, _e):
             """v.
 
             Parameters
@@ -253,8 +291,10 @@ class SuperDiffSDXLPipeline(DiffusionPipeline, ConfigMixin):
             _e :
                 _e
             """
-            _x / ((sigma**2 + 1) ** 0.5), t, encoder_hidden_states=_e
-        ).sample
+            return self.model(
+                _x / ((sigma**2 + 1) ** 0.5), t, encoder_hidden_states=_e
+            ).sample
+
         embeds = torch.cat(embeddings)
         latent_input = latents
         vel = v(latent_input, embeds)
@@ -309,10 +349,15 @@ class SuperDiffSDXLPipeline(DiffusionPipeline, ConfigMixin):
             self.seed
         )  # Seed generator to create the initial latent noise
 
-        latents = torch.randn((batch_size, self.unet.in_channels, height // 8, width // 8),
-                              generator=self.generator, dtype=self.dtype, device=self.device,)
+        latents = torch.randn(
+            (batch_size, self.unet.in_channels, height // 8, width // 8),
+            generator=self.generator,
+            dtype=self.dtype,
+            device=self.device,
+        )
         prompt_embeds, added_cond_kwargs = self.prepare_prompt_input(
-            prompt_1, prompt_2, batch_size, height, width)
+            prompt_1, prompt_2, batch_size, height, width
+        )
 
         return {
             "latents": latents,
@@ -338,38 +383,68 @@ class SuperDiffSDXLPipeline(DiffusionPipeline, ConfigMixin):
         added_cond_kwargs = model_inputs["added_cond_kwargs"]
 
         t = torch.tensor(1.0)
-        dt = 1.0/self.num_inference_steps
+        dt = 1.0 / self.num_inference_steps
         train_number_steps = 1000
-        latents = latents * (sigma(t)**2+1)**0.5
+        latents = latents * (sigma(t) ** 2 + 1) ** 0.5
         with torch.no_grad():
             for i in tqdm(range(self.num_inference_steps)):
                 latent_model_input = torch.cat([latents] * 3)
                 sigma_t = sigma(t)
-                dsigma = sigma(t-dt) - sigma_t
-                latent_model_input /= (sigma_t**2+1)**0.5
+                dsigma = sigma(t - dt) - sigma_t
+                latent_model_input /= (sigma_t**2 + 1) ** 0.5
                 with torch.no_grad():
-                    noise_pred = self.unet(latent_model_input, t*train_number_steps, encoder_hidden_states=prompt_embeds,
-                                           added_cond_kwargs=added_cond_kwargs, return_dict=False)[0]
-
-                noise_pred_uncond, noise_pred_text_o, noise_pred_text_b = noise_pred.chunk(
-                    3)
+                    noise_pred = self.unet(
+                        latent_model_input,
+                        t * train_number_steps,
+                        encoder_hidden_states=prompt_embeds,
+                        added_cond_kwargs=added_cond_kwargs,
+                        return_dict=False,
+                    )[0]
+
+                (
+                    noise_pred_uncond,
+                    noise_pred_text_o,
+                    noise_pred_text_b,
+                ) = noise_pred.chunk(3)
 
                 # noise = torch.sqrt(2*torch.abs(dsigma)*sigma_t)*torch.randn_like(latents)
-                noise = torch.sqrt(2*torch.abs(dsigma)*sigma_t)*torch.empty_like(
-                    latents, device=self.device).normal_(generator=self.generator)
-
-                dx_ind = 2*dsigma*(noise_pred_uncond + self.guidance_scale *
-                                   (noise_pred_text_b - noise_pred_uncond)) + noise
-                kappa = (torch.abs(dsigma)*(noise_pred_text_b-noise_pred_text_o)*(noise_pred_text_b+noise_pred_text_o)
-                         ).sum((1, 2, 3))-(dx_ind*((noise_pred_text_o-noise_pred_text_b))).sum((1, 2, 3))
-                kappa /= 2*dsigma*self.guidance_scale * \
-                    ((noise_pred_text_o-noise_pred_text_b)**2).sum((1, 2, 3))
-                noise_pred = noise_pred_uncond + self.guidance_scale * \
-                    ((noise_pred_text_b - noise_pred_uncond) +
-                     kappa[:, None, None, None]*(noise_pred_text_o-noise_pred_text_b))
+                noise = torch.sqrt(2 * torch.abs(dsigma) * sigma_t) * torch.empty_like(
+                    latents, device=self.device
+                ).normal_(generator=self.generator)
+
+                dx_ind = (
+                    2
+                    * dsigma
+                    * (
+                        noise_pred_uncond
+                        + self.guidance_scale *
+                        (noise_pred_text_b - noise_pred_uncond)
+                    )
+                    + noise
+                )
+                kappa = (
+                    torch.abs(dsigma)
+                    * (noise_pred_text_b - noise_pred_text_o)
+                    * (noise_pred_text_b + noise_pred_text_o)
+                ).sum((1, 2, 3)) - (
+                    dx_ind * ((noise_pred_text_o - noise_pred_text_b))
+                ).sum(
+                    (1, 2, 3)
+                )
+                kappa /= (
+                    2
+                    * dsigma
+                    * self.guidance_scale
+                    * ((noise_pred_text_o - noise_pred_text_b) ** 2).sum((1, 2, 3))
+                )
+                noise_pred = noise_pred_uncond + self.guidance_scale * (
+                    (noise_pred_text_b - noise_pred_uncond)
+                    + kappa[:, None, None, None]
+                    * (noise_pred_text_o - noise_pred_text_b)
+                )
 
                 if i < self.num_inference_steps - 1:
-                    latents += 2*dsigma * noise_pred + noise
+                    latents += 2 * dsigma * noise_pred + noise
                 else:
                     latents += dsigma * noise_pred
 
@@ -389,7 +464,7 @@ class SuperDiffSDXLPipeline(DiffusionPipeline, ConfigMixin):
         Callable
 
         """
-        latents = latents/self.vae.config.scaling_factor
+        latents = latents / self.vae.config.scaling_factor
         latents = latents.to(torch.float32)
         with torch.no_grad():
             image = self.vae.decode(latents, return_dict=False)[0]