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)