5. fastwam 模型 video expert pre dit过程

代码

原始代码

python 复制代码

    def pre_dit(
        self,
        x: torch.Tensor,
        timestep: torch.Tensor,
        context: torch.Tensor,
        context_mask: Optional[torch.Tensor] = None,
        action: Optional[torch.Tensor] = None,
        fuse_vae_embedding_in_latents: bool = False,
        control_camera_latents_input: Optional[torch.Tensor] = None,
    ) -> Dict[str, Any]:
        x, timestep, context_mask = self._validate_forward_inputs(
            x=x,
            timestep=timestep,
            context=context,
            context_mask=context_mask,
            action=action,
        )

        batch_size = x.shape[0]
        patch_h = int(self.patch_size[1])
        patch_w = int(self.patch_size[2])
        if x.shape[3] % patch_h != 0 or x.shape[4] % patch_w != 0:
            raise ValueError(
                "Latent spatial shape must be divisible by DiT patch size, "
                f"got HxW=({x.shape[3]}, {x.shape[4]}), patch=({patch_h}, {patch_w})"
            )
        tokens_per_frame = (x.shape[3] // patch_h) * (x.shape[4] // patch_w)

        if self.seperated_timestep and fuse_vae_embedding_in_latents:
            if not hasattr(self, "patch_size") or len(self.patch_size) < 3:
                raise ValueError(f"Invalid dit.patch_size: {getattr(self, 'patch_size', None)}")
            
            token_timesteps = torch.ones(
                (batch_size, x.shape[2], tokens_per_frame),
                dtype=timestep.dtype,
                device=timestep.device,
            ) * timestep.view(batch_size, 1, 1)
            token_timesteps[:, 0, :] = 0
            token_timesteps = token_timesteps.reshape(batch_size, -1)
            token_t_emb = sinusoidal_embedding_1d(self.freq_dim, token_timesteps.reshape(-1))
            t = self.time_embedding(token_t_emb).reshape(batch_size, -1, self.hidden_dim)
            t_mod = self.time_projection(t).unflatten(2, (6, self.hidden_dim))
        else:
            raise NotImplementedError("Only support seperated_timestep with fuse_vae_embedding_in_latents for now.")
            t = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, timestep))
            t_mod = self.time_projection(t).unflatten(1, (6, self.hidden_dim))

        print("patchify前的x:", x.shape)            
        x = self.patchify(x, control_camera_latents_input=control_camera_latents_input)
        print("patchify后的x:", x.shape)  
        f, h, w = x.shape[2:]

        print("text_embedding前的维度:", context.shape)
        context = self.text_embedding(context) # (B, L, dim)
        print("text_embedding后的维度:", context.shape)
        context_len = context.shape[1]
        if self.action_conditioned and action is not None:
            action_len = action.shape[1]
            action_emb = self.action_embedding(action) # (B, action_len, dim)
            action_pos_embed = sinusoidal_embedding_1d(self.hidden_dim, 
                torch.arange(action_len, device=action_emb.device)) # (action_len, dim)
            action_emb = action_emb + action_pos_embed.unsqueeze(0) # (B, action_len, dim)
            context = torch.cat([context, action_emb], dim=1) # (B, context_len + action_len, dim)

            # new mask
            num_temporal_groups = f - 1 # first latent frame do not attend to actions
            if num_temporal_groups <= 0:
                raise ValueError(
                    "Action-conditioned context mask requires at least 2 latent frames when `action` is provided."
                )
            assert action_emb.shape[1] % num_temporal_groups == 0, \
                f"Action embedding length {action_emb.shape[1]} must be divisible by number of temporal groups {num_temporal_groups}"
            # Each latent frame (from the 2nd one) attends to the corresponding group of action tokens
            action_group_mask = create_group_causal_attn_mask(
                num_temporal_groups=num_temporal_groups,
                num_query_per_group=tokens_per_frame,
                num_key_per_group=action_len // num_temporal_groups,
                mode=self.action_group_causal_mask_mode,
            ).to(context.device) # ((f-1)*tokens_per_frame, action_len)

            seq_len = f * h * w # query length
            final_context_mask = torch.zeros((batch_size, seq_len, context.shape[1]), dtype=torch.bool, device=context.device) # (B, seq_len, L + action_len)
            # all latent frames attend to text tokens
            final_context_mask[:, :, :context_len] = context_mask.unsqueeze(1).expand(-1, seq_len, -1) # (B, seq_len, L)
            # latent frames from the 2nd one attend to action tokens
            final_context_mask[:, tokens_per_frame:, context_len:] = action_group_mask.unsqueeze(0).expand(batch_size, -1, -1) # (B, seq_len, action_len)
            context_mask = final_context_mask
        elif self.action_conditioned and action is None:
            if f != 1:
                raise ValueError(
                    "Action-conditioned model requires `action` unless running single-frame text-only mode with num_latent_frames=1."
                )
            context_mask = context_mask.unsqueeze(1).expand(-1, f * h * w, -1) # (B, seq_len, L)
        else:
            context_mask = context_mask.unsqueeze(1).expand(-1, f * h * w, -1) # (B, seq_len, L)

        x_tokens = rearrange(x, "b c f h w -> b (f h w) c").contiguous()

        freqs = torch.cat([
            self.freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
            self.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
            self.freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
        ], dim=-1).reshape(f * h * w, 1, -1).to(x_tokens.device)

        return {
            "tokens": x_tokens,
            "freqs": freqs,
            "t": t,
            "t_mod": t_mod,
            "context": context,
            "context_mask": context_mask,
            "meta": {
                "grid_size": (f, h, w),
                "tokens_per_frame": tokens_per_frame,
                "batch_size": batch_size,
            },
        }

运行结果

bash 复制代码

--------before video_expert.pre_dit--------------
first_frame_latents: (1, 48, 1, 14, 28) torch.bfloat16 cuda:0
timestep_video: (1,) torch.bfloat16 cuda:0
context: (1, 129, 4096) torch.bfloat16 cuda:0
context_mask: (1, 129) torch.bool cuda:0
fuse_flag: True
patchify前的x: torch.Size([1, 48, 1, 14, 28])
patchify后的x: torch.Size([1, 3072, 1, 7, 14])
text_embedding前的维度: torch.Size([1, 129, 4096])
text_embedding后的维度: torch.Size([1, 129, 3072])
--------after video_expert.pre_dit--------------
[Processing] Key: tokens | Shape: torch.Size([1, 98, 3072]) | Dtype: torch.bfloat16
[Processing] Key: freqs | Shape: torch.Size([98, 1, 64]) | Dtype: torch.complex128
[Processing] Key: t | Shape: torch.Size([1, 98, 3072]) | Dtype: torch.bfloat16
[Processing] Key: t_mod | Shape: torch.Size([1, 98, 6, 3072]) | Dtype: torch.bfloat16
[Processing] Key: context | Shape: torch.Size([1, 129, 3072]) | Dtype: torch.bfloat16
[Processing] Key: context_mask | Shape: torch.Size([1, 98, 129]) | Dtype: torch.bool
[Processing] Key: meta | Type: <class 'dict'>

流程图

#mermaid-svg-U6NTxgAQtoBiEgXg{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-U6NTxgAQtoBiEgXg .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-U6NTxgAQtoBiEgXg .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-U6NTxgAQtoBiEgXg .error-icon{fill:#552222;}#mermaid-svg-U6NTxgAQtoBiEgXg .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-U6NTxgAQtoBiEgXg .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-U6NTxgAQtoBiEgXg .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-U6NTxgAQtoBiEgXg .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-U6NTxgAQtoBiEgXg .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-U6NTxgAQtoBiEgXg .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-U6NTxgAQtoBiEgXg .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-U6NTxgAQtoBiEgXg .marker{fill:#333333;stroke:#333333;}#mermaid-svg-U6NTxgAQtoBiEgXg .marker.cross{stroke:#333333;}#mermaid-svg-U6NTxgAQtoBiEgXg svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-U6NTxgAQtoBiEgXg p{margin:0;}#mermaid-svg-U6NTxgAQtoBiEgXg .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-U6NTxgAQtoBiEgXg .cluster-label text{fill:#333;}#mermaid-svg-U6NTxgAQtoBiEgXg .cluster-label span{color:#333;}#mermaid-svg-U6NTxgAQtoBiEgXg .cluster-label span p{background-color:transparent;}#mermaid-svg-U6NTxgAQtoBiEgXg .label text,#mermaid-svg-U6NTxgAQtoBiEgXg span{fill:#333;color:#333;}#mermaid-svg-U6NTxgAQtoBiEgXg .node rect,#mermaid-svg-U6NTxgAQtoBiEgXg .node circle,#mermaid-svg-U6NTxgAQtoBiEgXg .node ellipse,#mermaid-svg-U6NTxgAQtoBiEgXg .node polygon,#mermaid-svg-U6NTxgAQtoBiEgXg .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-U6NTxgAQtoBiEgXg .rough-node .label text,#mermaid-svg-U6NTxgAQtoBiEgXg .node .label text,#mermaid-svg-U6NTxgAQtoBiEgXg .image-shape .label,#mermaid-svg-U6NTxgAQtoBiEgXg .icon-shape .label{text-anchor:middle;}#mermaid-svg-U6NTxgAQtoBiEgXg .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-U6NTxgAQtoBiEgXg .rough-node .label,#mermaid-svg-U6NTxgAQtoBiEgXg .node .label,#mermaid-svg-U6NTxgAQtoBiEgXg .image-shape .label,#mermaid-svg-U6NTxgAQtoBiEgXg .icon-shape .label{text-align:center;}#mermaid-svg-U6NTxgAQtoBiEgXg .node.clickable{cursor:pointer;}#mermaid-svg-U6NTxgAQtoBiEgXg .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-U6NTxgAQtoBiEgXg .arrowheadPath{fill:#333333;}#mermaid-svg-U6NTxgAQtoBiEgXg .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-U6NTxgAQtoBiEgXg .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-U6NTxgAQtoBiEgXg .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-U6NTxgAQtoBiEgXg .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-U6NTxgAQtoBiEgXg .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-U6NTxgAQtoBiEgXg .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-U6NTxgAQtoBiEgXg .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-U6NTxgAQtoBiEgXg .cluster text{fill:#333;}#mermaid-svg-U6NTxgAQtoBiEgXg .cluster span{color:#333;}#mermaid-svg-U6NTxgAQtoBiEgXg 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-U6NTxgAQtoBiEgXg .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-U6NTxgAQtoBiEgXg rect.text{fill:none;stroke-width:0;}#mermaid-svg-U6NTxgAQtoBiEgXg .icon-shape,#mermaid-svg-U6NTxgAQtoBiEgXg .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-U6NTxgAQtoBiEgXg .icon-shape p,#mermaid-svg-U6NTxgAQtoBiEgXg .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-U6NTxgAQtoBiEgXg .icon-shape .label rect,#mermaid-svg-U6NTxgAQtoBiEgXg .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-U6NTxgAQtoBiEgXg .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-U6NTxgAQtoBiEgXg .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-U6NTxgAQtoBiEgXg :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} pre_dit inputs

x (first_frame_latents): (1,48,1,14,28) bf16 cuda

timestep: (1,) bf16 cuda

context: (1,129,4096) bf16 cuda

context_mask: (1,129) bool cuda

fuse_vae_embedding_in_latents: True
Validate shapes

_validate_forward_inputs
Compute tokens_per_frame

patch_size=(1,2,2)

H'=14,W'=28

patchify grid: h=7,w=14

tokens_per_frame = 7*14 = 98
Separated timestep (per-token)

seperated_timestep=True & fuse=True

Build token_timesteps:

shape (B,F,tokens_per_frame)=(1,1,98)

set frame0 timestep=0

flatten -> (1,98)

sinusoidal_embedding -> token_t_emb

(time_embedding)-> t: (1,98,3072)

(time_projection)-> t_mod: (1,98,6,3072)
Patchify video latents

Before: x (1,48,1,14,28)

After: x (1,3072,1,7,14)
Flatten to token sequence

x_tokens = rearrange

(b,c,f,h,w)->(b,f*h*w,c)

tokens: (1,98,3072)
Text embedding

text_embedding: 4096 -> 3072

context: (1,129,4096) -> (1,129,3072)
Expand context_mask to per-query

context_mask: (1,129)

-> (1,seq_len,129)

seq_len = f*h*w = 98

context_mask: (1,98,129)
Build rotary freqs

freqs: (seq_len,1,64)

= (98,1,64) complex128
pre_dit outputs dict
Return payload

(tokens, freqs, t, t_mod,

context, context_mask,

meta={grid_size=(1,7,14), tokens_per_frame=98, batch_size=1})