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
34 lines
1.1 KiB
34 lines
1.1 KiB
# This is a sample Python script.
|
|
|
|
# Press Shift+F10 to execute it or replace it with your code.
|
|
# Press Double Shift to search everywhere for classes, files, tool windows, actions, and settings.
|
|
|
|
|
|
def print_hi(name):
|
|
# Use a breakpoint in the code line below to debug your script.
|
|
print(f'Hi, {name}') # Press Ctrl+F8 to toggle the breakpoint.
|
|
|
|
|
|
# Press the green button in the gutter to run the script.
|
|
if __name__ == '__main__':
|
|
print_hi('PyCharm')
|
|
|
|
# See PyCharm help at https://www.jetbrains.com/help/pycharm/
|
|
import logging
|
|
import torch
|
|
from diffusers import StableDiffusionPipeline, DPMSolverMultistepSchedule
|
|
from towhee.operator import PyOperator
|
|
|
|
log = logging.getLogger(PyOperator)
|
|
|
|
class Generation(PyOperator):
|
|
def __init__(self,model_id='stabilityai/stable-diffusion-2-1'):
|
|
self._model_id=model_id
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
|
|
# pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
|
|
prompt = "an orange cat"
|
|
image = pipe(prompt).images[0]
|
|
|
|
|