import gradio as gr
import numpy as np
import random
import spaces
import torch
import time
from diffusers import DiffusionPipeline, AutoencoderTiny
from custom_pipeline import FluxWithCFGPipeline
from transformers import pipeline

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

# Initialize translator
translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")

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

# Translation function
def translate_to_english(text):
    if any(ord('가') <= ord(char) <= ord('힣') for char in text):
        translated = translator(text)[0]['translation_text']
        return translated
    return text

# Inference function
@spaces.GPU(duration=25)
def generate_image(prompt, seed=24, width=DEFAULT_WIDTH, height=DEFAULT_HEIGHT, randomize_seed=False, progress=gr.Progress(track_tqdm=True)):
    # Translate prompt if Korean
    english_prompt = translate_to_english(prompt)
    
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(int(float(seed)))

    img = pipe.generate_images( 
            prompt=english_prompt,
            width=width,
            height=height,
            num_inference_steps=DEFAULT_INFERENCE_STEPS,
            generator=generator
        )
    return img, seed

# Example prompts
examples = [
    "sexy woman & man , under wear, full body, sunday",
    "A glamorous young woman with long, wavy blonde hair and smokey eye makeup, posing in a luxury hotel room. She's wearing a sparkly gold cocktail dress and holding up a white card with 'Invite' written on it in elegant calligraphy. Soft, warm lighting creates a luxurious atmosphere. ", 
    "A fit male fitness influencer with short dark hair and stubble, standing shirtless in a modern gym. He has defined abs and arm muscles, and is holding a protein shake in one hand and a card that says 'Invite' in the other. Bright, clean lighting highlights his physique.",
    "A bohemian-style female travel blogger with sun-kissed skin and messy beach waves, sitting on a tropical beach at sunset. She's wearing a flowy white sundress and holding up a weathered postcard with 'Invite scrawled on it. Golden hour lighting bathes the scene in warm tones. ",
    "A trendy male fashion influencer with perfectly styled hair and designer stubble, posing on a city street. He's wearing a tailored suit and holding up a sleek black business card with 'Invite' printed in minimalist white font. The background shows blurred city lights, creating a chic urban atmosphere.",
    "A fresh-faced young female beauty guru with freckles and natural makeup, sitting at a vanity covered in cosmetics. She's wearing a pastel pink robe and holding up a makeup palette with 'Invite' written on it in lipstick. Soft, flattering lighting enhances her radiant complexion. ",
    "A stylish young woman with long, wavy ombre hair and winged eyeliner, posing in front of a neon-lit city skyline at night. She's wearing a sleek black leather jacket over a sparkly crop top and holding up a holographic business card that says 'Invite' in futuristic font. The card reflects the colorful neon lights, creating a cyberpunk aesthetic.",    
]    

css = """
footer {visibility: hidden;}
.container {max-width: 1200px; margin: auto; padding: 20px;}
.generate-box, .image-box {
    background-color: #f0f0f0;
    border-radius: 10px;
    padding: 20px;
    margin-bottom: 20px;
    height: 600px;  /* 고정된 높이 설정 */
    display: flex;
    flex-direction: column;
}
.image-box img {
    max-height: 100%;
    width: 100%;
    object-fit: contain;
}
.generate-box .row {display: flex; align-items: center; margin-bottom: 10px;}
.generate-box .row > * {margin-right: 10px;}
.generate-box .row > *:last-child {margin-right: 0;}
.advanced-options {background-color: #e0e0e0; border-radius