Spaces:
Running
on
Zero
Running
on
Zero
import torch | |
from diffusers import DiffusionPipeline | |
import gradio as gr | |
import os | |
import spaces | |
model_list = os.environ.get("MODELS").split(",") | |
lora_list = os.environ.get("LORAS") # Not in use | |
def generate(prompt, model): | |
pipe = DiffusionPipeline.from_pretrained( | |
model, | |
torch_dtype=torch.float16, | |
use_safetensors=True, | |
) | |
pipe.to('cuda') | |
negative_prompt = "nsfw, lowres, (bad), text, error, fewer, extra, missing, worst quality, jpeg artifacts, low quality, watermark, unfinished, displeasing, oldest, early, chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract]" | |
image = pipe( | |
prompt, | |
negative_prompt=negative_prompt, | |
width=832, | |
height=1216, | |
guidance_scale=7, | |
num_inference_steps=28 | |
).images[0] | |
return image | |
with gr.Blocks() as demo: | |
inp_prompt = gr.Textbox(label="Prompt") | |
inp_model = gr.Dropdown(model_list, label="Select a model") | |
out = gr.Image() | |
btn = gr.Button("Run") | |
btn.click(fn=generate, inputs=[inp_prompt, inp_model], outputs=out) | |
demo.launch() |