[BlockInfo] Lumina (#11227)

* block info

* device

* Make tensor int again

---------

Co-authored-by: Jedrzej Kosinski <kosinkadink1@gmail.com>
This commit is contained in:
Haoming 2025-12-16 09:57:28 +08:00 committed by GitHub
parent 41bcf0619d
commit fc4af86068
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -634,8 +634,11 @@ class NextDiT(nn.Module):
img, mask, img_size, cap_size, freqs_cis = self.patchify_and_embed(x, cap_feats, cap_mask, adaln_input, num_tokens, transformer_options=transformer_options)
freqs_cis = freqs_cis.to(img.device)
transformer_options["total_blocks"] = len(self.layers)
transformer_options["block_type"] = "double"
img_input = img
for i, layer in enumerate(self.layers):
transformer_options["block_index"] = i
img = layer(img, mask, freqs_cis, adaln_input, transformer_options=transformer_options)
if "double_block" in patches:
for p in patches["double_block"]: