Correct controlnet out of list error

src/diffusers/pipelines/controlnet/pipeline_controlnet.py

@@ -947,9 +947,9 @@ def __call__(
 
         # 7.1 Create tensor stating which controlnets to keep
         controlnet_keep = []
-        for i in range(num_inference_steps):
+        for i in range(len(timesteps)):
             keeps = [
-                1.0 - float(i / num_inference_steps < s or (i + 1) / num_inference_steps > e)
+                1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e)
                 for s, e in zip(control_guidance_start, control_guidance_end)
             ]
             controlnet_keep.append(keeps[0] if len(keeps) == 1 else keeps)

src/diffusers/pipelines/controlnet/pipeline_controlnet_img2img.py

@@ -1040,9 +1040,9 @@ def __call__(
 
         # 7.1 Create tensor stating which controlnets to keep
         controlnet_keep = []
-        for i in range(num_inference_steps):
+        for i in range(len(timesteps)):
             keeps = [
-                1.0 - float(i / num_inference_steps < s or (i + 1) / num_inference_steps > e)
+                1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e)
                 for s, e in zip(control_guidance_start, control_guidance_end)
             ]
             controlnet_keep.append(keeps[0] if len(keeps) == 1 else keeps)

src/diffusers/pipelines/controlnet/pipeline_controlnet_inpaint.py

@@ -1275,9 +1275,9 @@ def __call__(
 
         # 7.1 Create tensor stating which controlnets to keep
         controlnet_keep = []
-        for i in range(num_inference_steps):
+        for i in range(len(timesteps)):
             keeps = [
-                1.0 - float(i / num_inference_steps < s or (i + 1) / num_inference_steps > e)
+                1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e)
                 for s, e in zip(control_guidance_start, control_guidance_end)
             ]
             controlnet_keep.append(keeps[0] if len(keeps) == 1 else keeps)

tests/pipelines/controlnet/test_controlnet.py

@@ -46,7 +46,7 @@
     TEXT_TO_IMAGE_IMAGE_PARAMS,
     TEXT_TO_IMAGE_PARAMS,
 )
-from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin
+from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin, PipelineKarrasSchedulerTesterMixin
 
 
 enable_full_determinism()
@@ -97,7 +97,7 @@ def _test_stable_diffusion_compile(in_queue, out_queue, timeout):
     out_queue.join()
 
 
-class ControlNetPipelineFastTests(PipelineLatentTesterMixin, PipelineTesterMixin, unittest.TestCase):
+class ControlNetPipelineFastTests(PipelineLatentTesterMixin, PipelineTesterMixin, PipelineKarrasSchedulerTesterMixin, unittest.TestCase):
     pipeline_class = StableDiffusionControlNetPipeline
     params = TEXT_TO_IMAGE_PARAMS
     batch_params = TEXT_TO_IMAGE_BATCH_PARAMS
@@ -207,7 +207,7 @@ def test_inference_batch_single_identical(self):
         self._test_inference_batch_single_identical(expected_max_diff=2e-3)
 
 
-class StableDiffusionMultiControlNetPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
+class StableDiffusionMultiControlNetPipelineFastTests(PipelineTesterMixin, PipelineKarrasSchedulerTesterMixin, unittest.TestCase):
     pipeline_class = StableDiffusionControlNetPipeline
     params = TEXT_TO_IMAGE_PARAMS
     batch_params = TEXT_TO_IMAGE_BATCH_PARAMS

tests/pipelines/controlnet/test_controlnet_img2img.py

@@ -42,13 +42,13 @@
     TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
     TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
 )
-from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin
+from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin, PipelineKarrasSchedulerTesterMixin
 
 
 enable_full_determinism()
 
 
-class ControlNetImg2ImgPipelineFastTests(PipelineLatentTesterMixin, PipelineTesterMixin, unittest.TestCase):
+class ControlNetImg2ImgPipelineFastTests(PipelineLatentTesterMixin, PipelineTesterMixin, PipelineKarrasSchedulerTesterMixin, unittest.TestCase):
     pipeline_class = StableDiffusionControlNetImg2ImgPipeline
     params = TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
     batch_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
@@ -161,7 +161,7 @@ def test_inference_batch_single_identical(self):
         self._test_inference_batch_single_identical(expected_max_diff=2e-3)
 
 
-class StableDiffusionMultiControlNetPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
+class StableDiffusionMultiControlNetPipelineFastTests(PipelineTesterMixin, PipelineKarrasSchedulerTesterMixin, unittest.TestCase):
     pipeline_class = StableDiffusionControlNetImg2ImgPipeline
     params = TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
     batch_params = TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS

tests/pipelines/controlnet/test_controlnet_inpaint.py

@@ -42,13 +42,13 @@
     TEXT_GUIDED_IMAGE_INPAINTING_PARAMS,
     TEXT_TO_IMAGE_IMAGE_PARAMS,
 )
-from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin
+from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin, PipelineKarrasSchedulerTesterMixin
 
 
 enable_full_determinism()
 
 
-class ControlNetInpaintPipelineFastTests(PipelineLatentTesterMixin, PipelineTesterMixin, unittest.TestCase):
+class ControlNetInpaintPipelineFastTests(PipelineLatentTesterMixin, PipelineTesterMixin, PipelineKarrasSchedulerTesterMixin, unittest.TestCase):
     pipeline_class = StableDiffusionControlNetInpaintPipeline
     params = TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
     batch_params = TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
@@ -237,7 +237,7 @@ def get_dummy_components(self):
         return components
 
 
-class MultiControlNetInpaintPipelineFastTests(PipelineTesterMixin, unittest.TestCase):
+class MultiControlNetInpaintPipelineFastTests(PipelineTesterMixin,  PipelineKarrasSchedulerTesterMixin, unittest.TestCase):
     pipeline_class = StableDiffusionControlNetInpaintPipeline
     params = TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
     batch_params = TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS

tests/pipelines/test_pipelines_common.py

@@ -6,6 +6,7 @@
 import tempfile
 import unittest
 from typing import Callable, Union
+from jax._src.lax.lax import check_same_dtypes
 
 import numpy as np
 import PIL
@@ -17,6 +18,7 @@
 from diffusers.utils import logging
 from diffusers.utils.import_utils import is_accelerate_available, is_accelerate_version, is_xformers_available
 from diffusers.utils.testing_utils import require_torch, torch_device
+from diffusers.schedulers import KarrasDiffusionSchedulers
 
 
 def to_np(tensor):
@@ -26,6 +28,11 @@ def to_np(tensor):
     return tensor
 
 
+def check_same_shape(tensor_list):
+    shapes = [tensor.shape for tensor in tensor_list]
+    return all(shape == shapes[0] for shape in shapes[1:])
+
+
 class PipelineLatentTesterMixin:
     """
     This mixin is designed to be used with PipelineTesterMixin and unittest.TestCase classes.
@@ -155,6 +162,29 @@ def test_latents_input(self):
         self.assertLess(max_diff, 1e-4, "passing latents as image input generate different result from passing image")
 
 
+@require_torch
+class PipelineKarrasSchedulerTesterMixin:
+    """
+    This mixin is designed to be used with unittest.TestCase classes.
+    It provides a set of common tests for each PyTorch pipeline that makes use of KarrasDiffusionSchedulers
+    equivalence of dict and tuple outputs, etc.
+    """
+    def test_karras_schedulers_shape(self):
+        components = self.get_dummy_components()
+        pipe = self.pipeline_class(**components)
+        pipe.to(torch_device)
+        pipe.set_progress_bar_config(disable=None)
+        inputs = self.get_dummy_inputs(torch_device)
+
+        outputs = []
+        for scheduler_cls in KarrasDiffusionSchedulers:
+            pipe.scheduler = scheduler_cls.from_config(pipe.scheduler.config)
+            output = pipe(**inputs)[0]
+            outputs.append(output)
+
+        assert check_same_shape(outputs)
+
+
 @require_torch
 class PipelineTesterMixin:
     """