File tree 1 file changed +7
-4
lines changed 1 file changed +7
-4
lines changed Original file line number Diff line number Diff line change @@ -208,9 +208,12 @@ def route_callback(k_callback_values):
208
208
model_wrap_cfg = CFGDenoiser (self .model , threshold = threshold , warmup = max (0.8 * S ,S - 10 ))
209
209
model_wrap_cfg .prepare_to_sample (S , extra_conditioning_info = extra_conditioning_info )
210
210
211
- attention_map_token_ids = range (1 , extra_conditioning_info .tokens_count_including_eos_bos - 1 )
212
- attention_maps_saver = None if attention_maps_callback is None else AttentionMapSaver (token_ids = attention_map_token_ids , latents_shape = x .shape [- 2 :])
213
- if attention_maps_callback is not None :
211
+ # setup attention maps saving. checks for None are because there are multiple code paths to get here.
212
+ attention_maps_saver = None
213
+ if attention_maps_callback is not None and extra_conditioning_info is not None :
214
+ eos_token_index = extra_conditioning_info .tokens_count_including_eos_bos - 1
215
+ attention_map_token_ids = range (1 , eos_token_index )
216
+ attention_maps_saver = AttentionMapSaver (token_ids = attention_map_token_ids , latents_shape = x .shape [- 2 :])
214
217
model_wrap_cfg .invokeai_diffuser .setup_attention_map_saving (attention_maps_saver )
215
218
216
219
extra_args = {
@@ -226,7 +229,7 @@ def route_callback(k_callback_values):
226
229
),
227
230
None ,
228
231
)
229
- if attention_maps_callback is not None :
232
+ if attention_maps_saver is not None :
230
233
attention_maps_callback (attention_maps_saver )
231
234
return sampling_result
232
235
You can’t perform that action at this time.
0 commit comments