Spaces:
Running
on
Zero
Running
on
Zero
import gradio as gr | |
import spaces | |
import torch | |
from diffusers import DiffusionPipeline | |
model_name = 'UnfilteredAI/NSFW-gen-v2' | |
pipe = DiffusionPipeline.from_pretrained( | |
model_name, | |
torch_dtype=torch.float16 | |
) | |
pipe.to('cuda') | |
def generate(prompt, negative_prompt, num_inference_steps, guidance_scale, width, height, num_samples): | |
return pipe( | |
prompt, | |
negative_prompt=negative_prompt, | |
num_inference_steps=num_inference_steps, | |
guidance_scale=guidance_scale, | |
width=width, | |
height=height, | |
num_images_per_prompt=num_samples | |
).images | |
gr.Interface( | |
fn=generate, | |
inputs=[ | |
gr.Text(label="Prompt"), | |
gr.Text("", label="Negative Prompt"), | |
gr.Number(7, label="Number inference steps"), | |
gr.Number(3, label="Guidance scale"), | |
gr.Number(512, label="Width"), | |
gr.Number(512, label="Height"), | |
gr.Number(1, label="# images"), | |
], | |
outputs=gr.Gallery(), | |
).launch() |