From c490b5c2b9dd5f5979d186f7c913983a3c577855 Mon Sep 17 00:00:00 2001
From: Jingya HUANG <44135271+JingyaHuang@users.noreply.github.com>
Date: Tue, 26 Nov 2024 10:49:45 +0100
Subject: [PATCH] Emphasize the order of loading models in diffusion pipelines
 (#742)

quick fix
---
 optimum/neuron/modeling_diffusion.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/optimum/neuron/modeling_diffusion.py b/optimum/neuron/modeling_diffusion.py
index 915b98d6a..57081834e 100644
--- 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,