Switched to config.json

app.py

@@ -10,8 +10,8 @@ from huggingface_hub import hf_hub_download
 
 # pull files from hub
 token = os.environ.get("HF_TOKEN", None)
-yaml_file_path = hf_hub_download(repo_id="ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL",
-                                 filename="Prithvi_EO_V2_300M_TL_config.yaml", token=token)
+config_path = hf_hub_download(repo_id="ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL",
+                                 filename="config.json", token=token)
 checkpoint = hf_hub_download(repo_id="ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL",
                              filename='Prithvi_EO_V2_300M_TL.pt', token=token)
 model_def = hf_hub_download(repo_id="ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL",
@@ -67,7 +67,7 @@ def extract_rgb_imgs(input_img, rec_img, mask_img, channels, mean, std):
     return outputs
 
 
-def predict_on_images(data_files: list, yaml_file_path: str, checkpoint: str, mask_ratio: float = None):
+def predict_on_images(data_files: list, config_path: str, checkpoint: str, mask_ratio: float = None):
     try:
         data_files = [x.name for x in data_files]
         print('Path extracted from example')
@@ -77,18 +77,17 @@ def predict_on_images(data_files: list, yaml_file_path: str, checkpoint: str, ma
     # Get parameters --------
     print('This is the printout', data_files)
 
-    with open(yaml_file_path, 'r') as f:
-        config = yaml.safe_load(f)
+    with open(config_path, 'r') as f:
+        config = yaml.safe_load(f)['pretrained_cfg']
 
     batch_size = 8
-    bands = config['DATA']['BANDS']
+    bands = config['bands']
     num_frames = len(data_files)
-    mean = config['DATA']['MEAN']
-    std = config['DATA']['STD']
-    coords_encoding = config['MODEL']['COORDS_ENCODING']
-    img_size = config['DATA']['INPUT_SIZE'][-1]
-
-    mask_ratio = mask_ratio or config['DATA']['MASK_RATIO']
+    mean = config['mean']
+    std = config['std']
+    coords_encoding = config['coords_encoding']
+    img_size = config['img_size']
+    mask_ratio = mask_ratio or config['mask_ratio']
 
     assert num_frames <= 4, "Demo only supports up to four timestamps"
 
@@ -110,21 +109,12 @@ def predict_on_images(data_files: list, yaml_file_path: str, checkpoint: str, ma
 
     # Create model and load checkpoint -------------------------------------------------------------
 
-    model = PrithviMAE(img_size=config['DATA']['INPUT_SIZE'][-2:],
-                       patch_size=config['MODEL']['PATCH_SIZE'],
-                       num_frames=num_frames,
-                       in_chans=len(bands),
-                       embed_dim=config['MODEL']['EMBED_DIM'],
-                       depth=config['MODEL']['DEPTH'],
-                       num_heads=config['MODEL']['NUM_HEADS'],
-                       decoder_embed_dim=config['MODEL']['DECODER_EMBED_DIM'],
-                       decoder_depth=config['MODEL']['DECODER_DEPTH'],
-                       decoder_num_heads=config['MODEL']['DECODER_NUM_HEADS'],
-                       mlp_ratio=config['MODEL']['MLP_RATIO'],
-                       norm_layer=partial(torch.nn.LayerNorm, eps=1e-6),
-                       norm_pix_loss=config['MODEL']['NORM_PIX_LOSS'],
-                       coords_encoding=coords_encoding,
-                       coords_scale_learn=config['MODEL']['COORDS_SCALE_LEARN'])
+    config.update(
+        num_frames=num_frames,
+        coords_encoding=coords_encoding,
+    )
+
+    model = PrithviMAE(**config)
 
     total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
     print(f"\n--> Model has {total_params:,} parameters.\n")
@@ -196,7 +186,7 @@ def predict_on_images(data_files: list, yaml_file_path: str, checkpoint: str, ma
     return outputs
 
 
-run_inference = partial(predict_on_images, yaml_file_path=yaml_file_path,checkpoint=checkpoint)
+run_inference = partial(predict_on_images, config_path=config_path,checkpoint=checkpoint)
 
 with gr.Blocks() as demo: