import gradio as gr
from transformers import MarianMTModel, MarianTokenizer
import torch  # 추가된 부분
from hidiffusion import apply_hidiffusion
from diffusers import DiffusionPipeline, DDIMScheduler, AutoencoderKL
import spaces

# Load translation model
model_name = 'Helsinki-NLP/opus-mt-ko-en'
tokenizer = MarianTokenizer.from_pretrained(model_name)
model = MarianMTModel.from_pretrained(model_name)

# Define translation function
def translate_ko_to_en(text):
    translated = model.generate(**tokenizer(text, return_tensors="pt", padding=True))
    translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
    return translated_text

# Load image generation model
sd_model = "stabilityai/stable-diffusion-xl-base-1.0"
vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
scheduler = DDIMScheduler.from_pretrained(sd_model, subfolder="scheduler")
pipe = DiffusionPipeline.from_pretrained(sd_model, vae=vae, scheduler=scheduler, torch_dtype=torch.float16, use_safetensors=True, variant="fp16").to("cuda")
pipe.enable_vae_tiling()

# Define image generation function
@spaces.GPU
def run_hidiffusion(prompt, negative_prompt="", progress=gr.Progress(track_tqdm=True)):
    apply_hidiffusion(pipe)
    images = []
    for _ in range(2):
        result = pipe(prompt, guidance_scale=7.5, height=1024, width=1024, eta=1.0, negative_prompt=negative_prompt, num_inference_steps=25)
        images.append(result.images[0])
    return images

    
# Define prompt update function
def update_prompt(ages, styles, ethnicity, attire, hair, background_color):
    prompt_ko = f"{ages}, {styles}, {ethnicity}, {attire}, {hair}, {background_color}"
    prompt_en = translate_ko_to_en(prompt_ko)
    common_prompt = "full body visible, model pose, direct frontal gaze, msterpiece, realritics photo, High Quality, UHD"
    final_prompt = f"a {prompt_en}, {common_prompt}"
    return final_prompt

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# AI Human Model Generator")

    with gr.Tab("1024x1024픽셀"):
        with gr.Row():
            ages = gr.Dropdown(["young ages", "middle ages", "old ages", "highteen ages", "kid ages"], label="연령", value="young ages")   
            styles = gr.Dropdown(["Beautiful model woman", "Handsome model man", "Expert woman", "Expert man"], label="스타일", value="Beautiful model woman")            
            ethnicity = gr.Dropdown(["Asian ethnicity", "Caucasian ethnicity", "Black ethnicity"], label="인종", value="Asian ethnicity")
            attire = gr.Dropdown(["Business attire", "Nurse uniform attire", "Casual attire", "Golf wear attire", "T-shirts attire", "Coat attire"], label="의상", value="Business attire")
            hair = gr.Dropdown(["Long", "Curl", "Straight", "Short"], label="헤어 스타일", value="Long")
            background_color = gr.Dropdown(["White", "Blue", "Gray", "Black", "Green"], label="배경 스타일 컬러톤", value="White")
            
            negative_prompt_input = gr.Textbox(
                label="네가티브 프롬프트",
                value="(deformed, distorted, disfigured:1.3), poorly drawn, bad anatomy, wrong anatomy, extra limb, missing limb, floating limbs, (mutated hands and fingers:1.4), disconnected limbs, mutation, mutated, ugly, disgusting, blurry, amputation, (NSFW:1.25)",
                visible=False                
            )
        
        btn = gr.Button("Run")
        output_gallery = gr.Gallery(label="Results")
        btn.click(fn=lambda x, y, z, w, u, v: run_hidiffusion(update_prompt(x, y, z, w, u, v), negative_prompt_input.value), inputs=[ages, styles, ethnicity, attire, hair, background_color], outputs=[output_gallery])

demo.launch()