init

app.py

@@ -15,7 +15,7 @@ import numpy as np
 import PIL.Image
 
 from io import BytesIO
-
+import ugatit_test
 
 ORIGINAL_REPO_URL = 'https://github.com/taki0112/UGATIT'
 TITLE = 'taki0112/UGATIT'
@@ -26,6 +26,9 @@ ARTICLE = """
 
 """
 
+
+MODEL_REPO = 'hylee/UGATIT_model'
+
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser()
     parser.add_argument('--device', type=str, default='cpu')
@@ -41,13 +44,22 @@ def parse_args() -> argparse.Namespace:
     return parser.parse_args()
 
 
+def load_checkpoint():
+    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')
+    print(checkpoint_path)
+    return 'UGATIT_selfie2anime_lsgan_4resblock_6dis_1_1_10_10_1000_sn_smoothing'
+
 
 def run(
-    image
+    image,
+    checkpoint_dir: str,
 ) -> tuple[PIL.Image.Image]:
-    
 
-    return PIL.Image.open(image.name)
+    result = ugatit_test.main_test(image.name, checkpoint_dir)
+
+    return PIL.Image.open(result)
 
 
 def main():
@@ -55,7 +67,9 @@ def main():
 
     args = parse_args()
 
-    func = functools.partial(run)
+    checkpoint_dir = load_checkpoint()
+
+    func = functools.partial(run, checkpoint_dir=checkpoint_dir)
     func = functools.update_wrapper(func, run)
 
     

ugatit_test.py

@@ -8,7 +8,7 @@ from ugatit.utils import *
 
 class UgatitTest:
 
-    def __init__(self, sess):
+    def __init__(self, sess, checkpoint_dir):
         self.light = False
 
         if self.light:
@@ -18,7 +18,7 @@ class UgatitTest:
 
         self.sess = sess
         self.phase = 'test'
-        self.checkpoint_dir = '/home/hylee/cartoon/UGATIT/checkpoint'
+        self.checkpoint_dir = checkpoint_dir
         self.result_dir = 'results'
         self.log_dir = 'logs'
         self.dataset_name = 'selfie2anime'
@@ -57,8 +57,8 @@ class UgatitTest:
         self.img_size = 256
         self.img_ch = 3
 
-        self.sample_dir = os.path.join('/home/hylee/cartoon/UGATIT/samples', self.model_dir)
-        check_folder(self.sample_dir)
+        #self.sample_dir = os.path.join('/home/hylee/cartoon/UGATIT/samples', self.model_dir)
+        #check_folder(self.sample_dir)
 
         # self.trainA, self.trainB = prepare_data(dataset_name=self.dataset_name, size=self.img_size
         self.trainA_dataset = glob('./dataset/{}/*.*'.format(self.dataset_name + '/trainA'))
@@ -350,12 +350,12 @@ class UgatitTest:
 
 
 gan = None
-def main_test(img_path):
+def main_test(img_path, checkpoint_dir):
     # open session
     sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
     global gan
     if gan is None:
-        gan = UgatitTest(sess)
+        gan = UgatitTest(sess, checkpoint_dir)
         # build graph
         gan.build_model()
         # show network architecture