make style

modeling_glide.py

@@ -18,7 +18,14 @@ import numpy as np
 import torch
 
 import tqdm
-from diffusers import ClassifierFreeGuidanceScheduler, GlideDDIMScheduler, CLIPTextModel, DiffusionPipeline, GLIDETextToImageUNetModel, GLIDESuperResUNetModel
+from diffusers import (
+    ClassifierFreeGuidanceScheduler,
+    CLIPTextModel,
+    DiffusionPipeline,
+    GlideDDIMScheduler,
+    GLIDESuperResUNetModel,
+    GLIDETextToImageUNetModel,
+)
 from transformers import GPT2Tokenizer
 
 
@@ -46,12 +53,16 @@ class GLIDE(DiffusionPipeline):
         text_encoder: CLIPTextModel,
         tokenizer: GPT2Tokenizer,
         upscale_unet: GLIDESuperResUNetModel,
-        upscale_noise_scheduler: GlideDDIMScheduler
+        upscale_noise_scheduler: GlideDDIMScheduler,
     ):
         super().__init__()
         self.register_modules(
-            text_unet=text_unet, text_noise_scheduler=text_noise_scheduler, text_encoder=text_encoder, tokenizer=tokenizer,
-            upscale_unet=upscale_unet, upscale_noise_scheduler=upscale_noise_scheduler
+            text_unet=text_unet,
+            text_noise_scheduler=text_noise_scheduler,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            upscale_unet=upscale_unet,
+            upscale_noise_scheduler=upscale_noise_scheduler,
         )
 
     def q_posterior_mean_variance(self, scheduler, x_start, x_t, t):
@@ -67,9 +78,7 @@ class GLIDE(DiffusionPipeline):
             + _extract_into_tensor(scheduler.posterior_mean_coef2, t, x_t.shape) * x_t
         )
         posterior_variance = _extract_into_tensor(scheduler.posterior_variance, t, x_t.shape)
-        posterior_log_variance_clipped = _extract_into_tensor(
-            scheduler.posterior_log_variance_clipped, t, x_t.shape
-        )
+        posterior_log_variance_clipped = _extract_into_tensor(scheduler.posterior_log_variance_clipped, t, x_t.shape)
         assert (
             posterior_mean.shape[0]
             == posterior_variance.shape[0]
@@ -190,19 +199,30 @@ class GLIDE(DiffusionPipeline):
         # A value of 1.0 is sharper, but sometimes results in grainy artifacts.
         upsample_temp = 0.997
 
-        image = self.upscale_noise_scheduler.sample_noise(
-            (batch_size, 3, 256, 256), device=torch_device, generator=generator
-        ) * upsample_temp
+        image = (
+            self.upscale_noise_scheduler.sample_noise(
+                (batch_size, 3, 256, 256), device=torch_device, generator=generator
+            )
+            * upsample_temp
+        )
 
         num_timesteps = len(self.upscale_noise_scheduler)
-        for t in tqdm.tqdm(reversed(range(len(self.upscale_noise_scheduler))), total=len(self.upscale_noise_scheduler)):
+        for t in tqdm.tqdm(
+            reversed(range(len(self.upscale_noise_scheduler))), total=len(self.upscale_noise_scheduler)
+        ):
             # i) define coefficients for time step t
             clipped_image_coeff = 1 / torch.sqrt(self.upscale_noise_scheduler.get_alpha_prod(t))
             clipped_noise_coeff = torch.sqrt(1 / self.upscale_noise_scheduler.get_alpha_prod(t) - 1)
-            image_coeff = (1 - self.upscale_noise_scheduler.get_alpha_prod(t - 1)) * torch.sqrt(
-                self.upscale_noise_scheduler.get_alpha(t)) / (1 - self.upscale_noise_scheduler.get_alpha_prod(t))
-            clipped_coeff = torch.sqrt(self.upscale_noise_scheduler.get_alpha_prod(t - 1)) * self.upscale_noise_scheduler.get_beta(
-                t) / (1 - self.upscale_noise_scheduler.get_alpha_prod(t))
+            image_coeff = (
+                (1 - self.upscale_noise_scheduler.get_alpha_prod(t - 1))
+                * torch.sqrt(self.upscale_noise_scheduler.get_alpha(t))
+                / (1 - self.upscale_noise_scheduler.get_alpha_prod(t))
+            )
+            clipped_coeff = (
+                torch.sqrt(self.upscale_noise_scheduler.get_alpha_prod(t - 1))
+                * self.upscale_noise_scheduler.get_beta(t)
+                / (1 - self.upscale_noise_scheduler.get_alpha_prod(t))
+            )
 
             # ii) predict noise residual
             time_input = torch.tensor([t] * image.shape[0], device=torch_device)
@@ -216,8 +236,9 @@ class GLIDE(DiffusionPipeline):
             prev_image = clipped_coeff * pred_mean + image_coeff * image
 
             # iv) sample variance
-            prev_variance = self.upscale_noise_scheduler.sample_variance(t, prev_image.shape, device=torch_device,
-                                                                 generator=generator)
+            prev_variance = self.upscale_noise_scheduler.sample_variance(
+                t, prev_image.shape, device=torch_device, generator=generator
+            )
 
             # v) sample  x_{t-1} ~ N(prev_image, prev_variance)
             sampled_prev_image = prev_image + prev_variance