Update Time_TravelRephotography/projector.py

Time_TravelRephotography/projector.py

@@ -67,9 +67,13 @@ 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)
     return generator