Update src/layers_cache.py

src/layers_cache.py

@@ -146,7 +146,7 @@ class MultiSingleStreamBlockLoraProcessor(nn.Module):
                 start = i * scaled_cond_size + scaled_block_size
                 end = (i + 1) * scaled_cond_size + scaled_block_size
                 mask[start:end, start:end] = 0  # Diagonal blocks
-            mask = mask * -1e20
+            mask = mask * -1e10
             mask = mask.to(query.dtype)
 
             hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False, attn_mask=mask)
@@ -305,7 +305,7 @@ class MultiDoubleStreamBlockLoraProcessor(nn.Module):
                 start = i * scaled_cond_size + scaled_block_size
                 end = (i + 1) * scaled_cond_size + scaled_block_size
                 mask[start:end, start:end] = 0  # Diagonal blocks
-            mask = mask * -1e20
+            mask = mask * -1e10
             mask = mask.to(query.dtype)
 
             hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False, attn_mask=mask)