Method layer_forward

sat/dit_video_concat.py:674–786

(
        self,
        hidden_states,
        mask,
        *args,
        **kwargs,
    )

        return hidden_states

    def layer_forward(
        self,
        hidden_states,
        mask,
        *args,
        **kwargs,
    ):
        text_length = kwargs["text_length"]
        # hidden_states (b,(n_t+t*n_i),d)
        text_hidden_states = hidden_states[:, :text_length]  # (b,n,d)
        img_hidden_states = hidden_states[:, text_length:]  # (b,(t n),d)
        
        split_cond_flag = kwargs['split_cond_flag']
        layer = self.transformer.layers[kwargs["layer_id"]]
        adaLN_modulation = self.adaLN_modulations[kwargs["layer_id"]]

        (
            shift_msa,
            scale_msa,
            gate_msa,
            shift_mlp,
            scale_mlp,
            gate_mlp,
            text_shift_msa,
            text_scale_msa,
            text_gate_msa,
            text_shift_mlp,
            text_scale_mlp,
            text_gate_mlp,
        ) = adaLN_modulation(kwargs["emb"]).chunk(12, dim=1)  # * modulations for img and text

        gate_msa, gate_mlp, text_gate_msa, text_gate_mlp = (
            gate_msa.unsqueeze(1),
            gate_mlp.unsqueeze(1),
            text_gate_msa.unsqueeze(1),
            text_gate_mlp.unsqueeze(1),
        )

        # * prepare modulation for aug frames
        if split_cond_flag:
            cond_inds = kwargs["cond_inds"]
            pred_inds = [i for i in range(self.compressed_num_frames) if i not in cond_inds]
            assert cond_inds[-1] < pred_inds[0], f"cond frames must be ahead of pred frames, {cond_inds}, {pred_inds}"

            (
                aug_shift_msa,
                aug_scale_msa,
                aug_gate_msa,
                aug_shift_mlp,
                aug_scale_mlp,
                aug_gate_mlp,
            ) = adaLN_modulation(kwargs["aug_emb"]).chunk(12, dim=1)[: 6]  # * modulations for aug
            aug_gate_msa, aug_gate_mlp = aug_gate_msa.unsqueeze(1), aug_gate_mlp.unsqueeze(1)

        # self full attention (b,(t n),d)
        img_attention_input = layer.input_layernorm(img_hidden_states)
        text_attention_input = layer.input_layernorm(text_hidden_states)