Update app.py

app.py

@@ -121,9 +121,10 @@ def generate(input=DEFAULT_INPUT, filter_input="", negative_input=DEFAULT_NEGATI
 
     print(image_paths)
     
-    nsfw_prediction = repo_nsfw_classifier(**processor_nsfw_classifier(images=Image.open(image_paths[0]), return_tensors="pt")).logits
+    nsfw_prediction = repo_nsfw_classifier(**processor_nsfw_classifier(images=Image.open(image_paths[0]), return_tensors="pt")).logits.argmax(-1).item()
 
     print(nsfw_prediction)
+    print(repo_nsfw_classifier.config.id2label[nsfw_prediction])
 
     return image_paths, {item['label']: round(item['score'], 3) for item in nsfw_prediction}