From 1d67d090dad810a5fa297a91cbcaa485f9a3d6ab Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Mon, 6 Jan 2025 11:38:40 +0200
Subject: [PATCH] Fix preview

---
 hyvideo/diffusion/pipelines/pipeline_hunyuan_video.py | 10 ++++++++--
 latent_preview.py                                     |  2 +-
 nodes.py                                              |  3 ++-
 3 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/hyvideo/diffusion/pipelines/pipeline_hunyuan_video.py b/hyvideo/diffusion/pipelines/pipeline_hunyuan_video.py
index 4b4706e..02810dd 100644
--- a/hyvideo/diffusion/pipelines/pipeline_hunyuan_video.py
+++ b/hyvideo/diffusion/pipelines/pipeline_hunyuan_video.py
@@ -633,11 +633,12 @@ class HunyuanVideoPipeline(DiffusionPipeline):
         from ....latent_preview import prepare_callback
         callback = prepare_callback(self.transformer, num_inference_steps)
 
+        print(self.scheduler.sigmas)
+
         
         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
@@ -789,7 +790,12 @@ class HunyuanVideoPipeline(DiffusionPipeline):
                     if progress_bar is not None:
                         progress_bar.update()
                     if callback is not None:
-                        callback(i, latents.detach()[-1].permute(1,0,2,3), None, num_inference_steps)
+                        callback(
+                            i, 
+                            (latent_model_input - noise_pred).detach()[0].permute(1,0,2,3),
+                            None, 
+                            num_inference_steps
+                        )
                     else:
                         comfy_pbar.update(1)
 
diff --git a/latent_preview.py b/latent_preview.py
index 1a6d5dc..a79c627 100644
--- a/latent_preview.py
+++ b/latent_preview.py
@@ -20,7 +20,7 @@ class LatentPreviewer:
 
     def decode_latent_to_preview_image(self, preview_format, x0):
         preview_image = self.decode_latent_to_preview(x0)
-        return ("GIF", preview_image, MAX_PREVIEW_RESOLUTION)
+        return ("JPEG", preview_image, MAX_PREVIEW_RESOLUTION)
 
 class Latent2RGBPreviewer(LatentPreviewer):
     def __init__(self):
diff --git a/nodes.py b/nodes.py
index ec45313..856dd91 100644
--- a/nodes.py
+++ b/nodes.py
@@ -334,7 +334,7 @@ class HyVideoModelLoader:
             "prediction_type": 'flow_prediction'
         }
         scheduler = FlowMatchDiscreteScheduler.from_config(scheduler_config)
-        print(scheduler.config)
+        
         pipe = HunyuanVideoPipeline(
             transformer=transformer,
             scheduler=scheduler,
@@ -1164,6 +1164,7 @@ class HyVideoSampler:
 
         model["scheduler_config"]["flow_shift"] = flow_shift
         model["scheduler_config"]["algorithm_type"] = "sde-dpmsolver++"
+        #model["scheduler_config"]["use_beta_sigmas"] = True
         
         noise_scheduler = scheduler_mapping[scheduler].from_config(model["scheduler_config"])
         model["pipe"].scheduler = noise_scheduler