Improve text to image (#1)

- Improve text to image (681b7e4a4612b7467a146d325d0a17be155613bb)

text_to_image.py

@@ -1,8 +1,9 @@
 from transformers.tools.base import Tool, get_default_device
 from transformers.utils import is_accelerate_available, is_diffusers_available
+import torch
 
 if is_diffusers_available():
-    from diffusers import DiffusionPipeline
+    from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
 
 
 TEXT_TO_IMAGE_DESCRIPTION = (
@@ -34,13 +35,17 @@ class TextToImageTool(Tool):
             self.device = get_default_device()
 
         self.pipeline = DiffusionPipeline.from_pretrained(self.default_checkpoint)
+        self.pipeline.scheduler = DPMSolverMultistepScheduler.from_config(self.pipeline.scheduler.config)
         self.pipeline.to(self.device)
 
+        if self.device.type == "cuda":
+            self.pipeline.to(torch_dtype=torch.float16)
+
         self.is_initialized = True
 
-    def __call__(self, prompt):
+    def __call__(self, prompt, negative_prompt="low quality, bad quality, deformed, low resolution", added_prompt=" , highest quality, highly realistic, very high resolution"):
         if not self.is_initialized:
             self.setup()
 
-        return self.pipeline(prompt).images[0]
+        return self.pipeline(prompt + added_prompt, num_inference_steps=25).images[0]