Make hidream work with any latent resolution.

This commit is contained in:
comfyanonymous 2025-04-16 18:34:14 -04:00
parent 9899d187b1
commit 1fc00ba4b6

View File

@ -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]