Update app.py

app.py

@@ -1,27 +1,33 @@
 import gradio as gr
 from diffusers import DiffusionPipeline
 import random
+import spaces
 
 pipe = DiffusionPipeline.from_pretrained("UnfilteredAI/NSFW-gen-v2")
 
-def generate_image(prompt, negative_prompt="", height=1024, width=1024, seed=None):
-    generator = None
-    if seed is not None:
-        generator = torch.Generator(device="cuda").manual_seed(seed)
-    
-    images = pipe(prompt, negative_prompt=negative_prompt, height=height, width=width, generator=generator).images
-    return images[0]
+@spaces.GPU
+def generate(prompt, negative_prompt, num_inference_steps, guidance_scale, width, height, num_samples):
+    return pipe(
+        prompt=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
 
-# Define the Gradio interface
-inputs = [
-    gr.Textbox(label="Prompt", default="A fantasy landscape", lines=2),
-    gr.Textbox(label="Negative Prompt", default="", lines=2),
-    gr.Slider(minimum=256, maximum=1024, step=64, default=1024, label="Height"),
-    gr.Slider(minimum=256, maximum=1024, step=64, default=1024, label="Width"),
-    gr.Number(default=None, label="Seed", optional=True)
-]
 
-outputs = gr.Image(type="pil", label="Generated Image")
-
-c = gr.Interface(fn=generate_image, inputs=inputs, outputs=outputs, title="NSFW-Gen 2 Image Generator").launch()
-c.launch()
+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()