type hints

Author: sidthekidder

src/diffusers/pipelines/ddim/pipeline_ddim.py

@@ -15,6 +15,7 @@
 
 
 import warnings
+from typing import Optional
 
 import torch
 
@@ -28,7 +29,15 @@ def __init__(self, unet, scheduler):
         self.register_modules(unet=unet, scheduler=scheduler)
 
     @torch.no_grad()
-    def __call__(self, batch_size=1, generator=None, eta=0.0, num_inference_steps=50, output_type="pil", **kwargs):
+    def __call__(
+        self,
+        batch_size: int = 1,
+        generator: Optional[torch.Generator] = None,
+        eta: float = 0.0,
+        num_inference_steps: int = 50,
+        output_type: Optional[str] = "pil",
+        **kwargs,
+    ):
 
         if "torch_device" in kwargs:
             device = kwargs.pop("torch_device")

src/diffusers/pipelines/ddpm/pipeline_ddpm.py

@@ -15,6 +15,7 @@
 
 
 import warnings
+from typing import Optional
 
 import torch
 
@@ -28,7 +29,13 @@ def __init__(self, unet, scheduler):
         self.register_modules(unet=unet, scheduler=scheduler)
 
     @torch.no_grad()
-    def __call__(self, batch_size=1, generator=None, output_type="pil", **kwargs):
+    def __call__(
+        self,
+        batch_size: int = 1,
+        generator: Optional[torch.Generator] = None,
+        output_type: Optional[str] = "pil",
+        **kwargs,
+    ):
         if "torch_device" in kwargs:
             device = kwargs.pop("torch_device")
             warnings.warn(

src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py

@@ -10,13 +10,23 @@
 from transformers.configuration_utils import PretrainedConfig
 from transformers.modeling_outputs import BaseModelOutput
 from transformers.modeling_utils import PreTrainedModel
+from transformers.tokenization_utils import PreTrainedTokenizer
 from transformers.utils import logging
 
+from ...models import UNet2DModel, UNet2DConditionModel, AutoencoderKL, VQModel
 from ...pipeline_utils import DiffusionPipeline
+from ...schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler
 
 
 class LDMTextToImagePipeline(DiffusionPipeline):
-    def __init__(self, vqvae, bert, tokenizer, unet, scheduler):
+    def __init__(
+        self,
+        vqvae: Union[VQModel, AutoencoderKL],
+        bert: PreTrainedModel,
+        tokenizer: PreTrainedTokenizer,
+        unet: Union[UNet2DModel, UNet2DConditionModel],
+        scheduler: Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler],
+    ):
         super().__init__()
         scheduler = scheduler.set_format("pt")
         self.register_modules(vqvae=vqvae, bert=bert, tokenizer=tokenizer, unet=unet, scheduler=scheduler)
@@ -614,7 +624,7 @@ def custom_forward(*inputs):
 
 
 class LDMBertModel(LDMBertPreTrainedModel):
-    def __init__(self, config):
+    def __init__(self, config: LDMBertConfig):
         super().__init__(config)
         self.model = LDMBertEncoder(config)
         self.to_logits = nn.Linear(config.hidden_size, config.vocab_size)