#!/usr/bin/env python3
"""
HarukiMix AI Image Generator
Generates beautiful Japanese girl images using the harukimix model
"""

import gradio as gr
import torch
from diffusers import StableDiffusionXLPipeline
import os
from pathlib import Path

# Model configuration
MODEL_ID = "John6666/haruki-mix-v21-sdxl"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Initialize pipeline
print(f"Loading model: {MODEL_ID}")
print(f"Using device: {DEVICE}")

try:
    pipe = StableDiffusionXLPipeline.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        use_safetensors=True,
        variant="fp16" if torch.cuda.is_available() else None
    )
    pipe = pipe.to(DEVICE)
    print("✓ Model loaded successfully")
except Exception as e:
    print(f"Error loading model: {e}")
    pipe = None

def generate_image(
    prompt: str,
    negative_prompt: str = "",
    num_inference_steps: int = 30,
    guidance_scale: float = 7.5,
    height: int = 768,
    width: int = 768,
    seed: int = -1
) -> tuple:
    """
    Generate an image using harukimix model
    
    Args:
        prompt: Text description of the image to generate
        negative_prompt: Things to avoid in the image
        num_inference_steps: Number of denoising steps (higher = better quality but slower)
        guidance_scale: How much to follow the prompt (higher = more adherence)
        height: Image height (must be multiple of 8)
        width: Image width (must be multiple of 8)
        seed: Random seed for reproducibility (-1 for random)
    
    Returns:
        Tuple of (image, info_text)
    """
    if pipe is None:
        return None, "❌ Model failed to load. Please check your GPU/CPU resources."
    
    if not prompt.strip():
        return None, "❌ Please enter a prompt"
    
    try:
        # Set seed for reproducibility
        if seed >= 0:
            generator = torch.Generator(device=DEVICE).manual_seed(seed)
        else:
            generator = None
        
        # Generate image
        with torch.no_grad():
            result = pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                height=height,
                width=width,
                num_inference_steps=num_inference_steps,
                guidance_scale=guidance_scale,
                generator=generator
            )
        
        image = result.images[0]
        info = f"✓ Generated successfully!\n\nPrompt: {prompt}\nSteps: {num_inference_steps}, Scale: {guidance_scale}"
        
        return image, info
        
    except Exception as e:
        error_msg = f"❌ Error during generation: {str(e)}"
        print(error_msg)
        return None, error_msg

# Create Gradio interface
with gr.Blocks(title="HarukiMix AI Image Generator", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🎨 HarukiMix AI Image Generator
    
    Generate beautiful Japanese girl images using the harukimix model.
    
    **Note:** First generation may take 30-60 seconds as the model loads.
    Subsequent generations will be faster.
    """)
    
    with gr.Row():
        with gr.Column(scale=2):
            # Input section
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="e.g., beautiful japanese girl, long black hair, school uniform, smile, detailed face, masterpiece",
                lines=3,
                value="beautiful japanese girl, long black hair, school uniform, smile, detailed face, masterpiece"
            )
            
            negative_prompt = gr.Textbox(
                label="Negative Prompt (what to avoid)",
                placeholder="e.g., blurry, low quality, distorted, ugly",
                lines=2,
                value="blurry, low quality, distorted, ugly, bad anatomy"
            )
            
            with gr.Row():
                num_steps = gr.Slider(
                    label="Inference Steps",
                    minimum=10,
                    maximum=50,
                    value=30,
                    step=1,
                    info="Higher = better quality but slower"
                )
                
                guidance_scale = gr.Slider(
                    label="Guidance Scale",
                    minimum=1.0,
                    maximum=20.0,
                    value=7.5,
                    step=0.5,
                    info="How much to follow the prompt"
                )
            
            with gr.Row():
                height = gr.Slider(
                    label="Height",
                    minimum=512,
                    maximum=1024,
                    value=768,
                    step=64
                )
                
                width = gr.Slider(
                    label="Width",
                    minimum=512,
                    maximum=1024,
                    value=768,
                    step=64
                )
            
            seed = gr.Number(
                label="Seed (-1 for random)",
                value=-1,
                precision=0
            )
            
            generate_btn = gr.Button("🎨 Generate Image", variant="primary", scale=2)
        
        with gr.Column(scale=1):
            # Output section
            output_image = gr.Image(label="Generated Image", type="pil")
            output_info = gr.Textbox(label="Status", lines=4)
    
    # Connect button to generation function
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt, negative_prompt, num_steps, guidance_scale, height, width, seed],
        outputs=[output_image, output_info]
    )
    
    # Example prompts
    gr.Markdown("## 📝 Example Prompts")
    gr.Examples(
        examples=[
            [
                "beautiful japanese girl, long black hair, school uniform, smile, detailed face, masterpiece, best quality, 8k",
                "blurry, low quality, distorted, ugly"
            ],
            [
                "cute japanese girl, pink hair, kawaii style, big eyes, happy expression, detailed, high quality",
                "blurry, low quality, distorted"
            ],
            [
                "japanese woman, elegant kimono, traditional style, beautiful face, detailed, masterpiece",
                "blurry, low quality, distorted, modern"
            ],
            [
                "young japanese girl, casual clothes, natural lighting, smile, realistic, detailed, high quality",
                "blurry, low quality, distorted, ugly, bad anatomy"
            ]
        ],
        inputs=[prompt, negative_prompt],
        label="Try these prompts"
    )

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True
    )