import gradio as gr
from diffusers import StableDiffusionPipeline
import torch

# Load Stable Diffusion model
model_id = "CompVis/stable-diffusion-v1-4"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")

# Function to generate an image
def generate_image(prompt, style):
    styled_prompt = f"{style} {prompt}"
    image = pipe(styled_prompt).images[0]
    return image

# Define Gradio interface
styles = ["Realistic", "Cartoon", "Abstract", "Fantasy"]

with gr.Blocks() as demo:
    gr.Markdown("# Illustrations4Free: Generate AI Illustrations")
    
    with gr.Row():
        prompt_input = gr.Textbox(label="Enter Your Prompt", placeholder="Describe the image you'd like to generate.")
        style_dropdown = gr.Dropdown(choices=styles, label="Choose a Style", value="Realistic")
    
    generate_button = gr.Button("Generate Illustration")
    output_image = gr.Image(label="Generated Illustration")

    generate_button.click(
        fn=generate_image,
        inputs=[prompt_input, style_dropdown],
        outputs=output_image
    )

demo.launch()