# !/usr/bin/env python # -*- coding: UTF-8 -*- import os import torch import gc from PIL import Image import numpy as np import cv2 import comfy.model_management as mm import torch.nn.functional as F from comfy.utils import common_upscale,ProgressBar import folder_paths from diffusers import AutoencoderKLWan from .FlashVSR.vae import WanVAE from .FlashVSR.diffsynth.pipelines.flashvsr_full import TorchColorCorrectorWavelet from .FlashVSR.examples.WanVSR.infer_flashvsr_full import upscale_lq_video_bilinear,dup_first_frame_1cthw_simple from safetensors.torch import load_file cur_path = os.path.dirname(os.path.abspath(__file__)) def load_vae_model(vae,tcd_encoder,lite_vae, device="cuda"): vae_path=folder_paths.get_full_path("vae", vae) if vae != "none" else None tcd_encoder_path=folder_paths.get_full_path("FlashVSR", tcd_encoder) if tcd_encoder != "none" else None lite_vae_path=folder_paths.get_full_path("vae", lite_vae) if lite_vae != "none" else "none" if vae_path is not None: if vae_path.endswith(".safetensors"): config=AutoencoderKLWan.load_config(os.path.join(cur_path,"FlashVSR/config.json")) vae=AutoencoderKLWan.from_config(config).to(device,dtype=torch.bfloat16) vae_dict=load_file(vae_path,device="cpu") vae.load_state_dict(vae_dict,strict=False) del vae_dict vae.vae_mode="diffusers" else: from .FlashVSR.diffsynth import ModelManager,FlashVSRFullPipeline mm = ModelManager(torch_dtype=torch.bfloat16, device="cpu") mm.load_models([vae_path,]) pipe = FlashVSRFullPipeline.from_model_manager(mm, device="cuda") pipe.vae.model.encoder = None pipe.vae.model.conv1 = None pipe.enable_vram_management(num_persistent_param_in_dit=None,vae_only=True) pipe.load_models_to_device(["vae"]) vae=pipe.vae vae.vae_mode="normal" del pipe elif tcd_encoder_path is not None: from .FlashVSR.examples.WanVSR.utils.TCDecoder import build_tcdecoder vae = build_tcdecoder(new_channels=[512, 256, 128, 128], new_latent_channels=16+768) tcd_state=torch.load(tcd_encoder_path,weights_only=False,) if isinstance(tcd_state,dict) and "decoder" in tcd_state: tcd_state=tcd_state["decoder"] # 兼容 FlashVSR+ 的 TCDecoder.ckpt 包装结构 load_report=vae.load_state_dict(tcd_state, strict=False) total=len(vae.state_dict()) miss=len(load_report.missing_keys); unexp=len(load_report.unexpected_keys) print(f"[FlashVSR] TCDecoder load: matched {total-miss}/{total}, missing={miss}, unexpected={unexp}") if miss > total*0.05: print("[FlashVSR][WARN] TCDecoder 权重覆盖率过低，解码质量会严重异常！" "请确认 tcd_encoder 选择的是 FlashVSR+ 发布的 TCDecoder.ckpt") vae.vae_mode="tcd" elif lite_vae_path is not None: if "light" in lite_vae_path.lower() or "tae" in lite_vae_path.lower(): if os.path.basename(lite_vae_path).split(".")[0]=="lightvaew2_1": print("use lightvae decoder") vae = WanVAE(vae_path=lite_vae_path,dtype=torch.bfloat16,device=device,use_lightvae=True) vae.vae_mode="lightvae" elif os.path.basename(lite_vae_path).split(".")[0]=="taew2_1": from .FlashVSR.vae_tiny import WanVAE_tiny print("use vae_tiny decoder") vae = WanVAE_tiny(vae_path=lite_vae_path,dtype=torch.bfloat16,device=device,need_scaled=False) vae.vae_mode="tae" elif os.path.basename(lite_vae_path).split(".")[0]=="lighttaew2_1": from .FlashVSR.vae_tiny import WanVAE_tiny print("use vae_tiny light decoder") vae = WanVAE_tiny(vae_path=lite_vae_path,dtype=torch.bfloat16,device=device,need_scaled=True) vae.vae_mode="lighttae" else: raise ValueError(f"Unknown vae_name: {lite_vae_path},only support lightvae,tae,tae_tiny,lighttae_tiny") else: print("use upscale2x decoder") config=AutoencoderKLWan.load_config(os.path.join(cur_path,"FlashVSR/examples/config.json")) vae=AutoencoderKLWan.from_config(config).to(device,dtype=torch.bfloat16) vae_dict=load_file(lite_vae_path,device="cpu") vae.load_state_dict(vae_dict,strict=False) del vae_dict vae.vae_mode="upscale2x" else: raise ValueError("vae_path,tcd_encoder_path,lite_vae_path is None") if not vae.vae_mode=="lightvae": vae.eval() vae.requires_grad_(False) return vae def decode_latents( vae,latent,color_fix,fix_method,full_tiled, device): vae_mode=getattr(vae,"vae_mode","none") samples=latent["samples"] LQ_cur_idx=latent.get("LQ_cur_idx",0) pad_frames=latent.get("pad_frames",0) LQ=latent.get("LQ",None) if color_fix: ColorCorrector=TorchColorCorrectorWavelet(levels=5) pad_first_frame = True if "wavelet"== fix_method and color_fix else False if isinstance(samples,list) and isinstance(LQ_cur_idx,list) and vae_mode!="tcd": raise ValueError("when use long mode only support tcd_encoder") #print(samples.shape) #torch.Size([1, 16, 22, 160, 96]) if vae_mode=="none": samples=add_mean(samples) images=vae.decode(samples) #print(f"decode complete,images shape: {images.shape}") # torch.Size([1, 85, 1024, 768, 3]) images=images.permute(0,4,1,2,3) # torch.Size([1, 3, 85, 1024, 768]) else: if not vae_mode in ["normal","lightvae"]: vae.to(device) with torch.no_grad(): if isinstance(vae,AutoencoderKLWan): samples=add_mean(samples) if full_tiled: vae.enable_tiling() else: vae.disable_tiling() images=vae.decode(samples,return_dict=False)[0] images = map_neg1_1_to_0_1(images.cpu().float()) if vae_mode=="upscale2x": ch = images.shape[1] if ch == 3: upscale = 1 else: if ch % 3 == 0: upscale = round((ch // 3) ** 0.5) else: upscale = 1 images = F.pixel_shuffle(images.movedim(2, 1), upscale_factor=int(upscale)).movedim(1, 2) # pixel shuffle needs [..., C, H, W] format #torch.Size([1, 3, 77, 2048, 1536]) else: if vae_mode=="tcd": vae.clean_mem() if isinstance(samples,list): print("use tcd decoder to infer long mode") frames_total = [] for i,(cur,pre) in zip(samples,LQ_cur_idx): cur_LQ_frame = LQ[:,:,pre:cur,:,:] cur_frames = vae.decode_video(i.transpose(1, 2),parallel=False, show_progress_bar=False, cond=LQ[:,:,pre:cur,:,:].to(device)).transpose(1, 2).mul_(2).sub_(1) cur_frames = map_neg1_1_to_0_1(cur_frames.cpu().float()) #[21,8,8,8,8,8] #torch.Size([1, 3, 21, 1280, 768]) cur_LQ_frame = map_neg1_1_to_0_1(cur_LQ_frame.cpu().float()).to(device) if color_fix: try: if pad_first_frame: # 加帧 cur_frames = dup_first_frame_1cthw_simple(cur_frames) cur_LQ_frame=dup_first_frame_1cthw_simple(cur_LQ_frame) cur_frames = ColorCorrector( cur_frames.to(device=device), cur_LQ_frame, clip_range=(-1, 1), chunk_size=None, method=fix_method, ) if pad_first_frame: #减帧 cur_frames = cur_frames[:, :, 1:, :, :] # remove first frame except: print("color correction failed, return uncorrected images") frames_total.append(cur_frames.to('cpu')) images = torch.cat(frames_total, dim=2) if pad_frames>0: print(f"Remove{images.shape[2]} 's pad frames {pad_frames} --> original frames: {images.shape[2]-pad_frames}") images = images[:, :, :-pad_frames, :, :] #torch.Size([1, 3, 85, 1024, 768]) -> torch.Size([1, 3, 81, 1024, 768]) if 81 output -4 input + 8 return images.squeeze(0).permute(1,2,3,0).cpu() else: images = vae.decode_video(samples.transpose(1, 2),parallel=False, show_progress_bar=False, cond=LQ[:,:,:LQ_cur_idx,:,:]).transpose(1, 2).mul_(2).sub_(1) elif vae_mode=="normal": vae.clear_cache() images = vae.decode(samples, device=device, tiled=full_tiled, tile_size=(60, 104), tile_stride=(30, 52)) else: if isinstance(vae,WanVAE): if full_tiled: vae.use_tiling=True else: vae.use_tiling=False images=vae.decode(samples.squeeze(0)) images = map_neg1_1_to_0_1(images.cpu().float()) if not vae_mode=="lightvae": vae.to("cpu") torch.cuda.empty_cache() #print(f"decode complete,images shape: {images.shape}") if color_fix: LQ = map_neg1_1_to_0_1(LQ.cpu().float()) try: if pad_first_frame: images = dup_first_frame_1cthw_simple(images) LQ=dup_first_frame_1cthw_simple(LQ) if vae_mode=="upscale2x" and LQ.shape[-1]!=images.shape[-1]: scale_=int(images.shape[-1]/LQ.shape[-1]) LQ=upscale_lq_video_bilinear(LQ,scale_) images = ColorCorrector( images.to(device=device), LQ[:, :, :images.shape[2], :, :], clip_range=(-1, 1), chunk_size=16, method=fix_method ) if pad_first_frame: images = images[:, :, 1:, :, :] # remove first frame torch.Size([1, 3, 78, 1024, 768]) --> torch.Size([1, 3, 77, 1024, 768]) except: print("color correction failed, return uncorrected images") if pad_frames>0: print(f"Remove{images.shape[2]} 's pad frames {pad_frames} --> original frames: {images.shape[2]-pad_frames}") images = images[:, :, :-pad_frames, :, :] #torch.Size([1, 3, 85, 1024, 768]) -> torch.Size([1, 3, 81, 1024, 768]) if 81 output -4 input + 8 return images.squeeze(0).permute(1,2,3,0).cpu() def clear_comfyui_cache(): cf_models=mm.loaded_models() try: for pipe in cf_models: pipe.unpatch_model(device_to=torch.device("cpu")) print(f"Unpatching models.{pipe}") except: pass mm.soft_empty_cache() torch.cuda.empty_cache() max_gpu_memory = torch.cuda.max_memory_allocated() print(f"After Max GPU memory allocated: {max_gpu_memory / 1000 ** 3:.2f} GB") def add_mean(latents): vae_config={"latents_mean": [ -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921 ], "latents_std": [ 2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.916 ],} latents_mean = (torch.tensor(vae_config["latents_mean"]).view(1, 16, 1, 1, 1).to(latents.device, latents.dtype)) latents_std = 1.0 / torch.tensor(vae_config["latents_std"]).view(1, 16, 1, 1, 1).to(latents.device, latents.dtype) latents = latents / latents_std + latents_mean return latents def gc_cleanup(): gc.collect() torch.cuda.empty_cache() def tensor2cv(tensor_image): if len(tensor_image.shape)==4:# b hwc to hwc tensor_image=tensor_image.squeeze(0) if tensor_image.is_cuda: tensor_image = tensor_image.cpu() tensor_image=tensor_image.numpy() #反归一化 maxValue=tensor_image.max() tensor_image=tensor_image*255/maxValue img_cv2=np.uint8(tensor_image)#32 to uint8 img_cv2=cv2.cvtColor(img_cv2,cv2.COLOR_RGB2BGR) return img_cv2 def phi2narry(img): img = torch.from_numpy(np.array(img).astype(np.float32) / 255.0).unsqueeze(0) return img def tensor2image(tensor): tensor = tensor.cpu() image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy() image = Image.fromarray(image_np, mode='RGB') return image def tensor2pillist(tensor_in): d1, _, _, _ = tensor_in.size() if d1 == 1: img_list = [tensor2image(tensor_in)] else: tensor_list = torch.chunk(tensor_in, chunks=d1) img_list=[tensor2image(i) for i in tensor_list] return img_list def tensor2pillist_upscale(tensor_in,width,height): d1, _, _, _ = tensor_in.size() if d1 == 1: img_list = [nomarl_upscale(tensor_in,width,height)] else: tensor_list = torch.chunk(tensor_in, chunks=d1) img_list=[nomarl_upscale(i,width,height) for i in tensor_list] return img_list def tensor2list(tensor_in,width,height): if tensor_in is None: return None d1, _, _, _ = tensor_in.size() if d1 == 1: tensor_list = [tensor_upscale(tensor_in,width,height)] else: tensor_list_ = torch.chunk(tensor_in, chunks=d1) tensor_list=[tensor_upscale(i,width,height) for i in tensor_list_] return tensor_list def tensor_upscale(tensor, width, height): samples = tensor.movedim(-1, 1) samples = common_upscale(samples, width, height, "nearest-exact", "center") samples = samples.movedim(1, -1) return samples def nomarl_upscale(img, width, height): samples = img.movedim(-1, 1) img = common_upscale(samples, width, height, "nearest-exact", "center") samples = img.movedim(1, -1) img = tensor2image(samples) return img def cv2tensor(img,bgr2rgb=True): assert type(img) == np.ndarray, 'the img type is {}, but ndarry expected'.format(type(img)) if bgr2rgb: img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = torch.from_numpy(img.transpose((2, 0, 1))) return img.float().div(255).permute(1, 2, 0).unsqueeze(0) def images_generator(img_list: list, ): # get img size sizes = {} for image_ in img_list: if isinstance(image_, Image.Image): count = sizes.get(image_.size, 0) sizes[image_.size] = count + 1 elif isinstance(image_, np.ndarray): count = sizes.get(image_.shape[:2][::-1], 0) sizes[image_.shape[:2][::-1]] = count + 1 else: raise "unsupport image list,must be pil or cv2!!!" size = max(sizes.items(), key=lambda x: x[1])[0] yield size[0], size[1] # any to tensor def load_image(img_in): if isinstance(img_in, Image.Image): img_in = img_in.convert("RGB") i = np.array(img_in, dtype=np.float32) i = torch.from_numpy(i).div_(255) if i.shape[0] != size[1] or i.shape[1] != size[0]: i = torch.from_numpy(i).movedim(-1, 0).unsqueeze(0) i = common_upscale(i, size[0], size[1], "lanczos", "center") i = i.squeeze(0).movedim(0, -1).numpy() return i elif isinstance(img_in, np.ndarray): i = cv2.cvtColor(img_in, cv2.COLOR_BGR2RGB).astype(np.float32) i = torch.from_numpy(i).div_(255) print(i.shape) return i else: raise "unsupport image list,must be pil,cv2 or tensor!!!" total_images = len(img_list) processed_images = 0 pbar = ProgressBar(total_images) images = map(load_image, img_list) try: prev_image = next(images) while True: next_image = next(images) yield prev_image processed_images += 1 pbar.update_absolute(processed_images, total_images) prev_image = next_image except StopIteration: pass if prev_image is not None: yield prev_image def map_neg1_1_to_0_1(t): """ 接受 torch.Tensor 或可转 torch.Tensor 的输入。 可处理形状: H,W,C 或 B,H,W,C 或 B,T,H,W,C（会按元素处理）。 返回 float tensor，范围 0..1。 """ if not torch.is_tensor(t): t = torch.tensor(t) t = t.float() # map -1..1 -> 0..1 t = (t + 1.0) * 0.5 # 限幅到 [0,1] t = t.clamp(0.0, 1.0) # 保持在 cpu 端，调用方可决定是否转 device/dtype return t.cpu() def load_images_list(img_list: list, ): gen = images_generator(img_list) (width, height) = next(gen) images = torch.from_numpy(np.fromiter(gen, np.dtype((np.float32, (height, width, 3))))) if len(images) == 0: raise FileNotFoundError(f"No images could be loaded .") return images def get_video_files(directory, extensions=None): if extensions is None: extensions = ['webm', 'mp4', 'mkv', 'gif', 'mov'] extensions = [ext.lower() for ext in extensions] video_files = [] for root, dirs, files in os.walk(directory): for file in files: _, ext = os.path.splitext(file) ext = ext.lower()[1:] if ext in extensions: full_path = os.path.join(root, file) video_files.append(full_path) return video_files