Update app.py

app.py

@@ -1,24 +1,39 @@
-import requests
-from PIL import Image
 import gradio as gr
-from io import BytesIO
+import spaces
+import torch
+from diffusers import DiffusionPipeline
 
-pipeline = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1")
 
-def get_image(key, prompt, inference_steps, filter):
-  payload = {
-    "key": key,
-    "prompt": prompt,
-    "negative_prompt": "((out of frame)), ((extra fingers)), mutated hands, ((poorly drawn hands)), ((poorly drawn face)), (((mutation))), (((deformed))), (((tiling))), ((naked)), ((tile)), ((fleshpile)), ((ugly)), (((abstract))), blurry, ((bad anatomy)), ((bad proportions)), ((extra limbs)), cloned face, (((skinny))), glitchy, ((extra breasts)), ((double torso)), ((extra arms)), ((extra hands)), ((mangled fingers)), ((missing breasts)), (missing lips), ((ugly face)), ((fat)), ((extra legs)), anime",
-    "width": "512",
-    "height": "512",
-    "samples": "1",
-    "num_inference_steps": inference_steps,"safety_checker": filter,"enhance_prompt": "yes","guidance_scale": 7.5}
-  headers = {}
-  response = requests.request("POST", url, headers=headers, data=payload)
-  url1 = str(json.loads(response.text)['output'][0])
-  r = requests.get(url1)
-  i = Image.open(BytesIO(r.content))
-  return i
+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()