Questions about corrdiff shape_x and shape_y

[Liu-CWB]

The defualt shape_x and shape_y for corrdiff are equal grid. (shape_x = shape_y).

I set the non equal shape_x and shape_y in yaml file as following ,
img_shape_x: 320
img_shape_y: 336

The trainging can be well done.
However the torch tensor size become (channel, 320, 336) which the shape_x and shape_y mismatch in the torch tensor. It should be the (channel,336, 320).

Therefore, I check and modify some code as following,

For cwb.py under examples/weather/corrdiff/datasets/cwb.py

def image_shape(self):                                                              
        """Get the shape of the image (same for input and output)."""                                                         
        #return (self.img_shape_x, self.img_shape_y) 
        return (self.img_shape_y, self.img_shape_x) 

For img_utils.py under examples/weather/corrdiff/datasets/img_utils.py
line

if inp_or_tar == "inp":
        #img = np.reshape(img, (n_channels * (n_history + 1), img_shape_x, img_shape_y))
        img = np.reshape(img, (n_channels * (n_history + 1), img_shape_y, img_shape_x))
    elif inp_or_tar == "tar":
        #img = np.reshape(img, (n_channels, img_shape_x, img_shape_y))    You, 11/07/25 13:33 • Uncommitted changes
        img = np.reshape(img, (n_channels, img_shape_y, img_shape_x))

After modify these two python script, the tensor can match the shape_x and shape_y
img_shape_x: 320
img_shape_y: 336

And the torch tensor size become (channel, 336, 320)

Is this modification suitable ?

Thanks

[CharlelieLrt]

@Liu-CWB thank you for reporting this. We will open a PR to fix these issues related to non-square images. What type of model are you using: the SongUNet without positional embeddings or the SongUNetPosEmbd with positional embeddings? I believe in the latter there are multiple inconsistencies related to the case of non-square images. This will require more work to fix. In the meantime, if you are not using positional embeddings, I think the changes you proposed above should be sufficient for a workaround.

We will keep the issue open until we have better support for non-square images.

[Liu-CWB]

@CharlelieLrt
I use the default setting -->SongUNetPosEmbd with positional embeddings.
I need to update the corrdiff model with the new data (some data are non-square images) for CWA weather applications.
Therefore, I need to fix this issue.
Thank you for your help.