add HellMeme V2

hellomeme/models/__init__.py

@@ -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

hellomeme/models/hm_blocks.py

@@ -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)
 

hellomeme/models/hm_control.py

@@ -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

hellomeme/models/hm_denoising_3d.py

@@ -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,

hellomeme/models/hm_denoising_motion.py

@@ -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,

hellomeme/pipelines/pipline_hm_image.py

@@ -22,42 +22,44 @@ from diffusers.pipelines.stable_diffusion.pipeline_output import StableDiffusion
 from diffusers import StableDiffusionImg2ImgPipeline
 from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_timesteps, retrieve_latents
 
-from ..models import HMDenoising3D, HMControlNet, HMControlNet2
+from ..models import HMDenoising3D, HMControlNet, HMControlNet2, HMV2ControlNet, HMV2ControlNet2
 from ..models import HMReferenceAdapter
 
 class HMImagePipeline(StableDiffusionImg2ImgPipeline):
     def caryomitosis(self, **kwargs):
-        if not hasattr(self, "unet_ref"):
-            self.unet_ref = HMDenoising3D.from_unet2d(self.unet)
-            self.unet_ref.cpu()
         if not isinstance(self.unet, HMDenoising3D):
             unet = HMDenoising3D.from_unet2d(unet=self.unet)
             # todo: 不够优雅
             del self.unet
             self.unet = unet
             self.unet.cpu()
-        self.vae_decode = copy.deepcopy(self.vae)
 
-    def insert_hm_modules(self, dtype):
+    def insert_hm_modules(self, version, dtype):
         if isinstance(self.unet, HMDenoising3D):
-            hm_adapter = HMReferenceAdapter.from_pretrained('songkey/hm_reference')
+            if version == 'v1':
+                hm_adapter = HMReferenceAdapter.from_pretrained('songkey/hm_reference')
+            else:
+                hm_adapter = HMReferenceAdapter.from_pretrained('songkey/hm2_reference')
             self.unet.insert_reference_adapter(hm_adapter)
-            self.unet_ref.insert_reference_adapter(hm_adapter)
-            self.unet_ref.to(device='cpu', dtype=dtype).eval()
             self.unet.to(device='cpu', dtype=dtype).eval()
 
         if hasattr(self, "mp_control"):
             del self.mp_control
-        self.mp_control = HMControlNet.from_pretrained('songkey/hm_control')
+        if version == 'v1':
+            self.mp_control = HMControlNet.from_pretrained('songkey/hm_control')
+        else:
+            self.mp_control = HMV2ControlNet.from_pretrained('songkey/hm2_control')
         self.mp_control.to(device='cpu', dtype=dtype).eval()
 
         if hasattr(self, "mp_control2"):
             del self.mp_control2
-        self.mp_control2 = HMControlNet2.from_pretrained('songkey/hm_control2')
+        if version == 'v1':
+            self.mp_control2 = HMControlNet2.from_pretrained('songkey/hm_control2')
+        else:
+            self.mp_control2 = HMV2ControlNet2.from_pretrained('songkey/hm2_control2')
         self.mp_control2.to(device='cpu', dtype=dtype).eval()
 
         self.vae.to(device='cpu', dtype=dtype).eval()
-        self.vae_decode.to(device='cpu', dtype=dtype).eval()
         self.text_encoder.to(device='cpu', dtype=dtype).eval()
 
     @torch.no_grad()
@@ -230,15 +232,13 @@ class HMImagePipeline(StableDiffusionImg2ImgPipeline):
         )
 
         latent_model_input = torch.cat([torch.zeros_like(ref_latents), ref_latents]) if self.do_classifier_free_guidance else ref_latents
-        self.unet_ref.to(device=device)
-        cached_res = self.unet_ref(
+        self.unet.to(device=device)
+        cached_res = self.unet(
             latent_model_input.unsqueeze(2),
             0,
             encoder_hidden_states=prompt_embeds,
-            # control_hidden_states=control_latents,
             return_dict=False,
         )[1]
-        self.unet_ref.cpu()
 
         # 7.2 Optionally get Guidance Scale Embedding
         timestep_cond = None
@@ -304,14 +304,14 @@ class HMImagePipeline(StableDiffusionImg2ImgPipeline):
 
         self.unet.cpu()
 
-        self.vae_decode.to(device=device)
+        self.vae.to(device=device)
         if not output_type == "latent":
-            image = self.vae_decode.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[
+            image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[
                 0
             ]
         else:
             image = latents
-        self.vae_decode.cpu()
+        self.vae.cpu()
 
         do_denormalize = [True] * image.shape[0]
 

hellomeme/pipelines/pipline_hm_video.py

@@ -21,54 +21,58 @@ from diffusers.utils.torch_utils import randn_tensor
 from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_timesteps, retrieve_latents
 from diffusers import StableDiffusionImg2ImgPipeline, MotionAdapter, EulerDiscreteScheduler
 
-from ..models import HMDenoising3D, HMDenoisingMotion, HMControlNet, HMControlNet2
+from ..models import HMDenoising3D, HMDenoisingMotion, HMControlNet, HMControlNet2, HMV2ControlNet, HMV2ControlNet2
 from ..models import HMReferenceAdapter
 from ..utils import dicts_to_device, cat_dicts, cat_dicts_ref
 
 class HMVideoPipeline(StableDiffusionImg2ImgPipeline):
-    def caryomitosis(self, patch_frames=12, **kwargs):
+    def caryomitosis(self, version, **kwargs):
         if hasattr(self, "unet_ref"):
             del self.unet_ref
-        self.unet_ref = HMDenoising3D.from_unet2d(self.unet)
-        self.unet_ref.cpu()
         self.unet_pre = HMDenoising3D.from_unet2d(self.unet)
         self.unet_pre.cpu()
 
         self.num_frames = 12
-        adapter = MotionAdapter.from_pretrained("songkey/hm_animatediff_frame12", torch_dtype=torch.float16)
+        if version == 'v1':
+            adapter = MotionAdapter.from_pretrained("songkey/hm_animatediff_frame12", torch_dtype=torch.float16)
+        else:
+            adapter = MotionAdapter.from_pretrained("songkey/hm2_animatediff_frame12", torch_dtype=torch.float16)
         unet = HMDenoisingMotion.from_unet2d(unet=self.unet, motion_adapter=adapter, load_weights=True)
         # todo: 不够优雅
         del self.unet
         self.unet = unet
 
-        self.vae_decode = copy.deepcopy(self.vae)
-
-    def insert_hm_modules(self, dtype):
-        hm_adapter = HMReferenceAdapter.from_pretrained('songkey/hm_reference')
+    def insert_hm_modules(self, version, dtype):
+        if version == 'v1':
+            hm_adapter = HMReferenceAdapter.from_pretrained('songkey/hm_reference')
+        else:
+            hm_adapter = HMReferenceAdapter.from_pretrained('songkey/hm2_reference')
         if isinstance(self.unet, HMDenoisingMotion):
             self.unet.insert_reference_adapter(hm_adapter)
             self.unet.to(device='cpu', dtype=dtype).eval()
 
-        if hasattr(self, "unet_ref"):
-            self.unet_ref.insert_reference_adapter(hm_adapter)
-            self.unet_ref.to(device='cpu', dtype=dtype).eval()
-
         if hasattr(self, "unet_pre"):
             self.unet_pre.insert_reference_adapter(hm_adapter)
             self.unet_pre.to(device='cpu', dtype=dtype).eval()
 
         if hasattr(self, "mp_control"):
             del self.mp_control
-        self.mp_control = HMControlNet.from_pretrained('songkey/hm_control')
+        if version == 'v1':
+            self.mp_control = HMControlNet.from_pretrained('songkey/hm_control')
+        else:
+            self.mp_control = HMV2ControlNet.from_pretrained('songkey/hm2_control')
+
         self.mp_control.to(device='cpu', dtype=dtype).eval()
 
         if hasattr(self, "mp_control2"):
             del self.mp_control2
-        self.mp_control2 = HMControlNet2.from_pretrained('songkey/hm_control2')
+        if version == 'v1':
+            self.mp_control2 = HMControlNet2.from_pretrained('songkey/hm_control2')
+        else:
+            self.mp_control2 = HMV2ControlNet2.from_pretrained('songkey/hm2_control2')
         self.mp_control2.to(device='cpu', dtype=dtype).eval()
 
         self.vae.to(device='cpu', dtype=dtype).eval()
-        self.vae_decode.to(device='cpu', dtype=dtype).eval()
         self.text_encoder.to(device='cpu', dtype=dtype).eval()
 
     @torch.no_grad()
@@ -218,15 +222,14 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline):
                 guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
             ).to(device=device, dtype=prompt_embeds.dtype)
 
-        self.unet_ref.to(device=device)
-        cached_res = self.unet_ref(
+        self.unet_pre.to(device=device)
+        cached_res = self.unet_pre(
             torch.cat([torch.zeros_like(ref_latents), ref_latents], dim=0) if
             self.do_classifier_free_guidance else ref_latents,
             0,
             encoder_hidden_states=prompt_embeds,
             return_dict=False,
         )[1]
-        self.unet_ref.cpu()
 
         pre_latents = []
         control_latents = []
@@ -234,7 +237,6 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline):
         base_noise = randn_tensor([batch_size, c, 1, h, w], dtype=prompt_embeds.dtype, generator=generator).to(device=device)
         self.mp_control.to(device=device)
         self.mp_control2.to(device=device)
-        self.unet_pre.to(device=device)
         with self.progress_bar(total=raw_video_len) as progress_bar:
             for idx in range(raw_video_len):
                 condition = drive_params['condition'][:, :, idx:idx+1].clone().to(device=device)
@@ -313,7 +315,6 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline):
         self.mp_control.cpu()
         self.mp_control2.cpu()
         self.unet_pre.cpu()
-        self.unet_ref.cpu()
 
         chunk_size = self.num_frames
         chunk_overlap = min(max(0, chunk_overlap), chunk_size // 2 - 1)
@@ -421,15 +422,15 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline):
         self.unet.cpu()
         latents_res = latents / self.vae.config.scaling_factor
 
-        self.vae_decode.to(device=device)
+        self.vae.to(device=device)
         res_frames = []
         with self.progress_bar(total=raw_video_len) as progress_bar:
             for i in range(raw_video_len):
-                ret_tensor = self.vae_decode.decode(latents_res[:, :, i, :, :].to(device=device),
+                ret_tensor = self.vae.decode(latents_res[:, :, i, :, :].to(device=device),
                                              return_dict=False, generator=generator)[0]
                 ret_image = self.image_processor.postprocess(ret_tensor, output_type=output_type)
                 res_frames.append(ret_image)
                 progress_bar.set_description(f"VAE DECODE")
                 progress_bar.update()
-        self.vae_decode.cpu()
+        self.vae.cpu()
         return res_frames, latents_res