import gradio as gr
from diffusers import AutoPipelineForText2Image
import torch

# Initialize the Imagen 4 model
pipeline = AutoPipelineForText2Image.from_pretrained(
    "google/imagen-4",
    torch_dtype=torch.float16,
    variant="fp16",
).to("cuda")

def generate_image(prompt, negative_prompt="", steps=30, guidance_scale=7.5):
    """Generate image from text prompt using Imagen 4"""
    with torch.inference_mode():
        result = pipeline(
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=steps,
            guidance_scale=guidance_scale
        )
        return result.images[0]

# Custom theme for modern look
custom_theme = gr.themes.Soft(
    primary_hue="blue",
    secondary_hue="indigo",
    neutral_hue="slate",
    font=gr.themes.GoogleFont("Inter"),
    text_size="lg",
    spacing_size="lg",
    radius_size="md"
).set(
    button_primary_background_fill="*primary_600",
    button_primary_background_fill_hover="*primary_700",
    block_title_text_weight="600",
)

# Gradio 6 app
with gr.Blocks() as demo:
    gr.Markdown(
        """
        # 🎨 Imagen 4 Image Generator
        **Built with [anycoder](https://huggingface.co/spaces/akhaliq/anycoder)**
        Generate high-quality images from text prompts using Google's Imagen 4 model.
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Describe the image you want to generate...",
                lines=3
            )
            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                placeholder="Things you don't want in the image...",
                lines=2
            )
            
            with gr.Accordion("Advanced Settings", open=False):
                steps = gr.Slider(
                    label="Inference Steps",
                    minimum=10,
                    maximum=100,
                    value=30,
                    step=1
                )
                guidance_scale = gr.Slider(
                    label="Guidance Scale",
                    minimum=1.0,
                    maximum=20.0,
                    value=7.5,
                    step=0.1
                )
            
            generate_btn = gr.Button("Generate Image", variant="primary")
        
        with gr.Column(scale=1):
            output_image = gr.Image(
                label="Generated Image",
                height=512,
                width=512
            )
            with gr.Row():
                download_btn = gr.Button("Download")
                share_btn = gr.Button("Share", variant="secondary")
    
    # Event handlers
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt, negative_prompt, steps, guidance_scale],
        outputs=output_image,
        api_visibility="public"
    )
    
    # Example prompts
    gr.Examples(
        examples=[
            ["A majestic lion in the savannah at sunset, photorealistic"],
            ["Cyberpunk cityscape at night with neon lights"],
            ["Cute anime-style cat wearing a spacesuit"],
        ],
        inputs=prompt,
        label="Try these examples!"
    )

# Launch with modern theme and settings
demo.launch(
    theme=custom_theme,
    footer_links=[
        {"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"},
        {"label": "Imagen 4 Model", "url": "https://huggingface.co/google/imagen-4"}
    ],
    allowed_paths=["."]
)