import gradio as gr
import torch
from diffusers import NewbiePipeline
import numpy as np
from PIL import Image
import io

def generate_image(prompt, steps=28, height=1024, width=1024):
    """Generate an image using the NewBie model"""
    try:
        # Load pipeline if not already loaded
        if not hasattr(generate_image, 'pipe'):
            model_id = "NewBie-AI/NewBie-image-Exp0.1"
            
            # Load pipeline
            generate_image.pipe = NewbiePipeline.from_pretrained(
                model_id,
                torch_dtype=torch.bfloat16,
            ).to("cuda")
        
        # Generate image
        image = generate_image.pipe(
            prompt,
            height=height,
            width=width,
            num_inference_steps=steps,
        ).images[0]
        
        return image
        
    except Exception as e:
        error_msg = f"Error generating image: {str(e)}"
        raise gr.Error(error_msg)

# Create a modern Gradio 6 application
with gr.Blocks() as demo:
    gr.Markdown(
        """
        # 🎨 NewBie Image Generator
        
        **Built with [anycoder](https://huggingface.co/spaces/akhaliq/anycoder)**
        
        Generate beautiful anime-style images using the NewBie AI model. Simply enter your prompt and click generate!
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Configuration")
            prompt_input = gr.Textbox(
                label="Prompt",
                placeholder="Describe the image you want to generate...",
                lines=2,
                max_lines=5,
                value="1girl, beautiful anime style, masterpiece"
            )
            
            steps_slider = gr.Slider(
                minimum=1,
                maximum=50,
                value=28,
                step=1,
                label="Inference Steps"
            )
            
            with gr.Row():
                height_input = gr.Number(
                    label="Height",
                    value=1024,
                )
                width_input = gr.Number(
                    label="Width",
                    value=1024,
            )
            
        with gr.Column(scale=2):
            gr.Markdown("### Generated Image")
            output_image = gr.Image(
                label="Generated Image",
                height=400,
                width=400
            )
    
    generate_btn = gr.Button(
        "✨ Generate Image",
        variant="primary",
        size="lg"
    )
    
    # Example prompts for users
    examples = gr.Examples(
        examples=[
            ["1girl, beautiful anime style, masterpiece, detailed face"],
            ["cute cat wearing a hat, cartoon style, vibrant colors"],
            ["fantasy landscape with mountains and rivers, epic scale"],
            ["portrait of a warrior with armor, dynamic lighting"],
            ["sci-fi cityscape at night, neon lights, futuristic"],
        ],
        inputs=[prompt_input],
    )
    
    # Status message
    status_text = gr.Textbox(
        label="Status",
        interactive=False,
        value="Ready to generate! Enter a prompt and click the button."
            ]
    
    # Event handling with Gradio 6 syntax
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt_input, steps_slider, height_input, width_input],
        outputs=[output_image],
        api_visibility="public"
    )
    
    # Add some additional functionality
    with gr.Accordion("Advanced Settings", open=False):
        guidance_scale_info = gr.Markdown(
            "Model: NewBie-AI/NewBie-image-Exp0.1\nInference Steps: 28\nResolution: 1024x1024"
        )
    
    # Info section
    gr.Markdown(
        """
        ---
        **Tips for better results:**
        - Be descriptive and specific in your prompts
        - Include style keywords (anime, realistic, cartoon, etc.)")
        )
    
    # Footer with additional info
    gr.Markdown(
        """
        <div style='text-align: center; margin-top: 20px; color: #666;'>
        This application uses the NewBie diffusion model to generate high-quality images from text descriptions.
        </div>
        """
    )

# Launch the application with modern Gradio 6 theme
demo.launch(
    theme=gr.themes.Soft(
        primary_hue="indigo",
        secondary_hue="slate",
        font=gr.themes.GoogleFont("Inter"),
        text_size="md",
        spacing_size="md",
        radius_size="md"
    ).set(
        button_primary_background_fill="*primary_600",
        button_primary_background_fill_hover="*primary_700",
        block_title_text_weight="600",
    ),
    footer_links=[
        {"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder",
        {"label": "NewBie Model", "url": "https://huggingface.co/NewBie-AI/NewBie-image-Exp0.1"
    ]
)