Update src/pipelines/pipeline_echo_mimic.py

src/pipelines/pipeline_echo_mimic.py

@@ -34,6 +34,7 @@ from transformers import CLIPImageProcessor
 from src.models.mutual_self_attention import ReferenceAttentionControl
 from src.pipelines.context import get_context_scheduler
 from src.pipelines.utils import get_tensor_interpolation_method
+from src.utils.step_func import origin_by_velocity_and_sample, psuedo_velocity_wrt_noisy_and_timestep, get_alpha
 
 @dataclass
 class Audio2VideoPipelineOutput(BaseOutput):
@@ -417,9 +418,9 @@ class Audio2VideoPipeline(DiffusionPipeline):
             generator
         )
         # print(video_length, latents.shape)
-        face_locator_tensor = self.face_locator(face_mask_tensor)
-        uc_face_locator_tensor = torch.zeros_like(face_locator_tensor)
-        face_locator_tensor = torch.cat([uc_face_locator_tensor, face_locator_tensor], dim=0)
+        c_face_locator_tensor = self.face_locator(face_mask_tensor)
+        uc_face_locator_tensor = torch.zeros_like(c_face_locator_tensor)
+        face_locator_tensor = torch.cat([uc_face_locator_tensor, c_face_locator_tensor], dim=0)
         # Prepare extra step kwargs.
         extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
 
@@ -474,7 +475,7 @@ class Audio2VideoPipeline(DiffusionPipeline):
                         encoder_hidden_states=None,
                         return_dict=False,
                     )
-                    reference_control_reader.update(reference_control_writer, do_classifier_free_guidance=True)
+                    reference_control_reader.update(reference_control_writer, do_classifier_free_guidance=do_classifier_free_guidance)
 
 
                 num_context_batches = math.ceil(len(context_queue) / context_batch_size)
@@ -498,8 +499,8 @@ class Audio2VideoPipeline(DiffusionPipeline):
                         .to(device)
                         .repeat(2 if do_classifier_free_guidance else 1, 1, 1, 1, 1)
                     )
-                    audio_latents = torch.cat([audio_fea_final[:, c] for c in new_context]).to(device)
-                    audio_latents = torch.cat([torch.zeros_like(audio_latents), audio_latents], 0)
+                    c_audio_latents = torch.cat([audio_fea_final[:, c] for c in new_context]).to(device)
+                    audio_latents = torch.cat([torch.zeros_like(c_audio_latents), c_audio_latents], 0)
 
                     latent_model_input = self.scheduler.scale_model_input(
                         latent_model_input, t
@@ -508,11 +509,15 @@ class Audio2VideoPipeline(DiffusionPipeline):
                         latent_model_input,
                         t,
                         encoder_hidden_states=None,
-                        audio_cond_fea=audio_latents,
-                        face_musk_fea=face_locator_tensor,
+                        audio_cond_fea=audio_latents if do_classifier_free_guidance else c_audio_latents,
+                        face_musk_fea=face_locator_tensor if do_classifier_free_guidance else c_face_locator_tensor,
                         return_dict=False,
                     )[0]
 
+                    alphas_cumprod = self.scheduler.alphas_cumprod.to(latent_model_input.device)
+                    x_pred = origin_by_velocity_and_sample(pred, latent_model_input, alphas_cumprod, t)
+                    pred = psuedo_velocity_wrt_noisy_and_timestep(latent_model_input, x_pred, alphas_cumprod, t, torch.ones_like(t) * (-1))
+
                     for j, c in enumerate(new_context):
                         noise_pred[:, :, c] = noise_pred[:, :, c] + pred
                         counter[:, :, c] = counter[:, :, c] + 1
@@ -523,6 +528,8 @@ class Audio2VideoPipeline(DiffusionPipeline):
                     noise_pred = noise_pred_uncond + guidance_scale * (
                         noise_pred_text - noise_pred_uncond
                     )
+                else:
+                    noise_pred = noise_pred / counter
 
                 latents = self.scheduler.step(
                     noise_pred, t, latents, **extra_step_kwargs
@@ -583,4 +590,4 @@ class Audio2VideoPipeline(DiffusionPipeline):
         smoothed_tensor = torch.cat(
             [tensor[:, :, 0:1, :, :], internal_frames, tensor[:, :, -1:, :, :]], dim=2)
 
-        return smoothed_tensor
+        return smoothed_tensor