Update app.py

app.py

@@ -7,10 +7,11 @@ 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 = create_model('./models/cldm_v21_512_latctrl_coltrans.yaml').cuda()
 ckpt = hf_hub_download(repo_id="xywwww/scene_diffusion", filename="checkpoints/epoch=25-step=112553.ckpt")
 print(ckpt)
 model.load_state_dict(load_state_dict(ckpt, location='cuda'), strict=False)
+model = model.cuda()
 # model = load_model_checkpoint(model, ckpt)
 ddim_sampler = DDIMSampler(model)