import gradio as gr
import torch
from rudalle.pipelines import generate_images, show
from rudalle import get_rudalle_model, get_tokenizer, get_vae
from rudalle.utils import seed_everything
import random
import numpy as np
from PIL import Image

# Load model
device = "cuda" if torch.cuda.is_available() else "cpu"
model = get_rudalle_model('Malevich', pretrained=True, fp16=True, device=device)
vae = get_vae().to(device)
tokenizer = get_tokenizer()

def generate_images(prompt, negative_prompt=""):
    if not prompt.strip():
        prompt = "красивый пейзаж"  # default prompt in Russian if empty
    
    # Generate 4 images with random seeds
    images = []
    for _ in range(4):
        # Use random seed for each image to ensure variety
        seed = random.randint(0, 1000000)
        seed_everything(seed)
        
        with torch.no_grad():
            pil_images = generate_images(prompt, tokenizer, model, vae, top_k=512, top_p=0.995, images_num=1, negative_prompt=negative_prompt)
            if pil_images and len(pil_images) > 0:
                images.append(pil_images[0])
            else:
                # Create a placeholder image if generation fails
                img = Image.new('RGB', (512, 512), color='gray')
                images.append(img)
    
    return images

def select_image(images, evt: gr.SelectData):
    """Return the selected image for closer look"""
    if images and evt.index < len(images):
        return images[evt.index]
    return None

with gr.Blocks() as demo:
    gr.Markdown("# 🎨 ruDALL-E Malevich Image Generator")
    gr.HTML('<p style="text-align: center; margin-bottom: 20px;">Built with <a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank">anycoder</a></p>')
    
    with gr.Row():
        with gr.Column(scale=3):
            prompt_input = gr.Textbox(
                label="Prompt", 
                placeholder="Enter your prompt here (can be empty)...",
                lines=2,
                value=""
            )
            negative_prompt_input = gr.Textbox(
                label="Negative Prompt", 
                placeholder="What you don't want to see...",
                lines=2,
                value=""
            )
            generate_btn = gr.Button("Generate Images", variant="primary", size="lg")
        
        with gr.Column(scale=1):
            gr.Markdown("""
            ### Instructions:
            - Enter a prompt (Russian works best)
            - Leave empty for random landscapes
            - Negative prompt is optional
            - Click Generate to create 4 images
            - Click any image to see it larger
            """)
    
    with gr.Row():
        gallery = gr.Gallery(
            label="Generated Images",
            columns=2,
            rows=2,
            height="auto",
            object_fit="cover",
            show_label=True,
            allow_preview=True,
            interactive=True
        )
    
    with gr.Row():
        selected_image = gr.Image(
            label="Selected Image (Click an image above to view)",
            height=512,
            width=512,
            interactive=False
        )
    
    # Generate button event
    generate_btn.click(
        fn=generate_images,
        inputs=[prompt_input, negative_prompt_input],
        outputs=gallery,
        api_visibility="public"
    )
    
    # Gallery selection event
    gallery.select(
        fn=select_image,
        inputs=[gallery],
        outputs=[selected_image],
        api_visibility="private"
    )
    
    # Generate on enter key for prompt inputs
    prompt_input.submit(
        fn=generate_images,
        inputs=[prompt_input, negative_prompt_input],
        outputs=gallery
    )
    
    negative_prompt_input.submit(
        fn=generate_images,
        inputs=[prompt_input, negative_prompt_input],
        outputs=gallery
    )

# Launch with modern theme
demo.launch(
    theme=gr.themes.Soft(
        primary_hue="indigo",
        secondary_hue="purple",
        neutral_hue="slate",
        font=gr.themes.GoogleFont("Inter"),
        text_size="lg",
        spacing_size="lg",
        radius_size="md"
    ),
    footer_links=[{"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"}]
)