Update app.py

app.py

@@ -9,7 +9,10 @@ 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")
+
+# Load the NSFW classifier
+classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection",device=torch.device('cuda'))
+NSFW_THRESHOLD = 0.85
 
 # Define the function to generate the sketch
 @spaces.GPU
@@ -18,7 +21,18 @@ def generate_sketch(prompt, num_inference_steps, guidance_scale):
                  num_inference_steps=num_inference_steps, 
                  guidance_scale=guidance_scale,
                 ).images[0]
-    print(classifier(image))
+    
+
+    # Classify the image for NSFW content
+    classification = classifier(image)
+
+    print(classification)
+
+    # Check the classification results
+    for result in classification:
+        if result['label'] == 'nsfw' and result['score'] > NSFW_THRESHOLD: 
+            return "Inappropriate content detected. Please try another prompt."
+    
     image_path = "generated_sketch.png"
     
     image.save(image_path)