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;
}
"""


with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange") as app:
    gr.Markdown("# 🎭 Fashion AI Studio")
    gr.Markdown("Generate fashion images and try on virtual clothing using AI")
    
    with gr.Tabs():
        # Virtual Try-On 탭
        with gr.TabItem("👔 Virtual Try-On"):
            with gr.Row():
                with gr.Column():
                    prompt_input = gr.Textbox(
                        label="Style Description", 
                        placeholder="Describe the desired style (e.g., 'person wearing elegant dress')"
                    )
                    with gr.Row():
                        with gr.Group():
                            structure_image = gr.Image(
                                label="Your Photo (Full-body)", 
                                type="filepath"
                            )
                            gr.Markdown("*Upload a clear, well-lit full-body photo*")
                            depth_strength = gr.Slider(
                                minimum=0, 
                                maximum=50, 
                                value=15, 
                                label="Fitting Strength"
                            )
                        with gr.Group():
                            style_image = gr.Image(
                                label="Clothing Item", 
                                type="filepath"
                            )
                            gr.Markdown("*Upload the clothing item you want to try on*")
                            style_strength = gr.Slider(
                                minimum=0, 
                                maximum=1, 
                                value=0.5, 
                                label="Style Transfer Strength"
                            )
                    tryon_btn = gr.Button("Generate Try-On")
                    
                    gr.Examples(
                        examples=examples,
                        inputs=[prompt_input, structure_image, style_image, depth_strength, style_strength],
                        outputs=[output_image],
                        fn=generate_image,
                        cache_examples=True
                    )
                
                with gr.Column():
                    output_image.render()

        # Fashion Generation 탭
        with gr.TabItem("👗 Fashion Generation"):
            with gr.Column():
                # 모드 선택
                with gr.Group():
                    mode = gr.Radio(
                        choices=["Generate Model", "Generate Clothes"],
                        label="Generation Mode",
                        value="Generate Model"
                    )
                
                # 프롬프트 입력
                prompt = gr.TextArea(
                    label="✍️ Fashion Description (한글 또는 영어)",
                    placeholder="패션 모델이나 의류를 설명하세요...",
                    lines=5
                )
                
                # 결과 이미지
                result = gr.Image(label="Generated Fashion")
                
                generate_button = gr.Button("🚀 Generate Fashion")
                
                # 고급 설정 아코디언
                with gr.Accordion("🎨 Advanced Options", open=False):
                    with gr.Row():
                        cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=20, value=7.0)
                        steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=30)
                        lora_scale = gr.Slider(label="LoRA Scale", minimum=0, maximum=1, value=0.85)
                    
                    with gr.Row():
                        width = gr.Slider(label="Width", minimum=256, maximum=1536, value=512)
                        height = gr.Slider(label="Height", minimum=256, maximum=1536, value=768)
                    
                    with gr.Row():
                        randomize_seed = gr.Checkbox(True, label="Randomize seed")
                        seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, value=42)

                # 예제 탭
                with gr.Tabs():
                    with gr.TabItem("Model Examples"):
                        gr.Examples(examples=example_model_prompts, inputs=prompt)
                    with gr.TabItem("Clothes Examples"):
                        gr.Examples(examples=example_clothes_prompts, inputs=prompt)

    # 이벤트 핸들러
    tryon_btn.click(
        fn=generate_image,
        inputs=[prompt_input, structure_image, style_image, depth_strength, style_strength],
        outputs=[output_image]
    )
    
    generate_button.click(
        generate_fashion,
        inputs=[prompt, mode, cfg_scale, steps, randomize_seed, seed, width, height, lora_scale],
        outputs=[result, seed]
    )

if __name__ == "__main__":
    app.launch(share=True)