diff --git a/comfy/ldm/hidream_o1/model.py b/comfy/ldm/hidream_o1/model.py index a223e706f..6f8c77482 100644 --- a/comfy/ldm/hidream_o1/model.py +++ b/comfy/ldm/hidream_o1/model.py @@ -148,6 +148,12 @@ class HiDreamO1Transformer(nn.Module): raise ValueError("HiDreamO1Transformer requires input_ids and position_ids in conditioning") B, _, H, W = x.shape + + if H % self.patch_size!=0 or w% self.patch_size!=0: + raise ValueError( + f"Input dimensions ({H},{W}) must be divisible" + f"by patch size{self.patch_size}" + ) h_p, w_p = H // self.patch_size, W // self.patch_size tgt_image_len = h_p * w_p