stable-diffusion
copied
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
Readme
Files and versions
18 lines
612 B
18 lines
612 B
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
|
|
|
|
|