import torch
from diffusers import StableDiffusionPipeline
from PIL import Image
import gradio as gr

# Load the pre-trained Stable Diffusion model
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16
)
pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")

# Logo generation function
def generate_logo(description):
    prompt = f"A high-quality, professional logo: {description}, sharp, modern, HD"
    image = pipe(
        prompt,
        height=512,
        width=512,
        num_inference_steps=50,
        guidance_scale=7.5
    ).images[0]
    return image.resize((1024, 1024), resample=Image.LANCZOS)

# Gradio interface
iface = gr.Interface(
    fn=generate_logo,
    inputs=gr.Textbox(label="Logo Description"),
    outputs=gr.Image(type="pil", label="Generated Logo"),
    title="AI Logo Generator",
    description="Enter a description to generate a high-quality logo"
)

iface.launch(share=True)