What needs to be done to support 2.0

[amirothman]

Hi,

Firstly, thank you very much for the useful repo. I was trying to fine-tune with stable diffusion 2.0 and got the following error:

RuntimeError: Error(s) in loading state_dict for LatentDiffusion:
        size mismatch for model.diffusion_model.input_blocks.1.1.proj_in.weight: copying a param with shape torch.Size([320, 320]) from checkpoint, the shape
in current model is torch.Size([320, 320, 1, 1]).
        size mismatch for model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn2.to_k.weight: copying a param with shape torch.Size([320, 1024]) fr
om checkpoint, the shape in current model is torch.Size([320, 768]).

Tried with 512-base-ema.ckpt

What needs to be done, so that the trainer can support v2? Some pointers would be awesome so I could create the pull request :)

[victorchall]

ED2.0 will support SD2 checkpoints, it's not something I'll be adding here.

It's mostly adding the YAMLs for the new models, but the issue is SD2 is bigger and will not fit into even 24GB. A rewrite of the entire trainer is in progress.

[amirothman]

Awesome! Thanks for the response.