zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-10-23 05:58:07 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Feat: update pre-commit rev (#432)

Browse Source
This commit is contained in:
Nripesh Niketan
2024-02-11 19:23:27 +04:00
committed by GitHub
parent f45a1ab83c
commit f1ef378a58
3 changed files with 9 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
stable_diffusion/stable_diffusion/model_io.py
View File

@@ -186,9 +186,11 @@ def load_unet(key: str = _DEFAULT_MODEL, float16: bool = False):
out_channels=config["out_channels"],
block_out_channels=config["block_out_channels"],
layers_per_block=[config["layers_per_block"]] * n_blocks,
num_attention_heads=[config["attention_head_dim"]] * n_blocks
if isinstance(config["attention_head_dim"], int)
else config["attention_head_dim"],
num_attention_heads=(
[config["attention_head_dim"]] * n_blocks
if isinstance(config["attention_head_dim"], int)
else config["attention_head_dim"]
),
cross_attention_dim=[config["cross_attention_dim"]] * n_blocks,
norm_num_groups=config["norm_num_groups"],
)