diff --git a/.DS_Store b/.DS_Store
new file mode 100644
index 0000000..d332226
Binary files /dev/null and b/.DS_Store differ
diff --git a/README.md b/README.md
index 3101075..4750d2e 100644
--- a/README.md
+++ b/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"
+
-pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
+
-#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'.
-
+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.
+
-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.
-.
\ No newline at end of file
+
\ No newline at end of file
diff --git a/Stable_diff_text2image.py b/Stable_diff_text2image.py
deleted file mode 100644
index 1e7df6f..0000000
--- a/Stable_diff_text2image.py
+++ /dev/null
@@ -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
-
-
diff --git a/__init__.py b/__init__.py
index f9659f9..10b379e 100644
--- a/__init__.py
+++ b/__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)
\ No newline at end of file
diff --git a/a corgi.png b/a corgi.png
deleted file mode 100644
index 2980d10..0000000
Binary files a/a corgi.png and /dev/null differ
diff --git a/an orange cat.png b/an orange cat.png
deleted file mode 100644
index 5b030d0..0000000
Binary files a/an orange cat.png and /dev/null differ
diff --git a/an_orange_cat.png b/an_orange_cat.png
new file mode 100644
index 0000000..b63acd5
Binary files /dev/null and b/an_orange_cat.png differ
diff --git a/stable_diffusion.py b/stable_diffusion.py
new file mode 100644
index 0000000..c6137cb
--- /dev/null
+++ b/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