diffusers
diffusers copied to clipboard
Avoid creating tensor in CosmosAttnProcessor2_0
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.
thanks! would you like to share a PR so we can test it out on our end?