univa/models/modeling_univa_denoise_tower.py

from univa.models.configuration_univa_denoise_tower import UnivaDenoiseTowerConfig
from transformers.modeling_utils import PreTrainedModel

from typing import Any, Dict, Optional, Tuple, Union
import torch
from torch import nn
import numpy as np
from diffusers import FluxTransformer2DModel, SD3Transformer2DModel
from diffusers.utils import is_torch_version
from diffusers.models.modeling_outputs import Transformer2DModelOutput
from torch.nn.utils.rnn import pad_sequence


class UnivaDenoiseTower(PreTrainedModel):
    config_class = UnivaDenoiseTowerConfig
    base_model_prefix = "model"

    def __init__(self, config: UnivaDenoiseTowerConfig):
        super().__init__(config)
        self.config = config
        if config.denoiser_type == "flux":
            self.denoiser = FluxTransformer2DModel.from_config(config.denoiser_config)
        elif config.denoiser_type == "sd3":
            self.denoiser = SD3Transformer2DModel.from_config(config.denoiser_config)
        else:
            raise ValueError(f"Unknown denoiser type: {config.denoiser_type}")

        self._init_denoise_projector()
        self._init_vae_projector()
        self._init_siglip_projector()

    def _init_denoise_projector(self):
        """Initialize the denoise_projector for multi model input."""
        if self.config.denoise_projector_type == "mlp2x_gelu":
            self.denoise_projector = nn.Sequential(
                nn.Linear(
                    self.config.input_hidden_size,
                    self.config.output_hidden_size * 3,
                ),
                nn.SiLU(),
                nn.Linear(
                    self.config.output_hidden_size * 3, self.config.output_hidden_size
                ),
            )
        else:
            raise ValueError(
                f"Unknown denoise_projector_type: {self.config.denoise_projector_type}"
            )

    def _init_vae_projector(self):
        """Initialize the denoise_projector for multi model input."""
        if self.config.vae_projector_type == "mlp2x_gelu":
            self.vae_projector = nn.Sequential(
                nn.Linear(
                    self.config.vae_input_hidden_size,
                    # 2 * self.config.output_hidden_size,
                    3072,  # HARDCODE, x_embedder from flux
                ),
                nn.SiLU(),
                nn.Linear(
                    # 2 * self.config.output_hidden_size, 
                    3072,  # HARDCODE, x_embedder from flux
                    self.config.output_hidden_size
                ),
            )
        # elif self.config.vae_projector_type == "linear":
        # self.vae_projector = nn.Sequential(
        #     nn.Linear(
        #         self.config.vae_input_hidden_size,
        #         self.config.output_hidden_size,
        #     ),
        # )
        else:
            raise ValueError(
                f"Unknown vae_projector_type: {self.config.vae_projector_type}"
            )

    def _init_siglip_projector(self):
        """Initialize the denoise_projector for multi model input."""
        self.siglip_projector = nn.Sequential(
            nn.Linear(
                1152,  # HARDCODE, out from siglip
                4096 * 3,  # HARDCODE
            ),
            nn.SiLU(),
            nn.Linear(
                4096 * 3,  # HARDCODE
                4096,  # HARDCODE, context_embedder from flux
            ),
        )

    def _init_convnext_projector(self):
        """Initialize the denoise_projector for multi model input."""
        self.convnext_projector = nn.Sequential(
            nn.Linear(
                1152,  # HARDCODE, out from convnext
                4096 * 3,  # HARDCODE
            ),
            nn.SiLU(),
            nn.Linear(
                4096 * 3,  # HARDCODE
                4096,  # HARDCODE, context_embedder from flux
            ),
        )
        
    @staticmethod
    def _insert_image_feats(
            encoder_h, img_feats, img_pos, 
            output_hidden_size, vae_projector
            ):
        """
        encoder_h: Tensor[B, L, D]
        img_feats: list of B lists: 第 i 个元素是一个 list，长度 = len(img_pos[i])，
                其内第 k 项是一个 Tensor[Nik, D]
        img_pos:   list of B lists: 第 i 个元素是个位置列表 [p_i0, p_i1, ...]
                len(img_pos[i]) == len(img_feats[i])
        returns:   Tensor[B, L + Nmax, D]，在各自位置插入完后，按最长插入数 pad 右侧
        """
        B, L, D = encoder_h.shape
        device = encoder_h.device

        # —— 1. 每个样本先把多组 feats concat 成一条“插入流”，同时 expand positions
        flat_feats = []
        flat_pos   = []
        for feats_list, pos_list in zip(img_feats, img_pos):
            assert len(feats_list) == len(pos_list)
            # feats_list = [Tensor[N0,D], Tensor[N1,D], ...]
            # pos_list   = [p0,      p1,       ...]
            # concat 所有要插入的 tokens
            if len(feats_list) == 0:
                # 没有插入
                concat_f = torch.empty(0, output_hidden_size, device=device)
                pos_expanded = torch.empty(0, dtype=torch.long, device=device)
            else:
                concat_f = torch.cat(feats_list, dim=0)    # [Ni_total, D]
                concat_f = vae_projector(concat_f)
                # 对应位置也 expand 成同样长度
                # eg. feats_list[0].shape[0] 个 p0， feats_list[1].shape[0] 个 p1，…
                # ATTENTION p-1
                pos_expanded = torch.cat([
                    torch.full((f.shape[0],), p-1, dtype=torch.long, device=device)
                    for f, p in zip(feats_list, pos_list)
                ], dim=0)                                   # [Ni_total]
            flat_feats.append(concat_f)
            flat_pos.append(pos_expanded)

        # —— 2. pad 到同一个长度 Nmax
        padded_feats = pad_sequence(flat_feats, batch_first=True)    # [B, Nmax, D]
        pos_pad = pad_sequence(flat_pos, batch_first=True, padding_value=L)

        # —— 3. 准备所有 token 的“排序键”（sort‐key）
        # 原 token j 的 key = 2*j
        orig_key = (torch.arange(L, device=device) * 2).unsqueeze(0).expand(B, -1)       # [B, L]
        # 插入 token 的 key = 2*pos + 1
        ins_key  = pos_pad * 2 + 1                                                      # [B, Nmax]

        # —— 4. 拼接、一次性排序 + gather
        all_keys   = torch.cat([orig_key,    ins_key],    dim=1)                        # [B, L+Nmax]
        all_feats  = torch.cat([encoder_h, padded_feats], dim=1)                        # [B, L+Nmax, D]
        sort_idx   = all_keys.argsort(dim=1)                                            # [B, L+Nmax]
        new_seq    = all_feats.gather(1, sort_idx.unsqueeze(-1).expand(-1, -1, D))      # [B, L+Nmax, D]

        return new_seq

    def forward(
        self,
        hidden_states: torch.Tensor,
        timestep: torch.Tensor,
        encoder_hidden_states: torch.Tensor,
        pooled_projections: torch.Tensor,
        **kwargs,
    ) -> torch.Tensor:
        # if encoder_hidden_states is not None:
        #     encoder_hidden_states = self.denoise_projector(encoder_hidden_states)
        if self.config.denoiser_type == "flux":
            prefix_prompt_embeds = kwargs.pop("prefix_prompt_embeds", None)
            
            if encoder_hidden_states is not None:
                if prefix_prompt_embeds is not None:
                    encoder_hidden_states = torch.concat(
                        [encoder_hidden_states, prefix_prompt_embeds], dim=1
                    )
            else:
                assert prefix_prompt_embeds is not None
                encoder_hidden_states = prefix_prompt_embeds
            txt_ids = torch.zeros(encoder_hidden_states.shape[1], 3).to(
                hidden_states.device, dtype=hidden_states.dtype
            )

            joint_attention_kwargs = kwargs.pop('joint_attention_kwargs', None)
            # if joint_attention_kwargs is not None and 'attention_mask' in joint_attention_kwargs:
            #     attention_mask = joint_attention_kwargs['attention_mask']
            # else:
            #     attention_mask = torch.full(
            #         (hidden_states.shape[0], 1, hidden_states.shape[1]), 
            #         True, dtype=torch.bool, device=hidden_states.device
            #         )
                
            enc_attention_mask = kwargs.pop('enc_attention_mask', None)
            # if enc_attention_mask is None:
            #     enc_attention_mask = torch.full(
            #         (encoder_hidden_states.shape[0], 1, encoder_hidden_states.shape[1]), 
            #         True, dtype=torch.bool, device=encoder_hidden_states.device
            #         )
            # else:
            #     enc_attention_mask = enc_attention_mask.unsqueeze(1)
                    
            # attention_mask = torch.concat([enc_attention_mask, attention_mask], dim=-1)
            # attention_mask = attention_mask.unsqueeze(1)

            # joint_attention_kwargs['attention_mask'] = attention_mask
            # kwargs['joint_attention_kwargs'] = joint_attention_kwargs

            # print(f'hidden_states.shape, {hidden_states.shape}, encoder_hidden_states.shape, {encoder_hidden_states.shape}')
            # return self.fixed_flux_forward(
            return self.denoiser(
                hidden_states=hidden_states,
                timestep=timestep, # Note: timestep is in [0, 1]. It has been scaled by 1000 in the training script.
                encoder_hidden_states=encoder_hidden_states,
                pooled_projections=pooled_projections,
                txt_ids=txt_ids,
                **kwargs,
            )[0]

        elif self.config.denoiser_type == "sd3":
            prefix_prompt_embeds = kwargs.pop("prefix_prompt_embeds", None)
            if prefix_prompt_embeds is not None:
                encoder_hidden_states = torch.concat(
                    [prefix_prompt_embeds, encoder_hidden_states], dim=1
                )

            return self.denoiser(
                hidden_states=hidden_states,
                timestep=timestep,
                encoder_hidden_states=encoder_hidden_states,
                pooled_projections=pooled_projections,
                **kwargs,
            )[0]



    def fixed_flux_forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: torch.Tensor = None,
        pooled_projections: torch.Tensor = None,
        timestep: torch.LongTensor = None,
        img_ids: torch.Tensor = None,
        txt_ids: torch.Tensor = None,
        guidance: torch.Tensor = None,
        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
        controlnet_block_samples=None,
        controlnet_single_block_samples=None,
        return_dict: bool = True,
        controlnet_blocks_repeat: bool = False,
    ) -> Union[torch.FloatTensor, Transformer2DModelOutput]:
        """
        The [`FluxTransformer2DModel`] forward method.

        Args:
            hidden_states (`torch.FloatTensor` of shape `(batch size, channel, height, width)`):
                Input `hidden_states`.
            encoder_hidden_states (`torch.FloatTensor` of shape `(batch size, sequence_len, embed_dims)`):
                Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
            pooled_projections (`torch.FloatTensor` of shape `(batch_size, projection_dim)`): Embeddings projected
                from the embeddings of input conditions.
            timestep ( `torch.LongTensor`):
                Used to indicate denoising step.
            block_controlnet_hidden_states: (`list` of `torch.Tensor`):
                A list of tensors that if specified are added to the residuals of transformer blocks.
            joint_attention_kwargs (`dict`, *optional*):
                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
                `self.processor` in
                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
                tuple.

        Returns:
            If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
            `tuple` where the first element is the sample tensor.
        """

        hidden_states = self.denoiser.x_embedder(hidden_states)

        timestep = timestep.to(hidden_states.dtype) * 1000
        if guidance is not None:
            guidance = guidance.to(hidden_states.dtype) * 1000
        else:
            guidance = None

        temb = (
            self.denoiser.time_text_embed(timestep, pooled_projections)
            if guidance is None
            else self.denoiser.time_text_embed(timestep, guidance, pooled_projections)
        )
        encoder_hidden_states = self.denoiser.context_embedder(encoder_hidden_states)

        if txt_ids.ndim == 3:
            txt_ids = txt_ids[0]
        if img_ids.ndim == 3:
            img_ids = img_ids[0]

        ids = torch.cat((txt_ids, img_ids), dim=0)
        image_rotary_emb = self.denoiser.pos_embed(ids)
        if joint_attention_kwargs is not None and "ip_adapter_image_embeds" in joint_attention_kwargs:
            ip_adapter_image_embeds = joint_attention_kwargs.pop("ip_adapter_image_embeds")
            ip_hidden_states = self.denoiser.encoder_hid_proj(ip_adapter_image_embeds)
            joint_attention_kwargs.update({"ip_hidden_states": ip_hidden_states})

        for index_block, block in enumerate(self.denoiser.transformer_blocks):
            if torch.is_grad_enabled() and self.denoiser.gradient_checkpointing:

                def create_custom_forward(module, return_dict=None):
                    def custom_forward(*inputs):
                        if return_dict is not None:
                            return module(*inputs, return_dict=return_dict)
                        else:
                            return module(*inputs)

                    return custom_forward

                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
                encoder_hidden_states, hidden_states = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(block),
                    hidden_states,
                    encoder_hidden_states,
                    temb,
                    image_rotary_emb,
                    joint_attention_kwargs,  # add this line
                    **ckpt_kwargs,
                )

            else:
                encoder_hidden_states, hidden_states = block(
                    hidden_states=hidden_states,
                    encoder_hidden_states=encoder_hidden_states,
                    temb=temb,
                    image_rotary_emb=image_rotary_emb,
                    joint_attention_kwargs=joint_attention_kwargs,
                )

            # controlnet residual
            if controlnet_block_samples is not None:
                interval_control = len(self.denoiser.transformer_blocks) / len(controlnet_block_samples)
                interval_control = int(np.ceil(interval_control))
                # For Xlabs ControlNet.
                if controlnet_blocks_repeat:
                    hidden_states = (
                        hidden_states + controlnet_block_samples[index_block % len(controlnet_block_samples)]
                    )
                else:
                    hidden_states = hidden_states + controlnet_block_samples[index_block // interval_control]
        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)

        for index_block, block in enumerate(self.denoiser.single_transformer_blocks):
            if torch.is_grad_enabled() and self.denoiser.gradient_checkpointing:

                def create_custom_forward(module, return_dict=None):
                    def custom_forward(*inputs):
                        if return_dict is not None:
                            return module(*inputs, return_dict=return_dict)
                        else:
                            return module(*inputs)

                    return custom_forward

                ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
                hidden_states = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(block),
                    hidden_states,
                    temb,
                    image_rotary_emb,
                    joint_attention_kwargs, 
                    **ckpt_kwargs,
                )

            else:
                hidden_states = block(
                    hidden_states=hidden_states,
                    temb=temb,
                    image_rotary_emb=image_rotary_emb,
                    joint_attention_kwargs=joint_attention_kwargs,
                )

            # controlnet residual
            if controlnet_single_block_samples is not None:
                interval_control = len(self.denoiser.single_transformer_blocks) / len(controlnet_single_block_samples)
                interval_control = int(np.ceil(interval_control))
                hidden_states[:, encoder_hidden_states.shape[1] :, ...] = (
                    hidden_states[:, encoder_hidden_states.shape[1] :, ...]
                    + controlnet_single_block_samples[index_block // interval_control]
                )

        hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...]

        hidden_states = self.denoiser.norm_out(hidden_states, temb)
        output = self.denoiser.proj_out(hidden_states)


        if not return_dict:
            return (output,)

        return Transformer2DModelOutput(sample=output)