From fbbfc818ea3576e5629abc405c828c79c58c912a Mon Sep 17 00:00:00 2001
From: Bubbliiiing <47347516+bubbliiiing@users.noreply.github.com>
Date: Fri, 12 Jul 2024 17:47:47 +0800
Subject: [PATCH] Update pipeline (#50)

---
 easyanimate/pipeline/pipeline_easyanimate.py         | 7 +++++++
 easyanimate/pipeline/pipeline_easyanimate_inpaint.py | 7 +++++++
 2 files changed, 14 insertions(+)

diff --git a/easyanimate/pipeline/pipeline_easyanimate.py b/easyanimate/pipeline/pipeline_easyanimate.py
index ef53107..dc7aacf 100644
--- a/easyanimate/pipeline/pipeline_easyanimate.py
+++ b/easyanimate/pipeline/pipeline_easyanimate.py
@@ -633,6 +633,7 @@ class EasyAnimatePipeline(DiffusionPipeline):
         callback_steps: int = 1,
         clean_caption: bool = True,
         max_sequence_length: int = 120,
+        comfyui_progressbar: bool = False,
         **kwargs,
     ) -> Union[EasyAnimatePipelineOutput, Tuple]:
         """
@@ -783,6 +784,9 @@ class EasyAnimatePipeline(DiffusionPipeline):
         # 7. Denoising loop
         num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
 
+        if comfyui_progressbar:
+            from comfy.utils import ProgressBar
+            pbar = ProgressBar(num_inference_steps)
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):
                 latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
@@ -834,6 +838,9 @@ class EasyAnimatePipeline(DiffusionPipeline):
                         step_idx = i // getattr(self.scheduler, "order", 1)
                         callback(step_idx, t, latents)
 
+                if comfyui_progressbar:
+                    pbar.update(1)
+
         # Post-processing
         video = self.decode_latents(latents)
 
diff --git a/easyanimate/pipeline/pipeline_easyanimate_inpaint.py b/easyanimate/pipeline/pipeline_easyanimate_inpaint.py
index cf80f91..322f256 100644
--- a/easyanimate/pipeline/pipeline_easyanimate_inpaint.py
+++ b/easyanimate/pipeline/pipeline_easyanimate_inpaint.py
@@ -779,6 +779,8 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline):
         max_sequence_length: int = 120,
         clip_image: Image = None,
         clip_apply_ratio: float = 0.50,
+        comfyui_progressbar: bool = False,
+        **kwargs,
     ) -> Union[EasyAnimatePipelineOutput, Tuple]:
         """
         Function invoked when calling the pipeline for generation.
@@ -1061,6 +1063,9 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline):
         # 10. Denoising loop
         num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
         self._num_timesteps = len(timesteps)
+        if comfyui_progressbar:
+            from comfy.utils import ProgressBar
+            pbar = ProgressBar(num_inference_steps)
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):
                 latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
@@ -1130,6 +1135,8 @@ class EasyAnimateInpaintPipeline(DiffusionPipeline):
                         step_idx = i // getattr(self.scheduler, "order", 1)
                         callback(step_idx, t, latents)
 
+                if comfyui_progressbar:
+                    pbar.update(1)
         gc.collect()
         torch.cuda.empty_cache()
         torch.cuda.ipc_collect()