init

app.py

@@ -45,11 +45,32 @@ def parse_args() -> argparse.Namespace:
 
 
 def load_checkpoint():
+    dir = 'UGATIT_selfie2anime_lsgan_4resblock_6dis_1_1_10_10_1000_sn_smoothing'
     checkpoint_path = huggingface_hub.hf_hub_download(MODEL_REPO,
-                                                'UGATIT_selfie2anime_lsgan_4resblock_6dis_1_1_10_10_1000_sn_smoothing/checkpoint',
-                                                cache_dir='UGATIT_selfie2anime_lsgan_4resblock_6dis_1_1_10_10_1000_sn_smoothing')
+                                                dir+'/checkpoint',
+                                                cache_dir=dir,
+                                                force_filename='checkpoint')
     print(checkpoint_path)
-    return 'UGATIT_selfie2anime_lsgan_4resblock_6dis_1_1_10_10_1000_sn_smoothing'
+
+    checkpoint_path = huggingface_hub.hf_hub_download(MODEL_REPO,
+                                                      dir + '/UGATIT.model-1000000.index',
+                                                      cache_dir=dir,
+                                                      force_filename='UGATIT.model-1000000.index')
+    print(checkpoint_path)
+
+    checkpoint_path = huggingface_hub.hf_hub_download(MODEL_REPO,
+                                                      dir + '/UGATIT.model-1000000.meta',
+                                                      cache_dir=dir,
+                                                      force_filename='UGATIT.model-1000000.meta')
+    print(checkpoint_path)
+
+    checkpoint_path = huggingface_hub.hf_hub_download(MODEL_REPO,
+                                                      dir + '/UGATIT.model-1000000.data-00000-of-00001',
+                                                      cache_dir=dir,
+                                                      force_filename='UGATIT.model-1000000.data-00000-of-00001')
+    print(checkpoint_path)
+
+    return dir
 
 
 def run(