Inconsistency in x_mask Logic Between T2IFinalLayer and STDiTBlock Leading to Potential Information Leakage

[fanzhongyi]

Thank you for this impressive open-source work! I've noticed an inconsistency in the application logic of x_mask between the T2IFinalLayer and the STDiTBlock. Since this operation is adapted from PixArt, I believe the underlying design logic (i.e., isolation) is sound. However, the implementation in T2IFinalLayer leads to the temporal variable x being overwritten, resulting in information leakage from t to t0.

I am unsure if this behavior is intentional or if it’s a typo that has led to a bug. Additionally, I'm unclear about the implications of this issue. While it only appears in the final layer, it is also where the leakage occurs closest to the output. I look forward to your insights on this matter. Thank you very much!