From 1fc00ba4b6576ed5910a88caa47866774ee6d0ca Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Wed, 16 Apr 2025 18:34:14 -0400
Subject: [PATCH] Make hidream work with any latent resolution.

---
 comfy/ldm/hidream/model.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/comfy/ldm/hidream/model.py b/comfy/ldm/hidream/model.py
index 39c67a19..fcb5a9c5 100644
--- a/comfy/ldm/hidream/model.py
+++ b/comfy/ldm/hidream/model.py
@@ -13,6 +13,7 @@ from comfy.ldm.flux.layers import LastLayer
 
 from comfy.ldm.modules.attention import optimized_attention
 import comfy.model_management
+import comfy.ldm.common_dit
 
 
 # Copied from https://github.com/black-forest-labs/flux/blob/main/src/flux/modules/layers.py
@@ -701,7 +702,8 @@ class HiDreamImageTransformer2DModel(nn.Module):
         control = None,
         transformer_options = {},
     ) -> torch.Tensor:
-        hidden_states = x
+        bs, c, h, w = x.shape
+        hidden_states = comfy.ldm.common_dit.pad_to_patch_size(x, (self.patch_size, self.patch_size))
         timesteps = t
         pooled_embeds = y
         T5_encoder_hidden_states = context
@@ -794,4 +796,4 @@ class HiDreamImageTransformer2DModel(nn.Module):
         hidden_states = hidden_states[:, :image_tokens_seq_len, ...]
         output = self.final_layer(hidden_states, adaln_input)
         output = self.unpatchify(output, img_sizes)
-        return -output
+        return -output[:, :, :h, :w]