app.py

import torch
from diffusers import FluxPipeline
import gradio as gr
import spaces

# 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")

# Define the function to generate the sketch
@spaces.GPU
def generate_sketch(prompt, num_inference_steps, guidance_scale):
    image = pipe(prompt, 
                 num_inference_steps=num_inference_steps, 
                 guidance_scale=guidance_scale,
                ).images[0]
    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="image",
    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()