From eb9277f574c6ba9da8494afade631e7a8553402a Mon Sep 17 00:00:00 2001
From: Angelos Katharopoulos <a_katharopoulos@apple.com>
Date: Mon, 2 Dec 2024 13:15:50 -0800
Subject: [PATCH] Allow loading from diffusers ckpt (#1117)

---
 flux/flux/model.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/flux/flux/model.py b/flux/flux/model.py
index 18ea70b0..d8ad9d9b 100644
--- a/flux/flux/model.py
+++ b/flux/flux/model.py
@@ -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"]: