sd

main.py → 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

__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
+def text2image(*args,**kwargs):
+    return StableDiffusion(*args,**kwargs)