Update app.py

app.py

@@ -100,8 +100,8 @@ def generate(input=DEFAULT_INPUT, filter_input="", negative_input=DEFAULT_NEGATI
     
     print(steps, guidance)
 
-    #repo_nsfw_classifier.to(DEVICE)
-    #processor_nsfw_classifier.to(DEVICE)
+    repo_nsfw_classifier.to(DEVICE)
+    processor_nsfw_classifier.to(DEVICE)
     repo.to(DEVICE)
     
     parameters  = {
@@ -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.argmax(-1).item()
+    nsfw_prediction = repo_nsfw_classifier(**processor_nsfw_classifier(images=Image.open(image_paths[0]), return_tensors="pt")).logits
 
-    print(nsfw_prediction)
+    print(nsfw_prediction.argmax(-1))
+    print(nsfw_prediction.argmax(-1).item())
     print(repo_nsfw_classifier.config.id2label[nsfw_prediction])
 
     return image_paths, {item['label']: round(item['score'], 3) for item in nsfw_prediction}