[Type hint] PNDM pipeline

src/diffusers/models/unet_2d.py

@@ -120,7 +120,6 @@ def __init__(
     def forward(
         self, sample: torch.FloatTensor, timestep: Union[torch.Tensor, float, int]
     ) -> Dict[str, torch.FloatTensor]:
-
         # 0. center input if necessary
         if self.config.center_input_sample:
             sample = 2 * sample - 1.0

src/diffusers/models/unet_2d_condition.py

@@ -121,7 +121,6 @@ def forward(
         timestep: Union[torch.Tensor, float, int],
         encoder_hidden_states: torch.Tensor,
     ) -> Dict[str, torch.FloatTensor]:
-
         # 0. center input if necessary
         if self.config.center_input_sample:
             sample = 2 * sample - 1.0
@@ -145,7 +144,6 @@ def forward(
         # 3. down
         down_block_res_samples = (sample,)
         for downsample_block in self.down_blocks:
-
             if hasattr(downsample_block, "attentions") and downsample_block.attentions is not None:
                 sample, res_samples = downsample_block(
                     hidden_states=sample, temb=emb, encoder_hidden_states=encoder_hidden_states
@@ -160,7 +158,6 @@ def forward(
 
         # 5. up
         for upsample_block in self.up_blocks:
-
             res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
             down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
 

src/diffusers/models/unet_blocks.py

@@ -947,7 +947,6 @@ def __init__(
 
     def forward(self, hidden_states, res_hidden_states_tuple, temb=None):
         for resnet, attn in zip(self.resnets, self.attentions):
-
             # pop res hidden states
             res_hidden_states = res_hidden_states_tuple[-1]
             res_hidden_states_tuple = res_hidden_states_tuple[:-1]
@@ -1027,7 +1026,6 @@ def __init__(
 
     def forward(self, hidden_states, res_hidden_states_tuple, temb=None, encoder_hidden_states=None):
         for resnet, attn in zip(self.resnets, self.attentions):
-
             # pop res hidden states
             res_hidden_states = res_hidden_states_tuple[-1]
             res_hidden_states_tuple = res_hidden_states_tuple[:-1]
@@ -1091,7 +1089,6 @@ def __init__(
 
     def forward(self, hidden_states, res_hidden_states_tuple, temb=None):
         for resnet in self.resnets:
-
             # pop res hidden states
             res_hidden_states = res_hidden_states_tuple[-1]
             res_hidden_states_tuple = res_hidden_states_tuple[:-1]

src/diffusers/pipeline_utils.py

@@ -55,7 +55,6 @@
 
 
 class DiffusionPipeline(ConfigMixin):
-
     config_name = "model_index.json"
 
     def register_modules(self, **kwargs):

src/diffusers/pipelines/ddim/pipeline_ddim.py

@@ -29,7 +29,6 @@ def __init__(self, unet, scheduler):
 
     @torch.no_grad()
     def __call__(self, batch_size=1, generator=None, eta=0.0, num_inference_steps=50, output_type="pil", **kwargs):
-
         if "torch_device" in kwargs:
             device = kwargs.pop("torch_device")
             warnings.warn(

src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py

@@ -630,7 +630,6 @@ def forward(
         output_hidden_states=None,
         return_dict=None,
     ):
-
         outputs = self.model(
             input_ids,
             attention_mask=attention_mask,

src/diffusers/pipelines/pndm/pipeline_pndm.py

@@ -15,20 +15,33 @@
 
 
 import warnings
+from typing import Optional
 
 import torch
 
+from ...models import UNet2DModel
 from ...pipeline_utils import DiffusionPipeline
+from ...schedulers import PNDMScheduler
 
 
 class PNDMPipeline(DiffusionPipeline):
-    def __init__(self, unet, scheduler):
+    unet: UNet2DModel
+    scheduler: PNDMScheduler
+
+    def __init__(self, unet: UNet2DModel, scheduler: PNDMScheduler):
         super().__init__()
         scheduler = scheduler.set_format("pt")
         self.register_modules(unet=unet, scheduler=scheduler)
 
     @torch.no_grad()
-    def __call__(self, batch_size=1, generator=None, num_inference_steps=50, output_type="pil", **kwargs):
+    def __call__(
+        self,
+        batch_size: int = 1,
+        num_inference_steps: int = 50,
+        generator: Optional[torch.Generator] = None,
+        output_type: Optional[str] = "pil",
+        **kwargs,
+    ):
         # For more information on the sampling method you can take a look at Algorithm 2 of
         # the official paper: https://arxiv.org/pdf/2202.09778.pdf
 

src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py

@@ -59,7 +59,6 @@ def __call__(
         generator: Optional[torch.Generator] = None,
         output_type: Optional[str] = "pil",
     ):
-
         if isinstance(prompt, str):
             batch_size = 1
         elif isinstance(prompt, list):

src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py

@@ -73,7 +73,6 @@ def __call__(
         generator: Optional[torch.Generator] = None,
         output_type: Optional[str] = "pil",
     ):
-
         if isinstance(prompt, str):
             batch_size = 1
         elif isinstance(prompt, list):

src/diffusers/schedulers/scheduling_ddim.py

@@ -62,7 +62,6 @@ def __init__(
         set_alpha_to_one=True,
         tensor_format="pt",
     ):
-
         if beta_schedule == "linear":
             self.betas = np.linspace(beta_start, beta_end, num_train_timesteps, dtype=np.float32)
         elif beta_schedule == "scaled_linear":

src/diffusers/schedulers/scheduling_ddpm.py

@@ -60,7 +60,6 @@ def __init__(
         clip_sample=True,
         tensor_format="pt",
     ):
-
         if trained_betas is not None:
             self.betas = np.asarray(trained_betas)
         elif beta_schedule == "linear":

src/diffusers/schedulers/scheduling_pndm.py

@@ -58,7 +58,6 @@ def __init__(
         tensor_format="pt",
         skip_prk_steps=False,
     ):
-
         if beta_schedule == "linear":
             self.betas = np.linspace(beta_start, beta_end, num_train_timesteps, dtype=np.float32)
         elif beta_schedule == "scaled_linear":

src/diffusers/schedulers/scheduling_sde_vp.py

@@ -26,7 +26,6 @@
 class ScoreSdeVpScheduler(SchedulerMixin, ConfigMixin):
     @register_to_config
     def __init__(self, num_train_timesteps=2000, beta_min=0.1, beta_max=20, sampling_eps=1e-3, tensor_format="np"):
-
         self.sigmas = None
         self.discrete_sigmas = None
         self.timesteps = None

src/diffusers/schedulers/scheduling_utils.py

@@ -21,7 +21,6 @@
 
 
 class SchedulerMixin:
-
     config_name = SCHEDULER_CONFIG_NAME
     ignore_for_config = ["tensor_format"]
 

src/diffusers/utils/logging.py

@@ -65,17 +65,14 @@ def _get_default_logging_level():
 
 
 def _get_library_name() -> str:
-
     return __name__.split(".")[0]
 
 
 def _get_library_root_logger() -> logging.Logger:
-
     return logging.getLogger(_get_library_name())
 
 
 def _configure_library_root_logger() -> None:
-
     global _default_handler
 
     with _lock:
@@ -93,7 +90,6 @@ def _configure_library_root_logger() -> None:
 
 
 def _reset_library_root_logger() -> None:
-
     global _default_handler
 
     with _lock: