fix typing in kandinsky

Author: hlky

src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2.py

@@ -17,7 +17,7 @@
 import torch
 
 from ...models import UNet2DConditionModel, VQModel
-from ...schedulers import DDPMScheduler
+from ...schedulers import DDPMScheduler, UnCLIPScheduler
 from ...utils import deprecate, logging, replace_example_docstring
 from ...utils.torch_utils import randn_tensor
 from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput
@@ -83,7 +83,7 @@ class KandinskyV22Pipeline(DiffusionPipeline):
     def __init__(
         self,
         unet: UNet2DConditionModel,
-        scheduler: DDPMScheduler,
+        scheduler: Union[DDPMScheduler, UnCLIPScheduler],
         movq: VQModel,
     ):
         super().__init__()

src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_combined.py

@@ -140,13 +140,13 @@ class KandinskyV22CombinedPipeline(DiffusionPipeline):
     def __init__(
         self,
         unet: UNet2DConditionModel,
-        scheduler: DDPMScheduler,
+        scheduler: Union[DDPMScheduler, UnCLIPScheduler],
         movq: VQModel,
         prior_prior: PriorTransformer,
         prior_image_encoder: CLIPVisionModelWithProjection,
         prior_text_encoder: CLIPTextModelWithProjection,
         prior_tokenizer: CLIPTokenizer,
-        prior_scheduler: UnCLIPScheduler,
+        prior_scheduler: Union[DDPMScheduler, UnCLIPScheduler],
         prior_image_processor: CLIPImageProcessor,
     ):
         super().__init__()

src/diffusers/pipelines/kandinsky2_2/pipeline_kandinsky2_2_prior.py

@@ -5,7 +5,7 @@
 from transformers import CLIPImageProcessor, CLIPTextModelWithProjection, CLIPTokenizer, CLIPVisionModelWithProjection
 
 from ...models import PriorTransformer
-from ...schedulers import UnCLIPScheduler
+from ...schedulers import DDPMScheduler, UnCLIPScheduler
 from ...utils import (
     logging,
     replace_example_docstring,
@@ -114,7 +114,7 @@ def __init__(
         image_encoder: CLIPVisionModelWithProjection,
         text_encoder: CLIPTextModelWithProjection,
         tokenizer: CLIPTokenizer,
-        scheduler: UnCLIPScheduler,
+        scheduler: Union[DDPMScheduler, UnCLIPScheduler],
         image_processor: CLIPImageProcessor,
     ):
         super().__init__()