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

# 핵심 문제: 각 노이징 제거 단계가 분산을 도입함
def denoise_step(x_t, model, t):
    noise_pred = model(x_t, t)
    # 이 샘플링은 확률성을 도입함
    x_t_minus_1 = scheduler.step(noise_pred, t, x_t).prev_sample
    return x_t_minus_1  # 미묘한 변화가 프레임 간에 누적됨

class IdentityEncoder(nn.Module):
    def __init__(self, embed_dim=768):
        super().__init__()
        self.face_encoder = FaceRecognitionBackbone()  # 사전학습된 얼굴 모델
        self.projection = nn.Linear(512, embed_dim)
        self.identity_bank = nn.Parameter(torch.randn(32, embed_dim))
 
    def encode_identity(self, reference_frame):
        # 참조 프레임에서 정체성 특징 추출
        face_features = self.face_encoder(reference_frame)
        identity_embed = self.projection(face_features)
 
        # 학습된 정체성 토큰과 상호 주의
        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):
        # 프레임 내 표준 공간 주의
        x = self.spatial_attn(x, x, x)[0] + x
 
        # 프레임 간 시간적 주의
        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)
 
        # 얼굴 영역을 사용한 정체성 특화 주의
        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):
        # 참조 이미지에서 정체성 인코딩
        identity_embeds = []
        for ref in reference_images:
            identity_embeds.append(self.identity_encoder(ref))
 
        # 견고한 정체성을 위해 여러 참조 통합
        identity_tokens = torch.stack(identity_embeds).mean(dim=0)
 
        # 정체성 조건화를 통한 생성
        video = self.model.generate(
            prompt=prompt,
            num_frames=num_frames,
            cross_attention_kwargs={
                "identity_tokens": identity_tokens,
                "identity_strength": 0.8  # 일관성과 창의성의 균형
            }
        )
        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):
    # 참조에 대한 프레임별 정체성 매칭
    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_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)
            # 혼동을 방지하기 위한 분리자 추가
            char_tokens = torch.cat([char_tokens, self.character_separator])
            all_tokens.append(char_tokens)
        return torch.cat(all_tokens, dim=0)