Update app.py

app.py

@@ -33,6 +33,8 @@ footer {
 }
 '''
 
+repo_nsfw_classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection")
+
 vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
 controlnet = ControlNetModel.from_pretrained("MakiPan/controlnet-encoded-hands-130k", torch_dtype=torch.float16)
 
@@ -125,11 +127,15 @@ def generate(input=DEFAULT_INPUT, filter_input="", negative_input=DEFAULT_NEGATI
 
     print(image_paths)
 
-    classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection")(Image.open(image_paths[0]))
+    nsfw_prediction = repo_nsfw_classifier(Image.open(image_paths[0]))
+
+    print(nsfw_prediction)
+
+    nsfw_format = [f"{index['label']}: {index['score']:.3%}" for index in nsfw_prediction]
     
-    print(classifier)
+    print(nsfw_format)
     
-    return image_paths, [f"{index['label']}: {index['score']:.3%}" for index in classifier]
+    return image_paths, nsfw_format
 
 def cloud():
     print("[CLOUD] | Space maintained.")
@@ -156,9 +162,9 @@ with gr.Blocks(css=css) as main:
 
     with gr.Column():
         images = gr.Gallery(columns=1, label="Image")
-        classifier = gr.Label()
+        nsfw_classifier = gr.Label()
             
-    submit.click(generate, inputs=[input, filter_input, negative_input, model, height, width, steps, guidance, number, seed], outputs=[images, classifier], queue=False)
+    submit.click(generate, inputs=[input, filter_input, negative_input, model, height, width, steps, guidance, number, seed], outputs=[images, nsfw_classifier], queue=False)
     maintain.click(cloud, inputs=[], outputs=[], queue=False)
 
 main.launch(show_api=True)