add HellMeme V2

2024-12-08 22:13:35 +08:00
parent 24dd59172e
commit c9ca4a689e
16 changed files with 229 additions and 5415 deletions
@@ -9,6 +9,6 @@
 """

 from .hm_denoising_motion import HMDenoisingMotion
-from .hm_control import HMControlNet, HMControlNet2
+from .hm_control import HMControlNet, HMControlNet2, HMV2ControlNet, HMV2ControlNet2
 from .hm_blocks import HMReferenceAdapter
 from .hm_denoising_3d import HMDenoising3D
@@ -50,28 +50,29 @@ class HMReferenceAdapter(ModelMixin, ConfigMixin):
                         "CrossAttnUpBlock2D",
                         "CrossAttnUpBlock2D"),
                 num_attention_heads: Optional[Union[int, Tuple[int]]] = 8,
+                 version='v1'
                 ):
        super().__init__()

-        self.reference_modules_down = nn.ModuleList([])
-
        if isinstance(num_attention_heads, int):
            num_attention_heads = (num_attention_heads,) * len(down_block_types)

-        for i, down_block_type in enumerate(down_block_types):
-            output_channel = block_out_channels[i]
+        if version == 'v1':
+            self.reference_modules_down = nn.ModuleList([])
+            for i, down_block_type in enumerate(down_block_types):
+                output_channel = block_out_channels[i]

-            self.reference_modules_down.append(
-                SKReferenceAttention(
-                    in_channels=output_channel,
-                    num_attention_heads=num_attention_heads[i],
+                self.reference_modules_down.append(
+                    SKReferenceAttention(
+                        in_channels=output_channel,
+                        num_attention_heads=num_attention_heads[i],
+                    )
                )
-            )

-        self.reference_modules_mid = SKReferenceAttention(
-            in_channels=block_out_channels[-1],
-            num_attention_heads=num_attention_heads[-1],
-        )
+            self.reference_modules_mid = SKReferenceAttention(
+                in_channels=block_out_channels[-1],
+                num_attention_heads=num_attention_heads[-1],
+            )

        self.reference_modules_up = nn.ModuleList([])

@@ -326,7 +327,10 @@ class InsertReferenceAdapter(object):
        self.reference_modules_up = None

    def insert_reference_adapter(self, adapter: HMReferenceAdapter):
-        self.reference_modules_down = copy.deepcopy(adapter.reference_modules_down)
-        self.reference_modules_mid = copy.deepcopy(adapter.reference_modules_mid)
-        self.reference_modules_up = copy.deepcopy(adapter.reference_modules_up)
+        if hasattr(adapter, "reference_modules_down"):
+            self.reference_modules_down = copy.deepcopy(adapter.reference_modules_down)
+        if hasattr(adapter, "reference_modules_mid"):
+            self.reference_modules_mid = copy.deepcopy(adapter.reference_modules_mid)
+        if hasattr(adapter, "reference_modules_up"):
+            self.reference_modules_up = copy.deepcopy(adapter.reference_modules_up)

@@ -24,7 +24,7 @@ class HMControlNet(ModelMixin, ConfigMixin):
    @register_to_config
    def __init__(
            self,
-            embedding_channels: int,
+            embedding_channels: int = 1280,
            input_channels: int = 3,
            scale_factor: int = 8,
            cross_attention_dim: int = 320,
@@ -91,7 +91,7 @@ class HMControlNet2(ModelMixin, ConfigMixin):
    @register_to_config
    def __init__(
            self,
-            embedding_channels: int,
+            embedding_channels: int = 1280,
            input_channels: int = 3,
            scale_factor: int = 8,
            cross_attention_dim: int = 320,
@@ -143,4 +143,129 @@ class HMControlNet2(ModelMixin, ConfigMixin):
            ret_dict[f'down2_{idx}'] = rearrange(embedding,
                                                     "(b f) c h w -> b c f h w",
                                                     f=video_length)
+        return ret_dict
+
+class HMV2ControlNet(ModelMixin, ConfigMixin):
+
+    def __init__(
+            self,
+            embedding_channels: int = 1280,
+            input_channels: int = 3,
+            scale_factor: int = 4,
+            cross_attention_dim: int = 320,
+            block_out_channels: Tuple[int] = (128, 640, 1280, 1280, 1280),
+    ):
+        super().__init__()
+        self.scale_factor = scale_factor
+        self.embedding_channels = embedding_channels
+        self.cross_attention_dim = cross_attention_dim
+        self.conv_in = nn.Conv2d(input_channels, block_out_channels[0], kernel_size=3, padding=1, bias=False)
+
+        self.exp_embedding = Timesteps(cross_attention_dim, True, 0)
+        self.exp_proj = TimestepEmbedding(cross_attention_dim, cross_attention_dim)
+        self.face_proj = TimestepEmbedding(1024, cross_attention_dim)
+
+        self.blocks_down = nn.ModuleList([])
+        for i in range(1, len(block_out_channels)):
+            channel_in = block_out_channels[i - 1]
+            channel_out = block_out_channels[i]
+            self.blocks_down.append(
+                SKCrossAttention(
+                    channel_in=channel_in,
+                    channel_out=channel_out,
+                    cross_attention_dim=cross_attention_dim,
+                    num_positional_embeddings=64,
+                    num_positional_embeddings_hidden=1024
+                )
+            )
+
+    def forward(self, condition, drive_coeff, face_parts):
+        bs, _, video_length, h, w = condition.shape
+
+        conditioning = rearrange(condition, "b c f h w -> (b f) c h w")
+        conditioning = F.interpolate(conditioning,
+                                     size=(h // self.scale_factor, w // self.scale_factor),
+                                     mode='bilinear',
+                                     align_corners=False)
+        embedding = self.conv_in(conditioning)
+
+        drive_coeff = rearrange(drive_coeff * 500., "b f c -> (b f c)")
+        drive_embedding = self.exp_embedding(drive_coeff).to(dtype=embedding.dtype)
+        drive_embedding = rearrange(drive_embedding, "(b f c) d -> b f c d",
+                                    b=bs, f=video_length, d=self.cross_attention_dim)
+
+        face_parts = rearrange(face_parts, "b f c d -> (b f) c d")
+        face_embedding = self.face_proj(face_parts)
+
+        face_embedding = rearrange(face_embedding, "(b f) c d -> b f c d", f=video_length)
+
+        drive_embedding = torch.cat([face_embedding, drive_embedding], dim=2)
+        drive_embedding = rearrange(drive_embedding, "b f c d -> (b f) c d")
+
+        drive_embedding = self.exp_proj(drive_embedding)
+
+        ret_dict = {}
+        for idx, block in enumerate(self.blocks_down):
+            embedding = block(embedding, drive_embedding)
+            ret_dict[f'up_v2_{len(self.blocks_down) - idx - 1}'] = rearrange(embedding,
+                                                                          "(b f) c h w -> b c f h w",
+                                                                          f=video_length)
+        return ret_dict
+
+class HMV2ControlNet2(ModelMixin, ConfigMixin):
+    def __init__(
+            self,
+            embedding_channels: int = 1280,
+            input_channels: int = 3,
+            scale_factor: int = 4,
+            cross_attention_dim: int = 320,
+            block_out_channels: Tuple[int] = (128, 640, 1280, 1280, 1280),
+    ):
+        super().__init__()
+        self.scale_factor = scale_factor
+        self.embedding_channels = embedding_channels
+        self.cross_attention_dim = cross_attention_dim
+        self.conv_in = nn.Conv2d(input_channels, block_out_channels[0], kernel_size=3, padding=1, bias=False)
+
+        self.exp_embedding = Timesteps(cross_attention_dim, True, 0)
+        self.emo_proj = TimestepEmbedding(cross_attention_dim, cross_attention_dim)
+        # self.face_proj = TimestepEmbedding(1024, cross_attention_dim)
+
+        self.blocks_down = nn.ModuleList([])
+        for i in range(1, len(block_out_channels)):
+            channel_in = block_out_channels[i-1]
+            channel_out = block_out_channels[i]
+            self.blocks_down.append(
+                SKCrossAttention(
+                    channel_in=channel_in,
+                    channel_out=channel_out,
+                    cross_attention_dim=cross_attention_dim,
+                    num_positional_embeddings=64,
+                    num_positional_embeddings_hidden=1024
+                )
+            )
+
+    def forward(self, condition, emo_embedding):
+        bs, _, video_length, h, w = condition.shape
+
+        conditioning = rearrange(condition, "b c f h w -> (b f) c h w")
+        conditioning = F.interpolate(conditioning,
+                                   size=(h // self.scale_factor, w // self.scale_factor),
+                                   mode='bilinear',
+                                   align_corners=False)
+        embedding = self.conv_in(conditioning)
+
+        drive_coeff = rearrange(emo_embedding * 20., "b f c -> (b f c)")
+        drive_embedding = self.exp_embedding(drive_coeff).to(dtype=embedding.dtype)
+        drive_embedding = rearrange(drive_embedding, "(b f c) d -> b f c d",
+                                    b=bs, f=video_length, d=self.cross_attention_dim)
+        drive_embedding = rearrange(drive_embedding, "b f c d -> (b f) c d")
+        drive_embedding = self.emo_proj(drive_embedding)
+
+        ret_dict = {}
+        for idx, block in enumerate(self.blocks_down):
+            embedding = block(embedding, drive_embedding)
+            ret_dict[f'up2_v2_{len(self.blocks_down)-idx-1}'] = rearrange(embedding,
+                                                     "(b f) c h w -> b c f h w",
+                                                     f=video_length)
        return ret_dict
@@ -251,6 +251,10 @@ class HMDenoising3D(UNet2DConditionModel, CopyWeights, InsertReferenceAdapter):

            if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
                res_cache[f"up_{i}"] = sample.clone()
+                if not control_hidden_states is None and f'up_v2_{i}' in control_hidden_states:
+                    sample += rearrange(control_hidden_states[f'up_v2_{i}'], "b c f h w -> (b f) c h w")
+                if not control_hidden_states is None and f'up2_v2_{i}' in control_hidden_states:
+                    sample += rearrange(control_hidden_states[f'up2_v2_{i}'], "b c f h w -> (b f) c h w")
                if hasattr(self, "reference_modules_up") and not reference_hidden_states is None and f'up_{i}' in reference_hidden_states:
                    sample = self.reference_modules_up[i-1](sample, reference_hidden_states[f'up_{i}'], num_frames)

@@ -265,6 +269,10 @@ class HMDenoising3D(UNet2DConditionModel, CopyWeights, InsertReferenceAdapter):
                    encoder_attention_mask=encoder_attention_mask,
                )
            else:
+                if not control_hidden_states is None and f'up_v2_{i}' in control_hidden_states:
+                    sample += rearrange(control_hidden_states[f'up_v2_{i}'], "b c f h w -> (b f) c h w")
+                if not control_hidden_states is None and f'up2_v2_{i}' in control_hidden_states:
+                    sample += rearrange(control_hidden_states[f'up2_v2_{i}'], "b c f h w -> (b f) c h w")
                sample = upsample_block(
                    hidden_states=sample,
                    temb=emb,
@@ -202,6 +202,10 @@ class HMDenoisingMotion(UNetMotionModel, InsertReferenceAdapter):

            if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
                if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
+                    if not control_hidden_states is None and f'up_v2_{i}' in control_hidden_states:
+                        sample += rearrange(control_hidden_states[f'up_v2_{i}'], "b c f h w -> (b f) c h w")
+                    if not control_hidden_states is None and f'up2_v2_{i}' in control_hidden_states:
+                        sample += rearrange(control_hidden_states[f'up2_v2_{i}'], "b c f h w -> (b f) c h w")
                    if hasattr(self,
                               "reference_modules_up") and not reference_hidden_states is None and f'up_{i}' in reference_hidden_states:
                        sample = self.reference_modules_up[i - 1](sample, reference_hidden_states[f'up_{i}'],
@@ -218,6 +222,10 @@ class HMDenoisingMotion(UNetMotionModel, InsertReferenceAdapter):
                    cross_attention_kwargs=cross_attention_kwargs,
                )
            else:
+                if not control_hidden_states is None and f'up_v2_{i}' in control_hidden_states:
+                    sample += rearrange(control_hidden_states[f'up_v2_{i}'], "b c f h w -> (b f) c h w")
+                if not control_hidden_states is None and f'up2_v2_{i}' in control_hidden_states:
+                    sample += rearrange(control_hidden_states[f'up2_v2_{i}'], "b c f h w -> (b f) c h w")
                sample = upsample_block(
                    hidden_states=sample,
                    temb=emb,