Update app.py

app.py

@@ -1,4 +1,331 @@
+# import streamlit as st
+
+# x = st.slider('Select a value')
+# st.write(x, 'squared is', x * x)
+
+
 import streamlit as st
+import random
+from pathlib import Path
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from huggingface_hub import hf_hub_download, snapshot_download
+import tarfile
+import os
+import sys
+import yaml
+
+st.title("PrithviWxC Model Inference")
+
+st.write("Setting up environment...")
+
+# Set up torch backends and seeds
+torch.jit.enable_onednn_fusion(True)
+if torch.cuda.is_available():
+    st.write(f"Using device: {torch.cuda.get_device_name()}")
+    torch.backends.cudnn.benchmark = True
+    torch.backends.cudnn.deterministic = True
+
+random.seed(42)
+if torch.cuda.is_available():
+    torch.cuda.manual_seed(42)
+torch.manual_seed(42)
+np.random.seed(42)
+
+# Set device
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+else:
+    device = torch.device("cpu")
+
+st.write(f"Using device: {device}")
+
+# Download and extract PrithviWxC module
+st.write("Downloading and setting up PrithviWxC module...")
+
+module_tar_path = hf_hub_download(
+    repo_id="Prithvi-WxC/prithvi.wxc.2300m.v1",
+    filename="PrithviWxC.tar.gz",
+    local_dir=".",
+    force_download=True
+)
+
+with tarfile.open(module_tar_path, "r:gz") as tar:
+    tar.extractall(path=".")
+
+# Add the module path to sys.path
+sys.path.append(os.path.abspath("./PrithviWxC"))
+
+st.write("PrithviWxC module imported successfully.")
+
+# Now import the module
+from PrithviWxC.dataloaders.merra2 import Merra2Dataset, input_scalers, output_scalers, static_input_scalers, preproc
+from PrithviWxC.model import PrithviWxC
+
+# Variables and times
+surface_vars = [
+    "EFLUX",
+    "GWETROOT",
+    "HFLUX",
+    "LAI",
+    "LWGAB",
+    "LWGEM",
+    "LWTUP",
+    "PS",
+    "QV2M",
+    "SLP",
+    "SWGNT",
+    "SWTNT",
+    "T2M",
+    "TQI",
+    "TQL",
+    "TQV",
+    "TS",
+    "U10M",
+    "V10M",
+    "Z0M",
+]
+static_surface_vars = ["FRACI", "FRLAND", "FROCEAN", "PHIS"]
+vertical_vars = ["CLOUD", "H", "OMEGA", "PL", "QI", "QL", "QV", "T", "U", "V"]
+levels = [
+    34.0,
+    39.0,
+    41.0,
+    43.0,
+    44.0,
+    45.0,
+    48.0,
+    51.0,
+    53.0,
+    56.0,
+    63.0,
+    68.0,
+    71.0,
+    72.0,
+]
+padding = {"level": [0, 0], "lat": [0, -1], "lon": [0, 0]}
+
+st.write("Setting up dataset parameters...")
+
+# User inputs for lead times and input times
+lead_time = st.number_input("Lead Time (hours)", min_value=1, max_value=24, value=6)
+input_time = st.number_input("Input Time Difference (hours)", min_value=-24, max_value=0, value=-6)
+
+lead_times = [lead_time]  # This variable can be changed to change the task
+input_times = [input_time]  # This variable can be changed to change the task
+
+# Data file
+time_range = ("2020-01-01T00:00:00", "2020-01-01T23:59:59")
+
+st.write("Downloading data files...")
+
+surf_dir = Path("./merra-2")
+snapshot_download(
+    repo_id="Prithvi-WxC/prithvi.wxc.2300m.v1",
+    allow_patterns="merra-2/MERRA2_sfc_2020010[1].nc",
+    local_dir=".",
+    force_download=True,
+)
+
+vert_dir = Path("./merra-2")
+snapshot_download(
+    repo_id="Prithvi-WxC/prithvi.wxc.2300m.v1",
+    allow_patterns="merra-2/MERRA_pres_2020010[1].nc",
+    local_dir=".",
+    force_download=True,
+)
+
+# Climatology
+surf_clim_dir = Path("./climatology")
+snapshot_download(
+    repo_id="Prithvi-WxC/prithvi.wxc.2300m.v1",
+    allow_patterns="climatology/climate_surface_doy00[1]*.nc",
+    local_dir=".",
+    force_download=True,
+)
+
+vert_clim_dir = Path("./climatology")
+snapshot_download(
+    repo_id="Prithvi-WxC/prithvi.wxc.2300m.v1",
+    allow_patterns="climatology/climate_vertical_doy00[1]*.nc",
+    local_dir=".",
+    force_download=True,
+)
+
+st.write("Setting positional encoding...")
+
+positional_encoding = "fourier"
+
+st.write("Initializing dataset...")
+
+dataset = Merra2Dataset(
+    time_range=time_range,
+    lead_times=lead_times,
+    input_times=input_times,
+    data_path_surface=surf_dir,
+    data_path_vertical=vert_dir,
+    climatology_path_surface=surf_clim_dir,
+    climatology_path_vertical=vert_clim_dir,
+    surface_vars=surface_vars,
+    static_surface_vars=static_surface_vars,
+    vertical_vars=vertical_vars,
+    levels=levels,
+    positional_encoding=positional_encoding,
+)
+
+assert len(dataset) > 0, "There doesn't seem to be any valid data."
+
+st.write("Loading scalers...")
+
+surf_in_scal_path = Path("./climatology/musigma_surface.nc")
+hf_hub_download(
+    repo_id="Prithvi-WxC/prithvi.wxc.2300m.v1",
+    filename=f"climatology/{surf_in_scal_path.name}",
+    local_dir=".",
+    force_download=True,
+)
+
+vert_in_scal_path = Path("./climatology/musigma_vertical.nc")
+hf_hub_download(
+    repo_id="Prithvi-WxC/prithvi.wxc.2300m.v1",
+    filename=f"climatology/{vert_in_scal_path.name}",
+    local_dir=".",
+    force_download=True,
+)
+
+surf_out_scal_path = Path("./climatology/anomaly_variance_surface.nc")
+hf_hub_download(
+    repo_id="Prithvi-WxC/prithvi.wxc.2300m.v1",
+    filename=f"climatology/{surf_out_scal_path.name}",
+    local_dir=".",
+    force_download=True,
+)
+
+vert_out_scal_path = Path("./climatology/anomaly_variance_vertical.nc")
+hf_hub_download(
+    repo_id="Prithvi-WxC/prithvi.wxc.2300m.v1",
+    filename=f"climatology/{vert_out_scal_path.name}",
+    local_dir=".",
+    force_download=True,
+)
+
+in_mu, in_sig = input_scalers(
+    surface_vars,
+    vertical_vars,
+    levels,
+    surf_in_scal_path,
+    vert_in_scal_path,
+)
+
+output_sig = output_scalers(
+    surface_vars,
+    vertical_vars,
+    levels,
+    surf_out_scal_path,
+    vert_out_scal_path,
+)
+
+static_mu, static_sig = static_input_scalers(
+    surf_in_scal_path,
+    static_surface_vars,
+)
+
+st.write("Setting up model...")
+
+residual = "climate"
+masking_mode = "local"
+decoder_shifting = True
+masking_ratio = 0.99
+
+# Load model config
+hf_hub_download(
+    repo_id="Prithvi-WxC/prithvi.wxc.2300m.v1",
+    filename="config.yaml",
+    local_dir=".",
+    force_download=True,
+)
+
+with open("./config.yaml", "r") as f:
+    config = yaml.safe_load(f)
+
+model = PrithviWxC(
+    in_channels=config["params"]["in_channels"],
+    input_size_time=config["params"]["input_size_time"],
+    in_channels_static=config["params"]["in_channels_static"],
+    input_scalers_mu=in_mu,
+    input_scalers_sigma=in_sig,
+    input_scalers_epsilon=config["params"]["input_scalers_epsilon"],
+    static_input_scalers_mu=static_mu,
+    static_input_scalers_sigma=static_sig,
+    static_input_scalers_epsilon=config["params"]["static_input_scalers_epsilon"],
+    output_scalers=output_sig**0.5,
+    n_lats_px=config["params"]["n_lats_px"],
+    n_lons_px=config["params"]["n_lons_px"],
+    patch_size_px=config["params"]["patch_size_px"],
+    mask_unit_size_px=config["params"]["mask_unit_size_px"],
+    mask_ratio_inputs=masking_ratio,
+    embed_dim=config["params"]["embed_dim"],
+    n_blocks_encoder=config["params"]["n_blocks_encoder"],
+    n_blocks_decoder=config["params"]["n_blocks_decoder"],
+    mlp_multiplier=config["params"]["mlp_multiplier"],
+    n_heads=config["params"]["n_heads"],
+    dropout=config["params"]["dropout"],
+    drop_path=config["params"]["drop_path"],
+    parameter_dropout=config["params"]["parameter_dropout"],
+    residual=residual,
+    masking_mode=masking_mode,
+    decoder_shifting=decoder_shifting,
+    positional_encoding=positional_encoding,
+    checkpoint_encoder=[],
+    checkpoint_decoder=[],
+)
+
+st.write("Loading model weights...")
+
+weights_path = Path("./weights/prithvi.wxc.2300m.v1.pt")
+hf_hub_download(
+    repo_id="Prithvi-WxC/prithvi.wxc.2300m.v1",
+    filename=weights_path.name,
+    local_dir="./weights",
+    force_download=True,
+)
+
+state_dict = torch.load(weights_path, map_location=device)
+if "model_state" in state_dict:
+    state_dict = state_dict["model_state"]
+model.load_state_dict(state_dict, strict=True)
+
+model = model.to(device)
+
+st.write("Model loaded and ready.")
+
+if st.button("Run Inference"):
+    st.write("Running inference...")
+
+    data = next(iter(dataset))
+    batch = preproc([data], padding)
+
+    for k, v in batch.items():
+        if isinstance(v, torch.Tensor):
+            batch[k] = v.to(device)
+
+    with torch.no_grad():
+        model.eval()
+        out = model(batch)
+
+    st.write("Inference completed. Generating plot...")
+
+    t2m = out[0, 12].cpu().numpy()
+
+    lat = np.linspace(-90, 90, out.shape[-2])
+    lon = np.linspace(-180, 180, out.shape[-1])
+    X, Y = np.meshgrid(lon, lat)
+
+    fig, ax = plt.subplots()
+    cs = ax.contourf(X, Y, t2m, 100)
+    ax.set_aspect("equal")
+    plt.colorbar(cs)
+    st.pyplot(fig)
 
-x = st.slider('Select a value')
-st.write(x, 'squared is', x * x)
+    st.write("Plot generated.")