Allow loading from diffusers ckpt (#1117)

This commit is contained in:
Angelos Katharopoulos 2024-12-02 13:15:50 -08:00 committed by Billel Mokeddem
parent 7ee0a55705
commit a0e7965a4f

View File

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