diffusers icon indicating copy to clipboard operation
diffusers copied to clipboard

Avoid creating tensor in CosmosAttnProcessor2_0

Open chenxiao111222 opened this issue 8 months ago • 1 comments

https://github.com/huggingface/diffusers/blob/42077e6c734df2fc7bbed373abceab99635500ad/src/diffusers/models/transformers/transformer_cosmos.py#L188-L193

        # 4. Prepare for GQA
        query_idx = query.size(3)
        key_idx = key.size(3)
        value_idx = value.size(3)
        key = key.repeat_interleave(query_idx // key_idx, dim=3)
        value = value.repeat_interleave(query_idx // value_idx, dim=3)

Speedup ~10% here in Cosmos2TextToImagePipeline and Cosmos2VideoToWorldPipeline.

chenxiao111222 avatar Jun 20 '25 08:06 chenxiao111222

thanks! would you like to share a PR so we can test it out on our end?

yiyixuxu avatar Jun 21 '25 02:06 yiyixuxu