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