7. fastwam 模型 _predict_action_noise_with_cache部分

代码

python 复制代码

 for step_t_action, step_delta_action in zip(infer_timesteps_action, infer_deltas_action):
            timestep_action = step_t_action.unsqueeze(0).to(dtype=latents_action.dtype, device=self.device)

            pred_action_posi = self._predict_action_noise_with_cache(
                latents_action=latents_action,
                timestep_action=timestep_action,
                context=context,
                context_mask=context_mask,
                video_kv_cache=video_kv_cache,
                attention_mask=attention_mask,
                video_seq_len=video_seq_len,
            )
            pred_action = pred_action_posi

            latents_action = self.infer_action_scheduler.step(pred_action, step_delta_action, latents_action)

python 复制代码

@torch.no_grad()
    def _predict_action_noise_with_cache(
        self,
        latents_action: torch.Tensor,
        timestep_action: torch.Tensor,
        context: torch.Tensor,
        context_mask: torch.Tensor,
        video_kv_cache: list[dict[str, torch.Tensor]],
        attention_mask: torch.Tensor,
        video_seq_len: int,
    ) -> torch.Tensor:
        if self._pre_action_test == 0:     
            print("before action_pre")
            print("latents_action.shape:", latents_action.shape)
            print("timestep_action.shape:", timestep_action.shape)
            print("context.shape:", context.shape)
            print("context_mask.shape:", context_mask.shape)

        action_pre = self.action_expert.pre_dit(
            action_tokens=latents_action,
            timestep=timestep_action,
            context=context,
            context_mask=context_mask,
        )
        if self._pre_action_test == 0:  
            print("after action_pre")
            print("action_pre.tokens.shape:", action_pre["tokens"].shape)
            print("action_pre.freqs.shape:", action_pre["freqs"].shape)
            print("action_pre.t_mod.shape:", action_pre["t_mod"].shape)
            print("action_pre.context.shape:", action_pre["context"].shape)
            print("action_pre.context_mask.shape:", action_pre["context_mask"].shape)
            print("attention_mask.shape:", attention_mask.shape)
            print("video_seq_len:", video_seq_len)

        action_tokens = self.mot.forward_action_with_video_cache(
            action_tokens=action_pre["tokens"],
            action_freqs=action_pre["freqs"],
            action_t_mod=action_pre["t_mod"],
            action_context_payload={
                "context": action_pre["context"],
                "mask": action_pre["context_mask"],
            },
            video_kv_cache=video_kv_cache,
            attention_mask=attention_mask,
            video_seq_len=video_seq_len,
        )
        if self._pre_action_test == 0:  
            print("after forward_action_with_video_cache")
            print("action_tokens.shape:", action_tokens.shape)
        pred_action = self.action_expert.post_dit(action_tokens, action_pre)
        if self._pre_action_test == 0:  
            print("after post_dit")
            print("pred_action.shape:", pred_action.shape)
            self._pre_action_test += 1
        
        return pred_action

输出

bash 复制代码

before action_pre
latents_action.shape: torch.Size([1, 32, 7])
timestep_action.shape: torch.Size([1])
context.shape: torch.Size([1, 129, 4096])
context_mask.shape: torch.Size([1, 129])
after action_pre
action_pre.tokens.shape: torch.Size([1, 32, 1024])
action_pre.freqs.shape: torch.Size([32, 1, 64])
action_pre.t_mod.shape: torch.Size([1, 6, 1024])
action_pre.context.shape: torch.Size([1, 129, 1024])
action_pre.context_mask.shape: torch.Size([1, 32, 129])
attention_mask.shape: torch.Size([130, 130])
video_seq_len: 98
after forward_action_with_video_cache
action_tokens.shape: torch.Size([1, 32, 1024])
after post_dit
pred_action.shape: torch.Size([1, 32, 7])

流程图

#mermaid-svg-va4e42StKQVmeGPv{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-va4e42StKQVmeGPv .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-va4e42StKQVmeGPv .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-va4e42StKQVmeGPv .error-icon{fill:#552222;}#mermaid-svg-va4e42StKQVmeGPv .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-va4e42StKQVmeGPv .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-va4e42StKQVmeGPv .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-va4e42StKQVmeGPv .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-va4e42StKQVmeGPv .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-va4e42StKQVmeGPv .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-va4e42StKQVmeGPv .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-va4e42StKQVmeGPv .marker{fill:#333333;stroke:#333333;}#mermaid-svg-va4e42StKQVmeGPv .marker.cross{stroke:#333333;}#mermaid-svg-va4e42StKQVmeGPv svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-va4e42StKQVmeGPv p{margin:0;}#mermaid-svg-va4e42StKQVmeGPv .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-va4e42StKQVmeGPv .cluster-label text{fill:#333;}#mermaid-svg-va4e42StKQVmeGPv .cluster-label span{color:#333;}#mermaid-svg-va4e42StKQVmeGPv .cluster-label span p{background-color:transparent;}#mermaid-svg-va4e42StKQVmeGPv .label text,#mermaid-svg-va4e42StKQVmeGPv span{fill:#333;color:#333;}#mermaid-svg-va4e42StKQVmeGPv .node rect,#mermaid-svg-va4e42StKQVmeGPv .node circle,#mermaid-svg-va4e42StKQVmeGPv .node ellipse,#mermaid-svg-va4e42StKQVmeGPv .node polygon,#mermaid-svg-va4e42StKQVmeGPv .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-va4e42StKQVmeGPv .rough-node .label text,#mermaid-svg-va4e42StKQVmeGPv .node .label text,#mermaid-svg-va4e42StKQVmeGPv .image-shape .label,#mermaid-svg-va4e42StKQVmeGPv .icon-shape .label{text-anchor:middle;}#mermaid-svg-va4e42StKQVmeGPv .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-va4e42StKQVmeGPv .rough-node .label,#mermaid-svg-va4e42StKQVmeGPv .node .label,#mermaid-svg-va4e42StKQVmeGPv .image-shape .label,#mermaid-svg-va4e42StKQVmeGPv .icon-shape .label{text-align:center;}#mermaid-svg-va4e42StKQVmeGPv .node.clickable{cursor:pointer;}#mermaid-svg-va4e42StKQVmeGPv .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-va4e42StKQVmeGPv .arrowheadPath{fill:#333333;}#mermaid-svg-va4e42StKQVmeGPv .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-va4e42StKQVmeGPv .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-va4e42StKQVmeGPv .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-va4e42StKQVmeGPv .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-va4e42StKQVmeGPv .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-va4e42StKQVmeGPv .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-va4e42StKQVmeGPv .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-va4e42StKQVmeGPv .cluster text{fill:#333;}#mermaid-svg-va4e42StKQVmeGPv .cluster span{color:#333;}#mermaid-svg-va4e42StKQVmeGPv div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-va4e42StKQVmeGPv .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-va4e42StKQVmeGPv rect.text{fill:none;stroke-width:0;}#mermaid-svg-va4e42StKQVmeGPv .icon-shape,#mermaid-svg-va4e42StKQVmeGPv .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-va4e42StKQVmeGPv .icon-shape p,#mermaid-svg-va4e42StKQVmeGPv .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-va4e42StKQVmeGPv .icon-shape .label rect,#mermaid-svg-va4e42StKQVmeGPv .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-va4e42StKQVmeGPv .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-va4e42StKQVmeGPv .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-va4e42StKQVmeGPv :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} action_expert.post_dit
mot.forward_action_with_video_cache
action_expert.pre_dit
输入
latents_action

1, 32, 7

timestep_action

1

context

1, 129, 4096

context_mask

1, 129

video_kv_cache

List $num_layers$

each: k/v $1, 98, H\*Dh$
attention_mask

130, 130

= 98 video + 32 action
video_seq_len

98
action_expert.pre_dit

action_tokens: $1, 32, 7$

timestep: $1$

context: $1, 129, 4096$

context_mask: $1, 129$
action_pre.tokens

1, 32, 1024

action_pre.freqs

32, 1, 64

action_pre.t_mod

1, 6, 1024

action_pre.context

1, 129, 1024

action_pre.context_mask

1, 32, 129

action_context_payload

context: $1, 129, 1024$

mask: $1, 32, 129$
forward_action_with_video_cache

video_seq_len: 98

total_seq: 130
action_attention_mask

attention_mask $98:130, :130$

32, 130

q_action

1, 32, H\*Dh

k_cat = $video_k, action_k$

1, 130, H\*Dh

v_cat = $video_v, action_v$

1, 130, H\*Dh

mixed attention

action queries attend to video + action

output: $1, 32, H\*Dh$
action expert post blocks

self-attn output proj + cross-attn + FFN

output: $1, 32, 1024$
action_tokens

1, 32, 1024

action_expert.post_dit

input tokens: $1, 32, 1024$
pred_action

1, 32, 7