zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-08-30 02:53:41 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Allow loading from diffusers ckpt

Browse Source
This commit is contained in:
Angelos Katharopoulos 2024-11-22 20:52:50 -08:00
parent 042280ce50
commit 516d0e3af0
1 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
flux/flux/model.py
View File

@ -85,6 +85,8 @@ class Flux(nn.Module):
def sanitize(self, weights): def sanitize(self, weights):
new_weights = {} new_weights = {}
for k, w in weights.items(): for k, w in weights.items():
if k.startswith("model.diffusion_model."):
k = k[22:]
if k.endswith(".scale"): if k.endswith(".scale"):
k = k[:-6] + ".weight" k = k[:-6] + ".weight"
for seq in ["img_mlp", "txt_mlp", "adaLN_modulation"]: for seq in ["img_mlp", "txt_mlp", "adaLN_modulation"]: