企业官网建设流程全解析

代码

• 原始代码

defpre_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])ifx.shape[3]%patch_h!=0orx.shape[4]%patch_w!=0:raiseValueError("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)ifself.seperated_timestepandfuse_vae_embedding_in_latents:ifnothasattr(self,"patch_size")orlen(self.patch_size)<3:raiseValueError(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,:]=0token_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:raiseNotImplementedError("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]ifself.action_conditionedandactionisnotNone: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 masknum_temporal_groups=f-1# first latent frame do not attend to actionsifnum_temporal_groups<=0:raiseValueError("Action-conditioned context mask requires at least 2 latent frames when `action` is provided.")assertaction_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 tokensaction_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 lengthfinal_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 tokensfinal_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 tokensfinal_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_maskelifself.action_conditionedandactionisNone:iff!=1:raiseValueError("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,},}

• 运行结果

--------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'>

流程图