GrayShine commited on
Commit
355bbe3
1 Parent(s): 7f1217d

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -195,7 +195,7 @@ def init_model():
195
  else:
196
  raise ValueError("xformers is not available. Make sure it is installed correctly")
197
  ckpt_path = args.ckpt
198
- state_dict = find_model(ckpt_path)
199
  model.load_state_dict(state_dict)
200
  print('loading succeed')
201
  model.eval() # important!
 
195
  else:
196
  raise ValueError("xformers is not available. Make sure it is installed correctly")
197
  ckpt_path = args.ckpt
198
+ state_dict = state_dict = torch.load(ckpt_path, map_location=lambda storage, loc: storage)['ema']
199
  model.load_state_dict(state_dict)
200
  print('loading succeed')
201
  model.eval() # important!