How to replace the attention module of the unet ( instead of the attn_processor) with custom attention class?

[dingbang777]

How to replace the attention module of the unet ( instead of the attn_processor) with custom attention class?
I see usual case that use 'unet.set_attn_processor' to change the attn_processor, but what if I want to change the attention class partly in the unet? GPT tells me to do this but it the modification doesn't make effect.