Update app.py

app.py

@@ -15,8 +15,10 @@ class model:
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  if self.model_name != model_name:
  self.model_name = model_name
- model_ckpt_path = f"CELL-E_2-Image_Prediction/models/{model_name}.ckpt"
- model_config_path = f"CELL-E_2-Image_Prediction/models/{model_name}.yaml"
+ model_ckpt_path = hf_hub_download(repo_id=f"HuangLab/{model_name}", filename="model.ckpt")
+ config_config_path = hf_hub_download(repo_id=f"HuangLab/{model_name}", filename="config.yaml")
+ hf_hub_download(repo_id=f"HuangLab/{model_name}", filename="nucleus_vqgan.yaml")
+ hf_hub_download(repo_id=f"HuangLab/{model_name}", filename="threshold_vqgan.yaml")
 
  # Load model config and set ckpt_path if not provided in config
  config = OmegaConf.load(model_config_path)