jiangyzy commited on
Commit
30bc91f
1 Parent(s): 595554b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +7 -7
app.py CHANGED
@@ -195,15 +195,15 @@ 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
 
202
- # ckpt = torch.load(model_path, map_location="cpu")
203
- # model.load_state_dict(ckpt)
204
- # del ckpt
205
  model = model.to(device)
206
- sampler = DDIMSampler(model, device=device)
207
 
208
  # load demo
209
  demo = gr.Blocks()
 
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
 
202
+ ckpt = torch.load(model_path, map_location="cpu")
203
+ model.load_state_dict(ckpt)
204
+ del ckpt
205
  model = model.to(device)
206
+ sampler = None
207
 
208
  # load demo
209
  demo = gr.Blocks()