298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-05-10 06:25:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

Make hidream work with any latent resolution.

Browse Source
This commit is contained in:
comfyanonymous 2025-04-16 18:34:14 -04:00
parent 9899d187b1
commit 1fc00ba4b6
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
comfy/ldm/hidream/model.py
View File

@ -13,6 +13,7 @@ from comfy.ldm.flux.layers import LastLayer
from comfy.ldm.modules.attention import optimized_attention from comfy.ldm.modules.attention import optimized_attention
import comfy.model_management 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 # 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, control = None,
transformer_options = {}, transformer_options = {},
) -> torch.Tensor: ) -> 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 timesteps = t
pooled_embeds = y pooled_embeds = y
T5_encoder_hidden_states = context T5_encoder_hidden_states = context
@ -794,4 +796,4 @@ class HiDreamImageTransformer2DModel(nn.Module):
hidden_states = hidden_states[:, :image_tokens_seq_len, ...] hidden_states = hidden_states[:, :image_tokens_seq_len, ...]
output = self.final_layer(hidden_states, adaln_input) output = self.final_layer(hidden_states, adaln_input)
output = self.unpatchify(output, img_sizes) output = self.unpatchify(output, img_sizes)
return -output return -output[:, :, :h, :w]