optimize inter-frame continuity

hellomeme/pipelines/pipline_hm_video.py

@@ -287,15 +287,18 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline):
                     control_latent_input.update(control_latent2)
 
                 scheduler = EulerDiscreteScheduler(
-                    num_train_timesteps=1000,
-                    beta_start=0.00085,
-                    beta_end=0.012,
-                    beta_schedule="scaled_linear",
+                                num_train_timesteps=1000,
+                                beta_start=0.00085,
+                                beta_end=0.012,
+                                beta_schedule="scaled_linear",
+                            )
+
+                tmp_timesteps, tmp_steps = retrieve_timesteps(
+                    scheduler, 10, device, timesteps, 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])
+                    pred_latent = base_noise * scheduler.init_noise_sigma
                 else:
                     pred_latent = scheduler.add_noise(pred_latent, base_noise, tmp_timesteps[:1])
 
@@ -353,23 +356,19 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline):
 
         drive_idx_chunks[-1] = list(range(drive_idx_chunks[-1][0], paded_video_len))
 
-        scheduler = EulerDiscreteScheduler(
-            num_train_timesteps=1000,
-            beta_start=0.00085,
-            beta_end=0.012,
-            beta_schedule="scaled_linear",
-        )
-
         # 5. set timesteps
-        timesteps, num_inference_steps = retrieve_timesteps(scheduler, num_inference_steps, device, None, sigmas)
+        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)
 
         # 8. Denoising loop
-        num_warmup_steps = len(timesteps) - num_inference_steps * scheduler.order
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.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 = scheduler.add_noise(pre_latents, base_noise, latent_timestep)
+        latents = self.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:
@@ -381,7 +380,7 @@ 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 = scheduler.scale_model_input(latent_model_input_all, t)
+                latent_model_input_all = self.scheduler.scale_model_input(latent_model_input_all, t)
                 for cidx, chunk in enumerate(drive_idx_chunks):
                     if self.interrupt:
                         continue
@@ -413,7 +412,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 = scheduler.step(noise_pred.cpu(), t, latents,
+                latents = self.scheduler.step(noise_pred.cpu(), t, latents,
                                               **extra_step_kwargs, return_dict=False)[0]
 
                 if callback_on_step_end is not None:
@@ -427,11 +426,11 @@ 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) % scheduler.order == 0):
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.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(scheduler, "order", 1)
+                        step_idx = i // getattr(self.scheduler, "order", 1)
                         callback(step_idx, t, latents)
         self.unet.cpu()
         latents_res = latents / self.vae.config.scaling_factor

hellomeme/utils.py

@@ -308,7 +308,7 @@ def crop_and_resize(frames, landmarks, save_size=512, crop=True):
         all_tl, all_br = np.min(landmarks, axis=1), np.max(landmarks, axis=1)
         mean_wh = np.mean(all_br - all_tl, axis=0)
         tl, br = np.min(all_tl, axis=0), np.max(all_br, axis=0)
-        new_size = min(max(mean_wh) * 2.2, min(H, W) - 1)
+        new_size = min(max(mean_wh) * 2.5, min(H, W) - 1)
         fcenter = (tl + br) * 0.5
         ftl = fcenter - new_size * 0.5
         ftl[1] -= mean_wh[1] * 0.2