philipp-zettl's picture
Update app.py
8c7013a verified
raw
history blame
1 kB
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')
@spaces.GPU
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()