Spaces:
Sleeping
Sleeping
xxie
commited on
Commit
•
a6d435a
1
Parent(s):
198dde7
test
Browse files
demo.py
CHANGED
|
@@ -40,11 +40,6 @@ class DemoRunner:
|
|
| 40 |
cfg.model.model_name, cfg.model.predict_binary = 'diff-ho-attn', False # stage 2 does not predict segmentation
|
| 41 |
model_stage2 = CrossAttenHODiffusionModel(**cfg.model)
|
| 42 |
|
| 43 |
-
# Load from checkpoint
|
| 44 |
-
# ckpt_file1 = os.path.join(cfg.run.code_dir_abs, f'outputs/{cfg.run.stage1_name}/single/checkpoint-latest.pth')
|
| 45 |
-
# self.load_checkpoint(ckpt_file1, model_stage1)
|
| 46 |
-
# ckpt_file2 = os.path.join(cfg.run.code_dir_abs, f'outputs/{cfg.run.stage2_name}/single/checkpoint-latest.pth')
|
| 47 |
-
# self.load_checkpoint(ckpt_file2, model_stage2)
|
| 48 |
# Load ckpt from hf
|
| 49 |
ckpt_file1 = hf_hub_download("xiexh20/HDM-models", f'{cfg.run.stage1_name}.pth')
|
| 50 |
self.load_checkpoint(ckpt_file1, model_stage1)
|
|
|
|
| 40 |
cfg.model.model_name, cfg.model.predict_binary = 'diff-ho-attn', False # stage 2 does not predict segmentation
|
| 41 |
model_stage2 = CrossAttenHODiffusionModel(**cfg.model)
|
| 42 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 43 |
# Load ckpt from hf
|
| 44 |
ckpt_file1 = hf_hub_download("xiexh20/HDM-models", f'{cfg.run.stage1_name}.pth')
|
| 45 |
self.load_checkpoint(ckpt_file1, model_stage1)
|