Skip to content

Commit c38e7d6

Browse files
authored
block info (#10841)
1 parent eaf68c9 commit c38e7d6

File tree

1 file changed

+6
-0
lines changed

1 file changed

+6
-0
lines changed

comfy/ldm/flux/model.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -171,7 +171,10 @@ def forward_orig(
171171
pe = None
172172

173173
blocks_replace = patches_replace.get("dit", {})
174+
transformer_options["total_blocks"] = len(self.double_blocks)
175+
transformer_options["block_type"] = "double"
174176
for i, block in enumerate(self.double_blocks):
177+
transformer_options["block_index"] = i
175178
if ("double_block", i) in blocks_replace:
176179
def block_wrap(args):
177180
out = {}
@@ -215,7 +218,10 @@ def block_wrap(args):
215218
if self.params.global_modulation:
216219
vec, _ = self.single_stream_modulation(vec_orig)
217220

221+
transformer_options["total_blocks"] = len(self.single_blocks)
222+
transformer_options["block_type"] = "single"
218223
for i, block in enumerate(self.single_blocks):
224+
transformer_options["block_index"] = i
219225
if ("single_block", i) in blocks_replace:
220226
def block_wrap(args):
221227
out = {}

0 commit comments

Comments
 (0)