Support VHS animated previews

This commit is contained in:
kijai
2025-01-07 00:39:09 +02:00
parent 3acd21bec3
commit 11ee15a91e
4 changed files with 19 additions and 96 deletions
@@ -174,6 +174,7 @@ class HunyuanVideoPipeline(DiffusionPipeline):
self,
transformer: HYVideoDiffusionTransformer,
scheduler: KarrasDiffusionSchedulers,
comfy_model = None,
progress_bar_config: Dict[str, Any] = None,
base_dtype = torch.bfloat16,
):
@@ -187,6 +188,7 @@ class HunyuanVideoPipeline(DiffusionPipeline):
self._progress_bar_config.update(progress_bar_config)
self.base_dtype = base_dtype
self.comfy_model = comfy_model
# ==========================================================================================
self.register_modules(
@@ -630,8 +632,8 @@ class HunyuanVideoPipeline(DiffusionPipeline):
self._num_timesteps = len(timesteps)
# 8. Preview callback
from ....latent_preview import prepare_callback
callback = prepare_callback(self.transformer, num_inference_steps)
from latent_preview import prepare_callback
callback = prepare_callback(self.comfy_model, num_inference_steps)
#print(self.scheduler.sigmas)
-85
View File
@@ -1,85 +0,0 @@
import torch
from PIL import Image
from comfy.cli_args import args, LatentPreviewMethod
import comfy.model_management
import comfy.utils
MAX_PREVIEW_RESOLUTION = args.preview_size
def preview_to_image(latent_image):
latents_ubyte = (((latent_image + 1.0) / 2.0).clamp(0, 1) # change scale from -1..1 to 0..1
.mul(0xFF) # to 0..255
).to(device="cpu", dtype=torch.uint8, non_blocking=comfy.model_management.device_supports_non_blocking(latent_image.device))
return Image.fromarray(latents_ubyte.numpy())
class LatentPreviewer:
def decode_latent_to_preview(self, x0):
pass
def decode_latent_to_preview_image(self, preview_format, x0):
preview_image = self.decode_latent_to_preview(x0)
return ("JPEG", preview_image, MAX_PREVIEW_RESOLUTION)
class Latent2RGBPreviewer(LatentPreviewer):
def __init__(self):
latent_rgb_factors = [[-0.0395, -0.0331, 0.0445],
[ 0.0696, 0.0795, 0.0518],
[ 0.0135, -0.0945, -0.0282],
[ 0.0108, -0.0250, -0.0765],
[-0.0209, 0.0032, 0.0224],
[-0.0804, -0.0254, -0.0639],
[-0.0991, 0.0271, -0.0669],
[-0.0646, -0.0422, -0.0400],
[-0.0696, -0.0595, -0.0894],
[-0.0799, -0.0208, -0.0375],
[ 0.1166, 0.1627, 0.0962],
[ 0.1165, 0.0432, 0.0407],
[-0.2315, -0.1920, -0.1355],
[-0.0270, 0.0401, -0.0821],
[-0.0616, -0.0997, -0.0727],
[ 0.0249, -0.0469, -0.1703]]
self.latent_rgb_factors = torch.tensor(latent_rgb_factors, device="cpu").transpose(0, 1)
self.latent_rgb_factors_bias = torch.tensor([0.0259, -0.0192, -0.0761], device="cpu")
def decode_latent_to_preview(self, x0):
self.latent_rgb_factors = self.latent_rgb_factors.to(dtype=x0.dtype, device=x0.device)
if self.latent_rgb_factors_bias is not None:
self.latent_rgb_factors_bias = self.latent_rgb_factors_bias.to(dtype=x0.dtype, device=x0.device)
latent_image = torch.nn.functional.linear(x0[0].permute(1, 2, 0), self.latent_rgb_factors,
bias=self.latent_rgb_factors_bias)
return preview_to_image(latent_image)
def get_previewer():
previewer = None
method = args.preview_method
if method != LatentPreviewMethod.NoPreviews:
# TODO previewer method
if method == LatentPreviewMethod.Auto:
method = LatentPreviewMethod.Latent2RGB
if previewer is None:
previewer = Latent2RGBPreviewer()
return previewer
def prepare_callback(model, steps, x0_output_dict=None):
preview_format = "JPEG"
if preview_format not in ["JPEG", "PNG"]:
preview_format = "JPEG"
previewer = get_previewer()
pbar = comfy.utils.ProgressBar(steps)
def callback(step, x0, x, total_steps):
if x0_output_dict is not None:
x0_output_dict["x0"] = x0
preview_bytes = None
if previewer:
preview_bytes = previewer.decode_latent_to_preview_image(preview_format, x0)
pbar.update_absolute(step + 1, total_steps, preview_bytes)
return callback
+6 -3
View File
@@ -242,7 +242,7 @@ class HyVideoModelConfig:
def __init__(self, dtype):
self.unet_config = {}
self.unet_extra_config = {}
self.latent_format = comfy.latent_formats.LatentFormat()
self.latent_format = comfy.latent_formats.HunyuanVideo
self.latent_format.latent_channels = 16
self.manual_cast_dtype = dtype
self.sampling_settings = {"multiplier": 1.0}
@@ -332,7 +332,8 @@ class HyVideoModelLoader:
HyVideoModelConfig(base_dtype),
model_type=comfy.model_base.ModelType.FLOW,
device=device,
)
)
scheduler_config = {
"flow_shift": 9.0,
"reverse": True,
@@ -346,7 +347,8 @@ class HyVideoModelLoader:
transformer=transformer,
scheduler=scheduler,
progress_bar_config=None,
base_dtype=base_dtype
base_dtype=base_dtype,
comfy_model=comfy_model,
)
if not "torchao" in quantization:
@@ -362,6 +364,7 @@ class HyVideoModelLoader:
comfy_model.diffusion_model = transformer
patcher = comfy.model_patcher.ModelPatcher(comfy_model, device, offload_device)
pipe.comfy_model = patcher
del sd
gc.collect()
+9 -6
View File
@@ -90,6 +90,7 @@ class HyVideoInverseSampler:
CATEGORY = "HunyuanVideoWrapper"
def process(self, model, hyvid_embeds, flow_shift, steps, embedded_guidance_scale, seed, samples, gamma, start_step, end_step, gamma_trend, force_offload, interpolation_curve=None):
comfy_model_patcher = model
model = model.model
device = mm.get_torch_device()
offload_device = mm.unet_offload_device()
@@ -191,8 +192,8 @@ class HyVideoInverseSampler:
num_warmup_steps = len(timesteps) - steps * pipeline.scheduler.order
self._num_timesteps = len(timesteps)
from .latent_preview import prepare_callback
callback = prepare_callback(transformer, steps)
from latent_preview import prepare_callback
callback = prepare_callback(comfy_model_patcher, steps)
from comfy.utils import ProgressBar
from tqdm import tqdm
@@ -307,6 +308,7 @@ class HyVideoReSampler:
def process(self, model, hyvid_embeds, flow_shift, steps, embedded_guidance_scale,
samples, inversed_latents, force_offload, start_step, end_step, eta_base, eta_trend, interpolation_curve=None, feta_args=None):
comfy_model_patcher = model
model = model.model
device = mm.get_torch_device()
offload_device = mm.unet_offload_device()
@@ -377,8 +379,8 @@ class HyVideoReSampler:
# 7. Denoising loop
self._num_timesteps = len(timesteps)
from .latent_preview import prepare_callback
callback = prepare_callback(transformer, steps)
from latent_preview import prepare_callback
callback = prepare_callback(comfy_model_patcher, steps)
if feta_args is not None:
set_enhance_weight(feta_args["weight"])
@@ -512,6 +514,7 @@ class HyVideoPromptMixSampler:
def process(self, model, width, height, num_frames, hyvid_embeds, hyvid_embeds_2, flow_shift, steps, embedded_guidance_scale,
seed, force_offload, alpha, interpolation_curve=None, feta_args=None):
comfy_model_patcher = model
model = model.model
device = mm.get_torch_device()
offload_device = mm.unet_offload_device()
@@ -605,8 +608,8 @@ class HyVideoPromptMixSampler:
# 7. Denoising loop
self._num_timesteps = len(timesteps)
from .latent_preview import prepare_callback
callback = prepare_callback(transformer, steps)
from latent_preview import prepare_callback
callback = prepare_callback(comfy_model_patcher, steps)
from comfy.utils import ProgressBar
from tqdm import tqdm