Re-basin and Stable Diffusion Tensor Flow weights

[ogkalu2]

Hi Samuel. Thank you for being willing to look at this. Basically, i'm trying to see if it is possible to merge stable diffusion models i've finetuned with dreambooth with your method.

The first hurdle of course is that your implementation is not yet compatible with pytorch as far as i know. But the pytorch weights can be successfully converted to Tensorflow weights. This has been done before as well. I don't mind doing this if i have to.

The second hurdle will be successfully using your implementation on the TF models. I'm not sure how feasible this all is or how i would use your code on the SD TensorFlow weights.

[samuela]

Hi @ogkalu2! The exact framework that the model runs in is not super important. For example, we have used our JAX code to align the weights of two PyTorch models in the past. The only important part is that you can load the weights into Python/JAX and that you have a correct PermutationSpec for your model.

In general writing down the PermutationSpec will be the more challenging step. We are currently working on a pure-PyTorch version, including a "tracer" that can automatically generate the PermutationSpec for you but this is not ready for release just yet.

Would be very cool to see if this works on StableDiffusion models! Do let me know if you get it working!

[ogkalu2]

Hi @samuela Thanks for responding. Quite a bit's happened since the last comment.

I found a repo that had converted the code to pytorch. I also actually got down to writing the permutation spec for stable diffusion today. I think i'm on the right track but i'm not too sure. For instance, i'm not quite sure what the correct p_in and p_out values should be. Running it right now gives a couple different errors each time.

Sometimes i get something like

File "/content/drive/MyDrive/SD_rebasin/weight_matching.py", line 487, in weight_matching
A += w_a @ w_b.T
RuntimeError: mat1 and mat2 shapes cannot be multiplied (1280x11520 and 5760x1280)

or

"addmm_impl_cpu_" not implemented for 'Half'

Can you take a quick look here and see what you think I might be doing wrong ?

https://imgur.com/a/DMnb7P8

[samuela]

This looks to be an error in your permutation spec, I would try debugging what weight arrays that's occurring on.

[ogkalu2]

This looks to be an error in your permutation spec, I would try debugging what weight arrays that's occurring on.

I'll try that thanks.

I also get "RuntimeError: INDICES element is out of DATA bounds, id=256 axis_dim=256"

[ogkalu2]

This looks to be an error in your permutation spec, I would try debugging what weight arrays that's occurring on.

I've tried a couple different tests now removing everything besides a few lines of layers that were explicitly labelled in the state dict to see if that was the issue and i could build from that.

Example here https://imgur.com/a/IbAEfdP and https://imgur.com/a/o6SgMk8

But i still get those weird errors.

[ogkalu2]

@samuela
I finally got it to run the final permutation function for the few explicit blocks/layers (i'll build from there once i see everything else works) !

I feel so close but i've hit another wall on the apply permutation line. It's giving me "Key Erorr: Betas"

The good news is that i'm pretty sure i know what's happening here. Betas is the first key in Stable Diffusion's State Dict. It seems to be stuck trying to apply the permutation to Betas but betas wasn't defined in the permutation_spec list of layers or blocks to alter. More importantly, it's not the only key in the state dict i thought best left undisturbed.

Is there any way i can get the apply permutation function to skip keys that it doesn't have altered values for ?

[samuela]

Hmm, I'm not familiar with the stable diffusion architecture... is there a reason not to model permutations on the betas?

You could always add them to you PermutationSpec and just specify None for all the axes. That should make apply_permutation ignore them.

(Btw, since you're using the 3rd-party pytorch implementation some things may be different! That's a different codebase.)

[lopho]

You can skip betas, it's not a layer of the actual architecture but a stored parameter used for sampling and generating timestep embeddings. Same goes for alphas_cumprod, sqrt_alphas, and a dozen others.
Anything that isn't either in diffusionmodel, first_stage_model or cond_stage_model is such a parameter.
Actually you can probably skip first_stage_model (VAE) and cond_stage_model (CLIP) as well as those will probably be the same between two models. If not it would still make more sense to spec them separately as they are indeed separate models that only interact with each other manually, e.g. they have no connection in the model graph.

UNet is the interesting part, and this resides in diffusionmodel.model.*

EDIT:
to clarify, these are the keys you can skip

betas
alphas_cumprod
alphas_cumprod_prev
sqrt_alphas_cumprod
sqrt_one_minus_alphas_cumprod
log_one_minus_alphas_cumprod
sqrt_recip_alphas_cumprod
sqrt_recipm1_alphas_cumprod
posterior_variance
posterior_log_variance_clipped
posterior_mean_coef1
posterior_mean_coef2
model_ema.decay
model_ema.num_updates

and these are the ones you can skip if you retain the same VAE and text encoder

cond_stage_model.*
first_stage_model.*

[ogkalu2]

Hmm, I'm not familiar with the stable diffusion architecture... is there a reason not to model permutations on the betas?

You could always add them to you PermutationSpec and just specify None for all the axes. That should make apply_permutation ignore them.

(Btw, since you're using the 3rd-party pytorch implementation some things may be different! That's a different codebase.)

Mostly 3 reasons

• I can't tell the layer type of all the layers in the state_dict, only most of them. So i have to skip some to test it out first. Also some layers aren't really part of the architecture and will be the same between 2 models .

• A few correctly identified layers don't work well for some reason. I get an error like

File "/content/drive/MyDrive/SD_rebasin/weight_matching.py", line 293, in weight_matching
w_a = torch.moveaxis(w_a, axis, 0).reshape((n, -1))
RuntimeError: shape '[512, -1]' is invalid for input of size 768

• Strictly speaking, the Unet architecture begins with "model.diffusion_model.time_embed.0.weight" not betas so i thought i might skip it

For axes you mean the P_bgx and P_bgy values right ?. So None and None then ?

I have a question on that too. There are a lot of layers. I'm unsure how to correctly label them all. For the layers with 2 P values, do i just keep going sequentially ? I reach P_bg50 or so that way.

For the layers with only one , how would that work exactly ? It's a bit hard to tell when i need to go from say P_bg1 to P_bg2 and to P_bg3 and so on. The architecture is divided in 3 parts - The input blocks, the middle blocks and the output blocks. So I'm wondering, is it P_bg1, _bg2, P_bg3 for those set of blocks or something else ?

[ogkalu2]

You can skip betas, it's not a layer of the actual architecture but a stored parameter used for sampling and generating timestep embeddings. Same goes for alphas_cumprod, sqrt_alphas, and a dozen others. Anything that isn't either in diffusionmodel, first_stage_model or cond_stage_model is such a parameter. Actually you can probably skip first_stage_model (VAE) and cond_stage_model (CLIP) as well as those will probably be the same between two models. If not it would still make more sense to spec them separately as they are indeed separate models that only interact with each other manually, e.g. they have no connection in the model graph.

UNet is the interesting part, and this resides in diffusionmodel.model.*

EDIT: to clarify, these are the keys you can skip

betas
alphas_cumprod
alphas_cumprod_prev
sqrt_alphas_cumprod
sqrt_one_minus_alphas_cumprod
log_one_minus_alphas_cumprod
sqrt_recip_alphas_cumprod
sqrt_recipm1_alphas_cumprod
posterior_variance
posterior_log_variance_clipped
posterior_mean_coef1
posterior_mean_coef2
model_ema.decay
model_ema.num_updates

and these are the ones you can skip if you retain the same VAE and text encoder

cond_stage_model.*
first_stage_model.*

Thanks. Lots of dreambooth repos train the text encoder also now so i won't skip them.

I have some uncertainty on the type of certain layers.

The layers that have
emb_layers
proj_in
ptoj_out
transformer_blocks
skip_connection
self_attn
mid.attn

Do you have any idea ? Are they norm, conv, dense or neither ?

[lopho]

• emb_layers is SiLU + linear -> dense
• proj_in is Conv2d
• proj_out is Conv2d
• transformer_blocks is BasicTransformerBlock -> CrossAttention, LayerNorm, FeedForward
  • CrossAttention: all linear layers -> dense
  • FeedForward: GEGLU + linear -> dense
    • GEGLU: linear + gelu -> dense
• skip_connection: is either Identity or Conv2d (Identity has no weights, so going by key, its always Conv2d)
• self_attn: CLIPAttention -> dense
• mid.attn: AttnBlock -> Conv2d

Keep in mind that the different Attention layer types might be either all conv or dense, but they are not just sequentially chained.
Rather they do qkv queries and batch matrix mult on results. I'm not sure if that has any impact on the results of the permutations.

[ogkalu2]

• emb_layers is SiLU + linear -> dense

• proj_in is Conv2d

• proj_out is Conv2d

• transformer_blocks is BasicTransformerBlock -> CrossAttention, LayerNorm, FeedForward

  • CrossAttention: all linear layers -> dense

  • FeedForward: GEGLU + linear -> dense

    • GEGLU: linear + gelu -> dense

• skip_connection: is either Identity or Conv2d (Identity has no weights, so going by key, its always Conv2d)

• self_attn: CLIPAttention -> dense

• mid.attn: AttnBlock -> Conv2d

Keep in mind that the different Attention layer types might be either all conv or dense, but they are not just sequentially chained. Rather they do qkv queries and batch matrix mult on results. I'm not sure if that has any impact on the results of the permutations.

Thanks for the response. It's helped a lot. Yes i think i'm going to skip the attention layers, at least for the first go around.
As samuel suggested, Labeling the axes as none for betas worked on skipping it so i just have to do them all now.

Forgot to ask, what are the time_embed (i'm assuming dense now), model.out.0 and out.2 (norm and conv i think) and .op (i think conv ) ?

[lopho]

unet time_embed: linear, silu, linear, so its dense
unet out: GroupNorm, SiLU, Conv2d
op (which I guess you mean op of the unet downsampling block) is conv2d

It's all readily available in the implementation, so I suggest you read it yourself to get a better understanding.
unet: https://github.com/CompVis/stable-diffusion/blob/main/ldm/modules/diffusionmodules/openaimodel.py
unet attention: https://github.com/CompVis/stable-diffusion/blob/main/ldm/modules/attention.py
vae: https://github.com/CompVis/stable-diffusion/blob/main/ldm/models/autoencoder.py