From 783d1114594490567b6d3e861823a4d7abdf2234 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Mon, 6 Jan 2025 06:46:53 +0200
Subject: [PATCH] fix

---
 .../pipelines/pipeline_hunyuan_video.py       | 19 ++++++++++++++++---
 hyvideo/modules/__init__.py                   |  4 +---
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/hyvideo/diffusion/pipelines/pipeline_hunyuan_video.py b/hyvideo/diffusion/pipelines/pipeline_hunyuan_video.py
index 416e5c7..3896b69 100644
--- a/hyvideo/diffusion/pipelines/pipeline_hunyuan_video.py
+++ b/hyvideo/diffusion/pipelines/pipeline_hunyuan_video.py
@@ -29,6 +29,7 @@ from diffusers.utils import (
 )
 from diffusers.utils.torch_utils import randn_tensor
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.schedulers import DPMSolverMultistepScheduler
 
 from ...modules import HYVideoDiffusionTransformer
 from comfy.utils import ProgressBar
@@ -636,6 +637,7 @@ class HunyuanVideoPipeline(DiffusionPipeline):
         logger.info(f"Sampling {video_length} frames in {latents.shape[2]} latents at {width}x{height} with {len(timesteps)} inference steps")
         comfy_pbar = ProgressBar(len(timesteps))
         with self.progress_bar(total=len(timesteps)) as progress_bar:
+            old_pred_original_sample = None # for DPM-solver++
             for i, t in enumerate(timesteps):
                 if self.interrupt:
                     continue
@@ -764,9 +766,20 @@ class HunyuanVideoPipeline(DiffusionPipeline):
                         )
 
                 # compute the previous noisy sample x_t -> x_t-1
-                latents = self.scheduler.step(
-                    noise_pred, t, latents, **extra_step_kwargs, return_dict=False
-                )[0]
+                if not isinstance(self.scheduler, DPMSolverMultistepScheduler):
+                    latents = self.scheduler.step(
+                        noise_pred, t, latents, **extra_step_kwargs, return_dict=False
+                    )[0]
+                else:
+                    latents, old_pred_original_sample = self.scheduler.step(
+                        noise_pred,
+                        old_pred_original_sample,
+                        t,
+                        timesteps[i - 1] if i > 0 else None,
+                        latents,
+                        **extra_step_kwargs,
+                        return_dict=False,
+                    )
 
                 if callback_on_step_end is not None:
                     callback_kwargs = {}
diff --git a/hyvideo/modules/__init__.py b/hyvideo/modules/__init__.py
index a95842b..e2e7d99 100644
--- a/hyvideo/modules/__init__.py
+++ b/hyvideo/modules/__init__.py
@@ -1,3 +1 @@
-from .models import HYVideoDiffusionTransformer, HUNYUAN_VIDEO_CONFIG
-
-
+from .models import HYVideoDiffusionTransformer