import gradio as gr
from transformers import pipeline

# Load the model from Hugging Face using the pipeline
logo_pipeline = pipeline(model="Shakker-Labs/FLUX.1-dev-LoRA-Logo-Design")

# Function to generate logo from the prompt
def generate_logo(prompt):
    image = logo_pipeline(prompt)
    return image

# Gradio interface using Blocks
with gr.Blocks() as demo:
    gr.Markdown("## Logo Generator with Shakker-Labs/FLUX.1-dev-LoRA-Logo-Design")
    
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Logo Prompt", placeholder="Describe the logo you want to create...")
            submit_btn = gr.Button("Generate Logo")
        with gr.Column():
            output_image = gr.Image(label="Generated Logo")

    submit_btn.click(fn=generate_logo, inputs=prompt, outputs=output_image)

# Launch the Gradio app
demo.launch()