Update app.py

app.py

@@ -32,7 +32,7 @@ footer {
 }
 '''
 
-repo_nsfw_classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection")
+repo_nsfw_classifier = AutoModel.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")
@@ -98,7 +98,8 @@ def generate(input=DEFAULT_INPUT, filter_input="", negative_input=DEFAULT_NEGATI
         guidance = guidance_set
     
     print(steps, guidance)
-    
+
+    repo_nsfw_classifier.to(DEVICE)
     repo.to(DEVICE)
     
     parameters  = {
@@ -117,7 +118,7 @@ def generate(input=DEFAULT_INPUT, filter_input="", negative_input=DEFAULT_NEGATI
     image_paths = [save_image(img, seed) for img in images]
 
     print(image_paths)
-
+    
     nsfw_prediction = repo_nsfw_classifier(Image.open(image_paths[0]))
 
     print(nsfw_prediction)