Update app.py

app.py

@@ -1,5 +1,6 @@
 import torch
 from diffusers import FluxPipeline
+from transformers import pipeline
 import gradio as gr
 import spaces
 
@@ -8,6 +9,7 @@ pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=
 pipe.load_lora_weights("Shakker-Labs/FLUX.1-dev-LoRA-Children-Simple-Sketch", weight_name="FLUX-dev-lora-children-simple-sketch.safetensors")
 pipe.fuse_lora(lora_scale=1.5)
 pipe.to("cuda")
+classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection")
 
 # Define the function to generate the sketch
 @spaces.GPU
@@ -16,7 +18,9 @@ def generate_sketch(prompt, num_inference_steps, guidance_scale):
                  num_inference_steps=num_inference_steps, 
                  guidance_scale=guidance_scale,
                 ).images[0]
+    print(classifier(image))
     image_path = "generated_sketch.png"
+    
     image.save(image_path)
     return image_path