[Hi Dream] follow-up

src/diffusers/models/transformers/transformer_hidream_image.py

@@ -714,7 +714,8 @@ def forward(
         self,
         hidden_states: torch.Tensor,
         timesteps: torch.LongTensor = None,
-        encoder_hidden_states: torch.Tensor = None,
+        t5_encoder_hidden_states: torch.Tensor = None,
+        llama3_encoder_hidden_states: torch.Tensor = None,
         pooled_embeds: torch.Tensor = None,
         img_sizes: Optional[List[Tuple[int, int]]] = None,
         img_ids: Optional[torch.Tensor] = None,
@@ -773,9 +774,7 @@ def forward(
             )
         hidden_states = self.x_embedder(hidden_states)
 
-        T5_encoder_hidden_states = encoder_hidden_states[0]
-        encoder_hidden_states = encoder_hidden_states[-1]
-        encoder_hidden_states = [encoder_hidden_states[k] for k in self.config.llama_layers]
+        encoder_hidden_states = [llama3_encoder_hidden_states[k] for k in self.llama_layers]
 
         if self.caption_projection is not None:
             new_encoder_hidden_states = []
@@ -784,9 +783,9 @@ def forward(
                 enc_hidden_state = enc_hidden_state.view(batch_size, -1, hidden_states.shape[-1])
                 new_encoder_hidden_states.append(enc_hidden_state)
             encoder_hidden_states = new_encoder_hidden_states
-            T5_encoder_hidden_states = self.caption_projection[-1](T5_encoder_hidden_states)
-            T5_encoder_hidden_states = T5_encoder_hidden_states.view(batch_size, -1, hidden_states.shape[-1])
-            encoder_hidden_states.append(T5_encoder_hidden_states)
+            t5_encoder_hidden_states = self.caption_projection[-1](t5_encoder_hidden_states)
+            t5_encoder_hidden_states = t5_encoder_hidden_states.view(batch_size, -1, hidden_states.shape[-1])
+            encoder_hidden_states.append(t5_encoder_hidden_states)
 
         txt_ids = torch.zeros(
             batch_size,