videox_fun/models/flux2_transformer2d_control.py

# Modified from https://github.com/ali-vilab/VACE/blob/main/control/models/wan/wan_control.py
# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.

import glob
import inspect
import json
import os
from typing import Any, Dict, List, Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.loaders import FromOriginalModelMixin
from diffusers.models.attention_processor import Attention, AttentionProcessor
from diffusers.models.embeddings import (TimestepEmbedding, Timesteps,
                                         apply_rotary_emb,
                                         get_1d_rotary_pos_embed)
from diffusers.models.modeling_outputs import Transformer2DModelOutput
from diffusers.models.modeling_utils import ModelMixin
from diffusers.models.normalization import AdaLayerNormContinuous
from diffusers.utils import (USE_PEFT_BACKEND, is_torch_npu_available,
                             is_torch_version, logging, scale_lora_layers,
                             unscale_lora_layers)

from .flux2_transformer2d import (Flux2SingleTransformerBlock,
                                  Flux2Transformer2DModel,
                                  Flux2TransformerBlock)


class Flux2ControlTransformerBlock(Flux2TransformerBlock):
    def __init__(
        self, 
        dim: int,
        num_attention_heads: int,
        attention_head_dim: int,
        mlp_ratio: float = 3.0,
        eps: float = 1e-6,
        bias: bool = False,
        block_id=0
    ):
        super().__init__(dim, num_attention_heads, attention_head_dim, mlp_ratio, eps, bias)
        self.block_id = block_id
        if block_id == 0:
            self.before_proj = nn.Linear(dim, dim)
            nn.init.zeros_(self.before_proj.weight)
            nn.init.zeros_(self.before_proj.bias)
        self.after_proj = nn.Linear(dim, dim)
        nn.init.zeros_(self.after_proj.weight)
        nn.init.zeros_(self.after_proj.bias)

    def forward(self, c, x, **kwargs):
        if self.block_id == 0:
            c = self.before_proj(c) + x
            all_c = []
        else:
            all_c = list(torch.unbind(c))
            c = all_c.pop(-1)

        encoder_hidden_states, c = super().forward(c, **kwargs)
        c_skip = self.after_proj(c)
        all_c += [c_skip, c]
        c = torch.stack(all_c)
        return encoder_hidden_states, c
    
    
class BaseFlux2TransformerBlock(Flux2TransformerBlock):
    def __init__(
        self, 
        dim: int,
        num_attention_heads: int,
        attention_head_dim: int,
        mlp_ratio: float = 3.0,
        eps: float = 1e-6,
        bias: bool = False,
        block_id=0
    ):
        super().__init__(dim, num_attention_heads, attention_head_dim, mlp_ratio, eps, bias)
        self.block_id = block_id

    def forward(self, hidden_states, hints=None, context_scale=1.0, **kwargs):
        encoder_hidden_states, hidden_states = super().forward(hidden_states, **kwargs)
        if self.block_id is not None:
            hidden_states = hidden_states + hints[self.block_id] * context_scale
        return encoder_hidden_states, hidden_states


class Flux2ControlTransformer2DModel(Flux2Transformer2DModel):
    @register_to_config
    def __init__(
        self,
        control_layers=None,
        control_in_dim=None,
        patch_size: int = 1,
        in_channels: int = 128,
        out_channels: Optional[int] = None,
        num_layers: int = 8,
        num_single_layers: int = 48,
        attention_head_dim: int = 128,
        num_attention_heads: int = 48,
        joint_attention_dim: int = 15360,
        timestep_guidance_channels: int = 256,
        mlp_ratio: float = 3.0,
        axes_dims_rope: Tuple[int, ...] = (32, 32, 32, 32),
        rope_theta: int = 2000,
        eps: float = 1e-6,
    ):
        super().__init__(
            patch_size, in_channels, out_channels, num_layers, num_single_layers, attention_head_dim, 
            num_attention_heads, joint_attention_dim, timestep_guidance_channels, mlp_ratio, axes_dims_rope, 
            rope_theta, eps
        )

        self.control_layers = [i for i in range(0, self.num_layers, 2)] if control_layers is None else control_layers
        self.control_in_dim = self.in_dim if control_in_dim is None else control_in_dim

        assert 0 in self.control_layers
        self.control_layers_mapping = {i: n for n, i in enumerate(self.control_layers)}

        # blocks
        del self.transformer_blocks
        self.transformer_blocks = nn.ModuleList(
            [
                BaseFlux2TransformerBlock(
                    dim=self.inner_dim,
                    num_attention_heads=num_attention_heads,
                    attention_head_dim=attention_head_dim,
                    mlp_ratio=mlp_ratio,
                    eps=eps,
                    block_id=self.control_layers_mapping[i] if i in self.control_layers else None
                )
                for i in range(num_layers)
            ]
        )

        # control blocks
        self.control_transformer_blocks = nn.ModuleList(
            [
                Flux2ControlTransformerBlock(
                    dim=self.inner_dim,
                    num_attention_heads=num_attention_heads,
                    attention_head_dim=attention_head_dim,
                    mlp_ratio=mlp_ratio,
                    eps=eps,
                    block_id=i
                )
                for i in self.control_layers
            ]
        )

        # control patch embeddings
        self.control_img_in = nn.Linear(self.control_in_dim, self.inner_dim)

    def forward_control(
        self,
        x,
        control_context,
        kwargs
    ):
        # embeddings
        c = self.control_img_in(control_context)
        # Context Parallel
        if self.sp_world_size > 1:
            c = torch.chunk(c, self.sp_world_size, dim=1)[self.sp_world_rank]

        # arguments
        new_kwargs = dict(x=x)
        new_kwargs.update(kwargs)
        
        for block in self.control_transformer_blocks:
            if torch.is_grad_enabled() and self.gradient_checkpointing:
                def create_custom_forward(module, **static_kwargs):
                    def custom_forward(*inputs):
                        return module(*inputs, **static_kwargs)
                    return custom_forward
                ckpt_kwargs = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
                encoder_hidden_states, c = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(block, **new_kwargs),
                    c,
                    **ckpt_kwargs,
                )
            else:
                encoder_hidden_states, c = block(c, **new_kwargs)
            new_kwargs["encoder_hidden_states"] = encoder_hidden_states
 
        hints = torch.unbind(c)[:-1]
        return hints

    def forward(
        self,
        hidden_states: torch.Tensor,
        encoder_hidden_states: 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,
        control_context=None,
        control_context_scale=1.0,
        return_dict: bool = True,
    ):
        num_txt_tokens = encoder_hidden_states.shape[1]

        # 1. Calculate timestep embedding and modulation parameters
        timestep = timestep.to(hidden_states.dtype) * 1000
        guidance = guidance.to(hidden_states.dtype) * 1000

        temb = self.time_guidance_embed(timestep, guidance)

        double_stream_mod_img = self.double_stream_modulation_img(temb)
        double_stream_mod_txt = self.double_stream_modulation_txt(temb)
        single_stream_mod = self.single_stream_modulation(temb)[0]

        # 2. Input projection for image (hidden_states) and conditioning text (encoder_hidden_states)
        hidden_states = self.x_embedder(hidden_states)
        encoder_hidden_states = self.context_embedder(encoder_hidden_states)

        # 3. Calculate RoPE embeddings from image and text tokens
        # NOTE: the below logic means that we can't support batched inference with images of different resolutions or
        # text prompts of differents lengths. Is this a use case we want to support?
        if img_ids.ndim == 3:
            img_ids = img_ids[0]
        if txt_ids.ndim == 3:
            txt_ids = txt_ids[0]

        if is_torch_npu_available():
            freqs_cos_image, freqs_sin_image = self.pos_embed(img_ids.cpu())
            image_rotary_emb = (freqs_cos_image.npu(), freqs_sin_image.npu())
            freqs_cos_text, freqs_sin_text = self.pos_embed(txt_ids.cpu())
            text_rotary_emb = (freqs_cos_text.npu(), freqs_sin_text.npu())
        else:
            image_rotary_emb = self.pos_embed(img_ids)
            text_rotary_emb = self.pos_embed(txt_ids)
        concat_rotary_emb = (
            torch.cat([text_rotary_emb[0], image_rotary_emb[0]], dim=0),
            torch.cat([text_rotary_emb[1], image_rotary_emb[1]], dim=0),
        )

        # Arguments
        kwargs = dict(
            encoder_hidden_states=encoder_hidden_states,
            temb_mod_params_img=double_stream_mod_img,
            temb_mod_params_txt=double_stream_mod_txt,
            image_rotary_emb=concat_rotary_emb,
            joint_attention_kwargs=joint_attention_kwargs,
        )
        hints = self.forward_control(
            hidden_states, control_context, kwargs
        )

        for index_block, block in enumerate(self.transformer_blocks):
            # Arguments
            kwargs = dict(
                encoder_hidden_states=encoder_hidden_states,
                temb_mod_params_img=double_stream_mod_img,
                temb_mod_params_txt=double_stream_mod_txt,
                image_rotary_emb=concat_rotary_emb,
                joint_attention_kwargs=joint_attention_kwargs,
                hints=hints,
                context_scale=control_context_scale
            )
            if torch.is_grad_enabled() and self.gradient_checkpointing:
                def create_custom_forward(module, **static_kwargs):
                    def custom_forward(*inputs):
                        return module(*inputs, **static_kwargs)
                    return custom_forward

                ckpt_kwargs = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}

                encoder_hidden_states, hidden_states = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(block, **kwargs),
                    hidden_states,
                    **ckpt_kwargs,
                )
            else:
                encoder_hidden_states, hidden_states = block(hidden_states, **kwargs)

        for index_block, block in enumerate(self.single_transformer_blocks):
            if torch.is_grad_enabled() and self.gradient_checkpointing:
                def create_custom_forward(module):
                    def custom_forward(*inputs):
                        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,
                    single_stream_mod,
                    concat_rotary_emb,
                    joint_attention_kwargs,
                    **ckpt_kwargs,
                )
            else:
                encoder_hidden_states, hidden_states = block(
                    hidden_states=hidden_states,
                    encoder_hidden_states=encoder_hidden_states,
                    temb_mod_params=single_stream_mod,
                    image_rotary_emb=concat_rotary_emb,
                    joint_attention_kwargs=joint_attention_kwargs,
                )

        # 6. Output layers
        hidden_states = self.norm_out(hidden_states, temb)
        output = self.proj_out(hidden_states)

        if not return_dict:
            return (output,)

        return Transformer2DModelOutput(sample=output)