Update app.py
Browse files
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/
|
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/
|
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)
|