import gradio as gr
import numpy as np
import random
import torch

from PIL import Image as im

from helper.cond_encoder import CLIPEncoder
from helper.loader import Loader

from auto_encoder.models.variational_auto_encoder import VariationalAutoEncoder
from clip.models.ko_clip import KoCLIPWrapper
from diffusion_model.sampler.ddim import DDIM
from diffusion_model.models.latent_diffusion_model import LatentDiffusionModel
from diffusion_model.network.unet import Unet
from diffusion_model.network.unet_wrapper import UnetWrapper
from huggingface_hub import hf_hub_download

# import spaces #[uncomment to use ZeroGPU]

device = "cuda" if torch.cuda.is_available() else "cpu"
loader = Loader(device)
repo_id = "JuyeopDang/KoFace-Diffusion"
CONFIG_PATH = 'configs/composite_config.yaml'

if torch.cuda.is_available():
    torch_dtype = torch.float16
else:
    torch_dtype = torch.float32

def load_model_from_HF(model, repo_id, filename, is_ema=False):
    try:
        model_path = hf_hub_download(repo_id=repo_id, filename=filename)
    except Exception as e:
        print(f"파일 다운로드 또는 모델 로드 중 오류 발생: {e}")
    model_path = model_path[:-4]
    model = loader.model_load(model_path, model, is_ema=is_ema, print_dict=False)
    return model

examples = [
    ['Guidance Scale이 0인 경우, Text Condition을 이용하지 않고 임의의 얼굴을 생성합니다.', 0, im.open("./assets/Example0.webp")],
    ['동그란 얼굴에 풍성하고 살짝 긴 커트머리와 거의 나오지 않은 성대가 남성보다는 여성적인 분위기를 냅니다. 또렷한 눈과 입술이 인물의 섬세함을 더욱 부각시키고 지적으로 보이게 만듭니다.', 2, im.open("./assets/Example1.webp")],
    ['코가 큰 남자다.', 4, im.open("./assets/Example2.webp")],
    ['아줌마 파마같이 크고 부한 머리에 이목구비가 약간 여성스러워 보이는 얼굴이며, 적지않은 주름과 처진 눈에서 연륜이 느꺼진다. 큰 눈과 넓은 팔자 주름이 있는 고집스런 입매는 완고한 원칙주의자 같은 느낌이 있다.', 4, im.open("./assets/Example3.webp")],
    ['헤어 손질이 다소 미숙하여 세련된 느낌이 부족하다. 눈 끝이 올라가 있어 눈빛이 날카롭고 예민해 보인다. 전체적으로 마른 체격일 것으로 보이며, 업무 처리 능력은 뛰어나겠지만, 교우 관계는 원만하지 않을 수도 있다.', 7, im.open("./assets/Example4.webp")]
]


if __name__ == "__main__":
    vae = VariationalAutoEncoder(CONFIG_PATH)
    sampler = DDIM(CONFIG_PATH)
    clip = KoCLIPWrapper()
    cond_encoder = CLIPEncoder(clip, CONFIG_PATH)
    network = UnetWrapper(Unet, CONFIG_PATH, cond_encoder)
    dm = LatentDiffusionModel(network, sampler, vae)

    vae = load_model_from_HF(vae, repo_id, "composite_epoch2472.pth", False)
    clip = load_model_from_HF(clip, repo_id, "asian-composite-fine-tuned-koclip.pth", True)
    dm = load_model_from_HF(dm, repo_id, "asian-composite-clip-ldm.pth", True)

    def generate_image(y, gamma):
        images = dm.sample(2, y = y, gamma = gamma)
        images = images.permute(0, 2, 3, 1) 
        if type(images) is torch.Tensor:
          images = images.detach().cpu().numpy()
        images = np.clip(images / 2 + 0.5, 0, 1)
        return im.fromarray((images[0] * 255).astype(np.uint8))

    with gr.Blocks() as demo:
        gr.Markdown(
            """
            # KoFace AI - 한국인 몽타주 생성기
            * 몽타주를 설명하는 텍스트를 입력하고 '생성하기' 버튼을 눌러주세요.
            
            * **참고**: 현재 무료 티어를 이용하고 있어, **생성에 약 10분 정도의 시간**이 걸릴 수 있습니다. 
    
            * 이 AI는 Latent Diffusion Model을 직접 구현하여 만들었습니다. 자세한 코드나 실험 결과는 GitHub 저장소를 참고해 주세요.
    
            🔗 [GitHub 저장소 바로가기](https://github.com/Won-Seong/simple-latent-diffusion-model) | 📄 [사용한 데이터세트](https://www.aihub.or.kr/aihubdata/data/view.do?dataSetSn=618)
    
            
            """
        )
    
        with gr.Row():
            with gr.Column():
                text_input = gr.Textbox(
                    placeholder="몽타주를 설명하는 텍스트를 입력하세요.", 
                    label="Text Condition"
                )
                guidance_slider = gr.Slider(
                    0, 10, value=2, 
                    label="Guidance Scale", 
                    info="이 숫자가 크면 클수록 입력 텍스트를 더 강하게 이용하여 몽타주를 생성합니다."
                )
                submit_btn = gr.Button("생성하기")
    
            with gr.Column():
                image_output = gr.Image(label="생성된 몽타주")
    
        submit_btn.click(fn=generate_image, inputs=[text_input, guidance_slider], outputs=image_output)

        gr.Examples(examples, inputs = [text_input, guidance_slider, image_output])
    
    demo.launch()