skip scheduler

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, UnCLIPScheduler
+from ...schedulers import DDPMScheduler
 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: Union[DDPMScheduler, UnCLIPScheduler],
+        scheduler: DDPMScheduler,
         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: Union[DDPMScheduler, UnCLIPScheduler],
+        scheduler: DDPMScheduler,
         movq: VQModel,
         prior_prior: PriorTransformer,
         prior_image_encoder: CLIPVisionModelWithProjection,
         prior_text_encoder: CLIPTextModelWithProjection,
         prior_tokenizer: CLIPTokenizer,
-        prior_scheduler: Union[DDPMScheduler, UnCLIPScheduler],
+        prior_scheduler: 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 DDPMScheduler, UnCLIPScheduler
+from ...schedulers import UnCLIPScheduler
 from ...utils import (
     logging,
     replace_example_docstring,
@@ -114,7 +114,7 @@ def __init__(
         image_encoder: CLIPVisionModelWithProjection,
         text_encoder: CLIPTextModelWithProjection,
         tokenizer: CLIPTokenizer,
-        scheduler: Union[DDPMScheduler, UnCLIPScheduler],
+        scheduler: UnCLIPScheduler,
         image_processor: CLIPImageProcessor,
     ):
         super().__init__()

src/diffusers/pipelines/pipeline_utils.py

@@ -46,8 +46,7 @@
 from ..models.attention_processor import FusedAttnProcessor2_0
 from ..models.modeling_utils import _LOW_CPU_MEM_USAGE_DEFAULT, ModelMixin
 from ..quantizers.bitsandbytes.utils import _check_bnb_status
-from ..schedulers.scheduling_utils import SCHEDULER_CONFIG_NAME, SchedulerMixin
-from ..schedulers.scheduling_utils_flax import FlaxSchedulerMixin
+from ..schedulers.scheduling_utils import SCHEDULER_CONFIG_NAME
 from ..utils import (
     CONFIG_NAME,
     DEPRECATED_REVISION_ARGS,
@@ -839,6 +838,8 @@ def load_module(name, value):
         for key in init_dict.keys():
             if key not in passed_class_obj:
                 continue
+            if "scheduler" in key:
+                continue
 
             class_obj = passed_class_obj[key]
             _expected_class_types = []
@@ -849,16 +850,7 @@ def load_module(name, value):
                     _expected_class_types.append(expected_type.__name__)
 
             _is_valid_type = class_obj.__class__.__name__ in _expected_class_types
-            if (
-                isinstance(class_obj, SchedulerMixin) or isinstance(class_obj, FlaxSchedulerMixin)
-            ) and not _is_valid_type:
-                _requires_flow_match = any("FlowMatch" in class_type for class_type in _expected_class_types)
-                _is_flow_match = "FlowMatch" in class_obj.__class__.__name__
-                if _requires_flow_match and not _is_flow_match:
-                    logger.warning(f"Expected FlowMatch scheduler, got {class_obj.__class__.__name__}.")
-                elif not _requires_flow_match and _is_flow_match:
-                    logger.warning(f"Expected non-FlowMatch scheduler, got {class_obj.__class__.__name__}.")
-            elif not _is_valid_type:
+            if not _is_valid_type:
                 logger.warning(
                     f"Expected types for {key}: {_expected_class_types}, got {class_obj.__class__.__name__}."
                 )

src/diffusers/schedulers/scheduling_utils.py

@@ -46,7 +46,6 @@ class KarrasDiffusionSchedulers(Enum):
     UniPCMultistepScheduler = 13
     DPMSolverSDEScheduler = 14
     EDMEulerScheduler = 15
-    LCMScheduler = 16
 
 
 AysSchedules = {

tests/pipelines/stable_diffusion_2/test_stable_diffusion_latent_upscale.py

@@ -250,7 +250,6 @@ def test_karras_schedulers_shape(self):
             "KDPM2AncestralDiscreteScheduler",
             "DPMSolverSDEScheduler",
             "EDMEulerScheduler",
-            "LCMScheduler",
         ]
         components = self.get_dummy_components()
         pipe = self.pipeline_class(**components)

tests/pipelines/test_pipelines.py

@@ -47,8 +47,6 @@
     DPMSolverMultistepScheduler,
     EulerAncestralDiscreteScheduler,
     EulerDiscreteScheduler,
-    FlowMatchEulerDiscreteScheduler,
-    FluxPipeline,
     LMSDiscreteScheduler,
     ModelMixin,
     PNDMScheduler,
@@ -1814,32 +1812,6 @@ def test_wrong_model(self):
         assert "is of type" in str(error_context.exception)
         assert "but should be" in str(error_context.exception)
 
-    def test_wrong_model_scheduler_type(self):
-        scheduler = EulerDiscreteScheduler.from_pretrained("hf-internal-testing/tiny-flux-pipe", subfolder="scheduler")
-        with self.assertLogs(
-            logging.get_logger("diffusers.pipelines.pipeline_utils"), level="WARNING"
-        ) as warning_context:
-            _ = FluxPipeline.from_pretrained("hf-internal-testing/tiny-flux-pipe", scheduler=scheduler)
-
-        assert any("Expected" in message for message in warning_context.output)
-        assert any("scheduler" in message for message in warning_context.output)
-        assert any("EulerDiscreteScheduler" in message for message in warning_context.output)
-
-    def test_wrong_model_scheduler_enum(self):
-        scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
-            "hf-internal-testing/diffusers-stable-diffusion-tiny-all", subfolder="scheduler"
-        )
-        with self.assertLogs(
-            logging.get_logger("diffusers.pipelines.pipeline_utils"), level="WARNING"
-        ) as warning_context:
-            _ = StableDiffusionPipeline.from_pretrained(
-                "hf-internal-testing/diffusers-stable-diffusion-tiny-all", scheduler=scheduler
-            )
-
-        assert any("Expected" in message for message in warning_context.output)
-        assert any("scheduler" in message for message in warning_context.output)
-        assert any("FlowMatchEulerDiscreteScheduler" in message for message in warning_context.output)
-
 
 @slow
 @require_torch_gpu