move reshape logic

Author: hlky

src/diffusers/models/transformers/transformer_hidream_image.py

@@ -795,6 +795,21 @@ def forward(
         batch_size = hidden_states.shape[0]
         hidden_states_type = hidden_states.dtype
 
+        if hidden_states.shape[-2] != hidden_states.shape[-1]:
+            B, C, H, W = hidden_states.shape
+            patch_size = self.config.patch_size
+            pH, pW = H // patch_size, W // patch_size
+            out = torch.zeros(
+                (B, C, self.max_seq, patch_size * patch_size),
+                dtype=hidden_states.dtype,
+                device=hidden_states.device,
+            )
+            hidden_states = hidden_states.reshape(B, C, pH, patch_size, pW, patch_size)
+            hidden_states = hidden_states.permute(0, 1, 2, 4, 3, 5)
+            hidden_states = hidden_states.reshape(B, C, pH * pW, patch_size * patch_size)
+            out[:, :, 0 : pH * pW] = hidden_states
+            hidden_states = out
+
         # 0. time
         timesteps = self.expand_timesteps(timesteps, batch_size, hidden_states.device)
         timesteps = self.t_embedder(timesteps, hidden_states_type)

src/diffusers/pipelines/hidream_image/pipeline_hidream_image.py

@@ -666,21 +666,6 @@ def __call__(
                 # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
                 timestep = t.expand(latent_model_input.shape[0])
 
-                if latent_model_input.shape[-2] != latent_model_input.shape[-1]:
-                    B, C, H, W = latent_model_input.shape
-                    patch_size = self.transformer.config.patch_size
-                    pH, pW = H // patch_size, W // patch_size
-                    out = torch.zeros(
-                        (B, C, self.transformer.max_seq, patch_size * patch_size),
-                        dtype=latent_model_input.dtype,
-                        device=latent_model_input.device,
-                    )
-                    latent_model_input = latent_model_input.reshape(B, C, pH, patch_size, pW, patch_size)
-                    latent_model_input = latent_model_input.permute(0, 1, 2, 4, 3, 5)
-                    latent_model_input = latent_model_input.reshape(B, C, pH * pW, patch_size * patch_size)
-                    out[:, :, 0 : pH * pW] = latent_model_input
-                    latent_model_input = out
-
                 noise_pred = self.transformer(
                     hidden_states=latent_model_input,
                     timesteps=timestep,