Update app.py

app.py

@@ -1,3 +1,33 @@
 import gradio as gr
+import torch
+from PIL import Image
+from transformers import AutoModelForImageClassification, ViTImageProcessor
 
-gr.load("models/Falconsai/nsfw_image_detection").launch()
+# Load model
+model = AutoModelForImageClassification.from_pretrained("Falconsai/nsfw_image_detection")
+processor = ViTImageProcessor.from_pretrained('Falconsai/nsfw_image_detection')
+
+# Function to make predictions
+def classify_image(img):
+    with torch.no_grad():
+        inputs = processor(images=img, return_tensors="pt")
+        outputs = model(**inputs)
+        logits = outputs.logits
+
+    predicted_label = logits.argmax(-1).item()
+    return model.config.id2label[predicted_label]
+
+# Gradio Interface
+iface = gr.Interface(
+    fn=classify_image,
+    inputs="image",
+    outputs="text",
+    live=True,
+    interpretation="default"
+)
+
+# Launch the Gradio Interface
+iface.launch()
+
+
+#gr.load("models/Falconsai/nsfw_image_detection").launch()