Update pipeline.py

pipeline.py

@@ -777,6 +777,7 @@ class AnimateDiffPipeline(DiffusionPipeline, TextualInversionLoaderMixin, IPAdap
         init_image: Optional[PipelineImageInput] = None,
         init_image_strength: Optional[float] = 1.0,
         init_noise_correlation: Optional[float] = 0.0,
+        latent_mode: Optional[str] = "normal",
     ):
         r"""
         The call function to the pipeline for generation.
@@ -905,44 +906,47 @@ class AnimateDiffPipeline(DiffusionPipeline, TextualInversionLoaderMixin, IPAdap
         
         # 5. Prepare latent variables
         num_channels_latents = self.unet.config.in_channels
-        # latents = self.prepare_latents(
-        #     batch_size * num_videos_per_prompt,
-        #     num_channels_latents,
-        #     num_frames,
-        #     height,
-        #     width,
-        #     prompt_embeds.dtype,
-        #     device,
-        #     generator,
-        #     latents,
-        # )
-        latents = self.prepare_motion_latents(
-            batch_size * num_videos_per_prompt,
-            num_channels_latents,
-            num_frames,
-            height,
-            width,
-            prompt_embeds.dtype,
-            device,
-            generator,
-            latents,
-            x_velocity=x_velocity,
-            y_velocity=y_velocity,
-            scale_velocity=scale_velocity,
-        )
-        # latents, init_latents = self.prepare_correlated_latents(
-        #     init_image,
-        #     init_image_strength,
-        #     init_noise_correlation,
-        #     batch_size,
-        #     num_channels_latents,
-        #     num_frames,
-        #     height,
-        #     width,
-        #     prompt_embeds.dtype,
-        #     device,
-        #     generator,
-        # )
+        if(latent_mode == "normal"):
+            latents = self.prepare_latents(
+                batch_size * num_videos_per_prompt,
+                num_channels_latents,
+                num_frames,
+                height,
+                width,
+                prompt_embeds.dtype,
+                device,
+                generator,
+                latents,
+            )
+        elif(latent_mode == "motion"):
+            latents = self.prepare_motion_latents(
+                batch_size * num_videos_per_prompt,
+                num_channels_latents,
+                num_frames,
+                height,
+                width,
+                prompt_embeds.dtype,
+                device,
+                generator,
+                latents,
+                x_velocity=x_velocity,
+                y_velocity=y_velocity,
+                scale_velocity=scale_velocity,
+            )
+        elif(latent_mode == "correlated"):        
+            latents, init_latents = self.prepare_correlated_latents(
+                init_image,
+                init_image_strength,
+                init_noise_correlation,
+                batch_size,
+                num_channels_latents,
+                num_frames,
+                height,
+                width,
+                prompt_embeds.dtype,
+                device,
+                generator,
+            )
 
 
         # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline