import gradio as gr
from diffusers import StableDiffusionXLImg2ImgPipeline
import torch

# Load a lightweight pipeline that works well on CPU
def load_image_generator():
    try:
        model = StableDiffusionXLImg2ImgPipeline.from_pretrained(
            "stabilityai/stable-diffusion-2-1", 
            torch_dtype=torch.float16,
            variant="fp16",
            use_safetensors=True
        )
        # Ensure it runs on CPU
        #model = model.to("cpu")
        return model
    except Exception as e:
        print(f"Error loading model: {e}")
        return None

# Generate chatbot icon
def generate_chatbot_icon(
    prompt, 
    negative_prompt="low quality, bad composition, blurry", 
    num_inference_steps=20,
    guidance_scale=7.5,
    strength=0.75
):
    # Load the model
    model = load_image_generator()
    if model is None:
        return None
    
    # Default icon if no initial image
    default_init_image = torch.randn((1, 3, 512, 512))
    
    try:
        # Generate the image
        image = model(
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            strength=strength,
            image=default_init_image
        ).images[0]
        
        return image
    except Exception as e:
        print(f"Error generating image: {e}")
        return None

# Create Gradio interface
def create_gradio_interface():
    with gr.Blocks() as demo:
        gr.Markdown("# 🤖 Chatbot Icon Generator")
        
        with gr.Row():
            with gr.Column():
                # Prompt input
                prompt = gr.Textbox(
                    label="Icon Description", 
                    value="Cute minimalist chatbot avatar, clean design, friendly expression, cartoon style"
                )
                
                # Generate button
                generate_btn = gr.Button("Generate Icon")
                
            with gr.Column():
                # Output image
                output_image = gr.Image(label="Generated Chatbot Icon")
        
        # Connect generate button to function
        generate_btn.click(
            fn=generate_chatbot_icon, 
            inputs=[prompt], 
            outputs=[output_image]
        )
    
    return demo

# Launch the app
if __name__ == "__main__":
    demo = create_gradio_interface()
    demo.launch()