diff --git a/.gitignore b/.gitignore index d75870d..fe5defd 100644 --- a/.gitignore +++ b/.gitignore @@ -7,5 +7,4 @@ master_ip logs/ *.DS_Store .idea -*.pt tools/ \ No newline at end of file diff --git a/__init__.py b/__init__.py index f4e09a9..a7bf0b9 100644 --- a/__init__.py +++ b/__init__.py @@ -1,4 +1,7 @@ -from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS +from .nodes import NODE_CLASS_MAPPINGS as NODES_CLASS, NODE_DISPLAY_NAME_MAPPINGS as NODES_DISPLAY +from .nodes_rf_inversion import NODE_CLASS_MAPPINGS as NODE_CLASS_MAPPINGS_RF_INVERSION, NODE_DISPLAY_NAME_MAPPINGS as NODE_DISPLAY_NAME_MAPPINGS_RF_INVERSION +NODE_CLASS_MAPPINGS = {**NODES_CLASS, **NODE_CLASS_MAPPINGS_RF_INVERSION} +NODE_DISPLAY_NAME_MAPPINGS = {**NODES_DISPLAY, **NODE_DISPLAY_NAME_MAPPINGS_RF_INVERSION} __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file diff --git a/examples/hunhyuan_rf_inversion_testing_01.json b/examples/hunhyuan_rf_inversion_testing_01.json new file mode 100644 index 0000000..34d24f1 --- /dev/null +++ b/examples/hunhyuan_rf_inversion_testing_01.json @@ -0,0 +1,1126 @@ +{ + "last_node_id": 82, + "last_link_id": 154, + "nodes": [ + { + "id": 60, + "type": "VHS_LoadVideo", + "pos": [ + -528.7080688476562, + 58.57748794555664 + ], + "size": [ + 247.455078125, + 503.455078125 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [ + { + "name": "meta_batch", + "type": "VHS_BatchManager", + "link": null, + "shape": 7 + }, + { + "name": "vae", + "type": "VAE", + "link": null, + "shape": 7 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 75 + ], + "slot_index": 0 + }, + { + "name": "frame_count", + "type": "INT", + "links": [], + "slot_index": 1 + }, + { + "name": "audio", + "type": "AUDIO", + "links": null + }, + { + "name": "video_info", + "type": "VHS_VIDEOINFO", + "links": null + } + ], + "properties": { + "Node name for S&R": "VHS_LoadVideo" + }, + "widgets_values": { + "video": "wolf_interpolated.mp4", + "force_rate": 0, + "force_size": "Disabled", + "custom_width": 512, + "custom_height": 512, + "frame_load_cap": 49, + "skip_first_frames": 0, + "select_every_nth": 1, + "choose video to upload": "image", + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "frame_load_cap": 49, + "skip_first_frames": 0, + "force_rate": 0, + "filename": "wolf_interpolated.mp4", + "type": "input", + "format": "video/mp4", + "select_every_nth": 1 + }, + "muted": false + } + } + }, + { + "id": 7, + "type": "HyVideoVAELoader", + "pos": [ + -442.0136413574219, + -457.50823974609375 + ], + "size": [ + 466.5, + 82 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [ + { + "name": "compile_args", + "type": "COMPILEARGS", + "link": null, + "shape": 7 + } + ], + "outputs": [ + { + "name": "vae", + "type": "VAE", + "links": [ + 73, + 105 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "HyVideoVAELoader" + }, + "widgets_values": [ + "hyvid\\hunyuan_video_vae_bf16.safetensors", + "bf16" + ] + }, + { + "id": 71, + "type": "GetImageSizeAndCount", + "pos": [ + -66.3104019165039, + 135.95004272460938 + ], + "size": [ + 277.20001220703125, + 86 + ], + "flags": {}, + "order": 12, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 120 + } + ], + "outputs": [ + { + "name": "image", + "type": "IMAGE", + "links": [ + 117 + ], + "slot_index": 0 + }, + { + "name": "512 width", + "type": "INT", + "links": [], + "slot_index": 1 + }, + { + "name": "512 height", + "type": "INT", + "links": [], + "slot_index": 2 + }, + { + "name": "49 count", + "type": "INT", + "links": [], + "slot_index": 3 + } + ], + "properties": { + "Node name for S&R": "GetImageSizeAndCount" + }, + "widgets_values": [] + }, + { + "id": 72, + "type": "SetNode", + "pos": [ + 159.51576232910156, + 332.2814636230469 + ], + "size": [ + 210, + 58 + ], + "flags": { + "collapsed": true + }, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "link": 119 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 120 + ], + "slot_index": 0 + } + ], + "title": "Set_InputVideo", + "properties": { + "previousName": "InputVideo" + }, + "widgets_values": [ + "InputVideo" + ], + "color": "#2a363b", + "bgcolor": "#3f5159" + }, + { + "id": 81, + "type": "Reroute", + "pos": [ + 665.5242309570312, + -279.5227966308594 + ], + "size": [ + 75, + 26 + ], + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "", + "type": "*", + "link": 151 + } + ], + "outputs": [ + { + "name": "", + "type": "HYVIDEOMODEL", + "links": [ + 152, + 153 + ], + "slot_index": 0 + } + ], + "properties": { + "showOutputText": false, + "horizontal": false + } + }, + { + "id": 61, + "type": "HyVideoEncode", + "pos": [ + 269.31317138671875, + -177.6795196533203 + ], + "size": [ + 315, + 150 + ], + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "vae", + "type": "VAE", + "link": 73 + }, + { + "name": "image", + "type": "IMAGE", + "link": 117 + } + ], + "outputs": [ + { + "name": "samples", + "type": "LATENT", + "links": [ + 140, + 148 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "HyVideoEncode" + }, + "widgets_values": [ + true, + 64, + 256, + true + ] + }, + { + "id": 62, + "type": "ImageResizeKJ", + "pos": [ + -220.3394775390625, + 298.1787414550781 + ], + "size": [ + 315, + 266 + ], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "image", + "type": "IMAGE", + "link": 75 + }, + { + "name": "get_image_size", + "type": "IMAGE", + "link": null, + "shape": 7 + }, + { + "name": "width_input", + "type": "INT", + "link": null, + "widget": { + "name": "width_input" + }, + "shape": 7 + }, + { + "name": "height_input", + "type": "INT", + "link": null, + "widget": { + "name": "height_input" + }, + "shape": 7 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 119 + ], + "slot_index": 0 + }, + { + "name": "width", + "type": "INT", + "links": [], + "slot_index": 1 + }, + { + "name": "height", + "type": "INT", + "links": [], + "slot_index": 2 + } + ], + "properties": { + "Node name for S&R": "ImageResizeKJ" + }, + "widgets_values": [ + 512, + 512, + "lanczos", + false, + 16, + 0, + 0, + "disabled" + ] + }, + { + "id": 1, + "type": "HyVideoModelLoader", + "pos": [ + -473.4425964355469, + -289.11212158203125 + ], + "size": [ + 509.7506103515625, + 178 + ], + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "compile_args", + "type": "COMPILEARGS", + "link": 154, + "shape": 7 + }, + { + "name": "block_swap_args", + "type": "BLOCKSWAPARGS", + "link": null, + "shape": 7 + } + ], + "outputs": [ + { + "name": "model", + "type": "HYVIDEOMODEL", + "links": [ + 151 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "HyVideoModelLoader" + }, + "widgets_values": [ + "hyvideo\\hunyuan_video_720_fp8_e4m3fn.safetensors", + "bf16", + "fp8_e4m3fn_fast", + "offload_device", + "sageattn_varlen" + ] + }, + { + "id": 82, + "type": "HyVideoTorchCompileSettings", + "pos": [ + -965.2777099609375, + -357.0115051269531 + ], + "size": [ + 441, + 274 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "torch_compile_args", + "type": "COMPILEARGS", + "links": [ + 154 + ] + } + ], + "properties": { + "Node name for S&R": "HyVideoTorchCompileSettings" + }, + "widgets_values": [ + "inductor", + false, + "default", + false, + 64, + true, + true, + false, + false, + false + ] + }, + { + "id": 68, + "type": "ImageConcatMulti", + "pos": [ + 1430, + -440 + ], + "size": [ + 210, + 150 + ], + "flags": {}, + "order": 17, + "mode": 0, + "inputs": [ + { + "name": "image_1", + "type": "IMAGE", + "link": 121 + }, + { + "name": "image_2", + "type": "IMAGE", + "link": 113 + } + ], + "outputs": [ + { + "name": "images", + "type": "IMAGE", + "links": [ + 114 + ], + "slot_index": 0 + } + ], + "properties": {}, + "widgets_values": [ + 2, + "right", + false, + null + ] + }, + { + "id": 73, + "type": "GetNode", + "pos": [ + 1430, + -240 + ], + "size": [ + 210, + 58 + ], + "flags": { + "collapsed": true + }, + "order": 3, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 121 + ], + "slot_index": 0 + } + ], + "title": "Get_InputVideo", + "properties": {}, + "widgets_values": [ + "InputVideo" + ], + "color": "#2a363b", + "bgcolor": "#3f5159" + }, + { + "id": 66, + "type": "HyVideoDecode", + "pos": [ + 1030, + -450 + ], + "size": [ + 292.81866455078125, + 150 + ], + "flags": {}, + "order": 16, + "mode": 0, + "inputs": [ + { + "name": "vae", + "type": "VAE", + "link": 105 + }, + { + "name": "samples", + "type": "LATENT", + "link": 142 + } + ], + "outputs": [ + { + "name": "images", + "type": "IMAGE", + "links": [ + 113 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "HyVideoDecode" + }, + "widgets_values": [ + true, + 64, + 128, + true + ] + }, + { + "id": 16, + "type": "DownloadAndLoadHyVideoTextEncoder", + "pos": [ + -439.28472900390625, + -858.833984375 + ], + "size": [ + 429.6711730957031, + 178 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "hyvid_text_encoder", + "type": "HYVIDTEXTENCODER", + "links": [ + 102 + ] + } + ], + "properties": { + "Node name for S&R": "DownloadAndLoadHyVideoTextEncoder" + }, + "widgets_values": [ + "Kijai/llava-llama-3-8b-text-encoder-tokenizer", + "disabled", + "fp16", + false, + 2, + "disabled" + ] + }, + { + "id": 43, + "type": "CLIPLoader", + "pos": [ + -406.4485168457031, + -611.748046875 + ], + "size": [ + 369.7733154296875, + 86.28006744384766 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "CLIP", + "type": "CLIP", + "links": [ + 103 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "CLIPLoader" + }, + "widgets_values": [ + "clip_l.safetensors", + "stable_diffusion" + ] + }, + { + "id": 65, + "type": "HyVideoTextEncode", + "pos": [ + 535.5712890625, + -847.4974365234375 + ], + "size": [ + 459.8385009765625, + 345.2080383300781 + ], + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "text_encoders", + "type": "HYVIDTEXTENCODER", + "link": 102 + }, + { + "name": "custom_prompt_template", + "type": "PROMPT_TEMPLATE", + "link": null, + "shape": 7 + }, + { + "name": "clip_l", + "type": "CLIP", + "link": 103, + "shape": 7 + } + ], + "outputs": [ + { + "name": "hyvid_embeds", + "type": "HYVIDEMBEDS", + "links": [ + 139 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "HyVideoTextEncode" + }, + "widgets_values": [ + "Main content and theme: cinematic and photorealistic video of a wolf wearing sunglasses\nObject details:\na wolf\nsunglasses\nForest background\nThe video style is very cinematic and has dramatic lighting\nActions and movements: a wolf wearing sunglsses is turning it's head\nBackground and atmosphere:\nThe background is a peaceful forest, which creates a serene, calm atmosphere\nmasterpiece, best quality", + true, + "video" + ] + }, + { + "id": 70, + "type": "HyVideoEmptyTextEmbeds", + "pos": [ + 286.3083190917969, + 52.262840270996094 + ], + "size": [ + 302.1732177734375, + 26 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "hyvid_embeds", + "type": "HYVIDEMBEDS", + "links": [ + 147 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "HyVideoEmptyTextEmbeds" + }, + "widgets_values": [] + }, + { + "id": 80, + "type": "HyVideoInverseSampler", + "pos": [ + 694.1298217773438, + -4.059940338134766 + ], + "size": [ + 300.78759765625, + 522.9332885742188 + ], + "flags": {}, + "order": 14, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "HYVIDEOMODEL", + "link": 152 + }, + { + "name": "hyvid_embeds", + "type": "HYVIDEMBEDS", + "link": 147 + }, + { + "name": "samples", + "type": "LATENT", + "link": 148 + } + ], + "outputs": [ + { + "name": "samples", + "type": "LATENT", + "links": [ + 150 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "HyVideoInverseSampler" + }, + "widgets_values": [ + 30, + 0, + 1, + 1, + "fixed", + true, + 0.5 + ] + }, + { + "id": 69, + "type": "VHS_VideoCombine", + "pos": [ + 1788.1004638671875, + -439.3926696777344 + ], + "size": [ + 1098.79248046875, + 863.396240234375 + ], + "flags": {}, + "order": 18, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 114 + }, + { + "name": "audio", + "type": "AUDIO", + "link": null, + "shape": 7 + }, + { + "name": "meta_batch", + "type": "VHS_BatchManager", + "link": null, + "shape": 7 + }, + { + "name": "vae", + "type": "VAE", + "link": null, + "shape": 7 + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null + } + ], + "properties": { + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 24, + "loop_count": 0, + "filename_prefix": "HunyuanVideo", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 19, + "save_metadata": true, + "pingpong": false, + "save_output": false, + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "HunyuanVideo_00079.mp4", + "subfolder": "", + "type": "temp", + "format": "video/h264-mp4", + "frame_rate": 24 + }, + "muted": false + } + } + }, + { + "id": 78, + "type": "HyVideoReSampler", + "pos": [ + 1308.6800537109375, + -132.41998291015625 + ], + "size": [ + 326.2127380371094, + 568.9866943359375 + ], + "flags": {}, + "order": 15, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "HYVIDEOMODEL", + "link": 153 + }, + { + "name": "hyvid_embeds", + "type": "HYVIDEMBEDS", + "link": 139 + }, + { + "name": "samples", + "type": "LATENT", + "link": 140 + }, + { + "name": "inversed_latents", + "type": "LATENT", + "link": 150 + } + ], + "outputs": [ + { + "name": "samples", + "type": "LATENT", + "links": [ + 142 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "HyVideoReSampler" + }, + "widgets_values": [ + 30, + 6, + 1, + true, + 0, + 13, + 0.7000000000000001, + "linear_decrease" + ] + } + ], + "links": [ + [ + 73, + 7, + 0, + 61, + 0, + "VAE" + ], + [ + 75, + 60, + 0, + 62, + 0, + "IMAGE" + ], + [ + 102, + 16, + 0, + 65, + 0, + "HYVIDTEXTENCODER" + ], + [ + 103, + 43, + 0, + 65, + 2, + "CLIP" + ], + [ + 105, + 7, + 0, + 66, + 0, + "VAE" + ], + [ + 113, + 66, + 0, + 68, + 1, + "IMAGE" + ], + [ + 114, + 68, + 0, + 69, + 0, + "IMAGE" + ], + [ + 117, + 71, + 0, + 61, + 1, + "IMAGE" + ], + [ + 119, + 62, + 0, + 72, + 0, + "*" + ], + [ + 120, + 72, + 0, + 71, + 0, + "IMAGE" + ], + [ + 121, + 73, + 0, + 68, + 0, + "IMAGE" + ], + [ + 139, + 65, + 0, + 78, + 1, + "HYVIDEMBEDS" + ], + [ + 140, + 61, + 0, + 78, + 2, + "LATENT" + ], + [ + 142, + 78, + 0, + 66, + 1, + "LATENT" + ], + [ + 147, + 70, + 0, + 80, + 1, + "HYVIDEMBEDS" + ], + [ + 148, + 61, + 0, + 80, + 2, + "LATENT" + ], + [ + 150, + 80, + 0, + 78, + 3, + "LATENT" + ], + [ + 151, + 1, + 0, + 81, + 0, + "*" + ], + [ + 152, + 81, + 0, + 80, + 0, + "HYVIDEOMODEL" + ], + [ + 153, + 81, + 0, + 78, + 0, + "HYVIDEOMODEL" + ], + [ + 154, + 82, + 0, + 1, + 0, + "COMPILEARGS" + ] + ], + "groups": [], + "config": {}, + "extra": { + "ds": { + "scale": 0.5644739300537778, + "offset": [ + 1051.95782812112, + 1072.9689374346237 + ] + } + }, + "version": 0.4 +} \ No newline at end of file diff --git a/hunyuan_empty_prompt_embeds_dict.pt b/hunyuan_empty_prompt_embeds_dict.pt new file mode 100644 index 0000000..f6856fd Binary files /dev/null and b/hunyuan_empty_prompt_embeds_dict.pt differ diff --git a/hyvideo/diffusion/pipelines/pipeline_hunyuan_video.py b/hyvideo/diffusion/pipelines/pipeline_hunyuan_video.py index c7cf3f3..461ee45 100644 --- a/hyvideo/diffusion/pipelines/pipeline_hunyuan_video.py +++ b/hyvideo/diffusion/pipelines/pipeline_hunyuan_video.py @@ -509,6 +509,10 @@ class HunyuanVideoPipeline(DiffusionPipeline): num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order self._num_timesteps = len(timesteps) + # 8. Preview callback + from ....latent_preview import prepare_callback + callback = prepare_callback(self.transformer, num_inference_steps) + 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)) @@ -619,10 +623,10 @@ 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) + else: comfy_pbar.update(1) - if callback is not None and i % callback_steps == 0: - step_idx = i // getattr(self.scheduler, "order", 1) - callback(step_idx, t, latents) #latents = (latents / 2 + 0.5).clamp(0, 1).cpu() diff --git a/latent_preview.py b/latent_preview.py new file mode 100644 index 0000000..141c6b4 --- /dev/null +++ b/latent_preview.py @@ -0,0 +1,85 @@ +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 ("GIF", preview_image, MAX_PREVIEW_RESOLUTION) + +class Latent2RGBPreviewer(LatentPreviewer): + def __init__(self): + latent_rgb_factors = [[-0.41, -0.25, -0.26], + [-0.26, -0.49, -0.24], + [-0.37, -0.54, -0.3], + [-0.04, -0.29, -0.29], + [-0.52, -0.59, -0.39], + [-0.56, -0.6, -0.02], + [-0.53, -0.06, -0.48], + [-0.51, -0.28, -0.18], + [-0.59, -0.1, -0.33], + [-0.56, -0.54, -0.41], + [-0.61, -0.19, -0.5], + [-0.05, -0.25, -0.17], + [-0.23, -0.04, -0.22], + [-0.51, -0.56, -0.43], + [-0.13, -0.4, -0.05], + [-0.01, -0.01, -0.48]] + self.latent_rgb_factors = torch.tensor(latent_rgb_factors, device="cpu").transpose(0, 1) + self.latent_rgb_factors_bias = torch.tensor([0.138, 0.025, -0.299], 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 + diff --git a/nodes.py b/nodes.py index 6464793..a1f7043 100644 --- a/nodes.py +++ b/nodes.py @@ -855,7 +855,7 @@ class HyVideoSampler: return ({ "samples": out_latents },) - + #region VideoDecode class HyVideoDecode: @classmethod @@ -990,7 +990,7 @@ class HyVideoEncode: return ({"samples": latents},) -class CogVideoLatentPreview: +class HyVideoLatentPreview: @classmethod def INPUT_TYPES(s): return { @@ -1008,25 +1008,39 @@ class CogVideoLatentPreview: RETURN_TYPES = ("IMAGE", "STRING", ) RETURN_NAMES = ("images", "latent_rgb_factors",) FUNCTION = "sample" - CATEGORY = "PyramidFlowWrapper" + CATEGORY = "HunyuanVideoWrapper" def sample(self, samples, seed, min_val, max_val, r_bias, g_bias, b_bias): mm.soft_empty_cache() latents = samples["samples"].clone() print("in sample", latents.shape) - latents = latents.permute(0, 2, 1, 3, 4) # [batch_size, num_channels, num_frames, height, width] - - #[[0.0658900170023352, 0.04687556512203313, -0.056971557475649186], [-0.01265770449940036, -0.02814809569100843, -0.0768912512529372], [0.061456544746314665, 0.0005511617552452358, -0.0652574975291287], [-0.09020669168815276, -0.004755440180558637, -0.023763970904494294], [0.031766964513999865, -0.030959599938418375, 0.08654669098083616], [-0.005981764690055846, -0.08809119252349802, -0.06439852368217663], [-0.0212114426433989, 0.08894281999597677, 0.05155629477559985], [-0.013947446911030725, -0.08987475069900677, -0.08923124751217484], [-0.08235967967978511, 0.07268025379974379, 0.08830486164536037], [-0.08052049179735378, -0.050116143175332195, 0.02023752569687405], [-0.07607527759162447, 0.06827156419895981, 0.08678111754261035], [-0.04689089232553825, 0.017294986041038893, -0.10280492336438908], [-0.06105783150270304, 0.07311850680875913, 0.019995735372550075], [-0.09232589996527711, -0.012869815059053047, -0.04355587834255975], [-0.06679931010802251, 0.018399815879067458, 0.06802404982033876], [-0.013062632927118165, -0.04292991477896661, 0.07476243356192845]] - latent_rgb_factors =[[0.11945946736445662, 0.09919175788574555, -0.004832707433877734], [-0.0011977028264356232, 0.05496505130267682, 0.021321622433638193], [-0.014088548986590666, -0.008701477861945644, -0.020991313281459367], [0.03063921972519621, 0.12186477097625073, 0.0139593690235148], [0.0927403067854673, 0.030293187650929136, 0.05083134241694003], [0.0379112441305742, 0.04935199882777209, 0.058562766246777774], [0.017749911959153715, 0.008839453404921545, 0.036005638019226294], [0.10610119248526109, 0.02339855688237826, 0.057154257614084596], [0.1273639464837117, -0.010959856130713416, 0.043268631260428896], [-0.01873510946881321, 0.08220930648486932, 0.10613256772247093], [0.008429116376722327, 0.07623856561000408, 0.09295712117576727], [0.12938137079617007, 0.12360403483892413, 0.04478930933220116], [0.04565908794779364, 0.041064156741596365, -0.017695041535528512], [0.00019003240570281826, -0.013965147883381978, 0.05329669529635849], [0.08082391586738358, 0.11548306825496074, -0.021464170006615893], [-0.01517932393230994, -0.0057985555313003236, 0.07216646476618871]] + #latent_rgb_factors =[[-0.02531045419704009, -0.00504800612542497, 0.13293717293982546], [-0.03421835830845858, 0.13996708548892614, -0.07081038680118075], [0.011091819063647063, -0.03372949685846012, -0.0698232210116172], [-0.06276524604742019, -0.09322986677909442, 0.01826383612148913], [0.021290659938126788, -0.07719530444034409, -0.08247812477766273], [0.04401102991215147, -0.0026401932105894754, -0.01410913586718443], [0.08979717602613707, 0.05361221258740831, 0.11501425309699129], [0.04695121980405198, -0.13053491609675175, 0.05025986885867986], [-0.09704684176098193, 0.03397687417738002, -0.1105886644677771], [0.14694697234804935, -0.12316902186157716, 0.04210404546699645], [0.14432470831243552, -0.002580008133591355, -0.08490676947390643], [0.051502750076553944, -0.10071695490292451, -0.01786223610178095], [-0.12503276881774464, 0.08877830923879379, 0.1076584501927316], [-0.020191205513213406, -0.1493425056303128, -0.14289740371758308], [-0.06470138952271293, -0.07410426095060325, 0.00980804676890873], [0.11747671720735695, 0.10916082743849789, -0.12235599365235904]] + latent_rgb_factors = [[-0.41, -0.25, -0.26], + [-0.26, -0.49, -0.24], + [-0.37, -0.54, -0.3], + [-0.04, -0.29, -0.29], + [-0.52, -0.59, -0.39], + [-0.56, -0.6, -0.02], + [-0.53, -0.06, -0.48], + [-0.51, -0.28, -0.18], + [-0.59, -0.1, -0.33], + [-0.56, -0.54, -0.41], + [-0.61, -0.19, -0.5], + [-0.05, -0.25, -0.17], + [-0.23, -0.04, -0.22], + [-0.51, -0.56, -0.43], + [-0.13, -0.4, -0.05], + [-0.01, -0.01, -0.48]] + import random random.seed(seed) - latent_rgb_factors = [[random.uniform(min_val, max_val) for _ in range(3)] for _ in range(16)] + #latent_rgb_factors = [[random.uniform(min_val, max_val) for _ in range(3)] for _ in range(16)] out_factors = latent_rgb_factors print(latent_rgb_factors) - latent_rgb_factors_bias = [0.085, 0.137, 0.158] - #latent_rgb_factors_bias = [r_bias, g_bias, b_bias] + #latent_rgb_factors_bias = [0.138, 0.025, -0.299] + latent_rgb_factors_bias = [r_bias, g_bias, b_bias] latent_rgb_factors = torch.tensor(latent_rgb_factors, device=latents.device, dtype=latents.dtype).transpose(0, 1) latent_rgb_factors_bias = torch.tensor(latent_rgb_factors_bias, device=latents.device, dtype=latents.dtype) @@ -1063,7 +1077,8 @@ NODE_CLASS_MAPPINGS = { "HyVideoTorchCompileSettings": HyVideoTorchCompileSettings, "HyVideoSTG": HyVideoSTG, "HyVideoCustomPromptTemplate": HyVideoCustomPromptTemplate, -} + "HyVideoLatentPreview": HyVideoLatentPreview, + } NODE_DISPLAY_NAME_MAPPINGS = { "HyVideoSampler": "HunyuanVideo Sampler", "HyVideoDecode": "HunyuanVideo Decode", @@ -1076,4 +1091,5 @@ NODE_DISPLAY_NAME_MAPPINGS = { "HyVideoTorchCompileSettings": "HunyuanVideo Torch Compile Settings", "HyVideoSTG": "HunyuanVideo STG", "HyVideoCustomPromptTemplate": "HunyuanVideo Custom Prompt Template", + "HyVideoLatentPreview": "HunyuanVideo Latent Preview", } diff --git a/nodes_rf_inversion.py b/nodes_rf_inversion.py new file mode 100644 index 0000000..62b482d --- /dev/null +++ b/nodes_rf_inversion.py @@ -0,0 +1,437 @@ +#based on https://github.com/DarkMnDragon/rf-inversion-diffuser/blob/main/inversion_editing_cli.py +import torch +import gc +import os +from .utils import log, print_memory + +from .hyvideo.utils.data_utils import align_to +from diffusers.utils.torch_utils import randn_tensor +import comfy.model_management as mm +from .nodes import get_rotary_pos_embed + +script_directory = os.path.dirname(os.path.abspath(__file__)) + +def generate_eta_values( + timesteps, + start_step, + end_step, + eta, + eta_trend, +): + assert start_step < end_step and start_step >= 0 and end_step <= len(timesteps), "Invalid start_step and end_step" + # timesteps are monotonically decreasing, from 1.0 to 0.0 + print("eta timesteps", timesteps) + eta_values = [0.0] * (len(timesteps) - 1) + + if eta_trend == 'constant': + for i in range(start_step, end_step): + eta_values[i] = eta + elif eta_trend == 'linear_increase': + total_time = timesteps[start_step] - timesteps[end_step - 1] + for i in range(start_step, end_step): + eta_values[i] = eta * (timesteps[start_step] - timesteps[i]) / total_time + elif eta_trend == 'linear_decrease': + total_time = timesteps[start_step] - timesteps[end_step - 1] + for i in range(start_step, end_step): + eta_values[i] = eta * (timesteps[i] - timesteps[end_step - 1]) / total_time + else: + raise NotImplementedError(f"Unsupported eta_trend: {eta_trend}") + + return eta_values + +class HyVideoEmptyTextEmbeds: + @classmethod + def INPUT_TYPES(s): + return {"required": { + } + } + + RETURN_TYPES = ("HYVIDEMBEDS", ) + RETURN_NAMES = ("hyvid_embeds",) + FUNCTION = "process" + CATEGORY = "HunyuanVideoWrapper" + DESCRIPTION = "Empty Text Embeds for HunyuanVideoWrapper, to avoid having to encode prompts for inverse sampling" + + def process(self): + device = mm.text_encoder_device() + offload_device = mm.text_encoder_offload_device() + + prompt_embeds_dict = torch.load(os.path.join(script_directory, "hunyuan_empty_prompt_embeds_dict.pt")) + return (prompt_embeds_dict,) + +class HyVideoInverseSampler: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("HYVIDEOMODEL",), + "hyvid_embeds": ("HYVIDEMBEDS", ), + "samples": ("LATENT", {"tooltip": "init Latents to use for video2video process"} ), + "steps": ("INT", {"default": 30, "min": 1}), + "embedded_guidance_scale": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 30.0, "step": 0.01}), + "flow_shift": ("FLOAT", {"default": 1.0, "min": 1.0, "max": 30.0, "step": 0.01}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), + "force_offload": ("BOOLEAN", {"default": True}), + "gamma": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + }, + } + + RETURN_TYPES = ("LATENT",) + RETURN_NAMES = ("samples",) + FUNCTION = "process" + CATEGORY = "HunyuanVideoWrapper" + + def process(self, model, hyvid_embeds, flow_shift, steps, embedded_guidance_scale, seed, samples, gamma, force_offload): + device = mm.get_torch_device() + offload_device = mm.unet_offload_device() + dtype = model["dtype"] + transformer = model["pipe"].transformer + pipeline = model["pipe"] + + generator = torch.Generator(device=torch.device("cpu")).manual_seed(seed) + + latents = samples["samples"] if samples is not None else None + batch_size, num_channels_latents, latent_num_frames, latent_height, latent_width = latents.shape + height = latent_height * pipeline.vae_scale_factor + width = latent_width * pipeline.vae_scale_factor + num_frames = (latent_num_frames - 1) * 4 + 1 + + + if width <= 0 or height <= 0 or num_frames <= 0: + raise ValueError( + f"`height` and `width` and `video_length` must be positive integers, got height={height}, width={width}, video_length={num_frames}" + ) + if (num_frames - 1) % 4 != 0: + raise ValueError( + f"`video_length-1` must be a multiple of 4, got {num_frames}" + ) + + log.info( + f"Input (height, width, video_length) = ({height}, {width}, {num_frames})" + ) + + freqs_cos, freqs_sin = get_rotary_pos_embed(transformer, num_frames, height, width) + + pipeline.scheduler.shift = flow_shift + + if model["block_swap_args"] is not None: + for name, param in transformer.named_parameters(): + #print(name, param.data.device) + if "single" not in name and "double" not in name: + param.data = param.data.to(device) + + transformer.block_swap( + model["block_swap_args"]["double_blocks_to_swap"] - 1 , + model["block_swap_args"]["single_blocks_to_swap"] - 1, + offload_txt_in = model["block_swap_args"]["offload_txt_in"], + offload_img_in = model["block_swap_args"]["offload_img_in"], + ) + elif model["manual_offloading"]: + transformer.to(device) + + mm.unload_all_models() + mm.soft_empty_cache() + gc.collect() + + try: + torch.cuda.reset_peak_memory_stats(device) + except: + pass + + pipeline.scheduler.set_timesteps(steps, device=device) + timesteps = pipeline.scheduler.timesteps + timesteps = timesteps.flip(0) + print("timesteps", timesteps) + print("pipeline.scheduler.order", pipeline.scheduler.order) + print("len(timesteps)", len(timesteps)) + + latent_video_length = (num_frames - 1) // 4 + 1 + + # 5. Prepare latent variables + num_channels_latents = transformer.config.in_channels + + + latents = latents.to(device) + + shape = ( + 1, + num_channels_latents, + latent_video_length, + int(height) // pipeline.vae_scale_factor, + int(width) // pipeline.vae_scale_factor, + ) + noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + + frames_needed = noise.shape[1] + current_frames = latents.shape[1] + + if frames_needed > current_frames: + repeat_factor = frames_needed - current_frames + additional_frame = torch.randn((latents.size(0), repeat_factor, latents.size(2), latents.size(3), latents.size(4)), dtype=latents.dtype, device=latents.device) + latents = torch.cat((additional_frame, latents), dim=1) + self.additional_frames = repeat_factor + elif frames_needed < current_frames: + latents = latents[:, :frames_needed, :, :, :] + + + + # 7. Denoising loop + 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 comfy.utils import ProgressBar + from tqdm import tqdm + log.info(f"Sampling {num_frames} frames in {latents.shape[2]} latents at {width}x{height} with {len(timesteps)} inference steps") + comfy_pbar = ProgressBar(len(timesteps)) + with tqdm(total=len(timesteps)) as progress_bar: + for idx, (t, t_prev) in enumerate(zip(timesteps[:-1], timesteps[1:])): + latent_model_input = latents + + t_expand = t.repeat(latent_model_input.shape[0]) + guidance_expand = ( + torch.tensor( + [embedded_guidance_scale] * latent_model_input.shape[0], + dtype=torch.float32, + device=device, + ).to(pipeline.base_dtype) + * 1000.0 + if embedded_guidance_scale is not None + else None + ) + + # predict the noise residual + with torch.autocast( + device_type="cuda", dtype=pipeline.base_dtype, enabled=True + ): + noise_pred = transformer( # For an input image (129, 192, 336) (1, 256, 256) + latent_model_input, # [2, 16, 33, 24, 42] + t_expand, # [2] + text_states=hyvid_embeds["prompt_embeds"], # [2, 256, 4096] + text_mask=hyvid_embeds["attention_mask"], # [2, 256] + text_states_2=hyvid_embeds["prompt_embeds_2"], # [2, 768] + freqs_cos=freqs_cos, # [seqlen, head_dim] + freqs_sin=freqs_sin, # [seqlen, head_dim] + guidance=guidance_expand, + stg_block_idx=-1, + stg_mode=None, + return_dict=True, + )["x"] + sigma = t / 1000.0 + sigma_prev = t_prev / 1000.0 + target_noise_velocity = (noise - latents) / (1.0 - sigma) + interpolated_velocity = gamma * target_noise_velocity + (1 - gamma) * noise_pred + + latents = latents + (sigma_prev - sigma) * interpolated_velocity + + # compute the previous noisy sample x_t -> x_t-1 + #latents = pipeline.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + + + progress_bar.update() + if callback is not None: + print("callback", latents.shape) + callback(idx, latents.detach()[-1].permute(1,0,2,3), None, steps) + else: + comfy_pbar.update(1) + + + print_memory(device) + try: + torch.cuda.reset_peak_memory_stats(device) + except: + pass + + if force_offload: + if model["manual_offloading"]: + transformer.to(offload_device) + mm.soft_empty_cache() + gc.collect() + + return ({ + "samples": latents + },) + +class HyVideoReSampler: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model": ("HYVIDEOMODEL",), + "hyvid_embeds": ("HYVIDEMBEDS", ), + "samples": ("LATENT", {"tooltip": "init Latents to use for video2video process"} ), + "inversed_latents": ("LATENT", {"tooltip": "inversed latents from HyVideoInverseSampler"} ), + "steps": ("INT", {"default": 30, "min": 1}), + "embedded_guidance_scale": ("FLOAT", {"default": 6.0, "min": 0.0, "max": 30.0, "step": 0.01}), + "flow_shift": ("FLOAT", {"default": 1.0, "min": 1.0, "max": 30.0, "step": 0.01}), + "force_offload": ("BOOLEAN", {"default": True}), + "start_step": ("INT", {"default": 0, "min": 0}), + "end_step": ("INT", {"default": 18, "min": 0}), + "eta_base": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + "eta_trend": (['constant', 'linear_increase', 'linear_decrease'], {"default": "constant"}), + }, + } + + RETURN_TYPES = ("LATENT",) + RETURN_NAMES = ("samples",) + FUNCTION = "process" + CATEGORY = "HunyuanVideoWrapper" + + 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): + device = mm.get_torch_device() + offload_device = mm.unet_offload_device() + dtype = model["dtype"] + transformer = model["pipe"].transformer + pipeline = model["pipe"] + + target_latents = samples["samples"] + + batch_size, num_channels_latents, latent_num_frames, latent_height, latent_width = target_latents.shape + height = latent_height * pipeline.vae_scale_factor + width = latent_width * pipeline.vae_scale_factor + num_frames = (latent_num_frames - 1) * 4 + 1 + + if width <= 0 or height <= 0 or num_frames <= 0: + raise ValueError( + f"`height` and `width` and `video_length` must be positive integers, got height={height}, width={width}, video_length={num_frames}" + ) + if (num_frames - 1) % 4 != 0: + raise ValueError( + f"`video_length-1` must be a multiple of 4, got {num_frames}" + ) + + log.info( + f"Input (height, width, video_length) = ({height}, {width}, {num_frames})" + ) + + freqs_cos, freqs_sin = get_rotary_pos_embed(transformer, num_frames, height, width) + + pipeline.scheduler.shift = flow_shift + + if model["block_swap_args"] is not None: + for name, param in transformer.named_parameters(): + #print(name, param.data.device) + if "single" not in name and "double" not in name: + param.data = param.data.to(device) + + transformer.block_swap( + model["block_swap_args"]["double_blocks_to_swap"] - 1 , + model["block_swap_args"]["single_blocks_to_swap"] - 1, + offload_txt_in = model["block_swap_args"]["offload_txt_in"], + offload_img_in = model["block_swap_args"]["offload_img_in"], + ) + elif model["manual_offloading"]: + transformer.to(device) + + mm.unload_all_models() + mm.soft_empty_cache() + gc.collect() + + try: + torch.cuda.reset_peak_memory_stats(device) + except: + pass + + pipeline.scheduler.set_timesteps(steps, device=device) + timesteps = pipeline.scheduler.timesteps + + eta_values = generate_eta_values(timesteps / 1000, start_step, end_step, eta_base, eta_trend) + + + target_latents = target_latents.to(device) + latents = inversed_latents["samples"] + + # 7. Denoising loop + self._num_timesteps = len(timesteps) + + from .latent_preview import prepare_callback + callback = prepare_callback(transformer, steps) + + from comfy.utils import ProgressBar + from tqdm import tqdm + log.info(f"Sampling {num_frames} frames in {latents.shape[2]} latents at {width}x{height} with {len(timesteps)} inference steps") + comfy_pbar = ProgressBar(len(timesteps)) + + with tqdm(total=len(timesteps)) as progress_bar: + for idx, (t, t_prev) in enumerate(zip(timesteps[:-1], timesteps[1:])): + + latent_model_input = latents + + t_expand = t.repeat(latent_model_input.shape[0]) + guidance_expand = ( + torch.tensor( + [embedded_guidance_scale] * latent_model_input.shape[0], + dtype=torch.float32, + device=device, + ).to(pipeline.base_dtype) + * 1000.0 + if embedded_guidance_scale is not None + else None + ) + + # predict the noise residual + with torch.autocast( + device_type="cuda", dtype=pipeline.base_dtype, enabled=True + ): + noise_pred = transformer( # For an input image (129, 192, 336) (1, 256, 256) + latent_model_input, # [2, 16, 33, 24, 42] + t_expand, # [2] + text_states=hyvid_embeds["prompt_embeds"], # [2, 256, 4096] + text_mask=hyvid_embeds["attention_mask"], # [2, 256] + text_states_2=hyvid_embeds["prompt_embeds_2"], # [2, 768] + freqs_cos=freqs_cos, # [seqlen, head_dim] + freqs_sin=freqs_sin, # [seqlen, head_dim] + guidance=guidance_expand, + stg_block_idx=-1, + stg_mode=None, + return_dict=True, + )["x"] + sigma = t / 1000.0 + sigma_prev = t_prev / 1000.0 + noise_pred = noise_pred.to(torch.float32) + latents = latents.to(torch.float32) + target_latents = target_latents.to(torch.float32) + target_img_velocity = -(target_latents - latents) / sigma + + # interpolated velocity + eta = eta_values[idx] + interpolated_velocity = eta * target_img_velocity + (1 - eta) * noise_pred + latents = latents + (sigma_prev - sigma) * interpolated_velocity + + print(f"X_{sigma_prev:.3f} = X_{sigma:.3f} + {sigma_prev - sigma:.3f} * ({eta:.3f} * target_img_velocity + {1 - eta:.3f} * noise_pred)") + latents = latents.to(torch.bfloat16) + + if callback is not None: + callback(idx, latents.detach()[-1].permute(1,0,2,3), None, steps) + else: + comfy_pbar.update(1) + + print_memory(device) + try: + torch.cuda.reset_peak_memory_stats(device) + except: + pass + + if force_offload: + if model["manual_offloading"]: + transformer.to(offload_device) + mm.soft_empty_cache() + gc.collect() + + return ({ + "samples": latents + },) + +NODE_CLASS_MAPPINGS = { + "HyVideoInverseSampler": HyVideoInverseSampler, + "HyVideoReSampler": HyVideoReSampler, + "HyVideoEmptyTextEmbeds": HyVideoEmptyTextEmbeds +} +NODE_DISPLAY_NAME_MAPPINGS = { + "HyVideoInverseSampler": "HunyuanVideo Inverse Sampler", + "HyVideoReSampler": "HunyuanVideo ReSampler", + "HyVideoEmptyTextEmbeds": "HunyuanVideo Empty Text Embeds" +} \ No newline at end of file