Colab problem: The expanded size of the tensor (10) must match the existing size (11) at a non-singleton dimension 1.

[steinhaug]

I put up a quick colab hoping for this to work, but when running the demo I get an error:

Here is the Colab.

Running the test.py --epoch...

unet/config.json: 100% 743/743 [00:00<00:00, 2.91MB/s] diffusion_pytorch_model.safetensors: 100% 3.44G/3.44G [00:34<00:00, 99.9MB/s] Traceback (most recent call last): File "/content/DreamPose/test.py", line 48, in <module> unet = get_unet('CompVis/stable-diffusion-v1-4', "ebb811dd71cdc38a204ecbdd6ac5d580f529fd8c", resolution=512) File "/content/DreamPose/models/unet_dual_encoder.py", line 30, in get_unet unet.conv_in.weight[:, 4:] = torch.zeros(unet.conv_in.weight[:, 3:].shape) # new weights initialized to zero

RuntimeError: The expanded size of the tensor (10) must match the existing size (11) at non-singleton dimension 1. Target sizes: [320, 10, 3, 3]. Tensor sizes: [320, 11, 3, 3]

[xhqi0109]

+1

[isabellapu]

+1

[SaharaSheik]

I solved this by making the following changes to in this file models/unet_dual_encoder.py line 30
https://github.com/johannakarras/DreamPose/blob/5bf30b7df70cf6f2e0bb25556c6ff2cbf0f2b1bf/models/unet_dual_encoder.py#L30C1-L30C117

        unet.conv_in.weight[:, 4:] = torch.zeros(unet.conv_in.weight[:, 3:].shape) # new weights initialized to zero

change to

        unet.conv_in.weight[:, 4:] = torch.zeros(unet.conv_in.weight[:, 4:].shape) # new weights initialized to zero