Initialize latents correctly

animatediff/pipelines/pipeline_animation.py

@@ -328,17 +328,14 @@ class AnimationPipeline(DiffusionPipeline):
 
                 latents = latents.to(device)
             else:
+                latents = torch.randn(shape, generator=generator, device=rand_device, dtype=dtype).to(device)
+
                 # If init_latents is not None, repeat it for the entire batch
                 if init_latents is not None:
                     init_latents = init_latents.unsqueeze(0).repeat(batch_size, 1, 1, 1, 1)
                     for i in range(video_length):
-                        if init_latents is None:
-                            latents[:, :, i, :, :] = torch.randn(latents[:, :, i, :, :].shape, device=rand_device, dtype=dtype) * self.scheduler.init_noise_sigma
-                        else:
-                            init_alpha = (video_length - float(i)) / video_length / 30
-                            latents[:, :, i, :, :] = init_latents * init_alpha + latents[:, :, i, :, :] * (1 - init_alpha)
-                else:
-                    latents = torch.randn(shape, generator=generator, device=rand_device, dtype=dtype).to(device)
+                        init_alpha = (video_length - float(i)) / video_length / 30
+                        latents[:, :, i, :, :] = init_latents[:, :, i, :, :] * init_alpha + latents[:, :, i, :, :] * (1 - init_alpha)
 
         else:
             if latents.shape != shape:
@@ -353,6 +350,7 @@ class AnimationPipeline(DiffusionPipeline):
 
 
 
+
     @torch.no_grad()
     def __call__(
         self,