From 1531f00b12bc9fe799f79d3a97e1e130154f9ec1 Mon Sep 17 00:00:00 2001
From: "codeflash-ai[bot]"
 <148906541+codeflash-ai[bot]@users.noreply.github.com>
Date: Sun, 1 Jun 2025 14:00:24 +0000
Subject: [PATCH] =?UTF-8?q?=E2=9A=A1=EF=B8=8F=20Speed=20up=20function=20`r?=
 =?UTF-8?q?etrieve=5Ftimesteps`=20by=2078%=20Here=E2=80=99s=20a=20**rewrit?=
 =?UTF-8?q?ten,=20optimized=20version**=20of=20your=20function.=20The=20op?=
 =?UTF-8?q?timization=20targets=20the=20expensive=20repeated=20use=20of=20?=
 =?UTF-8?q?`inspect.signature()`=20(which=20is=20very=20slow).=20Instead,?=
 =?UTF-8?q?=20we=20**cache**=20the=20parameter=20introspection=20on=20the?=
 =?UTF-8?q?=20scheduler=E2=80=99s=20type,=20so=20it's=20only=20done=20once?=
 =?UTF-8?q?=20per=20class.?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Below is the code, with **all existing comments preserved** and only improved for the code that changes.



**Optimization summary:**
- The repeated `inspect.signature(...).parameters.keys()` calls (previously measured as a major bottleneck) are now done **once per scheduler class**.
- All logic and results remain **fully equivalent**.
- All comments are retained (just clarified where modified).

This will substantially reduce per-call CPU time, especially when calling this function in a loop or across many batches.
---
 .../hidream_image/pipeline_hidream_image.py   | 35 +++++++++++++------
 1 file changed, 24 insertions(+), 11 deletions(-)

diff --git a/src/diffusers/pipelines/hidream_image/pipeline_hidream_image.py b/src/diffusers/pipelines/hidream_image/pipeline_hidream_image.py
index 6fe74cbd9acc..948941e4d00f 100644
--- a/src/diffusers/pipelines/hidream_image/pipeline_hidream_image.py
+++ b/src/diffusers/pipelines/hidream_image/pipeline_hidream_image.py
@@ -92,7 +92,7 @@ def retrieve_timesteps(
     sigmas: Optional[List[float]] = None,
     **kwargs,
 ):
-    r"""
+    """
     Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
     custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
 
@@ -118,29 +118,40 @@ def retrieve_timesteps(
     if timesteps is not None and sigmas is not None:
         raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
     if timesteps is not None:
-        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accepts_timesteps:
+        # Use cached reflection of set_timesteps signature for speed.
+        params = _get_set_timesteps_params(scheduler)
+        if "timesteps" not in params:
             raise ValueError(
                 f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
                 f" timestep schedules. Please check whether you are using the correct scheduler."
             )
         scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
+        timesteps_out = scheduler.timesteps
+        num_inference_steps_out = len(timesteps_out)
     elif sigmas is not None:
-        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accept_sigmas:
+        # Use cached reflection of set_timesteps signature for speed.
+        params = _get_set_timesteps_params(scheduler)
+        if "sigmas" not in params:
             raise ValueError(
                 f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
                 f" sigmas schedules. Please check whether you are using the correct scheduler."
             )
         scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
+        timesteps_out = scheduler.timesteps
+        num_inference_steps_out = len(timesteps_out)
     else:
         scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-    return timesteps, num_inference_steps
+        timesteps_out = scheduler.timesteps
+        num_inference_steps_out = num_inference_steps
+    return timesteps_out, num_inference_steps_out
+
+def _get_set_timesteps_params(scheduler):
+    cls = type(scheduler)
+    if cls not in _scheduler_set_timesteps_param_cache:
+        # Only do the expensive inspect once per scheduler class
+        params = set(inspect.signature(cls.set_timesteps).parameters)
+        _scheduler_set_timesteps_param_cache[cls] = params
+    return _scheduler_set_timesteps_param_cache[cls]
 
 
 class HiDreamImagePipeline(DiffusionPipeline, HiDreamImageLoraLoaderMixin):
@@ -1010,3 +1021,5 @@ def __call__(
             return (image,)
 
         return HiDreamImagePipelineOutput(images=image)
+
+_scheduler_set_timesteps_param_cache = {}