Update app.py

app.py

@@ -1,3 +1,5 @@
+# Create a Gradio demo for the Multi-Head model
+
 import gradio as gr
 import torch
 import torch.nn as nn
@@ -42,7 +44,7 @@ MODEL_ID = "Ali7880/multihead-content-moderator"  # Change this!
 
 # Download model files
 checkpoint_path = hf_hub_download(MODEL_ID, "multihead_model.pt")
-checkpoint = torch.load(checkpoint_path, map_location='cpu')
+checkpoint = torch.load(checkpoint_path, map_location='cpu', weights_only=False)
 
 # Create and load model
 model = MultiHeadContentModerator(
@@ -88,8 +90,8 @@ def moderate_image(image):
     
     if flags:
         verdict = "❌ UNSAFE - " + ", ".join(flags)
+    else:
         verdict = f"✅ SAFE (Normal: {nsfw_results.get('normal', 0):.0%}, Safe: {violence_results.get('safe', 0):.0%})"
-        verdict = f"✅ SAFE (NSFW: {nsfw_results.get('safe', 0):.0%}, Violence: {violence_results.get('safe', 0):.0%})"
     
     return nsfw_results, violence_results, verdict
 
@@ -110,4 +112,4 @@ demo = gr.Interface(
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()