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

from helper.cond_encoder import CLIPEncoder

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

# import spaces #[uncomment to use ZeroGPU]

device = "cuda" if torch.cuda.is_available() else "cpu"

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

if __name__ == "__main__":
    from huggingface_hub import hf_hub_download
    CONFIG_PATH = 'configs/composite_config.yaml'
    
    repo_id = "JuyeopDang/KoFace-Diffusion"
    filename = "composite_epoch2472.pth" # 예: "pytorch_model.pt" 또는 "model.pt"
    vae = VariationalAutoEncoder(CONFIG_PATH)
    
    try:
        # 파일 다운로드
        # cache_dir을 지정하면 다운로드된 파일이 저장될 경로를 제어할 수 있습니다.
        # 기본적으로는 ~/.cache/huggingface/hub 에 저장됩니다.
        model_path = hf_hub_download(repo_id=repo_id, filename=filename)
        print(f"모델 가중치 파일이 다음 경로에 다운로드되었습니다: {model_path}")
    except Exception as e:
        print(f"파일 다운로드 또는 모델 로드 중 오류 발생: {e}")

    state_dict = torch.load(model_path, map_location='cuda')
    vae.load_state_dict(state_dict['model_state_dict'])
    print(vae)