AI 비디오의 캐릭터 일관성: 모델이 얼굴을 기억하는 방식

# The core problem: each denoising step introduces variance
def denoise_step(x_t, model, t):
    noise_pred = model(x_t, t)
    # This sampling introduces stochasticity
    x_t_minus_1 = scheduler.step(noise_pred, t, x_t).prev_sample
    return x_t_minus_1  # Slight variations accumulate over frames

class IdentityEncoder(nn.Module):
    def __init__(self, embed_dim=768):
        super().__init__()
        self.face_encoder = FaceRecognitionBackbone()  # Pre-trained face model
        self.projection = nn.Linear(512, embed_dim)
        self.identity_bank = nn.Parameter(torch.randn(32, embed_dim))
 
    def encode_identity(self, reference_frame):
        # Extract identity features from reference
        face_features = self.face_encoder(reference_frame)
        identity_embed = self.projection(face_features)
 
        # Cross-attend with learned identity tokens
        identity_tokens = self.cross_attention(
            query=self.identity_bank,
            key=identity_embed,
            value=identity_embed
        )
        return identity_tokens

class IdentityAwareAttention(nn.Module):
    def __init__(self, dim, num_heads=8):
        super().__init__()
        self.spatial_attn = nn.MultiheadAttention(dim, num_heads)
        self.temporal_attn = nn.MultiheadAttention(dim, num_heads)
        self.identity_attn = nn.MultiheadAttention(dim, num_heads)
 
    def forward(self, x, identity_tokens, face_masks):
        # Standard spatial attention within frames
        x = self.spatial_attn(x, x, x)[0] + x
 
        # Temporal attention across frames
        x = rearrange(x, '(b t) n d -> (b n) t d', t=num_frames)
        x = self.temporal_attn(x, x, x)[0] + x
        x = rearrange(x, '(b n) t d -> (b t) n d', n=num_patches)
 
        # Identity-specific attention using face regions
        face_tokens = x * face_masks.unsqueeze(-1)
        x = self.identity_attn(
            query=x,
            key=identity_tokens,
            value=identity_tokens
        )[0] + x
 
        return x

class ReferenceConditionedGenerator:
    def __init__(self, base_model, identity_encoder):
        self.model = base_model
        self.identity_encoder = identity_encoder
 
    def generate(self, prompt, reference_images, num_frames=120):
        # Encode identity from reference images
        identity_embeds = []
        for ref in reference_images:
            identity_embeds.append(self.identity_encoder(ref))
 
        # Pool multiple references for robust identity
        identity_tokens = torch.stack(identity_embeds).mean(dim=0)
 
        # Generate with identity conditioning
        video = self.model.generate(
            prompt=prompt,
            num_frames=num_frames,
            cross_attention_kwargs={
                "identity_tokens": identity_tokens,
                "identity_strength": 0.8  # Balances consistency vs creativity
            }
        )
        return video

L_identity = ||f(G(z, c)) - f(x_ref)||² + λ_temporal * Σ_t ||f(v_t) - f(v_{t+1})||²

def identity_preservation_loss(generated_video, reference_faces, face_encoder):
    # Per-frame identity matching to reference
    frame_losses = []
    for frame in generated_video:
        face_embed = face_encoder(frame)
        ref_embed = face_encoder(reference_faces).mean(dim=0)
        frame_losses.append(F.mse_loss(face_embed, ref_embed))
 
    reference_loss = torch.stack(frame_losses).mean()
 
    # Temporal consistency between adjacent frames
    temporal_losses = []
    for i in range(len(generated_video) - 1):
        curr_embed = face_encoder(generated_video[i])
        next_embed = face_encoder(generated_video[i + 1])
        temporal_losses.append(F.mse_loss(curr_embed, next_embed))
 
    temporal_loss = torch.stack(temporal_losses).mean()
 
    return reference_loss + 0.5 * temporal_loss

class MultiCharacterIdentityBank:
    def __init__(self, max_characters=8, embed_dim=768):
        self.banks = nn.ModuleList([
            IdentityBank(embed_dim) for _ in range(max_characters)
        ])
        self.character_separator = nn.Parameter(torch.randn(1, embed_dim))
 
    def encode_multiple(self, character_references):
        all_tokens = []
        for idx, refs in enumerate(character_references):
            char_tokens = self.banks[idx].encode(refs)
            # Add separator to prevent conflation
            char_tokens = torch.cat([char_tokens, self.character_separator])
            all_tokens.append(char_tokens)
        return torch.cat(all_tokens, dim=0)