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

Merge 282304a87d into 4b2a0df237

Browse Source
This commit is contained in:
Pranav 2025-06-10 21:20:21 -04:00 committed by GitHub
commit c79c611ea7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 12 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
stable_diffusion/stable_diffusion/model_io.py
View File

@ -130,7 +130,17 @@ def map_vae_weights(key, value):
if "upsamplers" in key:
key = key.replace("upsamplers.0.conv", "upsample")
# Map attention layers
# Map attention layers in SD-2-1-base:VAE
if "key" in key:
key = key.replace("key", "key_proj")
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 attention layers in SDXL Turbo
if "to_k" in key:
key = key.replace("to_k", "key_proj")
if "to_out.0" in key:
@ -140,6 +150,7 @@ def map_vae_weights(key, value):
if "to_v" in key:
key = key.replace("to_v", "value_proj")
# Map the mid block
if "mid_block.resnets.0" in key:
key = key.replace("mid_block.resnets.0", "mid_blocks.0")