diff --git a/dockerfiles/starlette/pytorch/Dockerfile.cpu b/dockerfiles/starlette/pytorch/Dockerfile.cpu
index 04d0a1db..012f93f7 100644
--- a/dockerfiles/starlette/pytorch/Dockerfile.cpu
+++ b/dockerfiles/starlette/pytorch/Dockerfile.cpu
@@ -5,7 +5,7 @@ COPY starlette_requirements.txt /tmp/requirements.txt
 RUN pip install  --no-cache-dir -r /tmp/requirements.txt && rm /tmp/requirements.txt
 
 # Think about a better solution -> base contaienr has pt 1.13. thats why need below 0.14
-RUN pip install  --no-cache-dir sentence_transformers torchvision~="0.14.0" diffusers=="0.8.1" accelerate=="0.14.0"
+RUN pip install  --no-cache-dir sentence_transformers torchvision~="0.14.0" diffusers=="0.9.0" accelerate=="0.14.0"
 
 # copy application
 COPY src/huggingface_inference_toolkit huggingface_inference_toolkit
diff --git a/dockerfiles/starlette/pytorch/Dockerfile.gpu b/dockerfiles/starlette/pytorch/Dockerfile.gpu
index 33e02056..74047e72 100644
--- a/dockerfiles/starlette/pytorch/Dockerfile.gpu
+++ b/dockerfiles/starlette/pytorch/Dockerfile.gpu
@@ -5,7 +5,7 @@ COPY starlette_requirements.txt /tmp/requirements.txt
 RUN pip install  --no-cache-dir -r /tmp/requirements.txt && rm /tmp/requirements.txt
 
 # Think about a better solution -> base contaienr has pt 1.13. thats why need below 0.14
-RUN pip install  --no-cache-dir sentence_transformers torchvision~="0.14.0" diffusers=="0.8.1" accelerate=="0.14.0"
+RUN pip install  --no-cache-dir sentence_transformers torchvision~="0.14.0" diffusers=="0.9.0" accelerate=="0.14.0"
 
 # copy application
 COPY src/huggingface_inference_toolkit huggingface_inference_toolkit
diff --git a/src/huggingface_inference_toolkit/diffusers_utils.py b/src/huggingface_inference_toolkit/diffusers_utils.py
index 3db96940..92e7ed06 100644
--- a/src/huggingface_inference_toolkit/diffusers_utils.py
+++ b/src/huggingface_inference_toolkit/diffusers_utils.py
@@ -13,7 +13,7 @@ def is_diffusers_available():
 if is_diffusers_available():
     import torch
 
-    from diffusers import StableDiffusionPipeline
+    from diffusers import DPMSolverMultistepScheduler, StableDiffusionPipeline
 
 
 def check_supported_pipeline(model_dir):
@@ -30,17 +30,39 @@ def check_supported_pipeline(model_dir):
 
 class DiffusersPipelineImageToText:
     def __init__(self, model_dir: str, device: str = None):  # needs "cuda" for GPU
+
         self.pipeline = StableDiffusionPipeline.from_pretrained(model_dir, torch_dtype=torch.float16)
+        # try to use DPMSolverMultistepScheduler
+        try:
+            self.pipeline.scheduler = DPMSolverMultistepScheduler.from_config(self.pipeline.scheduler.config)
+        except Exception:
+            pass
         self.pipeline.to(device)
 
-    def __call__(self, prompt, **kwargs):
-
-        if kwargs:
-            out = self.pipeline(prompt, **kwargs)
-        else:
-            out = self.pipeline(prompt)
+    def __call__(
+        self,
+        prompt,
+        num_inference_steps=25,
+        guidance_scale=7.5,
+        num_images_per_prompt=1,
+        height=None,
+        width=None,
+        negative_prompt=None,
+    ):
+        # TODO: add support for more images (Reason is correct output)
+        num_images_per_prompt = 1
+
+        # Call pipeline with parameters
+        out = self.pipeline(
+            prompt,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            num_images_per_prompt=num_images_per_prompt,
+            negative_prompt=negative_prompt,
+            height=height,
+            width=width,
+        )
 
-        # TODO: return more than 1 image if requested
         return out.images[0]