import torch
from diffusers import FluxPipeline
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")
pipe.fuse_lora(lora_scale=1.5)
pipe.to("cuda")

# Load the NSFW classifier
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,
                ).images[0]
    

    # Classify the image for NSFW content
    image_classification = image_classifier(image)

    print(image_classification)

    # Check the classification results
    for result in image_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)
    return image_path

# Gradio interface with sliders for num_inference_steps and guidance_scale
interface = gr.Interface(
    fn=generate_sketch,
    inputs=[
        "text",  # Prompt input
        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="auto",
    title="Kids Sketch Generator",
    description="Enter a text prompt and generate a fun sketch for kids with customizable inference steps and guidance scale."
)

# Launch the app
interface.launch()