72 lines
2.4 KiB
Python
72 lines
2.4 KiB
Python
import folder_paths
|
|
import torch
|
|
|
|
from comfy.utils import load_torch_file
|
|
import comfy.model_management as mm
|
|
device = mm.get_torch_device()
|
|
offload_device = mm.unet_offload_device()
|
|
|
|
class WanVideoAddFlashVSRInput:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {
|
|
"embeds": ("WANVIDIMAGE_EMBEDS",),
|
|
"images": ("IMAGE", {"tooltip": "Low-res video frames to enhance"}),
|
|
"strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01, "tooltip": "Strength to apply the FlashVSR latent"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("WANVIDIMAGE_EMBEDS",)
|
|
RETURN_NAMES = ("image_embeds",)
|
|
FUNCTION = "add"
|
|
CATEGORY = "WanVideoWrapper"
|
|
|
|
def add(self, embeds, images, strength):
|
|
updated = dict(embeds)
|
|
updated["flashvsr_LQ_images"] = images
|
|
updated["flashvsr_strength"] = strength
|
|
return (updated,)
|
|
|
|
|
|
class WanVideoFlashVSRDecoderLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model_name": (folder_paths.get_filename_list("vae"), {"tooltip": "These models are loaded from 'ComfyUI/models/vae'"}),
|
|
},
|
|
"optional": {
|
|
"precision": (["fp16", "fp32", "bf16"],
|
|
{"default": "bf16"}
|
|
),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("WANVAE",)
|
|
RETURN_NAMES = ("vae", )
|
|
FUNCTION = "loadmodel"
|
|
CATEGORY = "WanVideoWrapper"
|
|
DESCRIPTION = "Loads Wan VAE model from 'ComfyUI/models/vae'"
|
|
|
|
def loadmodel(self, model_name, precision):
|
|
from .TCDecoder import build_tcdecoder
|
|
dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[precision]
|
|
model_path = folder_paths.get_full_path("vae", model_name)
|
|
sd = load_torch_file(model_path, safe_load=True)
|
|
|
|
TCDecoder = build_tcdecoder(new_channels=[512, 256, 128, 128], new_latent_channels=16+768, dtype=dtype)
|
|
TCDecoder.load_state_dict(sd, strict=True)
|
|
TCDecoder.to(dtype)
|
|
|
|
return (TCDecoder,)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"WanVideoAddFlashVSRInput": WanVideoAddFlashVSRInput,
|
|
"WanVideoFlashVSRDecoderLoader": WanVideoFlashVSRDecoderLoader,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"WanVideoAddFlashVSRInput": "WanVideo Add FlashVSR Input",
|
|
"WanVideoFlashVSRDecoderLoader": "WanVideo FlashVSR Decoder Loader",
|
|
}
|