little change

hellomeme/__init__.py

@@ -10,4 +10,4 @@
 
 from .pipelines import (HMImagePipeline, HMVideoPipeline,
                         HM3ImagePipeline, HM3VideoPipeline,
-                        HM5ImagePipeline)
+                        HM5ImagePipeline, HM5VideoPipeline)

hellomeme/models/__init__.py

@@ -13,7 +13,7 @@ from .hm_control import (HMControlNet, HMControlNet2, HMV2ControlNet, HMV2Contro
                          HMV3ControlNet, HMControlNetBase, HM5ControlNetBase,
                          HM4SD15ControlProj, HM5SD15ControlProj)
 from .hm_adapters import (HMReferenceAdapter, HM3ReferenceAdapter, HM5ReferenceAdapter,
-                          HM3MotionAdapter, HMPipeline)
+                          HM3MotionAdapter, HM5MotionAdapter, HMPipeline)
 from .hm_denoising_3d import HMDenoising3D
 from .hm3_denoising_3d import HM3Denoising3D
 from .hm3_denoising_motion import HM3DenoisingMotion

hellomeme/models/hm3_denoising_3d.py

@@ -193,7 +193,9 @@ class HM3Denoising3D(UNet2DConditionModel, CopyWeights, InsertReferenceAdapter):
             if not control_hidden_states is None and f'down3_{idx}' in control_hidden_states:
                 sample += rearrange(control_hidden_states[f'down3_{idx}'], "b c f h w -> (b f) c h w")
             if hasattr(self, 'motion_down') and use_motion:
-                sample = self.motion_down[idx](sample, motion_pad_hidden_states[f'down_{idx}'], emb, num_frames)
+                sample = self.motion_down[idx](sample,
+                              None if motion_pad_hidden_states is None else motion_pad_hidden_states[f'down_{idx}'],
+                              emb, num_frames)
 
             down_block_res_samples += res_samples
 
@@ -251,7 +253,9 @@ class HM3Denoising3D(UNet2DConditionModel, CopyWeights, InsertReferenceAdapter):
             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](sample, reference_hidden_states[f'up_{i}'], num_frames)
             if hasattr(self, 'motion_up') and use_motion:
-                sample = self.motion_up[i](sample, motion_pad_hidden_states[f'up_{i}'], emb, num_frames)
+                sample = self.motion_up[i](sample,
+                              None if motion_pad_hidden_states is None else motion_pad_hidden_states[f'up_{i}'],
+                              emb, num_frames)
 
             if hasattr(upsample_block, "has_cross_attention") and upsample_block.has_cross_attention:
                 sample = upsample_block(

hellomeme/models/hm_adapters.py

@@ -42,7 +42,8 @@ from diffusers.models.lora import adjust_lora_scale_text_encoder
 from .hm_blocks import (SKReferenceAttentionV3,
                         SKReferenceAttentionV5,
                         SKReferenceAttention,
-                        SKMotionModule)
+                        SKMotionModule,
+                        SKMotionModuleV5)
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
@@ -183,6 +184,32 @@ class HM3MotionAdapter(ModelMixin, ConfigMixin):
             )
 
 
+class HM5MotionAdapter(ModelMixin, ConfigMixin):
+    @register_to_config
+    def __init__(self,  block_down_channels: Tuple[int] = (320, 640, 1280, 1280),
+                     block_up_channels: Tuple[int] = (1280, 1280, 1280, 640),
+                     num_attention_heads: int = 8):
+        super().__init__()
+        self.motion_down = nn.ModuleList([])
+
+        for i, in_channels in enumerate(block_down_channels):
+            self.motion_down.append(
+                SKMotionModuleV5(
+                    in_channels=in_channels,
+                    num_attention_heads=num_attention_heads,
+                )
+            )
+
+        self.motion_up = nn.ModuleList([])
+        for i, in_channels in enumerate(block_up_channels):
+            self.motion_up.append(
+                SKMotionModuleV5(
+                    in_channels=in_channels,
+                    num_attention_heads=num_attention_heads,
+                )
+            )
+
+
 class HMPipeline(StableDiffusionImg2ImgPipeline):
     @torch.no_grad()
     def load_lora_weights_sk(self, unet, pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],

hellomeme/models/hm_blocks.py

@@ -14,9 +14,12 @@ import torch.utils.checkpoint
 from typing import Tuple
 
 from einops import rearrange
-from diffusers.models.attention import Attention, FeedForward, BasicTransformerBlock
+from diffusers.models.attention import Attention, FeedForward
+from diffusers.models.unets.unet_motion_model import AnimateDiffTransformer3D
 from torchvision.models.resnet import BasicBlock
-from diffusers.models.embeddings import SinusoidalPositionalEmbedding, get_1d_sincos_pos_embed_from_grid, TimestepEmbedding
+from diffusers.models.embeddings import (SinusoidalPositionalEmbedding,
+                                         get_1d_sincos_pos_embed_from_grid,
+                                         TimestepEmbedding)
 
 from diffusers.utils import logging
 
@@ -725,6 +728,51 @@ class SKMotionModule(nn.Module):
         return hidden_states + self.proj(res_temp)
 
 
+class SKMotionModuleV5(nn.Module):
+    def __init__(
+            self,
+            in_channels: int,
+            num_attention_heads: int = 8,
+    ):
+        super().__init__()
+
+        self.attn1 = AnimateDiffTransformer3D(
+            in_channels=in_channels,
+            num_attention_heads=num_attention_heads,
+            attention_head_dim=in_channels // num_attention_heads,
+            dropout=0.0,
+            double_self_attention=True,
+            positional_embeddings='sinusoidal',
+            num_positional_embeddings=32,
+
+        )
+
+        self.attn_pad = AnimateDiffTransformer3D(
+            in_channels=in_channels,
+            num_attention_heads=num_attention_heads,
+            attention_head_dim=in_channels // num_attention_heads,
+            dropout=0.0,
+            double_self_attention=True,
+            positional_embeddings='sinusoidal',
+            num_positional_embeddings=32,
+        )
+
+
+    def forward(self, hidden_states, pad_states, temb, num_frames):
+        c, h, w = hidden_states.shape[-3:]
+        res_temp = hidden_states
+
+        if not pad_states is None:
+            res_temp = rearrange(res_temp, "(b f) c h w -> b c f h w", f=num_frames)
+            res_temp = torch.cat([pad_states[:,:,:1], res_temp, pad_states[:,:,-1:]], dim=2).contiguous()
+            res_temp = rearrange(res_temp, "b c f h w -> (b f) c h w")
+            res_temp = self.attn_pad(res_temp, num_frames=num_frames+2)
+            res_temp = rearrange(res_temp, "(b f) c h w -> b c f h w", f=num_frames+2).contiguous()
+            res_temp = rearrange(res_temp[:,:,1:-1], "b c f h w -> (b f) c h w", h=h, w=w) + hidden_states
+
+        res_temp = self.attn1(res_temp, num_frames=num_frames)
+        return res_temp
+
 # https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/controlnet.py
 def zero_module(module):
     for p in module.parameters():

hellomeme/pipelines/__init__.py

@@ -12,4 +12,5 @@ from .pipline_hm_image import HMImagePipeline
 from .pipline_hm_video import HMVideoPipeline
 from .pipline_hm3_image import HM3ImagePipeline
 from .pipline_hm3_video import HM3VideoPipeline
-from .pipline_hm5_image import HM5ImagePipeline
+from .pipline_hm5_image import HM5ImagePipeline
+from .pipline_hm5_video import HM5VideoPipeline

hellomeme/pipelines/pipline_hm5_video.py

@@ -0,0 +1,510 @@
+# coding: utf-8
+
+"""
+@File   : hm_pipline_image.py
+@Author : Songkey
+@Email  : songkey@pku.edu.cn
+@Date   : 5/30/2025
+@Desc   :
+adapted from: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py
+"""
+
+import copy
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import torch
+from einops import rearrange
+
+from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
+from diffusers.image_processor import PipelineImageInput
+from diffusers.utils import deprecate
+
+from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_timesteps, retrieve_latents
+from diffusers import DPMSolverMultistepScheduler
+from diffusers.utils.torch_utils import randn_tensor
+
+from ..models import (HM3Denoising3D, HM5MotionAdapter,
+                      HM5ReferenceAdapter, HMPipeline, HM5ControlNetBase, HM5SD15ControlProj)
+
+class HM5VideoPipeline(HMPipeline):
+    def caryomitosis(self, version='v5', modelscope=False, **kwargs):
+        if hasattr(self, "unet_ref"):
+            del self.unet_ref
+        self.unet_ref = HM3Denoising3D.from_unet2d(self.unet)
+        self.unet_ref.cpu()
+
+        self.num_frames = 12
+
+        if not isinstance(self.unet, HM3Denoising3D):
+            unet = HM3Denoising3D.from_unet2d(unet=self.unet)
+            # todo: 不够优雅
+            del self.unet
+            self.unet = unet
+            self.unet.cpu()
+
+        self.vae.cpu()
+        self.vae_decode = copy.deepcopy(self.vae)
+        self.text_encoder.cpu()
+        self.text_encoder_ref = copy.deepcopy(self.text_encoder)
+        self.safety_checker.cpu()
+
+    def insert_hm_modules(self, version='v5', dtype=torch.float16, modelscope=False):
+        self.version = version
+
+        if modelscope:
+            from modelscope import snapshot_download
+            hm_reference_dir = snapshot_download('songkey/hm5_reference')
+            hm_control_dir = snapshot_download('songkey/hm5_control_base')
+            hm_control_proj_dir = snapshot_download('songkey/hm5_control_proj')
+            hm_motion_dir = snapshot_download('songkey/hm5_motion')
+        else:
+            hm_reference_dir = 'songkey/hm5_reference'
+            hm_control_dir = 'songkey/hm5_control_base'
+            hm_control_proj_dir = 'songkey/hm5_control_proj'
+            hm_motion_dir = 'songkey/hm5_motion'
+
+        hm_adapter = HM5ReferenceAdapter.from_pretrained(hm_reference_dir)
+        motion_adapter = HM5MotionAdapter.from_pretrained(hm_motion_dir)
+
+        if isinstance(self.unet, HM3Denoising3D):
+            self.unet.insert_reference_adapter(hm_adapter)
+            self.unet.insert_reference_adapter(motion_adapter)
+            self.unet.to(device='cpu', dtype=dtype).eval()
+
+        if hasattr(self, "unet_ref"):
+            self.unet_ref.to(device='cpu', dtype=dtype).eval()
+
+        if hasattr(self, "mp_control"):
+            del self.mp_control
+        if hasattr(self, "mp_control_proj"):
+            del self.mp_control_proj
+
+        self.mp_control = HM5ControlNetBase.from_pretrained(hm_control_dir)
+        self.mp_control_proj = HM5SD15ControlProj.from_pretrained(hm_control_proj_dir)
+        self.mp_control_proj.to(device='cpu', dtype=dtype).eval()
+        self.mp_control.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()
+    def gen_ref_cache(self, latent, prompt_embeds, added_cond_kwargs, device, add_axis=False):
+        self.unet_ref.to(device=device)
+        latent = latent.to(device=device, dtype=prompt_embeds.dtype)
+        latent_model_input = torch.cat([torch.zeros_like(latent), latent]) if self.do_classifier_free_guidance else latent
+        cached_res = self.unet_ref(
+            latent_model_input.unsqueeze(2),
+            0,
+            encoder_hidden_states=prompt_embeds,
+            added_cond_kwargs=added_cond_kwargs,
+            return_dict=False,
+        )[1]
+        self.unet_ref.cpu()
+
+        if add_axis:
+            for k, v in cached_res.items():
+                cached_res[k] = rearrange(v, 'b c h w -> b c 1 h w')
+
+        return dicts_to_device([cached_res], device='cpu')[0]
+
+    @torch.no_grad()
+    def gen_video(self, indexes, pad_dict, base_noise, depth, control_dict, step, more_params):
+        device = more_params['device']
+        timesteps = more_params['timesteps']
+        timestep_cond = more_params['timestep_cond']
+        sigmas = more_params['sigmas']
+        prompt_embeds = more_params['prompt_embeds']
+        cached_res = more_params['cached_res']
+        added_cond_kwargs = more_params['added_cond_kwargs']
+        extra_step_kwargs = more_params['extra_step_kwargs']
+
+        num_indexes = len(indexes)
+        if num_indexes <= self.num_frames:
+            paded_indexes = indexes
+
+            paded_indexes_noise = list(range(len(indexes)))
+
+            paded_indexes = copy.deepcopy(paded_indexes)
+
+            drive_idx_chunks = [paded_indexes]
+            drive_idx_chunks_noise = [paded_indexes_noise]
+        else:
+            stride = self.num_frames - 1
+            paded_indexes = indexes
+            paded_indexes = copy.deepcopy(paded_indexes)
+
+            num_paded_indexes = len(paded_indexes)
+            paded_indexes_noise = list(range(num_paded_indexes))
+
+            drive_idx_chunks = [paded_indexes[i:min(i + self.num_frames, num_paded_indexes)] for i in
+                                range(0, num_paded_indexes, stride)]
+            drive_idx_chunks_noise = [paded_indexes_noise[i:min(i + self.num_frames, num_paded_indexes)] for i in
+                                range(0, num_paded_indexes, stride)]
+
+            if len(drive_idx_chunks[-1]) == 1:
+                drive_idx_chunks[-1] = drive_idx_chunks[:-1]
+                drive_idx_chunks_noise[-1] = drive_idx_chunks_noise[:-1]
+
+            margin_indexes = set()
+            for chunk in drive_idx_chunks:
+                margin_indexes.add(chunk[0])
+                margin_indexes.add(chunk[-1])
+            margin_indexes = sorted(list(margin_indexes))
+
+            self.gen_video(margin_indexes, pad_dict, base_noise, depth+1, control_dict, step, more_params)
+
+        ## generate latent
+        scheduler = DPMSolverMultistepScheduler(
+            num_train_timesteps=1000,
+            beta_start=0.00085,
+            beta_end=0.012,
+            beta_schedule="scaled_linear",
+            algorithm_type="sde-dpmsolver++",
+        )
+
+
+        timesteps, _ = retrieve_timesteps(scheduler, step, device, timesteps, sigmas)
+
+        N, C, _, H, W = base_noise.shape
+
+        latents = base_noise[:,:,:len(paded_indexes)] * scheduler.init_noise_sigma
+
+        self.unet.to(device=device)
+        with self.progress_bar(total=len(timesteps)) as progress_bar:
+            for i, t in enumerate(timesteps):
+                progress_bar.set_description(f"D{depth} GEN Frames")
+                progress_bar.update()
+
+                noise_pred = torch.zeros_like(latents)
+                noise_pred = torch.cat([noise_pred, noise_pred], dim=0) if \
+                    self.do_classifier_free_guidance else noise_pred
+                noise_pred_counter = torch.zeros([1, 1, noise_pred.size(2), 1, 1], dtype=prompt_embeds.dtype).cpu()
+
+                latent_model_input_all = torch.cat([latents] * 2, dim=0) if \
+                    self.do_classifier_free_guidance else latents
+
+                latent_model_input_all = scheduler.scale_model_input(latent_model_input_all, t)
+                for chunk, chunk_noise in zip(drive_idx_chunks, drive_idx_chunks_noise):
+                    if self.interrupt:
+                        continue
+
+                    control_latent = cat_dicts([control_dict[ix] for ix in chunk], dim=2)
+
+                    if chunk[0] in pad_dict and chunk[-1] in pad_dict:
+                        pad_latent = cat_dicts([pad_dict[chunk[0]], pad_dict[chunk[-1]]], dim=2)
+
+                    else:
+                        pad_latent = None
+
+                    latent_model_input = latent_model_input_all[:, :, chunk_noise].to(device=device)
+
+                    noise_pred_chunk = self.unet(
+                        latent_model_input,
+                        t,
+                        encoder_hidden_states=prompt_embeds,
+                        reference_hidden_states=dicts_to_device([cached_res], device=device)[0],
+                        control_hidden_states=dicts_to_device([control_latent], device=device)[0],
+                        motion_pad_hidden_states=None if pad_latent is None else dicts_to_device([pad_latent], device=device)[0],
+                        use_motion=True,
+                        timestep_cond=timestep_cond,
+                        cross_attention_kwargs=self.cross_attention_kwargs,
+                        added_cond_kwargs=added_cond_kwargs,
+                        return_dict=False,
+                    )[0].cpu()
+
+                    noise_pred[:, :, chunk_noise] = noise_pred[:, :, chunk_noise] + noise_pred_chunk
+                    noise_pred_counter[:, :, chunk_noise] = noise_pred_counter[:, :, chunk_noise] + 1
+
+                noise_pred = noise_pred / noise_pred_counter
+                if self.do_classifier_free_guidance:
+                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                    noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
+                latents = scheduler.step(noise_pred.cpu(), t, latents,
+                                              **extra_step_kwargs, return_dict=False)[0]
+        self.unet.cpu()
+
+        if depth == 0:
+            return latents[:,:,:num_indexes]
+        else:
+            for chunk, chunk_noise in zip(drive_idx_chunks, drive_idx_chunks_noise):
+                for idx, idx_noise in zip(chunk, chunk_noise):
+                    # if not idx in pad_dict:
+                        pad_dict[idx] = self.gen_ref_cache(latents[:,:,idx_noise], prompt_embeds,
+                                                           added_cond_kwargs, device, add_axis=True)
+
+    @torch.no_grad()
+    def __call__(
+            self,
+            prompt: Union[str, List[str]] = None,
+            image: PipelineImageInput = None,
+            chunk_overlap: int = 4,
+            drive_params: Dict[str, Any] = None,
+            strength: float = 0.8,
+            num_inference_steps: Optional[int] = 50,
+            timesteps: List[int] = None,
+            sigmas: List[float] = None,
+            guidance_scale: Optional[float] = 7.5,
+            negative_prompt: Optional[Union[str, List[str]]] = None,
+            eta: Optional[float] = 0.0,
+            generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+            prompt_embeds: Optional[torch.Tensor] = None,
+            negative_prompt_embeds: Optional[torch.Tensor] = None,
+            ip_adapter_image: Optional[PipelineImageInput] = None,
+            ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
+            output_type: Optional[str] = "pil",
+            device: Optional[str] = "cpu",
+            return_dict: bool = True,
+            cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+            clip_skip: int = None,
+            callback_on_step_end: Optional[
+                Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
+            ] = None,
+            callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+            **kwargs,
+    ):
+        callback = kwargs.pop("callback", None)
+        callback_steps = kwargs.pop("callback_steps", None)
+        num_images_per_prompt = 1
+
+        if callback is not None:
+            deprecate(
+                "callback",
+                "1.0.0",
+                "Passing `callback` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",
+            )
+        if callback_steps is not None:
+            deprecate(
+                "callback_steps",
+                "1.0.0",
+                "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider use `callback_on_step_end`",
+            )
+
+        if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
+            callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            strength,
+            callback_steps,
+            negative_prompt,
+            prompt_embeds,
+            negative_prompt_embeds,
+            ip_adapter_image,
+            ip_adapter_image_embeds,
+            callback_on_step_end_tensor_inputs,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._clip_skip = clip_skip
+        self._cross_attention_kwargs = cross_attention_kwargs
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        # device = self.device
+
+        # 3. Encode input prompt
+        text_encoder_lora_scale = (
+            self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None
+        )
+        self.text_encoder_ref.to(device=device)
+        prompt_embeds_ref, negative_prompt_embeds_ref = self.encode_prompt_sk(
+            self.text_encoder_ref,
+            prompt,
+            device,
+            num_images_per_prompt,
+            self.do_classifier_free_guidance,
+            negative_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            lora_scale=text_encoder_lora_scale,
+            clip_skip=self.clip_skip,
+        )
+        self.text_encoder_ref.cpu()
+
+        self.text_encoder.to(device=device)
+        prompt_embeds, negative_prompt_embeds = self.encode_prompt_sk(
+            self.text_encoder,
+            prompt,
+            device,
+            num_images_per_prompt,
+            self.do_classifier_free_guidance,
+            negative_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            lora_scale=text_encoder_lora_scale,
+            clip_skip=self.clip_skip,
+        )
+        self.text_encoder.cpu()
+
+        # For classifier free guidance, we need to do two forward passes.
+        # Here we concatenate the unconditional and text embeddings into a single batch
+        # to avoid doing two forward passes
+        if self.do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
+            prompt_embeds_ref = torch.cat([negative_prompt_embeds_ref, prompt_embeds_ref])
+
+        if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
+            image_embeds = self.prepare_ip_adapter_image_embeds(
+                ip_adapter_image,
+                ip_adapter_image_embeds,
+                device,
+                batch_size * num_images_per_prompt,
+                self.do_classifier_free_guidance,
+            )
+
+        # 4. Preprocess image
+        image = self.image_processor.preprocess(image).to(device=device, dtype=prompt_embeds.dtype)
+
+        raw_video_len = drive_params['condition'].size(2)
+
+        # 6. Prepare latent variables
+        self.vae.to(device=device)
+        ref_latents = [
+            retrieve_latents(self.vae.encode(image[i: i + 1].to(device=device)), generator=generator)
+            for i in range(batch_size)
+        ]
+        self.vae.cpu()
+
+        ref_latents = torch.cat(ref_latents, dim=0)
+        ref_latents = self.vae.config.scaling_factor * ref_latents
+        c, h, w = ref_latents.shape[1:]
+
+        # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 7.1 Add image embeds for IP-Adapter
+        added_cond_kwargs = (
+            {"image_embeds": image_embeds}
+            if ip_adapter_image is not None or ip_adapter_image_embeds is not None
+            else None
+        )
+
+        # 7.2 Optionally get Guidance Scale Embedding
+        timestep_cond = None
+        if self.unet.config.time_cond_proj_dim is not None:
+            guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(batch_size * num_images_per_prompt)
+            timestep_cond = self.get_guidance_scale_embedding(
+                guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
+            ).to(device=device, dtype=prompt_embeds.dtype)
+
+        input_noise = randn_tensor([batch_size, c, raw_video_len + self.num_frames + 1, h, w],
+                                   generator=generator, dtype=prompt_embeds.dtype)
+
+        cached_res = self.gen_ref_cache(ref_latents, prompt_embeds_ref, added_cond_kwargs, device)
+
+        control_latents = []
+        control_latent1_neg, control_latent2_neg = None, None
+        self.mp_control.to(device=device)
+        self.mp_control_proj.to(device=device)
+        with self.progress_bar(total=raw_video_len) as progress_bar:
+            for idx in range(raw_video_len):
+                progress_bar.set_description(f"GEN stage1")
+                progress_bar.update()
+
+                condition = drive_params['condition'][:, :, idx:idx+1].clone().to(device=device)
+
+                control_latent_input = {}
+                if 'drive_coeff' in drive_params:
+                    drive_coeff = drive_params['drive_coeff'][:, idx:idx+1].clone().to(device=device)
+                    face_parts = drive_params['face_parts'][:, idx:idx+1].clone().to(device=device)
+
+                    if control_latent1_neg is None and self.do_classifier_free_guidance:
+                        control_latent1_neg = self.mp_control(condition=torch.ones_like(condition)*-1,
+                                                    drive_coeff=torch.zeros_like(drive_coeff),
+                                                    face_parts=torch.zeros_like(face_parts))
+                        control_latent1_neg = self.mp_control_proj(control_latent1_neg)
+
+                    control_latent1 = self.mp_control(condition=condition,
+                                                    drive_coeff=drive_coeff,
+                                                    face_parts=face_parts)
+                    control_latent1 = self.mp_control_proj(control_latent1)
+                    if self.do_classifier_free_guidance:
+                        control_latent1 = cat_dicts([control_latent1_neg, control_latent1], dim=0)
+
+                    control_latent_input.update(control_latent1)
+                else:
+                    pd_fpg = drive_params['pd_fpg'][:, idx:idx+1].clone().to(device=device)
+
+                    if control_latent2_neg is None and self.do_classifier_free_guidance:
+                        control_latent2_neg = self.mp_control(condition=torch.ones_like(condition)*-1,
+                                emo_embedding=torch.zeros_like(pd_fpg))
+                        control_latent2_neg = self.mp_control_proj(control_latent2_neg)
+
+                    control_latent2 = self.mp_control(condition=condition, emo_embedding=pd_fpg)
+                    control_latent2 = self.mp_control_proj(control_latent2)
+                    if self.do_classifier_free_guidance:
+                        control_latent2 = cat_dicts([control_latent2_neg, control_latent2], dim=0)
+
+                    control_latent_input.update(control_latent2)
+                control_latents.append(dicts_to_device([control_latent_input], device='cpu')[0])
+
+        self.mp_control.cpu()
+        self.mp_control_proj.cpu()
+
+        more_params = dict(
+            timesteps=timesteps,
+            timestep_cond=timestep_cond,
+            sigmas=sigmas,
+            prompt_embeds=prompt_embeds,
+            cached_res=cached_res,
+            device=device,
+            added_cond_kwargs=added_cond_kwargs,
+            extra_step_kwargs=extra_step_kwargs,
+        )
+
+        indexes = list(range(raw_video_len))
+        latents = self.gen_video(indexes, dict(), input_noise, 0, control_latents, num_inference_steps, more_params)
+        latents_res = latents / self.vae_decode.config.scaling_factor
+
+        self.vae_decode.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),
+                                             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()
+        return res_frames, latents_res
+
+
+def merge_dicts(dictl, dictr, wl=0.5):
+    res = {}
+    for k in dictl.keys():
+        res[k] = dictl[k] * wl + dictr[k] * (1-wl)
+    return res
+
+
+def cat_dicts(dicts, dim=0):
+    res = {}
+    for k in dicts[0].keys():
+        res[k] = torch.cat([d[k].clone() for d in dicts], dim=dim)
+    return res
+
+
+def cat_dicts_ref(dicts):
+    res = {}
+    for k in dicts[0].keys():
+        res[k] = rearrange(torch.cat([d[k].clone().unsqueeze(1) for d in dicts], dim=1), "b f c h w -> (b f) c h w ")
+    return res
+
+
+def dicts_to_device(dicts, device):
+    ret = []
+    for d in dicts:
+        tmpd = {}
+        for k in d.keys():
+            tmpd[k] = d[k].clone().to(device)
+        ret.append(tmpd)
+    return ret

hellomeme/tools/sr.py

@@ -0,0 +1,394 @@
+# coding: utf-8
+
+"""
+@File   : sr.py
+@Author : Songkey
+@Email  : songkey@pku.edu.cn
+@Date   : 5/30/2025
+@Desc   : adapted from: https://github.com/xinntao/Real-ESRGAN
+"""
+
+import torch
+from torch import nn as nn
+from torch.nn import functional as F
+import cv2
+import numpy as np
+import math
+import os.path as osp
+
+def pixel_unshuffle(x, scale):
+    """ Pixel unshuffle.
+
+    Args:
+        x (Tensor): Input feature with shape (b, c, hh, hw).
+        scale (int): Downsample ratio.
+
+    Returns:
+        Tensor: the pixel unshuffled feature.
+    """
+    b, c, hh, hw = x.size()
+    out_channel = c * (scale**2)
+    assert hh % scale == 0 and hw % scale == 0
+    h = hh // scale
+    w = hw // scale
+    x_view = x.view(b, c, h, scale, w, scale)
+    return x_view.permute(0, 1, 3, 5, 2, 4).reshape(b, out_channel, h, w)
+
+def make_layer(basic_block, num_basic_block, **kwarg):
+    """Make layers by stacking the same blocks.
+
+    Args:
+        basic_block (nn.module): nn.module class for basic block.
+        num_basic_block (int): number of blocks.
+
+    Returns:
+        nn.Sequential: Stacked blocks in nn.Sequential.
+    """
+    layers = []
+    for _ in range(num_basic_block):
+        layers.append(basic_block(**kwarg))
+    return nn.Sequential(*layers)
+
+class ResidualDenseBlock(nn.Module):
+    """Residual Dense Block.
+
+    Used in RRDB block in ESRGAN.
+
+    Args:
+        num_feat (int): Channel number of intermediate features.
+        num_grow_ch (int): Channels for each growth.
+    """
+
+    def __init__(self, num_feat=64, num_grow_ch=32):
+        super(ResidualDenseBlock, self).__init__()
+        self.conv1 = nn.Conv2d(num_feat, num_grow_ch, 3, 1, 1)
+        self.conv2 = nn.Conv2d(num_feat + num_grow_ch, num_grow_ch, 3, 1, 1)
+        self.conv3 = nn.Conv2d(num_feat + 2 * num_grow_ch, num_grow_ch, 3, 1, 1)
+        self.conv4 = nn.Conv2d(num_feat + 3 * num_grow_ch, num_grow_ch, 3, 1, 1)
+        self.conv5 = nn.Conv2d(num_feat + 4 * num_grow_ch, num_feat, 3, 1, 1)
+
+        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+        # initialization
+        # default_init_weights([self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1)
+
+    def forward(self, x):
+        x1 = self.lrelu(self.conv1(x))
+        x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
+        x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
+        x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
+        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
+        # Empirically, we use 0.2 to scale the residual for better performance
+        return x5 * 0.2 + x
+
+
+class RRDB(nn.Module):
+    """Residual in Residual Dense Block.
+
+    Used in RRDB-Net in ESRGAN.
+
+    Args:
+        num_feat (int): Channel number of intermediate features.
+        num_grow_ch (int): Channels for each growth.
+    """
+
+    def __init__(self, num_feat, num_grow_ch=32):
+        super(RRDB, self).__init__()
+        self.rdb1 = ResidualDenseBlock(num_feat, num_grow_ch)
+        self.rdb2 = ResidualDenseBlock(num_feat, num_grow_ch)
+        self.rdb3 = ResidualDenseBlock(num_feat, num_grow_ch)
+
+    def forward(self, x):
+        out = self.rdb1(x)
+        out = self.rdb2(out)
+        out = self.rdb3(out)
+        # Empirically, we use 0.2 to scale the residual for better performance
+        return out * 0.2 + x
+
+class RRDBNet(nn.Module):
+    """Networks consisting of Residual in Residual Dense Block, which is used
+    in ESRGAN.
+
+    ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks.
+
+    We extend ESRGAN for scale x2 and scale x1.
+    Note: This is one option for scale 1, scale 2 in RRDBNet.
+    We first employ the pixel-unshuffle (an inverse operation of pixelshuffle to reduce the spatial size
+    and enlarge the channel size before feeding inputs into the main ESRGAN architecture.
+
+    Args:
+        num_in_ch (int): Channel number of inputs.
+        num_out_ch (int): Channel number of outputs.
+        num_feat (int): Channel number of intermediate features.
+            Default: 64
+        num_block (int): Block number in the trunk network. Defaults: 23
+        num_grow_ch (int): Channels for each growth. Default: 32.
+    """
+
+    def __init__(self, num_in_ch, num_out_ch, scale=4, num_feat=64, num_block=23, num_grow_ch=32):
+        super(RRDBNet, self).__init__()
+        self.scale = scale
+        if scale == 2:
+            num_in_ch = num_in_ch * 4
+        elif scale == 1:
+            num_in_ch = num_in_ch * 16
+        self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)
+        self.body = make_layer(RRDB, num_block, num_feat=num_feat, num_grow_ch=num_grow_ch)
+        self.conv_body = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        # upsample
+        self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
+
+        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+    def forward(self, x):
+        if self.scale == 2:
+            feat = pixel_unshuffle(x, scale=2)
+        elif self.scale == 1:
+            feat = pixel_unshuffle(x, scale=4)
+        else:
+            feat = x
+        feat = self.conv_first(feat)
+        body_feat = self.conv_body(self.body(feat))
+        feat = feat + body_feat
+        # upsample
+        feat = self.lrelu(self.conv_up1(F.interpolate(feat, scale_factor=2, mode='nearest')))
+        feat = self.lrelu(self.conv_up2(F.interpolate(feat, scale_factor=2, mode='nearest')))
+        out = self.conv_last(self.lrelu(self.conv_hr(feat)))
+        return out
+
+class RealESRGANer():
+    def __init__(self,
+                 scale,
+                 tile=0,
+                 tile_pad=10,
+                 pre_pad=10,
+                 half=True,
+                 device=None,
+                 gpu_id=None,
+                 modelscope=False):
+        self.scale = scale
+        self.tile_size = tile
+        self.tile_pad = tile_pad
+        self.pre_pad = pre_pad
+        self.mod_scale = None
+        self.half = half
+
+        # initialize model
+        if gpu_id:
+            self.device = torch.device(
+                f'cuda:{gpu_id}' if torch.cuda.is_available() else 'cpu') if device is None else device
+        else:
+            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') if device is None else device
+
+        if modelscope:
+            from modelscope import snapshot_download
+            model_path = osp.join(snapshot_download('songkey/ESRGAN'), 'RealESRGAN_x2plus.pth')
+        else:
+            from huggingface_hub import hf_hub_download
+            model_path = hf_hub_download('songkey/ESRGAN', filename='RealESRGAN_x2plus.pth')
+
+        loadnet = torch.load(model_path, map_location=torch.device('cpu'))
+
+        # prefer to use params_ema
+        if 'params_ema' in loadnet:
+            keyname = 'params_ema'
+        else:
+            keyname = 'params'
+        self.model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
+        self.model.load_state_dict(loadnet[keyname], strict=True)
+
+        self.model.eval()
+        self.model = self.model.to(self.device)
+        if self.half:
+            self.model = self.model.half()
+
+    def dni(self, net_a, net_b, dni_weight, key='params', loc='cpu'):
+        """Deep network interpolation.
+
+        ``Paper: Deep Network Interpolation for Continuous Imagery Effect Transition``
+        """
+        net_a = torch.load(net_a, map_location=torch.device(loc))
+        net_b = torch.load(net_b, map_location=torch.device(loc))
+        for k, v_a in net_a[key].items():
+            net_a[key][k] = dni_weight[0] * v_a + dni_weight[1] * net_b[key][k]
+        return net_a
+
+    def pre_process(self, img):
+        """Pre-process, such as pre-pad and mod pad, so that the images can be divisible
+        """
+        img = torch.from_numpy(np.transpose(img, (2, 0, 1))).float()
+        self.img = img.unsqueeze(0).to(self.device)
+        if self.half:
+            self.img = self.img.half()
+
+        # pre_pad
+        if self.pre_pad != 0:
+            self.img = F.pad(self.img, (0, self.pre_pad, 0, self.pre_pad), 'reflect')
+        # mod pad for divisible borders
+        if self.scale == 2:
+            self.mod_scale = 2
+        elif self.scale == 1:
+            self.mod_scale = 4
+        if self.mod_scale is not None:
+            self.mod_pad_h, self.mod_pad_w = 0, 0
+            _, _, h, w = self.img.size()
+            if (h % self.mod_scale != 0):
+                self.mod_pad_h = (self.mod_scale - h % self.mod_scale)
+            if (w % self.mod_scale != 0):
+                self.mod_pad_w = (self.mod_scale - w % self.mod_scale)
+            self.img = F.pad(self.img, (0, self.mod_pad_w, 0, self.mod_pad_h), 'reflect')
+
+    def process(self):
+        # model inference
+        self.output = self.model(self.img)
+
+    def tile_process(self):
+        """It will first crop input images to tiles, and then process each tile.
+        Finally, all the processed tiles are merged into one images.
+
+        Modified from: https://github.com/ata4/esrgan-launcher
+        """
+        batch, channel, height, width = self.img.shape
+        output_height = height * self.scale
+        output_width = width * self.scale
+        output_shape = (batch, channel, output_height, output_width)
+
+        # start with black image
+        self.output = self.img.new_zeros(output_shape)
+        tiles_x = math.ceil(width / self.tile_size)
+        tiles_y = math.ceil(height / self.tile_size)
+
+        # loop over all tiles
+        for y in range(tiles_y):
+            for x in range(tiles_x):
+                # extract tile from input image
+                ofs_x = x * self.tile_size
+                ofs_y = y * self.tile_size
+                # input tile area on total image
+                input_start_x = ofs_x
+                input_end_x = min(ofs_x + self.tile_size, width)
+                input_start_y = ofs_y
+                input_end_y = min(ofs_y + self.tile_size, height)
+
+                # input tile area on total image with padding
+                input_start_x_pad = max(input_start_x - self.tile_pad, 0)
+                input_end_x_pad = min(input_end_x + self.tile_pad, width)
+                input_start_y_pad = max(input_start_y - self.tile_pad, 0)
+                input_end_y_pad = min(input_end_y + self.tile_pad, height)
+
+                # input tile dimensions
+                input_tile_width = input_end_x - input_start_x
+                input_tile_height = input_end_y - input_start_y
+                tile_idx = y * tiles_x + x + 1
+                input_tile = self.img[:, :, input_start_y_pad:input_end_y_pad, input_start_x_pad:input_end_x_pad]
+
+                # upscale tile
+                try:
+                    with torch.no_grad():
+                        output_tile = self.model(input_tile)
+                except RuntimeError as error:
+                    print('Error', error)
+                print(f'\tTile {tile_idx}/{tiles_x * tiles_y}')
+
+                # output tile area on total image
+                output_start_x = input_start_x * self.scale
+                output_end_x = input_end_x * self.scale
+                output_start_y = input_start_y * self.scale
+                output_end_y = input_end_y * self.scale
+
+                # output tile area without padding
+                output_start_x_tile = (input_start_x - input_start_x_pad) * self.scale
+                output_end_x_tile = output_start_x_tile + input_tile_width * self.scale
+                output_start_y_tile = (input_start_y - input_start_y_pad) * self.scale
+                output_end_y_tile = output_start_y_tile + input_tile_height * self.scale
+
+                # put tile into output image
+                self.output[:, :, output_start_y:output_end_y,
+                            output_start_x:output_end_x] = output_tile[:, :, output_start_y_tile:output_end_y_tile,
+                                                                       output_start_x_tile:output_end_x_tile]
+
+    def post_process(self):
+        # remove extra pad
+        if self.mod_scale is not None:
+            _, _, h, w = self.output.size()
+            self.output = self.output[:, :, 0:h - self.mod_pad_h * self.scale, 0:w - self.mod_pad_w * self.scale]
+        # remove prepad
+        if self.pre_pad != 0:
+            _, _, h, w = self.output.size()
+            self.output = self.output[:, :, 0:h - self.pre_pad * self.scale, 0:w - self.pre_pad * self.scale]
+        return self.output
+
+    @torch.no_grad()
+    def enhance(self, img, outscale=None, alpha_upsampler='realesrgan'):
+        h_input, w_input = img.shape[0:2]
+        # img: numpy
+        img = img.astype(np.float32)
+        if np.max(img) > 256:  # 16-bit image
+            max_range = 65535
+            print('\tInput is a 16-bit image')
+        else:
+            max_range = 255
+        img = img / max_range
+        if len(img.shape) == 2:  # gray image
+            img_mode = 'L'
+            img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
+        elif img.shape[2] == 4:  # RGBA image with alpha channel
+            img_mode = 'RGBA'
+            alpha = img[:, :, 3]
+            img = img[:, :, 0:3]
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+            if alpha_upsampler == 'realesrgan':
+                alpha = cv2.cvtColor(alpha, cv2.COLOR_GRAY2RGB)
+        else:
+            img_mode = 'RGB'
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
+        # ------------------- process image (without the alpha channel) ------------------- #
+        self.pre_process(img)
+        if self.tile_size > 0:
+            self.tile_process()
+        else:
+            self.process()
+        output_img = self.post_process()
+        output_img = output_img.data.squeeze().float().cpu().clamp_(0, 1).numpy()
+        output_img = np.transpose(output_img[[2, 1, 0], :, :], (1, 2, 0))
+        if img_mode == 'L':
+            output_img = cv2.cvtColor(output_img, cv2.COLOR_BGR2GRAY)
+
+        # ------------------- process the alpha channel if necessary ------------------- #
+        if img_mode == 'RGBA':
+            if alpha_upsampler == 'realesrgan':
+                self.pre_process(alpha)
+                if self.tile_size > 0:
+                    self.tile_process()
+                else:
+                    self.process()
+                output_alpha = self.post_process()
+                output_alpha = output_alpha.data.squeeze().float().cpu().clamp_(0, 1).numpy()
+                output_alpha = np.transpose(output_alpha[[2, 1, 0], :, :], (1, 2, 0))
+                output_alpha = cv2.cvtColor(output_alpha, cv2.COLOR_BGR2GRAY)
+            else:  # use the cv2 resize for alpha channel
+                h, w = alpha.shape[0:2]
+                output_alpha = cv2.resize(alpha, (w * self.scale, h * self.scale), interpolation=cv2.INTER_LINEAR)
+
+            # merge the alpha channel
+            output_img = cv2.cvtColor(output_img, cv2.COLOR_BGR2BGRA)
+            output_img[:, :, 3] = output_alpha
+
+        # ------------------------------ return ------------------------------ #
+        if max_range == 65535:  # 16-bit image
+            output = (output_img * 65535.0).round().astype(np.uint16)
+        else:
+            output = (output_img * 255.0).round().astype(np.uint8)
+
+        if outscale is not None and outscale != float(self.scale):
+            output = cv2.resize(
+                output, (
+                    int(w_input * outscale),
+                    int(h_input * outscale),
+                ), interpolation=cv2.INTER_LANCZOS4)
+
+        return output, img_mode

meme.py

@@ -1,9 +1,7 @@
 import json
 import os
 import os.path as osp
-import random
 import sys
-import subprocess
 import importlib.metadata
 
 cur_dir = osp.dirname(osp.abspath(__file__))
@@ -11,7 +9,7 @@ cur_dir = osp.dirname(osp.abspath(__file__))
 installed_packages = [package.name for package in importlib.metadata.distributions()]
 
 REQUIRED = {
-    'diffusers':'0.31.0', 'transformers':'4.46.3', 'einops':'0.8.0', 'opencv-python':'4.10.0.84', 'tqdm':'4.67.0',
+    'diffusers':'0.33.1', 'transformers':'4.46.3', 'einops':'0.8.0', 'opencv-python':'4.10.0.84', 'tqdm':'4.67.0',
     'pillow':'10.2.0', 'onnxruntime-gpu':'1.18.1', 'onnx':'1.17.0', 'safetensors':'0.4.5',
     'accelerate':'1.1.1', 'peft':'0.13.2', 'modelscope':'', 'huggingface-hub':''
 }
@@ -54,7 +52,7 @@ from .hellomeme.utils import (get_drive_expression,
                               )
 from .hellomeme import (HMImagePipeline, HMVideoPipeline,
                         HM3ImagePipeline, HM3VideoPipeline,
-                        HM5ImagePipeline)
+                        HM5ImagePipeline, HM5VideoPipeline)
 
 config_path = osp.join(cur_dir, 'hellomeme', 'model_config.json')
 with open(config_path, 'r') as f:
@@ -169,7 +167,7 @@ class HMVideoPipelineLoader:
                 "checkpoint": (checkpoint_files, ),
                 "lora": (lora_files, ),
                 "vae": (vae_files, ),
-                "version": (['v1', 'v2', 'v3', 'v4'], ),
+                "version": (['v1', 'v2', 'v3', 'v4', 'v5'], ),
                 "stylize": (['x1', 'x2'], ),
                 "deployment": (['huggingface', 'modelscope'], ),
                 "lora_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1}),
@@ -192,6 +190,8 @@ class HMVideoPipelineLoader:
 
         if version == 'v3' or version == 'v4':
             pipeline = HM3VideoPipeline.from_pretrained(sd1_5_dir)
+        elif version == 'v5':
+            pipeline = HM5VideoPipeline.from_pretrained(sd1_5_dir)
         else:
             pipeline = HMVideoPipeline.from_pretrained(sd1_5_dir)
         pipeline.to(dtype=dtype)

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "comfyui_hellomeme"
 description = "This repository is the official implementation of the [a/HelloMeme](https://arxiv.org/pdf/2410.22901) ComfyUI interface"
-version = "1.3.3"
+version = "1.3.4"
 license = {file = "LICENSE"}
 
 [project.urls]