Update Time_TravelRephotography/projector.py

Time_TravelRephotography/projector.py

@@ -35,6 +35,9 @@ from models.degrade import (
     Downsample,
 )
 
+from huggingface_hub import hf_hub_download
+TOKEN = "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"
+
 
 def set_random_seed(seed: int):
     # FIXME (xuanluo): this setup still allows randomness somehow
@@ -67,11 +70,15 @@ def normalize(img: torch.Tensor, mean=0.5, std=0.5):
     return (img - mean) / std
 
 
-def create_generator(args: Namespace, device: torch.device):
-    generator = Generator(args.generator_size, 512, 8)
-    generator.load_state_dict(torch.load(args.ckpt)['g_ema'], strict=False)
+def create_generator(file_name: str, path:str,args: Namespace, device: torch.device):
+    path = hf_hub_download(f'{path}',
+                           f'{file_name}',
+                           use_auth_token=TOKEN)
+    with open(path, 'rb') as f:
+        generator = Generator(args.generator_size, 512, 8)
+        generator.load_state_dict(torch.load(f)['g_ema'], strict=False)
     generator.eval()
-    generator = generator.to(device)
+    generator.to(device)
     return generator