search

ml-explore / mlx-examples

Examples in the MLX framework
MIT License
6.3k stars 898 forks source link

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

Open pranav4501 opened 3 months ago

pranav4501 commented 3 months ago

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

angeloskath commented 2 months ago

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 commented 2 months ago

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