Update diffusers

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

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

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]