update

hellomeme/models/hm_blocks.py

@@ -136,7 +136,9 @@ class SKReferenceAttention(nn.Module):
 
     def forward(self, hidden_states, ref_stats, num_frames):
         h, w = hidden_states.shape[-2:]
-        ref_stats = ref_stats.repeat_interleave(num_frames, dim=0)
+
+        if ref_stats.shape[0] != hidden_states.shape[0]:
+            ref_stats = ref_stats.repeat_interleave(num_frames, dim=0)
         cat_stats = torch.cat([hidden_states, ref_stats], dim=-1)
 
         cat_stats = rearrange(cat_stats.contiguous(), "b c h w -> (b h) w c")

hellomeme/pipelines/pipline_hm_image.py

@@ -9,8 +9,8 @@
 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
-from einops import rearrange
 import torch
 
 from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
@@ -29,28 +29,36 @@ 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, device):
+    def insert_hm_modules(self, dtype):
         if isinstance(self.unet, HMDenoising3D):
             hm_adapter = HMReferenceAdapter.from_pretrained('songkey/hm_reference')
             self.unet.insert_reference_adapter(hm_adapter)
-            self.unet_ref = self.unet_ref.to(device=device, dtype=dtype).eval()
-            self.unet = self.unet.to(device=device, dtype=dtype).eval()
+            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')
-        self.mp_control = self.mp_control.to(device=device, dtype=dtype).eval()
+        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')
-        self.mp_control2 = self.mp_control2.to(device=device, dtype=dtype).eval()
+        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()
     def __call__(
@@ -71,6 +79,7 @@ class HMImagePipeline(StableDiffusionImg2ImgPipeline):
             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,
@@ -126,13 +135,14 @@ class HMImagePipeline(StableDiffusionImg2ImgPipeline):
         else:
             batch_size = prompt_embeds.shape[0]
 
-        device = self.device
+        # 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
         )
-        print('device:', device)
+
+        self.text_encoder.to(device=device)
         prompt_embeds, negative_prompt_embeds = self.encode_prompt(
             prompt,
             device,
@@ -144,6 +154,8 @@ class HMImagePipeline(StableDiffusionImg2ImgPipeline):
             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
@@ -169,10 +181,14 @@ class HMImagePipeline(StableDiffusionImg2ImgPipeline):
         timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device)
         latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)
 
+        # 6. Prepare reference latents
+        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
 
@@ -219,6 +235,7 @@ class HMImagePipeline(StableDiffusionImg2ImgPipeline):
             latent_model_input.unsqueeze(2),
             0,
             encoder_hidden_states=prompt_embeds,
+            # control_hidden_states=control_latents,
             return_dict=False,
         )[1]
         self.unet_ref.cpu()
@@ -237,6 +254,7 @@ class HMImagePipeline(StableDiffusionImg2ImgPipeline):
         # 8. Denoising loop
         num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
         self._num_timesteps = len(timesteps)
+        self.unet.to(device=device)
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):
                 if self.interrupt:
@@ -284,12 +302,16 @@ class HMImagePipeline(StableDiffusionImg2ImgPipeline):
                         step_idx = i // getattr(self.scheduler, "order", 1)
                         callback(step_idx, t, latents)
 
+        self.unet.cpu()
+
+        self.vae_decode.to(device=device)
         if not output_type == "latent":
-            image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[
+            image = self.vae_decode.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[
                 0
             ]
         else:
             image = latents
+        self.vae_decode.cpu()
 
         do_denormalize = [True] * image.shape[0]
 
@@ -298,4 +320,4 @@ class HMImagePipeline(StableDiffusionImg2ImgPipeline):
         # Offload all models
         self.maybe_free_model_hooks()
 
-        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=None)
+        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=None), latents.detach().cpu() / self.vae.config.scaling_factor

hellomeme/pipelines/pipline_hm_video.py

@@ -8,6 +8,7 @@
 @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
@@ -22,45 +23,53 @@ from diffusers import StableDiffusionImg2ImgPipeline, MotionAdapter, EulerDiscre
 
 from ..models import HMDenoising3D, HMDenoisingMotion, HMControlNet, HMControlNet2
 from ..models import HMReferenceAdapter
-from ..utils import dicts_to_device, cat_dicts
-
+from ..utils import dicts_to_device, cat_dicts, cat_dicts_ref
 
 class HMVideoPipeline(StableDiffusionImg2ImgPipeline):
     def caryomitosis(self, patch_frames=12, **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()
 
-        if patch_frames >= 14:
-            self.num_frames = 16
-            adapter = MotionAdapter.from_pretrained("songkey/hm_animatediff", torch_dtype=torch.float16)
-        else:
-            self.num_frames = 12
-            adapter = MotionAdapter.from_pretrained("songkey/hm_animatediff_frame12", torch_dtype=torch.float16)
+        self.num_frames = 12
+        adapter = MotionAdapter.from_pretrained("songkey/hm_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
 
-    def insert_hm_modules(self, dtype, device):
+        self.vae_decode = copy.deepcopy(self.vae)
+
+    def insert_hm_modules(self, dtype):
         hm_adapter = HMReferenceAdapter.from_pretrained('songkey/hm_reference')
         if isinstance(self.unet, HMDenoisingMotion):
             self.unet.insert_reference_adapter(hm_adapter)
-            self.unet = self.unet.to(device=device, dtype=dtype).eval()
+            self.unet.to(device='cpu', dtype=dtype).eval()
 
         if hasattr(self, "unet_ref"):
             self.unet_ref.insert_reference_adapter(hm_adapter)
-            self.unet_ref = self.unet_ref.to(device=device, dtype=dtype).eval()
+            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')
-        self.mp_control = self.mp_control.to(device=device, dtype=dtype).eval()
+        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')
-        self.mp_control2 = self.mp_control2.to(device=device, dtype=dtype).eval()
+        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()
     def __call__(
@@ -82,6 +91,7 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline):
             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,
@@ -137,12 +147,13 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline):
         else:
             batch_size = prompt_embeds.shape[0]
 
-        device = self.device
+        # 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.to(device=device)
         prompt_embeds, negative_prompt_embeds = self.encode_prompt(
             prompt,
             device,
@@ -154,6 +165,7 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline):
             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
@@ -172,20 +184,16 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline):
         # 4. Preprocess image
         image = self.image_processor.preprocess(image).to(device=device, dtype=prompt_embeds.dtype)
 
-        # 5. set timesteps
-        timesteps, num_inference_steps = retrieve_timesteps(
-            self.scheduler, num_inference_steps, device, timesteps, sigmas
-        )
-        timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device)
-        latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)
-
         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:]
@@ -201,14 +209,6 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline):
         )
 
         ref_latents = ref_latents.unsqueeze(2)
-        self.unet_ref.to(device=device)
-        cached_res = self.unet_ref(
-            torch.cat([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]
 
         # 7.2 Optionally get Guidance Scale Embedding
         timestep_cond = None
@@ -218,12 +218,23 @@ 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(
+            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 = []
         control_latent1_neg, control_latent2_neg = None, None
         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)
@@ -266,9 +277,7 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline):
                     beta_schedule="scaled_linear",
                 )
 
-                tmp_timesteps, _ = retrieve_timesteps(
-                    scheduler, 8, device, None, sigmas
-                )
+                tmp_timesteps, _ = retrieve_timesteps(scheduler, 8, device, None, sigmas)
                 if idx == 0:
                     pred_latent = scheduler.add_noise(ref_latents, base_noise, tmp_timesteps[:1])
                 else:
@@ -278,7 +287,7 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline):
                     latent_model_input = torch.cat([pred_latent, pred_latent], dim=0) if \
                         self.do_classifier_free_guidance else pred_latent
                     latent_model_input = scheduler.scale_model_input(latent_model_input, t)
-                    noise_pred_chunk = self.unet_ref(
+                    noise_pred_chunk = self.unet_pre(
                         latent_model_input,
                         t,
                         encoder_hidden_states=prompt_embeds,
@@ -303,6 +312,7 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline):
         pre_latents = torch.cat(pre_latents, dim=2)
         self.mp_control.cpu()
         self.mp_control2.cpu()
+        self.unet_pre.cpu()
         self.unet_ref.cpu()
 
         chunk_size = self.num_frames
@@ -328,20 +338,25 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline):
 
         drive_idx_chunks[-1] = list(range(drive_idx_chunks[-1][0], paded_video_len))
 
-        # 5. set timesteps
-        timesteps, num_inference_steps = retrieve_timesteps(
-            self.scheduler, num_inference_steps, device, None, sigmas
+        scheduler = EulerDiscreteScheduler(
+            num_train_timesteps=1000,
+            beta_start=0.00085,
+            beta_end=0.012,
+            beta_schedule="scaled_linear",
         )
-        timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device)
+
+        # 5. set timesteps
+        timesteps, num_inference_steps = retrieve_timesteps(scheduler, num_inference_steps, device, None, sigmas)
         latent_timestep = timesteps[:1].repeat(batch_size * num_images_per_prompt)
 
         # 8. Denoising loop
-        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        num_warmup_steps = len(timesteps) - num_inference_steps * scheduler.order
         self._num_timesteps = len(timesteps)
 
         base_noise = randn_tensor([batch_size, c, paded_video_len, h, w], dtype=prompt_embeds.dtype, generator=generator)
-        latents = self.scheduler.add_noise(pre_latents, base_noise, latent_timestep)
+        latents = scheduler.add_noise(pre_latents, base_noise, latent_timestep)
 
+        self.unet.to(device=device)
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):
                 noise_pred = torch.zeros([batch_size * 2 if self.do_classifier_free_guidance
@@ -351,12 +366,13 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline):
                 latent_model_input_all = torch.cat([latents.clone(), latents.clone()], dim=0) if \
                     self.do_classifier_free_guidance else latents.clone()
 
-                latent_model_input_all = self.scheduler.scale_model_input(latent_model_input_all, t)
+                latent_model_input_all = scheduler.scale_model_input(latent_model_input_all, t)
                 for cidx, chunk in enumerate(drive_idx_chunks):
                     if self.interrupt:
                         continue
 
                     control_latent = cat_dicts([control_latents[ix] for ix in chunk], dim=2)
+                    # cached_res = cat_dicts_ref([cached_latents[ix] for ix in chunk])
                     latent_model_input = latent_model_input_all[:, :, chunk].to(device=device)
 
                     # predict the noise residual
@@ -382,7 +398,7 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline):
                     noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
 
                 # compute the previous noisy sample x_t -> x_t-1
-                latents = self.scheduler.step(noise_pred.cpu(), t, latents,
+                latents = scheduler.step(noise_pred.cpu(), t, latents,
                                               **extra_step_kwargs, return_dict=False)[0]
 
                 if callback_on_step_end is not None:
@@ -396,22 +412,24 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline):
                     negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
 
                 # call the callback, if provided
-                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % scheduler.order == 0):
                     progress_bar.set_description(f"GEN stage2")
                     progress_bar.update()
                     if callback is not None and i % callback_steps == 0:
-                        step_idx = i // getattr(self.scheduler, "order", 1)
+                        step_idx = i // getattr(scheduler, "order", 1)
                         callback(step_idx, t, latents)
+        self.unet.cpu()
+        latents_res = latents / self.vae.config.scaling_factor
 
-        latents = latents / self.vae.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(latents[:, :, i, :, :].to(device=device),
+                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()
-        return res_frames
+        self.vae_decode.cpu()
+        return res_frames, latents_res

hellomeme/utils.py

@@ -18,10 +18,13 @@ import subprocess
 from einops import rearrange
 from .tools.utils import transform_points
 from .tools.hello_3dmm import get_project_points_rect
-from diffusers.pipelines.stable_diffusion.convert_from_ckpt import convert_ldm_unet_checkpoint
 
 from safetensors import safe_open
 
+def get_torch_device(gpu_id=-1):
+    if gpu_id < 0: return torch.device("cpu")
+    gpu_id = min(max(-1, gpu_id), torch.cuda.device_count())
+    return torch.device(f"cuda:{gpu_id}")
 
 def merge_dicts(dictl, dictr, wl=0.5):
     res = {}
@@ -36,6 +39,11 @@ def cat_dicts(dicts, dim=0):
         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 = []
@@ -55,11 +63,6 @@ def load_safetensors(model_path):
     return tensors
 
 
-def load_unet_from_safetensors(safetensors_path, unet_config):
-    original_stats = load_safetensors(safetensors_path)
-    return convert_ldm_unet_checkpoint(original_stats, unet_config)
-
-
 def image_preprocess(np_bgr, size, dtype=torch.float32):
     img_np = cv2.resize(np_bgr, size)
     return np_bgr_to_tensor(img_np, dtype)