C

app.py

@@ -11,6 +11,8 @@ import os
 os.environ["CUDA_VISIBLE_DEVICES"] = "7"
 # os.environ["OPENAI_LOGDIR"] = "./logs"
 # os.environ["MPI_DISABLED"] = "1"
+# os.environ.getattribute("HF_TOKEN")
+token = os.getenv("HF_TOKEN", None)
 import torch
 import torch.distributed as dist
 import torchvision.transforms as transforms
@@ -516,7 +518,7 @@ setattr(diffusion, "settings", settings)
 
 
 pretrained_dewarp_model = GeoTr_Seg_Inf()
-settings.env.seg_model_path = hf_hub_download(repo_id="hanquansanren/DvD", filename="seg.pth")
+settings.env.seg_model_path = hf_hub_download(repo_id="hanquansanren/DvD", filename="seg.pth", token=token)
 reload_segmodel(pretrained_dewarp_model.msk, settings.env.seg_model_path)
 # reload_model(pretrained_dewarp_model.GeoTr, settings.env.dewarping_model_path)
 pretrained_dewarp_model.to(dist_util.dev())
@@ -525,19 +527,19 @@ pretrained_dewarp_model.eval()
 if settings.env.use_line_mask:
     pretrained_line_seg_model = UNet(n_channels=3, n_classes=1)
     pretrained_seg_model = Seg()
-    settings.env.line_seg_model_path = hf_hub_download(repo_id="hanquansanren/DvD", filename="line_model2.pth")
+    settings.env.line_seg_model_path = hf_hub_download(repo_id="hanquansanren/DvD", filename="line_model2.pth", token=token)
     line_model_ckpt = dist_util.load_state_dict(settings.env.line_seg_model_path, map_location='cpu')['model']
     pretrained_line_seg_model.load_state_dict(line_model_ckpt, strict=True)    
     pretrained_line_seg_model.to(dist_util.dev())
     pretrained_line_seg_model.eval()
 
-    settings.env.new_seg_model_path = hf_hub_download(repo_id="hanquansanren/DvD", filename="seg_model.pth")
+    settings.env.new_seg_model_path = hf_hub_download(repo_id="hanquansanren/DvD", filename="seg_model.pth", token=token)
     seg_model_ckpt = dist_util.load_state_dict(settings.env.new_seg_model_path, map_location='cpu')['model']
     pretrained_seg_model.load_state_dict(seg_model_ckpt, strict=True)    
     pretrained_seg_model.to(dist_util.dev())
     pretrained_seg_model.eval()
 
-settings.env.model_path = hf_hub_download(repo_id="hanquansanren/DvD", filename="model1852000.pt")
+settings.env.model_path = hf_hub_download(repo_id="hanquansanren/DvD", filename="model1852000.pt", token=token)
 model.cpu().load_state_dict(dist_util.load_state_dict(settings.env.model_path, map_location="cpu"), strict=False)
 logger.log(f"Model loaded with {settings.env.model_path}")
 
@@ -564,4 +566,4 @@ if __name__ == '__main__':
 
 
 
-    
+