Update app.py

app.py

@@ -3,10 +3,14 @@ import torch
 from annotator.util import resize_image, HWC3
 from cldm.model import create_model, load_state_dict
 from cldm.ddim_hacked import DDIMSampler
+from huggingface_hub import hf_hub_download
 
 # Initialize the model and other components
+# config = "./models/cldm_v21_512_latctrl_coltrans.yaml'"
 model = create_model('./models/cldm_v21_512_latctrl_coltrans.yaml').cpu()
-model.load_state_dict(load_state_dict('xywwww/scene_diffusion/checkpoints/epoch=25-step=112553.ckpt', location='cuda'), strict=False)
+ckpt = hf_hub_download(repo_id="xywwww/scene_diffusion", filename="checkpoints/epoch=25-step=112553.ckpt")
+# model.load_state_dict(load_state_dict('xywwww/scene_diffusion/checkpoints/epoch=25-step=112553.ckpt', location='cuda'), strict=False)
+model = load_model_checkpoint(model, ckpt)
 model = model.cuda()
 ddim_sampler = DDIMSampler(model)