From b0067d298bf2b3d1884cd404e0cda489296e7ba1 Mon Sep 17 00:00:00 2001
From: Friday <3270939387@qq.com>
Date: Fri, 22 Sep 2023 21:06:06 +0800
Subject: [PATCH] sd

---
 main.py => Stable_diff_text2image.py | 4 ++--
 __init__.py                          | 7 +++----
 2 files changed, 5 insertions(+), 6 deletions(-)
 rename main.py => Stable_diff_text2image.py (81%)

diff --git a/main.py b/Stable_diff_text2image.py
similarity index 81%
rename from main.py
rename to Stable_diff_text2image.py
index b524b97..237cd0f 100644
--- a/main.py
+++ b/Stable_diff_text2image.py
@@ -6,11 +6,11 @@ from towhee.operator import PyOperator
 log = logging.getLogger(PyOperator)
 
 class StableDiffusion(PyOperator):
-    def __init__(self,model_id='stabilityai/stable-diffusion-2-1'):
+    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):
-        self.pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
         # pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
         image = pipe(prompt).images[0]
         return image
diff --git a/__init__.py b/__init__.py
index db18cbb..f9659f9 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,5 +1,4 @@
-from diffusers import StableDiffusionPipeline,DPMSolverMultistepSchedule
-from towhee.operator import PyOperator
+from Stable_diff_text2image import StableDiffusion
 
-def stablediffusion(*args,**kwargs):
-    return image
\ No newline at end of file
+def text2image(*args,**kwargs):
+    return StableDiffusion(*args,**kwargs)
\ No newline at end of file