Debug

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>

README.md

@@ -1,98 +1,68 @@
 
-# Title
+# Image generation using Stable Diffusion
 
-Text2image implementation through StableDiffusion
-
-
-.
-
-
-
-
-
-
-
-
-
-
-
-
-## Overview
-
-Stable Diffusion is a software library that provides efficient and accurate algorithms for solving diffusion equations numerically. It is designed to handle diffusion problems in various scientific and engineering fields, such as heat transfer, fluid dynamics, and chemical reactions.
-
-.
-
-
-
-
-## Author
-
-- CHANG XIAOYIN(3270939387@qq.com)
-
-.
-
-
-## Features
-
-- Efficient and scalable algorithms for solving large-scale diffusion problems
-
-- Visualization tools for analyzing and visualizing the simulated results
-
-- Numerical solution of diffusion equations with different boundary conditions
-
-- Support for various discretization schemes, including finite difference, finite element, and spectral methods
-
-.
+A text2image operator generates image given a text prompt.
+This operator is implemented with [Huggingface Diffusers](https://github.com/huggingface/diffusers).
 
 
 ## Code example
 
-```pythonscript
-import torch
+```python
+from towhee import pipe, ops
+
+pipe = (
+    pipe.input('prompt')
+        .map('prompt', 'image', ops.text2image.stable_diffusion())
+        .output('image')
+)
 
-from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
+image = pipe('an orange cat')
+image.save('an_orange_cat.png')
+```
 
-model_id = "stabilityai/stable-diffusion-2-1"
+<img src="./an_orange_cat.png" width="800px"/>
 
-pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
+<br />
 
-#pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
+## Factory Constructor
 
-prompt = "an orange cat"
+Create the operator via the following factory method:
 
-image = pipe(prompt).images[0] 
+***text2image.stable_diffusion(model_id='stabilityai/stable-diffusion-2-1', device=None)***
 
-image.save("orange_cat.png")
-```
-.
+**Parameters:**
 
+***model_id***: *str*
 
-## result picture
+The model id in string, defaults to 'stabilityai/stable-diffusion-2-1'.
 
-![App Screenshot](https://towhee.io/text2image/stable-diffusion/raw/branch/main/an%20orange%20cat.png)
+Supported model names: [pretrained diffuser models](https://huggingface.co/models?library=diffusers&sort=downloads)
 
-.
 
+***device***: *str*
 
-## Contributing
+The device to running model on, defaults to None.
+If None, it will automatically use cuda if gpu is available.
 
-f you find any issues or have suggestions for improvements, please submit them through the GitHub issue tracker. Contributions to the Stable Diffusion are welcome. You can fork the repository, make your changes, and submit a pull request.
+<br />
 
-Please ensure that your contributions adhere to the coding conventions and style guidelines outlined in the repository.
+## Interface
 
-.
+The operator takes a text prompt in string as input.
+It loads pretrained diffuser model and generates an image.
 
+***\_\_call\_\_(txt)***
 
-## License
+**Parameters:**
 
-Stable Diffusion is licensed under the MIT License. You are free to use, modify, and distribute the library in accordance with the terms of this license.
+***prompt***: *str*
 
-.
+​	The text in string.
 
+**Returns**:
 
-## Contacts
+*PIL.Image*
 
-If you have any questions or need further assistance, feel free to reach out to the development team at 3270939387@qq.com. We would be happy to assist you.
+​	The generated image.
 
-.
+<br />

Stable_diff_text2image.py

@@ -1,18 +0,0 @@
-import logging
-import torch
-from diffusers import StableDiffusionPipeline, DPMSolverMultistepSchedule
-from towhee.operator import PyOperator
-
-log = logging.getLogger(PyOperator)
-
-class StableDiffusion(PyOperator):
-    def __init__(self,model_id='stabilityai/stable-diffusion-2-1',pipe='StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32'):
-        self._model_id=model_id
-        self._pipe = pipe
-
-    def __call__(self, prompt:str):
-        # pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
-        image = pipe(prompt).images[0]
-        return image
-
-

__init__.py

@@ -1,4 +1,4 @@
-from Stable_diff_text2image import StableDiffusion
+from .stable_diffusion import StableDiffusion
 
-def text2image(*args,**kwargs):
+def stable_diffusion(*args,**kwargs):
     return StableDiffusion(*args,**kwargs)

stable_diffusion.py

@@ -0,0 +1,26 @@
+import logging
+import torch
+from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
+from towhee.operator import PyOperator
+
+log = logging.getLogger(PyOperator)
+
+class StableDiffusion(PyOperator):
+    def __init__(self,
+                 model_id: str ='stabilityai/stable-diffusion-2-1',
+                 device: str = None
+                 ):
+        if device is None:
+             device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self.device = device
+
+        if self.device == 'cpu':
+            self.pipe= StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
+        else:
+            self.pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
+            self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(self.pipe.scheduler.config)
+            self.pipe = self.pipe.to('cuda')
+        
+    def __call__(self, prompt: str):
+        image = self.pipe(prompt).images[0]
+        return image