Upload RotEx inference code, demos, documentation, and pretrained checkpoints

40440d6 verified 24 days ago

11.4 kB

"""
Advanced attention mechanisms for better semantic and style alignment.
Includes Cross-Attention, AdaIN, and Depth-Controlled Attention.
Author: https://github.com/datnguyen-tien204
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
from typing import Optional


class AdaptiveInstanceNorm2d(nn.Module):
    """
    Adaptive Instance Normalization (AdaIN) for style transfer.
    Aligns mean and variance of content features to match style features.
    """
    def __init__(self, num_features: int, style_dim: int):
        """
        Args:
            num_features: Number of feature channels
            style_dim: Dimension of style vector
        """
        super().__init__()
        self.num_features = num_features
        
        # Style modulation parameters
        self.style_scale = nn.Linear(style_dim, num_features)
        self.style_shift = nn.Linear(style_dim, num_features)
        
        # Initialize to identity transformation
        nn.init.constant_(self.style_scale.weight, 0)
        nn.init.constant_(self.style_scale.bias, 1)
        nn.init.constant_(self.style_shift.weight, 0)
        nn.init.constant_(self.style_shift.bias, 0)
    
    def forward(self, x: torch.Tensor, style: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: Input features [B, C, H, W]
            style: Style vector [B, style_dim]
        Returns:
            normalized: Style-adapted features [B, C, H, W]
        """
        # Instance normalization
        mean = x.mean(dim=[2, 3], keepdim=True)
        std = x.std(dim=[2, 3], keepdim=True) + 1e-5
        normalized = (x - mean) / std
        
        # Style modulation
        scale = self.style_scale(style).unsqueeze(-1).unsqueeze(-1)
        shift = self.style_shift(style).unsqueeze(-1).unsqueeze(-1)
        
        return normalized * scale + shift


class AdaIN1d(nn.Module):
    """
    AdaIN for 1D features (sequence data).
    """
    def __init__(self, num_features: int, style_dim: int):
        super().__init__()
        self.num_features = num_features

        self.norm = nn.LayerNorm(num_features, elementwise_affine=False)
        self.style_scale = nn.Linear(style_dim, num_features)
        self.style_shift = nn.Linear(style_dim, num_features)

        nn.init.constant_(self.style_scale.weight, 0)
        nn.init.constant_(self.style_scale.bias, 1)
        nn.init.constant_(self.style_shift.weight, 0)
        nn.init.constant_(self.style_shift.bias, 0)

    def forward(self, x: torch.Tensor, style: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: Input features [B, L, C]
            style: Style vector [B, style_dim] OR [B, L_s, style_dim]
        Returns:
            normalized: Style-adapted features [B, L, C]
        """
        if style.dim() == 3:
            style = style.mean(dim=1)

        normalized = self.norm(x)
        scale = self.style_scale(style).unsqueeze(1)
        shift = self.style_shift(style).unsqueeze(1)

        return normalized * scale + shift



class CrossAttention(nn.Module):
    """
    Cross-Attention mechanism for semantic alignment between content and style.
    """
    def __init__(self, query_dim: int, key_dim: int, num_heads: int = 8, 
                 dropout: float = 0.0):
        """
        Args:
            query_dim: Dimension of query features
            key_dim: Dimension of key/value features
            num_heads: Number of attention heads
            dropout: Dropout probability
        """
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = query_dim // num_heads
        assert query_dim % num_heads == 0, "query_dim must be divisible by num_heads"
        
        self.query = nn.Linear(query_dim, query_dim)
        self.key = nn.Linear(key_dim, query_dim)
        self.value = nn.Linear(key_dim, query_dim)
        
        self.proj = nn.Linear(query_dim, query_dim)
        self.dropout = nn.Dropout(dropout)
    
    def forward(self, query: torch.Tensor, key: torch.Tensor, 
                value: Optional[torch.Tensor] = None,
                mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        """
        Args:
            query: Query features [B, L_q, query_dim]
            key: Key features [B, L_k, key_dim]
            value: Value features [B, L_k, key_dim] (if None, uses key)
            mask: Attention mask [B, L_q, L_k] (optional)
        Returns:
            output: Attended features [B, L_q, query_dim]
        """
        if value is None:
            value = key
        
        B, L_q, _ = query.shape
        L_k = key.shape[1]
        
        # Linear projections
        q = self.query(query)
        k = self.key(key)
        v = self.value(value)
        
        # Reshape for multi-head attention
        q = rearrange(q, 'b l (h d) -> b h l d', h=self.num_heads)
        k = rearrange(k, 'b l (h d) -> b h l d', h=self.num_heads)
        v = rearrange(v, 'b l (h d) -> b h l d', h=self.num_heads)
        
        # Scaled dot-product attention
        scale = self.head_dim ** -0.5
        attn = torch.matmul(q, k.transpose(-2, -1)) * scale
        
        if mask is not None:
            attn = attn.masked_fill(mask.unsqueeze(1) == 0, float('-inf'))
        
        attn = F.softmax(attn, dim=-1)
        attn = self.dropout(attn)
        
        # Apply attention to values
        out = torch.matmul(attn, v)
        out = rearrange(out, 'b h l d -> b l (h d)')
        
        # Output projection
        out = self.proj(out)
        
        return out


class DepthControlledAttention(nn.Module):
    """
    Depth-controlled attention mechanism for structure-aware feature learning.
    Uses depth information to modulate attention patterns.
    """
    def __init__(self, feature_dim: int, depth_dim: int = 64, num_heads: int = 8):
        """
        Args:
            feature_dim: Dimension of input features
            depth_dim: Dimension of depth features
            num_heads: Number of attention heads
        """
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = feature_dim // num_heads
        
        # Process depth information
        self.depth_encoder = nn.Sequential(
            nn.Conv2d(1, depth_dim, kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(depth_dim, depth_dim, kernel_size=3, padding=1)
        )
        
        # Depth-modulated attention
        self.qkv = nn.Linear(feature_dim, feature_dim * 3)
        self.depth_gate = nn.Linear(depth_dim, feature_dim)
        self.proj = nn.Linear(feature_dim, feature_dim)
        
    def forward(self, x: torch.Tensor, depth: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: Input features [B, L, C]
            depth: Depth map [B, 1, H, W]
        Returns:
            output: Depth-modulated features [B, L, C]
        """
        B, L, C = x.shape
        
        # Encode depth
        depth_feat = self.depth_encoder(depth)  # [B, depth_dim, H, W]
        H, W = depth_feat.shape[-2:]
        depth_feat = rearrange(depth_feat, 'b c h w -> b (h w) c')
        
        # Generate Q, K, V
        qkv = self.qkv(x)
        q, k, v = rearrange(qkv, 'b l (three h d) -> three b h l d', 
                           three=3, h=self.num_heads)
        
        # Compute attention
        scale = self.head_dim ** -0.5
        attn = torch.matmul(q, k.transpose(-2, -1)) * scale
        
        # Modulate attention with depth
        if depth_feat.shape[1] == L:
            depth_gate = self.depth_gate(depth_feat)
            depth_gate = rearrange(depth_gate, 'b l (h d) -> b h l d', h=self.num_heads)
            depth_mask = torch.sigmoid(depth_gate.mean(dim=-1, keepdim=True))
            attn = attn * depth_mask
        
        attn = F.softmax(attn, dim=-1)
        
        # Apply attention
        out = torch.matmul(attn, v)
        out = rearrange(out, 'b h l d -> b l (h d)')
        out = self.proj(out)
        
        return out


class StyleAttentionBlock(nn.Module):
    """
    Combined attention block with AdaIN and Cross-Attention.
    """
    def __init__(self, feature_dim: int, style_dim: int, num_heads: int = 8):
        super().__init__()
        
        # Layer normalization
        self.norm1 = nn.LayerNorm(feature_dim)
        self.norm2 = nn.LayerNorm(feature_dim)
        
        # Cross-attention for style conditioning
        self.cross_attn = CrossAttention(feature_dim, style_dim, num_heads)
        
        # AdaIN for style transfer
        self.adain = AdaIN1d(feature_dim, style_dim)
        
        # Feed-forward network
        self.ffn = nn.Sequential(
            nn.Linear(feature_dim, feature_dim * 4),
            nn.GELU(),
            nn.Linear(feature_dim * 4, feature_dim)
        )
    
    def forward(self, x: torch.Tensor, style: torch.Tensor, 
                style_features: Optional[torch.Tensor] = None) -> torch.Tensor:
        """
        Args:
            x: Input features [B, L, C]
            style: Style vector [B, style_dim]
            style_features: Optional style feature sequence [B, L_s, style_dim]
        Returns:
            output: Style-conditioned features [B, L, C]
        """
        # Cross-attention with style features
        if style_features is not None:
            attn_out = self.cross_attn(self.norm1(x), style_features)
            x = x + attn_out
        
        # AdaIN style transfer
        x = self.adain(x, style)
        
        # Feed-forward
        x = x + self.ffn(self.norm2(x))
        
        return x


class MultiScaleAttention(nn.Module):
    """
    Multi-scale attention for capturing features at different resolutions.
    """
    def __init__(self, feature_dim: int, num_scales: int = 3, num_heads: int = 8):
        super().__init__()
        self.num_scales = num_scales
        
        self.attentions = nn.ModuleList([
            nn.MultiheadAttention(feature_dim, num_heads, batch_first=True)
            for _ in range(num_scales)
        ])
        
        self.fusion = nn.Linear(feature_dim * num_scales, feature_dim)
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: Input features [B, L, C]
        Returns:
            output: Multi-scale attended features [B, L, C]
        """
        B, L, C = x.shape
        
        outputs = []
        for i, attn in enumerate(self.attentions):
            # Downsample for different scales
            if i > 0:
                scale = 2 ** i
                # Simple downsampling by reshaping
                L_down = L // scale
                if L_down > 0:
                    x_down = x[:, ::scale, :]
                else:
                    x_down = x
                out, _ = attn(x_down, x_down, x_down)
                # Upsample back
                if out.shape[1] != L:
                    out = F.interpolate(out.transpose(1, 2), size=L, mode='linear', 
                                      align_corners=False).transpose(1, 2)
            else:
                out, _ = attn(x, x, x)
            
            outputs.append(out)
        
        # Concatenate and fuse
        multi_scale = torch.cat(outputs, dim=-1)
        fused = self.fusion(multi_scale)
        
        return fused

	"""
	Advanced attention mechanisms for better semantic and style alignment.
	Includes Cross-Attention, AdaIN, and Depth-Controlled Attention.
	Author: https://github.com/datnguyen-tien204
	"""
	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from einops import rearrange
	from typing import Optional


	class AdaptiveInstanceNorm2d(nn.Module):
	"""
	Adaptive Instance Normalization (AdaIN) for style transfer.
	Aligns mean and variance of content features to match style features.
	"""
	def __init__(self, num_features: int, style_dim: int):
	"""
	Args:
	num_features: Number of feature channels
	style_dim: Dimension of style vector
	"""
	super().__init__()
	self.num_features = num_features

	# Style modulation parameters
	self.style_scale = nn.Linear(style_dim, num_features)
	self.style_shift = nn.Linear(style_dim, num_features)

	# Initialize to identity transformation
	nn.init.constant_(self.style_scale.weight, 0)
	nn.init.constant_(self.style_scale.bias, 1)
	nn.init.constant_(self.style_shift.weight, 0)
	nn.init.constant_(self.style_shift.bias, 0)

	def forward(self, x: torch.Tensor, style: torch.Tensor) -> torch.Tensor:
	"""
	Args:
	x: Input features [B, C, H, W]
	style: Style vector [B, style_dim]
	Returns:
	normalized: Style-adapted features [B, C, H, W]
	"""
	# Instance normalization
	mean = x.mean(dim=[2, 3], keepdim=True)
	std = x.std(dim=[2, 3], keepdim=True) + 1e-5
	normalized = (x - mean) / std

	# Style modulation
	scale = self.style_scale(style).unsqueeze(-1).unsqueeze(-1)
	shift = self.style_shift(style).unsqueeze(-1).unsqueeze(-1)

	return normalized * scale + shift


	class AdaIN1d(nn.Module):
	"""
	AdaIN for 1D features (sequence data).
	"""
	def __init__(self, num_features: int, style_dim: int):
	super().__init__()
	self.num_features = num_features

	self.norm = nn.LayerNorm(num_features, elementwise_affine=False)
	self.style_scale = nn.Linear(style_dim, num_features)
	self.style_shift = nn.Linear(style_dim, num_features)

	nn.init.constant_(self.style_scale.weight, 0)
	nn.init.constant_(self.style_scale.bias, 1)
	nn.init.constant_(self.style_shift.weight, 0)
	nn.init.constant_(self.style_shift.bias, 0)

	def forward(self, x: torch.Tensor, style: torch.Tensor) -> torch.Tensor:
	"""
	Args:
	x: Input features [B, L, C]
	style: Style vector [B, style_dim] OR [B, L_s, style_dim]
	Returns:
	normalized: Style-adapted features [B, L, C]
	"""
	if style.dim() == 3:
	style = style.mean(dim=1)

	normalized = self.norm(x)
	scale = self.style_scale(style).unsqueeze(1)
	shift = self.style_shift(style).unsqueeze(1)

	return normalized * scale + shift



	class CrossAttention(nn.Module):
	"""
	Cross-Attention mechanism for semantic alignment between content and style.
	"""
	def __init__(self, query_dim: int, key_dim: int, num_heads: int = 8,
	dropout: float = 0.0):
	"""
	Args:
	query_dim: Dimension of query features
	key_dim: Dimension of key/value features
	num_heads: Number of attention heads
	dropout: Dropout probability
	"""
	super().__init__()
	self.num_heads = num_heads
	self.head_dim = query_dim // num_heads
	assert query_dim % num_heads == 0, "query_dim must be divisible by num_heads"

	self.query = nn.Linear(query_dim, query_dim)
	self.key = nn.Linear(key_dim, query_dim)
	self.value = nn.Linear(key_dim, query_dim)

	self.proj = nn.Linear(query_dim, query_dim)
	self.dropout = nn.Dropout(dropout)

	def forward(self, query: torch.Tensor, key: torch.Tensor,
	value: Optional[torch.Tensor] = None,
	mask: Optional[torch.Tensor] = None) -> torch.Tensor:
	"""
	Args:
	query: Query features [B, L_q, query_dim]
	key: Key features [B, L_k, key_dim]
	value: Value features [B, L_k, key_dim] (if None, uses key)
	mask: Attention mask [B, L_q, L_k] (optional)
	Returns:
	output: Attended features [B, L_q, query_dim]
	"""
	if value is None:
	value = key

	B, L_q, _ = query.shape
	L_k = key.shape[1]

	# Linear projections
	q = self.query(query)
	k = self.key(key)
	v = self.value(value)

	# Reshape for multi-head attention
	q = rearrange(q, 'b l (h d) -> b h l d', h=self.num_heads)
	k = rearrange(k, 'b l (h d) -> b h l d', h=self.num_heads)
	v = rearrange(v, 'b l (h d) -> b h l d', h=self.num_heads)

	# Scaled dot-product attention
	scale = self.head_dim ** -0.5
	attn = torch.matmul(q, k.transpose(-2, -1)) * scale

	if mask is not None:
	attn = attn.masked_fill(mask.unsqueeze(1) == 0, float('-inf'))

	attn = F.softmax(attn, dim=-1)
	attn = self.dropout(attn)

	# Apply attention to values
	out = torch.matmul(attn, v)
	out = rearrange(out, 'b h l d -> b l (h d)')

	# Output projection
	out = self.proj(out)

	return out


	class DepthControlledAttention(nn.Module):
	"""
	Depth-controlled attention mechanism for structure-aware feature learning.
	Uses depth information to modulate attention patterns.
	"""
	def __init__(self, feature_dim: int, depth_dim: int = 64, num_heads: int = 8):
	"""
	Args:
	feature_dim: Dimension of input features
	depth_dim: Dimension of depth features
	num_heads: Number of attention heads
	"""
	super().__init__()
	self.num_heads = num_heads
	self.head_dim = feature_dim // num_heads

	# Process depth information
	self.depth_encoder = nn.Sequential(
	nn.Conv2d(1, depth_dim, kernel_size=3, padding=1),
	nn.ReLU(inplace=True),
	nn.Conv2d(depth_dim, depth_dim, kernel_size=3, padding=1)
	)

	# Depth-modulated attention
	self.qkv = nn.Linear(feature_dim, feature_dim * 3)
	self.depth_gate = nn.Linear(depth_dim, feature_dim)
	self.proj = nn.Linear(feature_dim, feature_dim)

	def forward(self, x: torch.Tensor, depth: torch.Tensor) -> torch.Tensor:
	"""
	Args:
	x: Input features [B, L, C]
	depth: Depth map [B, 1, H, W]
	Returns:
	output: Depth-modulated features [B, L, C]
	"""
	B, L, C = x.shape

	# Encode depth
	depth_feat = self.depth_encoder(depth) # [B, depth_dim, H, W]
	H, W = depth_feat.shape[-2:]
	depth_feat = rearrange(depth_feat, 'b c h w -> b (h w) c')

	# Generate Q, K, V
	qkv = self.qkv(x)
	q, k, v = rearrange(qkv, 'b l (three h d) -> three b h l d',
	three=3, h=self.num_heads)

	# Compute attention
	scale = self.head_dim ** -0.5
	attn = torch.matmul(q, k.transpose(-2, -1)) * scale

	# Modulate attention with depth
	if depth_feat.shape[1] == L:
	depth_gate = self.depth_gate(depth_feat)
	depth_gate = rearrange(depth_gate, 'b l (h d) -> b h l d', h=self.num_heads)
	depth_mask = torch.sigmoid(depth_gate.mean(dim=-1, keepdim=True))
	attn = attn * depth_mask

	attn = F.softmax(attn, dim=-1)

	# Apply attention
	out = torch.matmul(attn, v)
	out = rearrange(out, 'b h l d -> b l (h d)')
	out = self.proj(out)

	return out


	class StyleAttentionBlock(nn.Module):
	"""
	Combined attention block with AdaIN and Cross-Attention.
	"""
	def __init__(self, feature_dim: int, style_dim: int, num_heads: int = 8):
	super().__init__()

	# Layer normalization
	self.norm1 = nn.LayerNorm(feature_dim)
	self.norm2 = nn.LayerNorm(feature_dim)

	# Cross-attention for style conditioning
	self.cross_attn = CrossAttention(feature_dim, style_dim, num_heads)

	# AdaIN for style transfer
	self.adain = AdaIN1d(feature_dim, style_dim)

	# Feed-forward network
	self.ffn = nn.Sequential(
	nn.Linear(feature_dim, feature_dim * 4),
	nn.GELU(),
	nn.Linear(feature_dim * 4, feature_dim)
	)

	def forward(self, x: torch.Tensor, style: torch.Tensor,
	style_features: Optional[torch.Tensor] = None) -> torch.Tensor:
	"""
	Args:
	x: Input features [B, L, C]
	style: Style vector [B, style_dim]
	style_features: Optional style feature sequence [B, L_s, style_dim]
	Returns:
	output: Style-conditioned features [B, L, C]
	"""
	# Cross-attention with style features
	if style_features is not None:
	attn_out = self.cross_attn(self.norm1(x), style_features)
	x = x + attn_out

	# AdaIN style transfer
	x = self.adain(x, style)

	# Feed-forward
	x = x + self.ffn(self.norm2(x))

	return x


	class MultiScaleAttention(nn.Module):
	"""
	Multi-scale attention for capturing features at different resolutions.
	"""
	def __init__(self, feature_dim: int, num_scales: int = 3, num_heads: int = 8):
	super().__init__()
	self.num_scales = num_scales

	self.attentions = nn.ModuleList([
	nn.MultiheadAttention(feature_dim, num_heads, batch_first=True)
	for _ in range(num_scales)
	])

	self.fusion = nn.Linear(feature_dim * num_scales, feature_dim)

	def forward(self, x: torch.Tensor) -> torch.Tensor:
	"""
	Args:
	x: Input features [B, L, C]
	Returns:
	output: Multi-scale attended features [B, L, C]
	"""
	B, L, C = x.shape

	outputs = []
	for i, attn in enumerate(self.attentions):
	# Downsample for different scales
	if i > 0:
	scale = 2 ** i
	# Simple downsampling by reshaping
	L_down = L // scale
	if L_down > 0:
	x_down = x[:, ::scale, :]
	else:
	x_down = x
	out, _ = attn(x_down, x_down, x_down)
	# Upsample back
	if out.shape[1] != L:
	out = F.interpolate(out.transpose(1, 2), size=L, mode='linear',
	align_corners=False).transpose(1, 2)
	else:
	out, _ = attn(x, x, x)

	outputs.append(out)

	# Concatenate and fuse
	multi_scale = torch.cat(outputs, dim=-1)
	fused = self.fusion(multi_scale)

	return fused