import gradio as gr
from diffusers import StableDiffusionPipeline

# Load the Hugging Face model
model_id = "kopyl/ui-icons-256"
pipeline = StableDiffusionPipeline.from_pretrained(model_id)
pipeline.to("cuda")  # Use GPU for faster inference, if available

def generate_icon(prompt):
    """Generate an icon from a text prompt."""
    result = pipeline(prompt).images[0]  # Generate the image
    return result

# Build Gradio Interface
with gr.Blocks() as app:
    gr.Markdown("## 🖌️ Icon Generator App")
    gr.Markdown("Generate 256x256 icons using the `kopyl/ui-icons-256` model.")
    
    with gr.Row():
        prompt_input = gr.Textbox(label="Enter a prompt", placeholder="E.g., 'a colorful app icon for a weather app'")
        generate_button = gr.Button("Generate Icon")
    
    output_image = gr.Image(label="Generated Icon")
    
    def on_click(prompt):
        return generate_icon(prompt)
    
    generate_button.click(on_click, inputs=prompt_input, outputs=output_image)

app.launch()  # Launch the app