Implementing StableFlow in other DiT Architecture

[willyawan16]

I am trying to implement your work into Cosmos Predict2 DiT architecture, but it seems that there are no changes to my image.
I have found cosmos' vital layers by using a set of 64 prompts for diverse objects and random seeds.
I am doing an image2image generation with initial prompt is "A dog at the beach" and the edited prompt ""A cat at the beach", no changes after generating as shown below

Input image

Edited image

Question(s)

For the self attention injection, it is written in the paper that parallel generation is done to selectively replacing the keys and value. Mean while in the code you just did it in one run and copying the keys and value directly from the first index, is there any explanation for this?

Very appreciate your work, looking forward to your reply!