Update app.py

app.py

@@ -1,3 +1,27 @@
 import gradio as gr
+from diffusers import DiffusionPipeline
+import random
+
+pipe = DiffusionPipeline.from_pretrained("models/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]
+
+# Define the Gradio interface
+inputs = [
+    gr.inputs.Textbox(label="Prompt", default="A fantasy landscape", lines=2),
+    gr.inputs.Textbox(label="Negative Prompt", default="", lines=2),
+    gr.inputs.Slider(minimum=256, maximum=1024, step=64, default=1024, label="Height"),
+    gr.inputs.Slider(minimum=256, maximum=1024, step=64, default=1024, label="Width"),
+    gr.inputs.Number(default=None, label="Seed", optional=True)
+]
+
+outputs = gr.outputs.Image(type="pil", label="Generated Image")
+
+gr.Interface(fn=generate_image, inputs=inputs, outputs=outputs, title="Stable Diffusion Image Generator").launch()
 
-gr.load("models/UnfilteredAI/NSFW-gen-v2").launch()