Spaces:
Runtime error
Runtime error
| import torch | |
| from diffusers import FluxPipeline | |
| from transformers import pipeline | |
| import gradio as gr | |
| import spaces | |
| device=torch.device('cuda') | |
| pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16) | |
| 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) | |
| NSFW_THRESHOLD = 0.3 | |
| # Define the function to generate the sketch | |
| def generate_sketch(prompt,style, num_inference_steps, guidance_scale): | |
| # Classify the text for NSFW content | |
| #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.") | |
| print(prompt) | |
| match style: | |
| case 'shou_xin': | |
| prompt = "shou_xin, " + prompt | |
| pipe.load_lora_weights("Datou1111/shou_xin", weight_name="shou_xin.safetensors") | |
| case 'sketched': | |
| prompt = "sketched style, " + prompt | |
| pipe.load_lora_weights("Shakker-Labs/FLUX.1-dev-LoRA-Children-Simple-Sketch", weight_name="FLUX-dev-lora-children-simple-sketch.safetensors") | |
| case 'sketch_paint': | |
| prompt = "Sketch paint, " + prompt | |
| pipe.load_lora_weights("strangerzonehf/Sketch-Paint", weight_name="Sketch-Paint.safetensors") | |
| case 'sketch_sized': | |
| prompt = "Sketch Sized, " + prompt | |
| pipe.load_lora_weights("strangerzonehf/Flux-Sketch-Sized-LoRA", weight_name="Flux-Sketch-Sized-LoRA.safetensors") | |
| case _: | |
| prompt = "shou_xin, " + prompt | |
| pipe.load_lora_weights("Datou1111/shou_xin", weight_name="shou_xin.safetensors") | |
| pipe.fuse_lora(lora_scale=1.5) | |
| 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 None,"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 image_path,None #gr.Image(value=image_path), gr.update(visible=False) | |
| # Gradio interface with sliders for num_inference_steps and guidance_scale | |
| interface = gr.Interface( | |
| fn=generate_sketch, | |
| inputs=[ | |
| "text", # Prompt input | |
| gr.Dropdown( | |
| ["sketched", "shou_xin","Sketch-Paint","Flux-Sketch-Sized-LoRA"],value='sketched', label="Style" | |
| ), | |
| 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(label="Generated Sketch"), | |
| gr.Textbox(label="Message") | |
| ], | |
| 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() |