LTT commited on
Commit
d90e62e
1 Parent(s): 661dc7f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -267,7 +267,7 @@ print('Loading custom white-background unet ...')
267
  if os.path.exists(infer_config.unet_path):
268
  unet_ckpt_path = infer_config.unet_path
269
  else:
270
- unet_ckpt_path = hf_hub_download(repo_id="LTT/PRM", filename="diffusion_pytorch_model.bin", repo_type="model")
271
  state_dict = torch.load(unet_ckpt_path, map_location='cpu')
272
  pipeline.unet.load_state_dict(state_dict, strict=True)
273
 
@@ -279,7 +279,7 @@ model = instantiate_from_config(model_config)
279
  if os.path.exists(infer_config.model_path):
280
  model_ckpt_path = infer_config.model_path
281
  else:
282
- model_ckpt_path = hf_hub_download(repo_id="LTT/PRM", filename="final_ckpt.ckpt", repo_type="model")
283
  state_dict = torch.load(model_ckpt_path, map_location='cpu')['state_dict']
284
  state_dict = {k[14:]: v for k, v in state_dict.items() if k.startswith('lrm_generator.')}
285
  model.load_state_dict(state_dict, strict=True)
 
267
  if os.path.exists(infer_config.unet_path):
268
  unet_ckpt_path = infer_config.unet_path
269
  else:
270
+ unet_ckpt_path = hf_hub_download(repo_id="LTT/xxx", filename="diffusion_pytorch_model.bin", repo_type="model")
271
  state_dict = torch.load(unet_ckpt_path, map_location='cpu')
272
  pipeline.unet.load_state_dict(state_dict, strict=True)
273
 
 
279
  if os.path.exists(infer_config.model_path):
280
  model_ckpt_path = infer_config.model_path
281
  else:
282
+ model_ckpt_path = hf_hub_download(repo_id="LTT/xxx", filename="final_ckpt.ckpt", repo_type="model")
283
  state_dict = torch.load(model_ckpt_path, map_location='cpu')['state_dict']
284
  state_dict = {k[14:]: v for k, v in state_dict.items() if k.startswith('lrm_generator.')}
285
  model.load_state_dict(state_dict, strict=True)