Update app.py

app.py

@@ -4,6 +4,9 @@ from transformers import pipeline
 import gradio as gr
 import spaces
 
+
+device=torch.device('cuda')
+
 # Load the model and LoRA weights
 pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
 pipe.load_lora_weights("Shakker-Labs/FLUX.1-dev-LoRA-Children-Simple-Sketch", weight_name="FLUX-dev-lora-children-simple-sketch.safetensors")
@@ -11,12 +14,22 @@ pipe.fuse_lora(lora_scale=1.5)
 pipe.to("cuda")
 
 # Load the NSFW classifier
-classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection",device=torch.device('cuda'))
+image_classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection",device=device)
+text_classifier = pipeline("sentiment-analysis", model="michellejieli/NSFW_text_classification",device=device)
+classifier("I see you’ve set aside this special time to humiliate yourself in public.")
 NSFW_THRESHOLD = 0.5
 
 # Define the function to generate the sketch
 @spaces.GPU
 def generate_sketch(prompt, num_inference_steps, guidance_scale):
+    # Classify the text for NSFW content
+    text_classification = text_classifier(prompt)
+    
+    # Check the classification results
+    for result in text_classification:
+        if result['label'] == 'nsfw' and result['score'] > NSFW_THRESHOLD: 
+            return "Inappropriate prompt detected. Please try another prompt."
+    
     image = pipe("sketched style, " + prompt, 
                  num_inference_steps=num_inference_steps, 
                  guidance_scale=guidance_scale,
@@ -24,12 +37,12 @@ def generate_sketch(prompt, num_inference_steps, guidance_scale):
     
 
     # Classify the image for NSFW content
-    classification = classifier(image)
+    image_classification = image_classifier(image)
 
-    print(classification)
+    print(image_classification)
 
     # Check the classification results
-    for result in classification:
+    for result in image_classification:
         if result['label'] == 'nsfw' and result['score'] > NSFW_THRESHOLD: 
             return "Inappropriate content detected. Please try another prompt."
     
@@ -46,7 +59,7 @@ interface = gr.Interface(
         gr.Slider(5, 50, value=24, step=1, label="Number of Inference Steps"),  # Slider for num_inference_steps
         gr.Slider(1.0, 10.0, value=3.5, step=0.1, label="Guidance Scale")  # Slider for guidance_scale
     ],
-    outputs="image",
+    outputs="auto",
     title="Kids Sketch Generator",
     description="Enter a text prompt and generate a fun sketch for kids with customizable inference steps and guidance scale."
 )