up

Avoid creating tensor in CosmosAttnProcessor2_0 (#11761 )
2025-12-08 05:24:20 +08:00 · 2025-07-10 22:23:04 +02:00 · 2025-06-21 11:53:58 +08:00
1 changed files with 9 additions and 3 deletions
--- a/src/diffusers/models/transformers/transformer_cosmos.py
+++ b/src/diffusers/models/transformers/transformer_cosmos.py
@@ -186,9 +186,15 @@ class CosmosAttnProcessor2_0:
            key = apply_rotary_emb(key, image_rotary_emb, use_real=True, use_real_unbind_dim=-2)

        # 4. Prepare for GQA
-        query_idx = torch.tensor(query.size(3), device=query.device)
-        key_idx = torch.tensor(key.size(3), device=key.device)
-        value_idx = torch.tensor(value.size(3), device=value.device)
+        if torch.onnx.is_in_onnx_export():
+            query_idx = torch.tensor(query.size(3), device=query.device)
+            key_idx = torch.tensor(key.size(3), device=key.device)
+            value_idx = torch.tensor(value.size(3), device=value.device)
+
+        else:
+            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)
Author	SHA1	Message	Date
yiyixuxu	fff9e31eb7	up	2025-07-10 22:23:04 +02:00
chenxiao	72e22e86c3	Avoid creating tensor in CosmosAttnProcessor2_0 (#11761 )	2025-06-21 11:53:58 +08:00