diff --git a/hellomeme/__init__.py b/hellomeme/__init__.py index 632a0f7..8c06cdc 100644 --- a/hellomeme/__init__.py +++ b/hellomeme/__init__.py @@ -8,5 +8,5 @@ @Desc : """ -from .piplines import HMVideoPipeline, HMImagePipeline +from .pipelines import HMVideoPipeline, HMImagePipeline from .models import HMDenoising3D, HMDenoisingMotion, HMControlNet, HMReferenceAdapter \ No newline at end of file diff --git a/hellomeme/piplines/__init__.py b/hellomeme/pipelines/__init__.py similarity index 100% rename from hellomeme/piplines/__init__.py rename to hellomeme/pipelines/__init__.py diff --git a/hellomeme/piplines/pipline_hm_image.py b/hellomeme/pipelines/pipline_hm_image.py similarity index 92% rename from hellomeme/piplines/pipline_hm_image.py rename to hellomeme/pipelines/pipline_hm_image.py index 062a524..1b16611 100644 --- a/hellomeme/piplines/pipline_hm_image.py +++ b/hellomeme/pipelines/pipline_hm_image.py @@ -170,27 +170,15 @@ class HMImagePipeline(StableDiffusionImg2ImgPipeline): ref_latents = torch.cat(ref_latents, dim=0) ref_latents = self.vae.config.scaling_factor * ref_latents - # 6. Prepare latent variables - # latents = self.prepare_latents( - # ref_latents, - # latent_timestep, - # batch_size, - # num_images_per_prompt, - # prompt_embeds.dtype, - # device, - # generator, - # ) + condition = drive_params['condition'].clone().to(device=device) + drive_coeff = drive_params['drive_coeff'].clone().to(device=device) + face_parts = drive_params['face_parts'].clone().to(device=device) - # condition, drive_coeff, face_parts if self.do_classifier_free_guidance: - drive_params['condition'] = torch.cat([torch.ones_like(drive_params['condition']) * -1, drive_params['condition']], dim=0) - drive_params['drive_coeff'] = torch.cat([torch.zeros_like(drive_params['drive_coeff']), drive_params['drive_coeff']], dim=0) - drive_params['face_parts'] = torch.cat([torch.zeros_like(drive_params['face_parts']), drive_params['face_parts']], dim=0) - - if hasattr(self, "mp_control"): - control_latents = self.mp_control(**drive_params) - else: - control_latents = None + condition = torch.cat([torch.ones_like(condition) * -1, condition], dim=0) + drive_coeff = torch.cat([torch.zeros_like(drive_coeff), drive_coeff], dim=0) + face_parts = torch.cat([torch.zeros_like(face_parts), face_parts], dim=0) + control_latents = self.mp_control(condition=condition, drive_coeff=drive_coeff, face_parts=face_parts) # 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) diff --git a/hellomeme/piplines/pipline_hm_video.py b/hellomeme/pipelines/pipline_hm_video.py similarity index 82% rename from hellomeme/piplines/pipline_hm_video.py rename to hellomeme/pipelines/pipline_hm_video.py index 344932c..6f502f5 100644 --- a/hellomeme/piplines/pipline_hm_video.py +++ b/hellomeme/pipelines/pipline_hm_video.py @@ -167,7 +167,7 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline): 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['drive_coeff'].size(1) + raw_video_len = drive_params['condition'].size(2) chunk_size = num_frames chunk_overlap = min(max(0, chunk_overlap), chunk_size // 2) chunk_stride = chunk_size - chunk_overlap @@ -210,7 +210,6 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline): self.do_classifier_free_guidance else ref_latents, 0, encoder_hidden_states=prompt_embeds, - # added_cond_kwargs=added_cond_kwargs, return_dict=False, )[1] # else: @@ -224,35 +223,26 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline): guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim ).to(device=device, dtype=prompt_embeds.dtype) - if self.do_classifier_free_guidance: - drive_params['condition'] = torch.cat( - [torch.ones_like(drive_params['condition']) * -1, drive_params['condition']], dim=0) - drive_params['drive_coeff'] = torch.cat( - [torch.zeros_like(drive_params['drive_coeff']), drive_params['drive_coeff']], dim=0) - drive_params['face_parts'] = torch.cat( - [torch.zeros_like(drive_params['face_parts']), drive_params['face_parts']], dim=0) + pre_latents = [] + control_latents = [] + control_latent_neg = None + base_noise = randn_tensor([batch_size, c, 1, h, w], dtype=prompt_embeds.dtype, generator=generator).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) + 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) - control_latents_chunks = [] - for cidx, chunk in enumerate(drive_idx_chunks): - condition_chunk = drive_params['condition'][:, :, chunk, :, :] - drive_coeff_chunk = drive_params['drive_coeff'][:, chunk, :] - face_parts_chunk = drive_params['face_parts'][:, chunk, :, :] + if control_latent_neg is None and self.do_classifier_free_guidance: + control_latent_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_latents_chunk = self.mp_control(condition=condition_chunk.to(device=device), - drive_coeff=drive_coeff_chunk.to(device=device), - face_parts=face_parts_chunk.to(device=device)) - control_latents_chunks.append(control_latents_chunk) - drive_idx_chunks[-1] = list(range(drive_idx_chunks[-1][0], paded_video_len)) - - with self.progress_bar(total=len(drive_idx_chunks)) as progress_bar: - chunk_latents = torch.zeros([batch_size, c, paded_video_len, h, w], dtype=prompt_embeds.dtype).cpu() - chunk_latents_counter = torch.zeros([1, 1, paded_video_len, 1, 1], dtype=prompt_embeds.dtype).cpu() - chunk_noise = randn_tensor([batch_size, c, 1, h, w], dtype=prompt_embeds.dtype, generator=generator).to( - device=device) - chunk_noise = chunk_noise.repeat_interleave(chunk_size, dim=2) - for cidx, chunk in enumerate(drive_idx_chunks): - control_latents_chunk = control_latents_chunks[cidx] + control_latent = self.mp_control(condition=condition, + drive_coeff=drive_coeff, + face_parts=face_parts) + if self.do_classifier_free_guidance: + control_latent = cat_dicts([control_latent_neg, control_latent], dim=0) scheduler = EulerDiscreteScheduler( num_train_timesteps=1000, @@ -264,18 +254,22 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline): tmp_timesteps, _ = retrieve_timesteps( scheduler, 8, device, None, sigmas ) - # chunk_latent = chunk_noise * self.scheduler.init_noise_sigma - chunk_latent = scheduler.add_noise(ref_latents, chunk_noise, tmp_timesteps[:1]) + if idx == 0: + pred_latent = scheduler.add_noise(ref_latents, base_noise, tmp_timesteps[:1]) + else: + pred_latent = scheduler.add_noise(pred_latent, base_noise, + tmp_timesteps[:1]) + for i, t in enumerate(tmp_timesteps): - latent_model_input = torch.cat([chunk_latent, chunk_latent], dim=0) if \ - self.do_classifier_free_guidance else chunk_latent + 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( latent_model_input, t, encoder_hidden_states=prompt_embeds, reference_hidden_states=cached_res, - control_hidden_states=control_latents_chunk, + control_hidden_states=control_latent, timestep_cond=timestep_cond, cross_attention_kwargs=self.cross_attention_kwargs, added_cond_kwargs=added_cond_kwargs, @@ -284,13 +278,17 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline): if self.do_classifier_free_guidance: noise_pred_uncond, noise_pred_text = noise_pred_chunk.chunk(2) noise_pred_chunk = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond) - chunk_latent = scheduler.step(noise_pred_chunk, t, chunk_latent, - **extra_step_kwargs, return_dict=False)[0] + pred_latent = scheduler.step(noise_pred_chunk, t, pred_latent, + **extra_step_kwargs, return_dict=False)[0] + pre_latents.append(pred_latent.cpu()) + control_latents.append(dicts_to_device([control_latent], device='cpu')[0]) - chunk_latents[:, :, chunk] = chunk_latents[:, :, chunk] + chunk_latent.cpu() - chunk_latents_counter[:, :, chunk] += 1 - progress_bar.update() - chunk_latents = chunk_latents / chunk_latents_counter + pre_latents = torch.cat(pre_latents, dim=2) + + if post_pad_num > 0: + pre_latents = torch.cat([pre_latents, + (pre_latents[:, :, -2:-1]).clone().repeat_interleave(post_pad_num, 2)], dim=2) + control_latents = control_latents + [control_latents[-1]] * post_pad_num # 5. set timesteps timesteps, num_inference_steps = retrieve_timesteps( @@ -304,9 +302,9 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline): self._num_timesteps = len(timesteps) base_noise = randn_tensor([batch_size, c, paded_video_len, h, w], dtype=prompt_embeds.dtype, generator=generator) - # ref_latents = ref_latents.repeat_interleave(paded_video_len, dim=2) + latents = self.scheduler.add_noise(pre_latents, base_noise, latent_timestep) - latents = self.scheduler.add_noise(chunk_latents.to(device=device), base_noise.to(device=device), latent_timestep).cpu() + drive_idx_chunks[-1] = list(range(drive_idx_chunks[-1][0], paded_video_len)) with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): @@ -322,7 +320,7 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline): if self.interrupt: continue - control_latents = control_latents_chunks[cidx] + control_latent = cat_dicts([control_latents[ix] for ix in chunk], dim=2) latent_model_input = latent_model_input_all[:, :, chunk].to(device=device) # predict the noise residual @@ -331,7 +329,7 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline): t, encoder_hidden_states=prompt_embeds, reference_hidden_states=cached_res, - control_hidden_states=control_latents, + control_hidden_states=dicts_to_device([control_latent], device=device)[0], timestep_cond=timestep_cond, cross_attention_kwargs=self.cross_attention_kwargs, added_cond_kwargs=added_cond_kwargs, @@ -373,4 +371,19 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline): res_frames = [self.vae.decode(latents[:,:,i,:,:].to(device=device), return_dict=False, generator=generator)[0] for i in range(raw_video_len)] res_frames = [self.image_processor.postprocess(frame, output_type=output_type) for frame in res_frames] - return res_frames \ No newline at end of file + return res_frames + +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 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 \ No newline at end of file diff --git a/hellomeme/utils.py b/hellomeme/utils.py index 20f3ca8..d3e65cb 100644 --- a/hellomeme/utils.py +++ b/hellomeme/utils.py @@ -146,7 +146,7 @@ def get_face_parts(images, landmarks, save_size=256): ret_list.append(np.stack(tmp_part_list, axis=0)) return np.stack(ret_list, axis=0).transpose(1, 0, 2, 3, 4) -def get_face_params(h3dmm, harkit_bs, frames, landmarks, save_size=512, align=True): +def get_face_params(h3dmm, harkit_bs, frames, landmarks, save_size=(512, 512), align=True): face_parts = [] for tmp_face_parts_list in get_face_parts(frames, landmarks, save_size=256): tmp_face_parts_list = [cv2.GaussianBlur(x, (9, 9), 0) for x in tmp_face_parts_list] @@ -223,22 +223,12 @@ def det_landmarks(face_aligner, frame_list, save_size=(512, 512), reset=False): face_aligner.reset_track() return save_frame_list, save_landmark_list -def get_drive_params(face_aligner, h3dmm, harkit_bs, frame_list, save_size, align): +def get_drive_params(face_aligner, h3dmm, harkit_bs, frame_list, save_size, align=True): frame_num = len(frame_list) frame_list, landmark_list = det_landmarks(face_aligner, frame_list, save_size=(512, 512), reset=False) assert len(frame_list) == frame_num return get_face_params(h3dmm, harkit_bs, frame_list, landmark_list, save_size=save_size, align=align) -def load_data_list(data_dir, post_fix='.pickle;.txt'): - post_fixs = post_fix.split(';') - ret_list = [] - for root, dirnames, filenames in os.walk(data_dir): - for name in filenames: - if os.path.splitext(name)[1].lower() in post_fixs and not name.startswith('.'): - data_path = os.path.join(root, name) - ret_list.append(data_path) - return ret_list - def crop_and_resize(frames, landmarks, save_size=512, crop=True): H, W = frames[0].shape[:2] if crop: @@ -376,4 +366,14 @@ def ff_change_fps(input_video, output_video, fps=15): output_video ] - subprocess.run(cmd) \ No newline at end of file + subprocess.run(cmd) + +def load_data_list(data_dir, post_fix='.pickle;.txt'): + post_fixs = post_fix.split(';') + ret_list = [] + for root, dirnames, filenames in os.walk(data_dir): + for name in filenames: + if os.path.splitext(name)[1] in post_fixs and not name.startswith('.'): + data_path = os.path.join(root, name) + ret_list.append(data_path) + return ret_list \ No newline at end of file