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

# 예시 프롬프트 정의
model_examples = [
    "professional fashion model wearing elegant black dress in studio lighting",
    "fashion model in casual street wear, urban background",
    "high fashion model in avant-garde outfit on runway"
]

clothes_examples = [
    "luxurious red evening gown with detailed embroidery",
    "casual denim jacket with vintage wash",
    "modern minimalist white blazer with clean lines"
]

@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 == "패션 모델 생성":
        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

def update_examples(mode):
    if mode == "패션 모델 생성":
        return gr.Examples(examples=model_examples)
    else:
        return gr.Examples(examples=clothes_examples)

with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange") as app:
    gr.Markdown("# 🎭 Fashion AI Studio")
    gr.Markdown("AI를 사용하여 패션 모델과 의류를 생성해보세요")
    
    with gr.Column():
        # 모드 선택
        mode = gr.Radio(
            choices=["패션 모델 생성", "패션 의류 생성"],
            label="생성 모드",
            value="패션 모델 생성"
        )
        
        # 프롬프트 입력
        prompt = gr.TextArea(
            label="✍️ 패션 설명 (한글 또는 영어)",
            placeholder="패션 모델이나 의류를 설명하세요...",
            lines=3
        )
        
        # 예시 섹션
        with gr.Column() as example_container:
            examples = gr.Examples(
                examples=model_examples,
                inputs=prompt,
                label="예시 프롬프트"
            )
        
        # 결과 이미지
        result = gr.Image(label="생성된 이미지")
        
        generate_button = gr.Button("🚀 이미지 생성")
        
        # 고급 설정
        with gr.Accordion("🎨 고급 설정", 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="시드 랜덤화")
                seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, value=42)

    # 이벤트 핸들러
    mode.change(
        fn=update_examples,
        inputs=[mode],
        outputs=[examples]
    )
    
    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)