add progress for musepose inference

musepose/pipelines/pipeline_pose2vid_long.py

@@ -3,6 +3,7 @@ import inspect
 import math
 from dataclasses import dataclass
 from typing import Callable, List, Optional, Union
+import gradio as gr
 
 import numpy as np
 import torch
@@ -53,6 +54,7 @@ class Pose2VideoPipeline(DiffusionPipeline):
         image_proj_model=None,
         tokenizer=None,
         text_encoder=None,
+        gradio_progress: gr.Progress = gr.Progress(),
     ):
         super().__init__()
 
@@ -77,6 +79,7 @@ class Pose2VideoPipeline(DiffusionPipeline):
             do_convert_rgb=True,
             do_normalize=False,
         )
+        self.gradio_progress = gradio_progress
 
     def enable_vae_slicing(self):
         self.vae.enable_slicing()
@@ -117,6 +120,8 @@ class Pose2VideoPipeline(DiffusionPipeline):
         video = []
         for frame_idx in tqdm(range(latents.shape[0])):
             video.append(self.vae.decode(latents[frame_idx : frame_idx + 1]).sample)
+            self.gradio_progress(frame_idx /latents.shape[0], "Writing Video..")
+
         video = torch.cat(video)
         video = rearrange(video, "(b f) c h w -> b c f h w", f=video_length)
         video = (video / 2 + 0.5).clamp(0, 1)
@@ -448,6 +453,8 @@ class Pose2VideoPipeline(DiffusionPipeline):
         num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):
+                self.gradio_progress(i/num_inference_steps, "Making Image Move..")
+
                 noise_pred = torch.zeros(
                     (
                         latents.shape[0] * (2 if do_classifier_free_guidance else 1),

musepose_inference.py

@@ -11,6 +11,7 @@ from transformers import CLIPVisionModelWithProjection
 import torch.nn.functional as F
 import gc
 from huggingface_hub import hf_hub_download
+import gradio as gr
 
 from musepose.models.pose_guider import PoseGuider
 from musepose.models.unet_2d_condition import UNet2DConditionModel
@@ -65,7 +66,8 @@ class MusePoseInference:
         seed: int,
         steps: int,
         fps: int,
-        skip: int
+        skip: int,
+        gradio_progress=gr.Progress()
     ):
         download_models(model_dir=self.model_dir)
         print(f"Model Paths: {self.musepose_model_paths}\n{self.image_gen_model_paths}\n{self.inference_config_path}")
@@ -144,6 +146,7 @@ class MusePoseInference:
             denoising_unet=self.denoising_unet,
             pose_guider=self.pose_guider,
             scheduler=scheduler,
+            gradio_progress=gradio_progress
         )
         self.pipe = self.pipe.to("cuda", dtype=weight_dtype)