Fix attention layers map for SD-2-1-Base

[pranav4501]

The safetensors are not mapped correctly for attention layers of SD-2-1-base.

[angeloskath]

These don't look correct given that they will result in value_proj_proj and query_proj_proj and to_k. Am I missing sth?

[pranav4501]

Yeah, I missed that. I didn't check for sdxl turbo. It should be fine now if we first check for query and then for to_q.
Also, why are we not using load_weights here? Doesn't it help with such mapping issues?

[angeloskath]

Hi @pranav4501 could you give an example of what wasn't working and what this fixes? I will take a closer look but I would appreciate sth like image2image.py is not working for this model or some weights and the corresponding model file sth that clearly shows the problem.

Thanks and sorry if I am putting extra work on you.

[pranav4501]

Hi @angeloskath , Sorry for the late response. You will have to make a couple changes to model_io.py of stable diffusion

1. Update this line(L-175) to

model.load_weights(list(weights),strict=True)

2. Comment out this line(L-262), Positions ids in CLIP is not implemented, I will try to look into it
3. Run The image generation

python3 txt2image.py "A photo of an astronaut riding a horse on Mars." --n_images 1 --steps 1 --model sd

4. You should encounter something like

ValueError: Received parameters not in model: encoder.mid_blocks.1.proj_attn.weight encoder.mid_blocks.1.value.bias decoder.mid_blocks.1.proj_attn.bias decoder.mid_blocks.1.value.bias