fix head transpose bug and torch compile bug

2025-03-13 22:30:39 +00:00
parent e45b41be79
commit 995992aab2
2 changed files with 2 additions and 3 deletions
@@ -172,7 +172,7 @@ def parallel_attention(q, k, v, img_q_len, img_kv_len, text_mask, mask_strategy=
                for i, head_idx in enumerate(heads):
                    hidden_states[:, head_idx:head_idx + 1, :, :] = processed_heads[:, i:i + 1, :, :]
                
-                hidden_states = hidden_states.transpose(1, 2)
+            hidden_states = hidden_states.transpose(1, 2)
    else:
        query = torch.cat([query, encoder_query], dim=1)
        key = torch.cat([key, encoder_key], dim=1)
@@ -45,5 +45,4 @@ CUDA_VISIBLE_DEVICES=1 torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_p
    --model_path $MODEL_BASE \
    --mask_strategy_file_path $mask_strategy_file_path \
    --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt \
-    --vae-sp \
-    --enable_torch_compile
+    --vae-sp