Delete ip_attention_processor_enhanced.py

ip_attention_processor_enhanced.py

@@ -1,321 +0,0 @@
-"""
-Enhanced IP-Adapter Attention Processor - Optimized for Maximum Face Preservation
-===================================================================================
-
-Improvements over base version:
-1. Adaptive scaling based on attention scores
-2. Multi-scale face feature integration
-3. Learnable blending weights per layer
-4. Face confidence-aware modulation
-5. Better gradient flow with skip connections
-
-Expected improvement: +2-3% additional face similarity
-
-Author: Pixagram Team
-License: MIT
-"""
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from typing import Optional, Dict
-from diffusers.models.attention_processor import AttnProcessor2_0
-
-
-class EnhancedIPAttnProcessor2_0(nn.Module):
-    """
-    Enhanced IP-Adapter attention with adaptive scaling and optimizations.
-    
-    Key improvements over base:
-    - Adaptive scale based on attention statistics
-    - Learnable per-layer blending weights
-    - Better numerical stability
-    - Optional face confidence modulation
-    
-    Args:
-        hidden_size: Attention layer hidden dimension
-        cross_attention_dim: Encoder hidden states dimension
-        scale: Base blending weight for face features
-        num_tokens: Number of face embedding tokens
-        adaptive_scale: Enable adaptive scaling (recommended)
-        learnable_scale: Make scale learnable per layer
-    """
-    
-    def __init__(
-        self,
-        hidden_size: int,
-        cross_attention_dim: Optional[int] = None,
-        scale: float = 1.0,
-        num_tokens: int = 4,
-        adaptive_scale: bool = True,
-        learnable_scale: bool = True
-    ):
-        super().__init__()
-        
-        if not hasattr(F, "scaled_dot_product_attention"):
-            raise ImportError("Requires PyTorch 2.0+")
-        
-        self.hidden_size = hidden_size
-        self.cross_attention_dim = cross_attention_dim or hidden_size
-        self.base_scale = scale
-        self.num_tokens = num_tokens
-        self.adaptive_scale = adaptive_scale
-        
-        # Dedicated K/V projections for face features
-        self.to_k_ip = nn.Linear(self.cross_attention_dim, hidden_size, bias=False)
-        self.to_v_ip = nn.Linear(self.cross_attention_dim, hidden_size, bias=False)
-        
-        # Learnable scale parameter (per layer)
-        if learnable_scale:
-            self.scale_param = nn.Parameter(torch.tensor(scale))
-        else:
-            self.register_buffer('scale_param', torch.tensor(scale))
-        
-        # Adaptive scaling module
-        if adaptive_scale:
-            self.adaptive_gate = nn.Sequential(
-                nn.Linear(hidden_size, hidden_size // 4),
-                nn.ReLU(),
-                nn.Linear(hidden_size // 4, 1),
-                nn.Sigmoid()
-            )
-        
-        # Better initialization
-        self._init_weights()
-    
-    def _init_weights(self):
-        """Xavier initialization for stable training."""
-        nn.init.xavier_uniform_(self.to_k_ip.weight)
-        nn.init.xavier_uniform_(self.to_v_ip.weight)
-        
-        if self.adaptive_scale:
-            for module in self.adaptive_gate:
-                if isinstance(module, nn.Linear):
-                    nn.init.xavier_uniform_(module.weight)
-                    if module.bias is not None:
-                        nn.init.zeros_(module.bias)
-    
-    def compute_adaptive_scale(
-        self, 
-        query: torch.Tensor, 
-        ip_key: torch.Tensor,
-        base_scale: float
-    ) -> torch.Tensor:
-        """
-        Compute adaptive scale based on query-key similarity.
-        Higher similarity = stronger face preservation.
-        """
-        # Compute mean query features
-        query_mean = query.mean(dim=(1, 2))  # [batch, head_dim * heads]
-        
-        # Pass through gating network
-        gate = self.adaptive_gate(query_mean)  # [batch, 1]
-        
-        # Modulate base scale
-        adaptive_scale = base_scale * (0.5 + gate)  # Range: [0.5*base, 1.5*base]
-        
-        return adaptive_scale.view(-1, 1, 1)  # [batch, 1, 1] for broadcasting
-    
-    def forward(
-        self,
-        attn,
-        hidden_states: torch.FloatTensor,
-        encoder_hidden_states: Optional[torch.FloatTensor] = None,
-        attention_mask: Optional[torch.FloatTensor] = None,
-        temb: Optional[torch.FloatTensor] = None,
-    ) -> torch.FloatTensor:
-        """Forward pass with adaptive face preservation."""
-        residual = hidden_states
-        
-        if attn.spatial_norm is not None:
-            hidden_states = attn.spatial_norm(hidden_states, temb)
-        
-        input_ndim = hidden_states.ndim
-        if input_ndim == 4:
-            batch_size, channel, height, width = hidden_states.shape
-            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
-        
-        batch_size, sequence_length, _ = (
-            hidden_states.shape if encoder_hidden_states is None 
-            else encoder_hidden_states.shape
-        )
-        
-        if attention_mask is not None:
-            attention_mask = attn.prepare_attention_mask(
-                attention_mask, sequence_length, batch_size
-            )
-            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
-        
-        if attn.group_norm is not None:
-            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
-        
-        query = attn.to_q(hidden_states)
-        
-        # Split text and face embeddings
-        if encoder_hidden_states is None:
-            encoder_hidden_states = hidden_states
-            ip_hidden_states = None
-        else:
-            end_pos = encoder_hidden_states.shape[1] - self.num_tokens
-            encoder_hidden_states, ip_hidden_states = (
-                encoder_hidden_states[:, :end_pos, :],
-                encoder_hidden_states[:, end_pos:, :]
-            )
-            if attn.norm_cross:
-                encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
-        
-        # Text attention
-        key = attn.to_k(encoder_hidden_states)
-        value = attn.to_v(encoder_hidden_states)
-        
-        inner_dim = key.shape[-1]
-        head_dim = inner_dim // attn.heads
-        
-        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-        
-        hidden_states = F.scaled_dot_product_attention(
-            query, key, value, 
-            attn_mask=attention_mask, 
-            dropout_p=0.0, 
-            is_causal=False
-        )
-        
-        hidden_states = hidden_states.transpose(1, 2).reshape(
-            batch_size, -1, attn.heads * head_dim
-        )
-        hidden_states = hidden_states.to(query.dtype)
-        
-        # Face attention with enhancements
-        if ip_hidden_states is not None:
-            # Dedicated K/V projections
-            ip_key = self.to_k_ip(ip_hidden_states)
-            ip_value = self.to_v_ip(ip_hidden_states)
-            
-            ip_key = ip_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-            ip_value = ip_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-            
-            # Face attention
-            ip_hidden_states = F.scaled_dot_product_attention(
-                query, ip_key, ip_value,
-                attn_mask=None,
-                dropout_p=0.0,
-                is_causal=False
-            )
-            
-            ip_hidden_states = ip_hidden_states.transpose(1, 2).reshape(
-                batch_size, -1, attn.heads * head_dim
-            )
-            ip_hidden_states = ip_hidden_states.to(query.dtype)
-            
-            # Compute effective scale
-            if self.adaptive_scale and self.training == False:  # Only in inference
-                try:
-                    adaptive_scale = self.compute_adaptive_scale(query, ip_key, self.scale_param.item())
-                    effective_scale = adaptive_scale
-                except:
-                    effective_scale = self.scale_param
-            else:
-                effective_scale = self.scale_param
-            
-            # Blend with adaptive scale
-            hidden_states = hidden_states + effective_scale * ip_hidden_states
-        
-        # Output projection
-        hidden_states = attn.to_out[0](hidden_states)
-        hidden_states = attn.to_out[1](hidden_states)
-        
-        if input_ndim == 4:
-            hidden_states = hidden_states.transpose(-1, -2).reshape(
-                batch_size, channel, height, width
-            )
-        
-        if attn.residual_connection:
-            hidden_states = hidden_states + residual
-        
-        hidden_states = hidden_states / attn.rescale_output_factor
-        
-        return hidden_states
-
-
-def setup_enhanced_ip_adapter_attention(
-    pipe,
-    ip_adapter_scale: float = 1.0,
-    num_tokens: int = 4,
-    device: str = "cuda",
-    dtype = torch.float16,
-    adaptive_scale: bool = True,
-    learnable_scale: bool = True
-) -> Dict[str, nn.Module]:
-    """
-    Setup enhanced IP-Adapter attention processors.
-    
-    Args:
-        pipe: Diffusers pipeline
-        ip_adapter_scale: Base face embedding strength
-        num_tokens: Number of face tokens
-        device: Device
-        dtype: Data type
-        adaptive_scale: Enable adaptive scaling
-        learnable_scale: Make scales learnable
-    
-    Returns:
-        Dict of attention processors
-    """
-    attn_procs = {}
-    
-    for name in pipe.unet.attn_processors.keys():
-        cross_attention_dim = None if name.endswith("attn1.processor") else pipe.unet.config.cross_attention_dim
-        
-        if name.startswith("mid_block"):
-            hidden_size = pipe.unet.config.block_out_channels[-1]
-        elif name.startswith("up_blocks"):
-            block_id = int(name[len("up_blocks.")])
-            hidden_size = list(reversed(pipe.unet.config.block_out_channels))[block_id]
-        elif name.startswith("down_blocks"):
-            block_id = int(name[len("down_blocks.")])
-            hidden_size = pipe.unet.config.block_out_channels[block_id]
-        else:
-            hidden_size = pipe.unet.config.block_out_channels[-1]
-        
-        if cross_attention_dim is None:
-            attn_procs[name] = AttnProcessor2_0()
-        else:
-            attn_procs[name] = EnhancedIPAttnProcessor2_0(
-                hidden_size=hidden_size,
-                cross_attention_dim=cross_attention_dim,
-                scale=ip_adapter_scale,
-                num_tokens=num_tokens,
-                adaptive_scale=adaptive_scale,
-                learnable_scale=learnable_scale
-            ).to(device, dtype=dtype)
-    
-    print(f"[OK] Enhanced attention processors created")
-    print(f"  - Total processors: {len(attn_procs)}")
-    print(f"  - Adaptive scaling: {adaptive_scale}")
-    print(f"  - Learnable scales: {learnable_scale}")
-    
-    return attn_procs
-
-
-# Backward compatibility
-IPAttnProcessor2_0 = EnhancedIPAttnProcessor2_0
-
-
-if __name__ == "__main__":
-    print("Testing Enhanced IP-Adapter Processor...")
-    
-    processor = EnhancedIPAttnProcessor2_0(
-        hidden_size=1280,
-        cross_attention_dim=2048,
-        scale=0.8,
-        num_tokens=4,
-        adaptive_scale=True,
-        learnable_scale=True
-    )
-    
-    print(f"\n[OK] Processor created successfully")
-    print(f"Parameters: {sum(p.numel() for p in processor.parameters()):,}")
-    print(f"Has adaptive scaling: {processor.adaptive_scale}")
-    print(f"Has learnable scale: {isinstance(processor.scale_param, nn.Parameter)}")