import spaces
import gradio as gr
import torch
from PIL import Image
from diffusers import DiffusionPipeline
import random
from transformers import pipeline

torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
torch.backends.cuda.matmul.allow_tf32 = True

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

# 기본 모델 및 LoRA 설정
base_model = "black-forest-labs/FLUX.1-dev"
model_lora_repo = "Motas/Flux_Fashion_Photography_Style"  # 패션 모델 LoRA
clothes_lora_repo = "prithivMLmods/Canopus-Clothing-Flux-LoRA"  # 의류 LoRA

pipe = DiffusionPipeline.from_pretrained(base_model, torch_dtype=torch.bfloat16)
pipe.to("cuda")

MAX_SEED = 2**32-1

# 예제 설정
example_model_prompts = [
    "professional fashion model, full body shot, standing pose, natural lighting, studio background, high fashion, elegant pose",
    "fashion model portrait, upper body, confident pose, fashion photography, neutral background, professional lighting",
    "stylish fashion model, three-quarter view, editorial pose, high-end fashion magazine style, minimal background"
]

example_clothes_prompts = [
    "luxury designer sweater, cashmere material, cream color, cable knit pattern, high-end fashion, product photography",
    "elegant business blazer, tailored fit, charcoal grey, premium wool fabric, professional wear",
    "modern streetwear hoodie, oversized fit, minimalist design, premium cotton, urban style"
]

@spaces.GPU()
def generate_fashion(prompt, mode, cfg_scale, steps, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
    # 한글 감지 및 번역
    def contains_korean(text):
        return any(ord('가') <= ord(char) <= ord('힣') for char in text)
    
    if contains_korean(prompt):
        translated = translator(prompt)[0]['translation_text']
        actual_prompt = translated
    else:
        actual_prompt = prompt

    # 모드에 따른 LoRA 및 트리거워드 설정
    if mode == "Generate Model":
        pipe.load_lora_weights(model_lora_repo)
        trigger_word = "fashion photography, professional model"
    else:
        pipe.load_lora_weights(clothes_lora_repo)
        trigger_word = "upper clothing, fashion item"

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device="cuda").manual_seed(seed)

    progress(0, "Starting fashion generation...")

    for i in range(1, steps + 1):
        if i % (steps // 10) == 0:
            progress(i / steps * 100, f"Processing step {i} of {steps}...")

    image = pipe(
        prompt=f"{actual_prompt} {trigger_word}",
        num_inference_steps=steps,
        guidance_scale=cfg_scale,
        width=width,
        height=height,
        generator=generator,
        joint_attention_kwargs={"scale": lora_scale},
    ).images[0]

    progress(100, "Completed!")
    return image, seed

# CSS 정의
custom_css = """
body {
    background-color: #f5f5f5;
    font-family: 'Arial', sans-serif;
}

.container {
    max-width: 1200px;
    margin: 0 auto;
    padding: 20px;
}

.header {
    text-align: center;
    color: #333;
    margin-bottom: 30px;
    font-size: 2.5em;
    text-transform: uppercase;
    letter-spacing: 2px;
}

.box-common {
    background-color: white;
    border-radius: 15px;
    padding: 25px;
    margin: 20px 0;
    box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1);
}

.mode-box {
    background-color: #fff;
    padding: 20px;
    border-radius: 10px;
    margin-bottom: 20px;
    border: 2px solid #eee;
}

.result-box {
    width: 90%;
    max-width: 1000px;
    margin: 20px auto;
}

.image-output {
    width: 100%;
    max-width: 800px;
    margin: 0 auto;
    display: block;
}

.prompt-box {
    width: 90%;
    max-width: 1000px;
    margin: 20px auto;
}

.generate-btn {
    background: linear-gradient(45deg, #ff6b6b, #ff8e8e) !important;
    color: white !important;
    padding: 15px 30px !important;
    border-radius: 8px !important;
    border: none !important;
    font-size: 1.1em !important;
    cursor: pointer !important;
    transition: all 0.3s ease !important;
    width: 250px !important;
    margin: 20px auto !important;
    display: block !important;
    text-transform: uppercase !important;
    letter-spacing: 1px !important;
}

.generate-btn:hover {
    background: linear-gradient(45deg, #ff8e8e, #ff6b6b) !important;
    transform: translateY(-2px) !important;
    box-shadow: 0 5px 15px rgba(255, 107, 107, 0.4) !important;
}

.accordion {
    width: 90%;
    max-width: 1000px;
    margin: 20px auto;
}

.examples-table {
    background-color: rgba(255, 255, 255, 0.95);
    border-radius: 8px;
    padding: 15px;
    margin-top: 20px;
}

.parameter-box {
    background-color: #f8f9fa;
    padding: 20px;
    border-radius: 8px;
    margin: 10px 0;
}
"""

# Gradio 인터페이스 구성
with gr.Blocks(css=custom_css, theme="Yntec/HaleyCH_Theme_Orange") as app:
    with gr.Column(elem_classes="container"):
        gr.Markdown("# 🎭 Fashion Generation Studio", elem_classes="header")
        
        # 모드 선택
        with gr.Group(elem_classes="mode-box box-common"):
            mode = gr.Radio(
                choices=["Generate Model", "Generate Clothes"],
                label="Generation Mode",
                value="Generate Model"
            )
        
        # 이미지 출력 영역
        with gr.Group(elem_classes="result-box box-common"):
            gr.Markdown("### 🖼️ Generated Fashion")
            result = gr.Image(label="Result", elem_classes="image-output")
        
        # 프롬프트 입력 영역
        with gr.Group(elem_classes="prompt-box box-common"):
            prompt = gr.TextArea(
                label="✍️ Fashion Description (한글 또는 영어)",
                placeholder="패션 모델이나 의류를 설명하세요...",
                lines=5
            )
            generate_button = gr.Button(
                "🚀 Generate Fashion",
                elem_classes="generate-btn"
            )

        # 예제 탭
        with gr.Tabs():
            with gr.TabItem("Model Examples"):
                gr.Examples(
                    examples=example_model_prompts,
                    inputs=prompt,
                    label="Model Prompts"
                )
            with gr.TabItem("Clothes Examples"):
                gr.Examples(
                    examples=example_clothes_prompts,
                    inputs=prompt,
                    label="Clothes Prompts"
                )
        
        # 고급 옵션
        with gr.Accordion("🎨 Advanced Options", open=False, elem_classes="accordion box-common"):
            with gr.Group(elem_classes="parameter-box"):
                gr.Markdown("### 🎛️ Generation Parameters")
                with gr.Row():
                    with gr.Column():
                        cfg_scale = gr.Slider(
                            label="CFG Scale",
                            minimum=1,
                            maximum=20,
                            step=0.5,
                            value=7.0
                        )
                        steps = gr.Slider(
                            label="Steps",
                            minimum=1,
                            maximum=100,
                            step=1,
                            value=30
                        )
                        lora_scale = gr.Slider(
                            label="LoRA Scale",
                            minimum=0,
                            maximum=1,
                            step=0.01,
                            value=0.85
                        )
            
            with gr.Group(elem_classes="parameter-box"):
                gr.Markdown("### 📐 Image Dimensions")
                with gr.Row():
                    width = gr.Slider(
                        label="Width",
                        minimum=256,
                        maximum=1536,
                        step=64,
                        value=512
                    )
                    height = gr.Slider(
                        label="Height",
                        minimum=256,
                        maximum=1536,
                        step=64,
                        value=768
                    )
            
            with gr.Group(elem_classes="parameter-box"):
                gr.Markdown("### 🎲 Seed Settings")
                with gr.Row():
                    randomize_seed = gr.Checkbox(
                        True,
                        label="Randomize seed"
                    )
                    seed = gr.Slider(
                        label="Seed",
                        minimum=0,
                        maximum=MAX_SEED,
                        step=1,
                        value=42
                    )

    # 이벤트 핸들러
    generate_button.click(
        generate_fashion,
        inputs=[prompt, mode, cfg_scale, steps, randomize_seed, seed, width, height, lora_scale],
        outputs=[result, seed]
    )

app.queue()
app.launch()