@@ -455,7 +455,7 @@ def forward(
455455 # 3. down
456456 down_block_res_samples = (sample ,)
457457 for downsample_block in self .down_blocks :
458- if hasattr (downsample_block , "attentions " ) and downsample_block .attentions is not None :
458+ if hasattr (downsample_block , "has_cross_attention " ) and downsample_block .has_cross_attention :
459459 sample , res_samples = downsample_block (
460460 hidden_states = sample ,
461461 temb = emb ,
@@ -481,7 +481,7 @@ def forward(
481481 if not is_final_block and forward_upsample_size :
482482 upsample_size = down_block_res_samples [- 1 ].shape [2 :]
483483
484- if hasattr (upsample_block , "attentions " ) and upsample_block .attentions is not None :
484+ if hasattr (upsample_block , "has_cross_attention " ) and upsample_block .has_cross_attention :
485485 sample = upsample_block (
486486 hidden_states = sample ,
487487 temb = emb ,
@@ -726,6 +726,7 @@ def __init__(
726726 resnets = []
727727 attentions = []
728728
729+ self .has_cross_attention = True
729730 self .attention_type = attention_type
730731 self .attn_num_head_channels = attn_num_head_channels
731732
@@ -924,6 +925,7 @@ def __init__(
924925 resnets = []
925926 attentions = []
926927
928+ self .has_cross_attention = True
927929 self .attention_type = attention_type
928930 self .attn_num_head_channels = attn_num_head_channels
929931
@@ -1043,6 +1045,7 @@ def __init__(
10431045 ):
10441046 super ().__init__ ()
10451047
1048+ self .has_cross_attention = True
10461049 self .attention_type = attention_type
10471050 self .attn_num_head_channels = attn_num_head_channels
10481051 resnet_groups = resnet_groups if resnet_groups is not None else min (in_channels // 4 , 32 )
0 commit comments