app.py

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

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

# import spaces #[uncomment to use ZeroGPU]

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

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

if __name__ == "__main__":
    from huggingface_hub import hf_hub_download
    CONFIG_PATH = 'configs/composite_config.yaml'

    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)
    
    print(dm)