Update app.py

app.py

@@ -15,7 +15,7 @@ pipe.to("cuda")
 
 # Load the NSFW classifier
 image_classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection",device=device)
-text_classifier = pipeline("text-classification", model="eliasalbouzidi/distilbert-nsfw-text-classifier",device=device)
+#text_classifier = pipeline("text-classification", model="eliasalbouzidi/distilbert-nsfw-text-classifier",device=device)
 
 NSFW_THRESHOLD = 0.5
 
@@ -23,13 +23,13 @@ NSFW_THRESHOLD = 0.5
 @spaces.GPU
 def generate_sketch(prompt, num_inference_steps, guidance_scale):
     # Classify the text for NSFW content
-    text_classification = text_classifier(prompt)
-    print(text_classification)
+    #text_classification = text_classifier(prompt)
+    #print(text_classification)
     
     # Check the classification results
-    for result in text_classification:
-        if result['label'] == 'nsfw' and result['score'] > NSFW_THRESHOLD: 
-            return gr.update(visible=False),gr.Text(value="Inappropriate prompt detected. Please try another prompt.")
+    #for result in text_classification:
+    #    if result['label'] == 'nsfw' and result['score'] > NSFW_THRESHOLD: 
+    #        return gr.update(visible=False),gr.Text(value="Inappropriate prompt detected. Please try another prompt.")
     
     image = pipe("sketched style, " + prompt, 
                  num_inference_steps=num_inference_steps, 
@@ -45,12 +45,12 @@ def generate_sketch(prompt, num_inference_steps, guidance_scale):
     # Check the classification results
     for result in image_classification:
         if result['label'] == 'nsfw' and result['score'] > NSFW_THRESHOLD: 
-           return gr.update(visible=False),gr.Text(value="Inappropriate content detected. Please try another prompt.")
+           return "Inappropriate content detected. Please try another prompt." #return gr.update(visible=False),gr.Text(value="Inappropriate content detected. Please try another prompt.")
     
     image_path = "generated_sketch.png"
     
     image.save(image_path)
-    return gr.Image(value=image_path), gr.update(visible=False)
+    return image_path #gr.Image(value=image_path), gr.update(visible=False)
 
 # Gradio interface with sliders for num_inference_steps and guidance_scale
 interface = gr.Interface(
@@ -60,7 +60,10 @@ 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=[gr.Image(), gr.Text()],
+    outputs=[
+        gr.Image(type="file", label="Generated Sketch"),  # For image output
+        gr.Textbox(label="Message")  # For message output (in case of NSFW content)
+    ],
     title="Kids Sketch Generator",
     description="Enter a text prompt and generate a fun sketch for kids with customizable inference steps and guidance scale."
 )