kijai
bb75cddd60
https://zhuang2002.github.io/FlashVSR/ This only implements the projection model and the VAE, which seems to be enough for upscaling. This does NOT implement any of the streaming and sparse attention code.
70 lines
2.2 KiB
Python
70 lines
2.2 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"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("WANVIDIMAGE_EMBEDS",)
|
|
RETURN_NAMES = ("image_embeds",)
|
|
FUNCTION = "add"
|
|
CATEGORY = "WanVideoWrapper"
|
|
|
|
def add(self, embeds, images):
|
|
updated = dict(embeds)
|
|
updated["flashvsr_LQ_images"] = images
|
|
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",
|
|
}
|