import gradio as gr
import numpy as np
import random
import spaces
import torch
import time
import os
from diffusers import DiffusionPipeline
from custom_pipeline import FLUXPipelineWithIntermediateOutputs
from transformers import pipeline

# 번역 모델 로드
translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")

# Constants
MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048
DEFAULT_WIDTH = 1024
DEFAULT_HEIGHT = 1024
DEFAULT_INFERENCE_STEPS = 1

# Device and model setup
dtype = torch.float16
pipe = FLUXPipelineWithIntermediateOutputs.from_pretrained(
    "black-forest-labs/FLUX.1-schnell", torch_dtype=dtype
).to("cuda")
torch.cuda.empty_cache()

# 한글 메뉴 이름 dictionary
korean_labels = {
    "Generated Image": "생성된 이미지",
    "Prompt": "프롬프트",
    "Enhance Image": "이미지 향상",
    "Advanced Options": "고급 옵션",
    "Seed": "시드",
    "Randomize Seed": "시드 무작위화",
    "Width": "너비",
    "Height": "높이",
    "Inference Steps": "추론 단계",
    "Inspiration Gallery": "영감 갤러리"
}

def translate_if_korean(text):
    if any('\u3131' <= char <= '\u3163' or '\uac00' <= char <= '\ud7a3' for char in text):
        return translator(text)[0]['translation_text']
    return text

# Inference function
@spaces.GPU(duration=25)
def generate_image(prompt, seed=42, width=DEFAULT_WIDTH, height=DEFAULT_HEIGHT, randomize_seed=False, num_inference_steps=DEFAULT_INFERENCE_STEPS):
    
    prompt = translate_if_korean(prompt)
    
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)

    start_time = time.time()

    # Only generate the last image in the sequence
    for img in pipe.generate_images(  
            prompt=prompt,
            guidance_scale=0,
            num_inference_steps=num_inference_steps,
            width=width,
            height=height,
            generator=generator
        ): 
        latency = f"처리 시간: {(time.time()-start_time):.2f} 초"    
        yield img, seed, latency

# Example prompts
examples = [
    "달에서 알에서 부화하는 작은 우주 비행사",
    "안녕하세요 세상이라고 쓰인 표지판을 들고 있는 고양이",
    "비너 슈니첼의 애니메이션 일러스트레이션",
    "하늘을 나는 자동차와 네온 불빛이 있는 미래적인 도시 풍경",
    "긴 갈색 웨이브 머리를 올려 묶고 안경을 쓴 젊은 여성의 사진. 그녀는 흰 피부에 눈과 입술을 강조한 은은한 화장을 했습니다. 그녀는 검은색 상의를 입었습니다. 배경은 도시 건물 외관으로 보이며, 햇빛이 그녀의 얼굴에 따뜻한 빛을 비추고 있습니다.",
    "스티브 잡스를 스타워즈 영화 캐릭터로 상상해보세요"
]

css = """
footer {
    visibility: hidden;
}
"""

# --- Gradio UI ---
with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange", css=css) as demo:
    with gr.Column(elem_id="app-container"):

        with gr.Row():
            with gr.Column(scale=3):
                result = gr.Image(label=korean_labels["Generated Image"], show_label=False, interactive=False)
            with gr.Column(scale=1):
                prompt = gr.Text(
                    label=korean_labels["Prompt"],
                    placeholder="생성하고 싶은 이미지를 설명하세요...",
                    lines=3,
                    show_label=False,
                    container=False,
                )
                enhanceBtn = gr.Button(f"🚀 {korean_labels['Enhance Image']}")

                with gr.Column(korean_labels["Advanced Options"]):
                    with gr.Row():
                        latency = gr.Text(show_label=False)
                    with gr.Row():
                        seed = gr.Number(label=korean_labels["Seed"], value=42, precision=0)
                        randomize_seed = gr.Checkbox(label=korean_labels["Randomize Seed"], value=False)
                    with gr.Row():
                        width = gr.Slider(label=korean_labels["Width"], minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=DEFAULT_WIDTH)
                        height = gr.Slider(label=korean_labels["Height"], minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=DEFAULT_HEIGHT)
                        num_inference_steps = gr.Slider(label=korean_labels["Inference Steps"], minimum=1, maximum=4, step=1, value=DEFAULT_INFERENCE_STEPS)

        with gr.Row():
            gr.Markdown(f"### 🌟 {korean_labels['Inspiration Gallery']}")
        with gr.Row():
            gr.Examples(
                examples=examples,
                fn=generate_image,
                inputs=[prompt],
                outputs=[result, seed],
                cache_examples="lazy" 
            )

    # Event handling - Trigger image generation on button click or input change
    enhanceBtn.click(
        fn=generate_image,
        inputs=[prompt, seed, width, height],
        outputs=[result, seed, latency],
        show_progress="hidden",
        show_api=False,
        queue=False
    )

    gr.on(
        triggers=[prompt.input, width.input, height.input, num_inference_steps.input],
        fn=generate_image,
        inputs=[prompt, seed, width, height, randomize_seed, num_inference_steps],
        outputs=[result, seed, latency],
        show_progress="hidden",
        show_api=False,
        trigger_mode="always_last",
        queue=False
    )

# Launch the app
demo.launch()