import gradio as gr
import cv2
import numpy as np
from PIL import Image
import torch
from AnimeGANv2.test import inference, load_model

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

generator = None

def get_generator():
    global generator
    if generator is None:
        generator = load_model(style="Hayao", device=DEVICE)
    return generator

def ghibli_transform(image):
    if image is None:
        return None, "Kein Bild"

    try:
        gen = get_generator()
        img_np = np.array(image)
        img_bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
        cartoon_bgr = inference.inference(gen, img_bgr, device=DEVICE)
        cartoon_rgb = cv2.cvtColor(cartoon_bgr, cv2.COLOR_BGR2RGB)
        result = Image.fromarray(cartoon_rgb)
        return result, "fertig"
    except Exception as e:
        return None, str(e)

css = """
.gradio-container { max-width: 1100px; margin: auto; padding: 1rem; }
h1 { color: #4CAF50; text-align: center; }
button { min-height: 50px; font-size: 1.1rem !important; }
"""

with gr.Blocks(css=css, theme=gr.themes.Soft(primary_hue="green")) as demo:
    gr.Markdown("# Foto → Studio Ghibli Style (Hayao)")
    
    with gr.Row():
        input_image = gr.Image(type="pil", label="Original", sources=["upload", "clipboard"], height=480)
        
        with gr.Column(min_width=280):
            btn = gr.Button("In Ghibli-Stil umwandeln", variant="primary")
        
        output_image = gr.Image(label="Ergebnis", height=480)

    status = gr.Textbox(label="Status", interactive=False)

    gr.Examples(
        examples=[
            "https://images.unsplash.com/photo-1506905925346-21bda4d32df4",
            "https://images.unsplash.com/photo-1543466835-00a7907e9de1",
            "https://images.unsplash.com/photo-1501785888041-af3ef285b470",
            "https://images.unsplash.com/photo-1470071459604-3b5ec3a7fe05",
        ],
        inputs=input_image
    )

    btn.click(
        fn=ghibli_transform,
        inputs=input_image,
        outputs=[output_image, status]
    )

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