Update app.py

app.py

@@ -7,7 +7,7 @@ import uuid
 import os
 
 from diffusers import StableDiffusionXLPipeline, StableDiffusion3Pipeline
-from transformers import AutoModel
+from transformers import AutoModelForImageClassification, ViTImageProcessor
 from PIL import Image
 
 # Pre-Initialize
@@ -32,7 +32,8 @@ footer {
 }
 '''
 
-repo_nsfw_classifier = AutoModel.from_pretrained("Falconsai/nsfw_image_detection")
+repo_nsfw_classifier = AutoModelForImageClassification.from_pretrained("Falconsai/nsfw_image_detection")
+processor_nsfw_classifier = ViTImageProcessor.from_pretrained("Falconsai/nsfw_image_detection")
 
 repo_default = StableDiffusionXLPipeline.from_pretrained("fluently/Fluently-XL-Final", torch_dtype=torch.float16, use_safetensors=True, add_watermarker=False)
 repo_default.load_lora_weights("ehristoforu/dalle-3-xl-v2", adapter_name="base")
@@ -99,7 +100,8 @@ def generate(input=DEFAULT_INPUT, filter_input="", negative_input=DEFAULT_NEGATI
     
     print(steps, guidance)
 
-    repo_nsfw_classifier.to(DEVICE)
+    #repo_nsfw_classifier.to(DEVICE)
+    #processor_nsfw_classifier.to(DEVICE)
     repo.to(DEVICE)
     
     parameters  = {
@@ -119,7 +121,7 @@ def generate(input=DEFAULT_INPUT, filter_input="", negative_input=DEFAULT_NEGATI
 
     print(image_paths)
     
-    nsfw_prediction = repo_nsfw_classifier(Image.open(image_paths[0]))
+    nsfw_prediction = repo_nsfw_classifier(**processor(images=Image.open(image_paths[0]), return_tensors="pt")).logits
 
     print(nsfw_prediction)