zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-08-09 18:36:38 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Pranav Veldurthi 2024-09-01 14:34:35 -04:00
parent 3c6e8b11af
commit 1fbc9361e4
1 changed files with 11 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
stable_diffusion/stable_diffusion/model_io.py
View File

@ -140,6 +140,17 @@ def map_vae_weights(key, value):
if "to_v" in key:
key = key.replace("to_v", "value_proj")
# Map attention layers in SD-2-1-base:VAE
if "key" in key:
key = key.replace("key", "to_k")
if "proj_attn" in key:
key = key.replace("proj_attn", "out_proj")
if "query" in key:
key = key.replace("query", "query_proj")
if "value" in key:
key = key.replace("value", "value_proj")
# Map the mid block
if "mid_block.resnets.0" in key:
key = key.replace("mid_block.resnets.0", "mid_blocks.0")