Update app.py

app.py

@@ -4,6 +4,7 @@ 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
@@ -15,29 +16,35 @@ 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
 
-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)
-    
+def load_model_from_HF(model, repo_id, filename, is_ema=False):
     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}")
+    model_path = model_path[:-4]
+    model = loader.model_load(model_path, model, is_ema=is_ema, print_dict=False)
+    return model
 
-    state_dict = torch.load(model_path, map_location=device)
-    vae.load_state_dict(state_dict['model_state_dict'])
-    print(vae)
+if __name__ == "__main__":
+    from huggingface_hub import hf_hub_download
+    CONFIG_PATH = 'configs/composite_config.yaml'
+
+    vae = VariationalAutoEncoder(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)