import gradio as gr
import torch
import spaces # Import spaces for ZeroGPU compatibility on Hugging Face
from models import infer, load_model

# Initialize model on startup
try:
    load_model()
except Exception as e:
    print(f"Warning: Model failed to load on startup (might load on first request): {e}")

# Define the inference function wrapper for Gradio
# Decorate with @spaces.GPU to enable ZeroGPU usage on Hugging Face Spaces
@spaces.GPU
def process_image(prompt, negative_prompt, height, width, steps, guidance_scale, seed):
    try:
        image, used_seed = infer(
            prompt=prompt,
            negative_prompt=negative_prompt,
            height=int(height),
            width=int(width),
            steps=int(steps),
            guidance_scale=float(guidance_scale),
            seed=int(seed)
        )
        return image, f"Seed used: {used_seed}"
    except Exception as e:
        return None, f"Error: {str(e)}"

# CSS for styling
css = """
.container { max_width: 1200px; margin: auto; }
.header-link { font-weight: bold; color: #3b82f6; text-decoration: none; }
.header-link:hover { text-decoration: underline; }
"""

# Create the Gradio Interface
with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
    
    # Header Section
    with gr.Row():
        with gr.Column():
            gr.Markdown(
                """
                # 🎨 Kandinsky 5.0 T2I Lite
                ### Text-to-Image Generation using Kandinsky 5.0 Lite SFT Diffusers
                [Built with anycoder](https://huggingface.co/spaces/akhaliq/anycoder)
                """
            )

    with gr.Row(elem_classes="container"):
        # Left Column: Inputs
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Describe the image you want to generate...",
                lines=3,
                value="A fluffy, expressive cat wearing a bright red hat with a soft, slightly textured fabric. The hat should look cozy and well-fitted on the cat’s head. On the front of the hat, add clean, bold white text that reads “SWEET”, clearly visible and neatly centered."
            )
            
            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                placeholder="What to avoid (e.g., blurry, bad anatomy)...",
                lines=2,
                value="low quality, bad quality, sketches, lowres, blurry"
            )
            
            with gr.Accordion("Advanced Settings", open=False):
                with gr.Row():
                    height = gr.Slider(label="Height", minimum=512, maximum=2048, step=64, value=1024)
                    width = gr.Slider(label="Width", minimum=512, maximum=2048, step=64, value=1024)
                
                steps = gr.Slider(label="Inference Steps", minimum=10, maximum=100, step=1, value=50)
                guidance_scale = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=20.0, step=0.1, value=3.5)
                seed = gr.Number(label="Seed", value=-1, precision=0, info="Set to -1 for random")

            run_btn = gr.Button("Generate Image", variant="primary", size="lg")

        # Right Column: Output
        with gr.Column(scale=1):
            output_image = gr.Image(label="Generated Image", type="pil")
            status_text = gr.Textbox(label="Status", interactive=False)

    # Event Handling
    run_btn.click(
        fn=process_image,
        inputs=[prompt, negative_prompt, height, width, steps, guidance_scale, seed],
        outputs=[output_image, status_text]
    )

    # Examples
    gr.Examples(
        examples=[
            [
                "A fluffy, expressive cat wearing a bright red hat with a soft, slightly textured fabric. The hat should look cozy and well-fitted on the cat’s head. On the front of the hat, add clean, bold white text that reads “SWEET”, clearly visible and neatly centered.",
                "low quality, bad quality",
                1024, 1024, 50, 3.5, -1
            ],
            [
                "A futuristic cityscape at sunset with flying cars and neon lights, cyberpunk style, highly detailed.",
                "blurry, low resolution, monochrome",
                1024, 1024, 30, 4.0, -1
            ],
            [
                "A serene landscape painting of a mountain lake with reflection, oil painting style, textured brushstrokes.",
                "photorealistic, vector art, watermark",
                1024, 1024, 40, 3.0, -1
            ]
        ],
        inputs=[prompt, negative_prompt, height, width, steps, guidance_scale, seed],
        fn=process_image,
        outputs=[output_image, status_text],
        cache_examples=False
    )

if __name__ == "__main__":
    demo.launch()