Emphasize the order of loading models in diffusion pipelines (#742)

quick fix
huggingface · Nov 26, 2024 · c490b5c · c490b5c
1 parent d78c7c7
commit c490b5c
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/optimum/neuron/modeling_diffusion.py b/optimum/neuron/modeling_diffusion.py
@@ -421,10 +421,11 @@ def load_model(
                 Whether to move manually the traced model to NeuronCore. It's only needed when `inline_weights_to_neff=False`, otherwise it is loaded automatically to a Neuron device.
         """
         submodels = {
-            "text_encoder": text_encoder_path,
-            "text_encoder_2": text_encoder_2_path,
+            # Load the UNet/Diffusion transformer first to avoid CPU OOM
             "unet": unet_path,
             "transformer": transformer_path,
+            "text_encoder": text_encoder_path,
+            "text_encoder_2": text_encoder_2_path,
             "vae_encoder": vae_encoder_path,
             "vae_decoder": vae_decoder_path,
             "controlnet": controlnet_paths,