jiangyzy commited on
Commit
ba01e37
1 Parent(s): 094b404

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -6
app.py CHANGED
@@ -195,12 +195,12 @@ def main(args):
195
  config = OmegaConf.load("configs/config_customnet.yaml")
196
  model = instantiate_from_config(config.model)
197
 
198
- # model_path='./customnet_v1.pt?download=true'
199
- # if not os.path.exists(model_path):
200
- # os.system(f'wget https://huggingface.co/TencentARC/CustomNet/resolve/main/customnet_v1.pt?download=true -P .')
201
- # ckpt = torch.load(model_path, map_location="cpu")
202
- # model.load_state_dict(ckpt)
203
- # del ckpt
204
 
205
  model = model.to(device)
206
  sampler = None
 
195
  config = OmegaConf.load("configs/config_customnet.yaml")
196
  model = instantiate_from_config(config.model)
197
 
198
+ model_path='./customnet_v1.pt?download=true'
199
+ if not os.path.exists(model_path):
200
+ os.system(f'wget https://huggingface.co/TencentARC/CustomNet/resolve/main/customnet_v1.pt?download=true -P .')
201
+ ckpt = torch.load(model_path, map_location="cpu")
202
+ model.load_state_dict(ckpt)
203
+ del ckpt
204
 
205
  model = model.to(device)
206
  sampler = None