Modularize wxc downstream tasks: downscaling, gravity

[romeokienzler]

from the conversation with @biancazadrozny and Johannes Schmude:

• Enable more complex embedding networks (convolutions; optionally also pixel shuffle)
• Enable more complex heads (convolutions; optionally also pixel shuffle)
• Optionally skip connections
• Optionally LoRA.

When we are referring to the "backbone", we are referring to these two objects: Encoder: https://github.com/NASA-IMPACT/Prithvi-WxC/blob/f9002006b413a3ffa968dc3f3edbe56afe81954b/PrithviWxC/model.py#L910 Decoder: https://github.com/NASA-IMPACT/Prithvi-WxC/blob/f9002006b413a3ffa968dc3f3edbe56afe81954b/PrithviWxC/model.py#L932

Downscaling:

The model is here: https://github.com/IBM/granite-wxc/blob/4032b17c31cf91bd44825a48dc2aa61f850623eb/granitewxc/models/finetune_model.py#L124 The backbone is here: https://github.com/IBM/granite-wxc/blob/4032b17c31cf91bd44825a48dc2aa61f850623eb/granitewxc/models/finetune_model.py#L45 If you read the forward function (https://github.com/IBM/granite-wxc/blob/4032b17c31cf91bd44825a48dc2aa61f850623eb/granitewxc/models/finetune_model.py#L248), you will find upscale and conv_after_backbone. I think those are the convolutions before/after the model.

Gravity waves:

Model is here: https://github.com/NASA-IMPACT/gravity-wave-finetuning/blob/b296edc4e6881cd32a5eab2f2cefa566acdff744/gravity_wave_model.py#L177 The U-Net seems to be defined via the encoder and decoder properties. I.e. the model is encoder-backbone-decoder, where the backbone is the pretrained part. (The naming is confusing). See here Encoder: https://github.com/NASA-IMPACT/gravity-wave-finetuning/blob/b296edc4e6881cd32a5eab2f2cefa566acdff744/gravity_wave_model.py#L305 Decoder: https://github.com/NASA-IMPACT/gravity-wave-finetuning/blob/b296edc4e6881cd32a5eab2f2cefa566acdff744/gravity_wave_model.py#L325 Since this is a U-Net, only enc4 hits the foundation model ...

[biancazadrozny]

@Joao-L-S-Almeida see above