diffsynth/models/qwen_image_connector.py

import torch
import torch.nn as nn
from typing import Optional, Tuple

from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
from transformers.models.qwen2.modeling_qwen2 import (
    Qwen2PreTrainedModel,
    Qwen2Attention,
    Qwen2MLP,
    Qwen2RMSNorm,
    Qwen2RotaryEmbedding,
    apply_rotary_pos_emb,
)
from transformers.integrations.sdpa_attention import sdpa_attention_forward
from torch.nn import functional as F


class MultiHeadRMSNorm(nn.Module):
    def __init__(self, dim, heads=1):
        super().__init__()
        self.scale = dim**0.5
        self.gamma = nn.Parameter(torch.ones(heads, 1, dim))

    def forward(self, x):
        return F.normalize(x, dim=-1) * self.gamma * self.scale


class Qwen2BidirectionalSdpaAttention(Qwen2Attention):
    """
    An SDPA-based attention that does NOT apply causal masking.
    Inherits from Qwen2Attention, but sets self.is_causal = False.
    """

    def __init__(self, config: Qwen2Config, layer_idx: Optional[int] = None):
        super().__init__(config, layer_idx)
        self.is_causal = False
        self.qk_norm = config.qk_norm
        if self.qk_norm:
            self.q_norm = MultiHeadRMSNorm(
                config.hidden_size // config.num_attention_heads,
                config.num_attention_heads,
            )
            self.k_norm = MultiHeadRMSNorm(
                config.hidden_size // config.num_attention_heads,
                config.num_key_value_heads,
            )

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
    ):
        input_shape = hidden_states.shape[:-1]
        hidden_shape = (*input_shape, -1, self.head_dim)

        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)

        if position_embeddings is not None:
            cos, sin = position_embeddings
            query_states, key_states = apply_rotary_pos_emb(
                query_states, key_states, cos, sin
            )

        if self.qk_norm:
            query_states = self.q_norm(query_states)
            key_states = self.k_norm(key_states)

        attn_output, attn_weights = sdpa_attention_forward(
            self,
            query_states,
            key_states,
            value_states,
            attention_mask=None,
            dropout=0.0 if not self.training else self.attention_dropout,
            scaling=self.scaling,
            is_causal=False,
        )

        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
        attn_output = self.o_proj(attn_output)
        return attn_output


class Qwen2EncoderLayer(nn.Module):
    def __init__(self, config: Qwen2Config, layer_idx: int):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.self_attn = Qwen2BidirectionalSdpaAttention(config, layer_idx)
        self.mlp = Qwen2MLP(config)

        self.input_layernorm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = Qwen2RMSNorm(
            config.hidden_size, eps=config.rms_norm_eps
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
    ):
                          
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)

        hidden_states = self.self_attn(
            hidden_states=hidden_states,
            position_embeddings=position_embeddings,
        )
        hidden_states = residual + hidden_states

                    
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states

        return hidden_states


class Qwen2Encoder(Qwen2PreTrainedModel):
    supports_gradient_checkpointing = True

    def __init__(self, config: Qwen2Config):
        super().__init__(config)
        self.layers = nn.ModuleList(
            [Qwen2EncoderLayer(config, i) for i in range(self.config.num_hidden_layers)]
        )
        if config.rope:
            self.rotary_emb = Qwen2RotaryEmbedding(config=config)
        else:
            self.rotary_emb = None
        if hasattr(config, "norm") and config.norm:
            self.norm = Qwen2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        else:
            self.norm = None
        self.gradient_checkpointing = True
        self.post_init()

    def _init_weights(self, module):
        std = self.config.initializer_range
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.bias is not None:
                module.bias.data.zero_()

    def forward(self, hidden_states):
        bsz, seq_len, _ = hidden_states.size()
        position_ids = torch.arange(seq_len, device=hidden_states.device).unsqueeze(0)

                                                            
        if self.rotary_emb is not None:
            position_embeddings = self.rotary_emb(hidden_states, position_ids)
        else:
            position_embeddings = None

        for layer in self.layers:
            if self.gradient_checkpointing and self.training:
                hidden_states = self._gradient_checkpointing_func(
                    layer.__call__,
                    hidden_states,
                    position_embeddings,
                )
            else:
                hidden_states = layer(
                    hidden_states,
                    position_embeddings=position_embeddings,
                )
        if self.norm:
            hidden_states = self.norm(hidden_states)
        return hidden_states


class RMSNorm(nn.Module):
    def __init__(self, dim, eps: float, elementwise_affine: bool = True):
        super().__init__()
        self.eps = eps
        if elementwise_affine:
            self.weight = nn.Parameter(torch.ones(dim))
        else:
            self.weight = None

    def forward(self, hidden_states):
        input_dtype = hidden_states.dtype
        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)

        if self.weight is not None:
                                                      
            if self.weight.dtype in [torch.float16, torch.bfloat16]:
                hidden_states = hidden_states.to(self.weight.dtype)
            hidden_states = hidden_states * self.weight
        else:
            hidden_states = hidden_states.to(input_dtype)

        return hidden_states


class QwenImageConnector(nn.Module):
    def __init__(
        self,
        input_dim: int=3584,
        output_dim: int=3584,
        num_layers: int=6,
        input_scale: float = 1.0,
    ):
        super().__init__()
        self.encoder = Qwen2Encoder(
            Qwen2Config(
                hidden_size=input_dim,
                intermediate_size=input_dim * 4,
                num_hidden_layers=num_layers,
                num_attention_heads=input_dim // 64,
                num_key_value_heads=input_dim // 64,
                initializer_range=0.014,
                use_cache=False,
                rope=True,
                qk_norm=True,
            ),
        )
        self.proj = nn.Sequential(
            nn.Linear(input_dim, output_dim),
            nn.GELU(approximate="tanh"),
            nn.Linear(output_dim, output_dim),
            RMSNorm(output_dim, eps=1e-5, elementwise_affine=True),
        )

                                                         
        with torch.no_grad():
            self.proj[-1].weight.fill_(input_scale)                                                     

    def forward(self, x):
        x = self.encoder(x)
        x = self.proj(x)
        return x