make from_flax work for controlnet (huggingface#3161)

fix from_flax Co-authored-by: Patrick von Platen <[email protected]>
pdoane · Apr 21, 2023 · bc0392a · bc0392a
1 parent 05d9bae
commit bc0392a
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/src/diffusers/models/modeling_pytorch_flax_utils.py b/src/diffusers/models/modeling_pytorch_flax_utils.py
@@ -110,6 +110,12 @@ def load_flax_weights_in_pytorch_model(pt_model, flax_state):
                     .replace("_1", ".1")
                     .replace("_2", ".2")
                     .replace("_3", ".3")
+                    .replace("_4", ".4")
+                    .replace("_5", ".5")
+                    .replace("_6", ".6")
+                    .replace("_7", ".7")
+                    .replace("_8", ".8")
+                    .replace("_9", ".9")
                 )
 
         flax_key = ".".join(flax_key_tuple_array)