Merge pull request #35 from smthemex/Pr

init

FlashVSR/diffsynth/models/wan_video_dit.py

@@ -38,6 +38,62 @@ except:
     print("**************导入block_sparse_attn失败,调用临时函数;Using wrapper block_sparse_attn ,need to install block_sparse_attn to get better performance*************")
     from ...examples.WanVSR.utils.utils import block_sparse_attn_func
 
+class BlockGPUManager:
+    def __init__(self, device="cuda",):
+        self.device = device
+        self.managed_modules = []
+        self.embedder_modules = []
+    
+    def setup_for_inference(self, transformer_model,):
+        self._collect_managed_modules(transformer_model)
+        self._initialize_embedder_modules()
+        return self
+
+    def _collect_managed_modules(self, transformer_model):
+        self.managed_modules = []
+        self.embedder_modules = []
+
+        for i, block in enumerate(transformer_model.blocks):
+            self.managed_modules.append(block)
+
+        if hasattr(transformer_model, 'time_embedding'):
+            self.embedder_modules.append(transformer_model.time_embedding)
+        
+        if hasattr(transformer_model, 'time_projection'):
+            self.embedder_modules.append(transformer_model.time_projection)
+        
+        if hasattr(transformer_model, 'text_embedding'):
+            self.embedder_modules.append(transformer_model.text_embedding)
+        
+        if hasattr(transformer_model, 'patch_embedding'):
+            self.embedder_modules.append(transformer_model.patch_embedding)
+        
+        if hasattr(transformer_model, 'head'):
+            self.embedder_modules.append(transformer_model.head)
+        
+        # if hasattr(transformer_model, 'LQ_proj_in'):
+        #     self.embedder_modules.append(transformer_model.LQ_proj_in)
+
+  
+    def _initialize_embedder_modules(self):
+        for module in self.embedder_modules:
+            if hasattr(module, 'to'):
+                module.to(self.device, non_blocking=True)
+        return self
+
+
+    def unload_all_blocks_to_cpu(self):
+        for module in self.managed_modules:
+            if hasattr(module, 'to'):
+                module.to('cpu')
+
+        for module in self.embedder_modules:
+            if hasattr(module, 'to'):
+                module.to('cpu')
+        
+        
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
 
 # ----------------------------
 # Local / window masks
@@ -609,6 +665,7 @@ class WanModel(torch.nn.Module):
                 local_num: Optional[int] = None,
                 is_full_block: bool = False,
                 causal_idx: Optional[int] = None,
+                gpu_manager=None,
                 **kwargs,
                 ):
         # time / text embeds
@@ -665,6 +722,16 @@ class WanModel(torch.nn.Module):
             if LQ_latents is not None and block_id < len(LQ_latents):
                 x += LQ_latents[block_id]
 
+            if gpu_manager is not None: # apple 1 block to save Vram
+                if block_id < len(self.blocks):
+                    module = gpu_manager.managed_modules[block_id]
+                    if hasattr(module, 'to'):
+                        module.to(gpu_manager.device)
+                if block_id > 0 and (block_id - 1) < len(self.blocks):
+                    prev_module = gpu_manager.managed_modules[block_id - 1]
+                    if hasattr(prev_module, 'to'):
+                        prev_module.to('cpu')
+
             if self.training and use_gradient_checkpointing:
                 if use_gradient_checkpointing_offload:
                     with torch.autograd.graph.save_on_cpu():

FlashVSR/diffsynth/pipelines/flashvsr_full.py

@@ -8,7 +8,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 import numpy as np
 from einops import rearrange
-from PIL import Image
+from ..models.wan_video_dit import BlockGPUManager
 from tqdm import tqdm
 # import pyfiglet
 
@@ -162,7 +162,7 @@ class FlashVSRFullPipeline(BasePipeline):
         self.width_division_factor = 16
         self.use_unified_sequence_parallel = False
         self.prompt_emb_posi = None
-        self.ColorCorrector = TorchColorCorrectorWavelet(levels=5)
+        #self.ColorCorrector = TorchColorCorrectorWavelet(levels=5)
         self.new_decoder=False
         self.VAE=None
         self.version="1.0"
@@ -178,60 +178,63 @@ class FlashVSRFullPipeline(BasePipeline):
                           ⚡FlashVSR
 """)
 
-    def enable_vram_management(self, num_persistent_param_in_dit=None):
+    def enable_vram_management(self, num_persistent_param_in_dit=None,vae_only=False):
         # 仅管理 dit / vae
-        dtype = next(iter(self.dit.parameters())).dtype
         from ..vram_management import enable_vram_management, AutoWrappedModule, AutoWrappedLinear
-        enable_vram_management(
-            self.dit,
-            module_map={
-                torch.nn.Linear: AutoWrappedLinear,
-                torch.nn.Conv3d: AutoWrappedModule,
-                torch.nn.LayerNorm: AutoWrappedModule,
-                RMSNorm: AutoWrappedModule,
-            },
-            module_config=dict(
-                offload_dtype=dtype,
-                offload_device="cpu",
-                onload_dtype=dtype,
-                onload_device=self.device,
-                computation_dtype=self.torch_dtype,
-                computation_device=self.device,
-            ),
-            max_num_param=num_persistent_param_in_dit,
-            overflow_module_config=dict(
-                offload_dtype=dtype,
-                offload_device="cpu",
-                onload_dtype=dtype,
-                onload_device="cpu",
-                computation_dtype=self.torch_dtype,
-                computation_device=self.device,
-            ),
-        )
-        dtype = next(iter(self.vae.parameters())).dtype
-        enable_vram_management(
-            self.vae,
-            module_map={
-                torch.nn.Linear: AutoWrappedLinear,
-                torch.nn.Conv2d: AutoWrappedModule,
-                RMS_norm: AutoWrappedModule,
-                CausalConv3d: AutoWrappedModule,
-                Upsample: AutoWrappedModule,
-                torch.nn.SiLU: AutoWrappedModule,
-                torch.nn.Dropout: AutoWrappedModule,
-            },
-            module_config=dict(
-                offload_dtype=dtype,
-                offload_device="cpu",
-                onload_dtype=dtype,
-                onload_device=self.device,
-                computation_dtype=self.torch_dtype,
-                computation_device=self.device,
-            ),
-        )
+        if not vae_only:
+            dtype = next(iter(self.dit.parameters())).dtype
+            
+            enable_vram_management(
+                self.dit,
+                module_map={
+                    torch.nn.Linear: AutoWrappedLinear,
+                    torch.nn.Conv3d: AutoWrappedModule,
+                    torch.nn.LayerNorm: AutoWrappedModule,
+                    RMSNorm: AutoWrappedModule,
+                },
+                module_config=dict(
+                    offload_dtype=dtype,
+                    offload_device="cpu",
+                    onload_dtype=dtype,
+                    onload_device=self.device,
+                    computation_dtype=self.torch_dtype,
+                    computation_device=self.device,
+                ),
+                max_num_param=num_persistent_param_in_dit,
+                overflow_module_config=dict(
+                    offload_dtype=dtype,
+                    offload_device="cpu",
+                    onload_dtype=dtype,
+                    onload_device="cpu",
+                    computation_dtype=self.torch_dtype,
+                    computation_device=self.device,
+                ),
+            )
+        if self.vae is not None:
+            dtype = next(iter(self.vae.parameters())).dtype
+            enable_vram_management(
+                self.vae,
+                module_map={
+                    torch.nn.Linear: AutoWrappedLinear,
+                    torch.nn.Conv2d: AutoWrappedModule,
+                    RMS_norm: AutoWrappedModule,
+                    CausalConv3d: AutoWrappedModule,
+                    Upsample: AutoWrappedModule,
+                    torch.nn.SiLU: AutoWrappedModule,
+                    torch.nn.Dropout: AutoWrappedModule,
+                },
+                module_config=dict(
+                    offload_dtype=dtype,
+                    offload_device="cpu",
+                    onload_dtype=dtype,
+                    onload_device=self.device,
+                    computation_dtype=self.torch_dtype,
+                    computation_device=self.device,
+                ),
+            )
         self.enable_cpu_offload()
 
-    def fetch_models(self, model_manager: ModelManager):
+    def fetch_models(self, model_manager: ModelManager):  
         self.dit = model_manager.fetch_model("wan_video_dit")
         self.vae = model_manager.fetch_model("wan_video_vae")
 
@@ -243,7 +246,6 @@ class FlashVSRFullPipeline(BasePipeline):
         pipe.fetch_models(model_manager)
         # 可选：统一序列并行入口（此处默认关闭）
         pipe.use_unified_sequence_parallel = False
-        
         return pipe
 
     def denoising_model(self):
@@ -381,7 +383,17 @@ class FlashVSRFullPipeline(BasePipeline):
         kv_ratio=3.0,
         local_range = 9,
         color_fix = True,
+        offload=False,
+
     ):
+        
+        
+        if not offload:
+            gpu_manager=None
+        else:
+            gpu_manager = BlockGPUManager(device="cuda")
+            gpu_manager.setup_for_inference(self.dit)
+
         # 只接受 cfg=1.0（与原代码一致）
         assert cfg_scale == 1.0, "cfg_scale must be 1.0"
 
@@ -416,10 +428,11 @@ class FlashVSRFullPipeline(BasePipeline):
         if hasattr(self.dit, "LQ_proj_in"):
             self.dit.LQ_proj_in.clear_cache()
 
-
         latents_total = []
-        self.vae.clear_cache()
-
+        if self.vae is not None:
+            self.vae.clear_cache()
+        LQ_pre_idx = 0
+        LQ_cur_idx = 0
         with torch.no_grad():
             for cur_process_idx in tqdm(range(process_total_num)):
                 torch.cuda.synchronize()
@@ -430,6 +443,11 @@ class FlashVSRFullPipeline(BasePipeline):
                     pre_cache_v = [None] * len(self.dit.blocks)
                     LQ_latents = None
                     inner_loop_num = 7
+                    if gpu_manager is not None: # apple 1 block to save Vram
+                        if hasattr(self.denoising_model().LQ_proj_in, 'to'):
+                            #print("move LQ_proj_in to gpu for the first 7 frames")
+                            self.denoising_model().LQ_proj_in.to(gpu_manager.device)
+                      
                     for inner_idx in range(inner_loop_num):
                         cur = self.denoising_model().LQ_proj_in.stream_forward(
                             LQ_video[:, :, max(0, inner_idx*4-3):(inner_idx+1)*4-3, :, :]
@@ -441,10 +459,18 @@ class FlashVSRFullPipeline(BasePipeline):
                         else:
                             for layer_idx in range(len(LQ_latents)):
                                 LQ_latents[layer_idx] = torch.cat([LQ_latents[layer_idx], cur[layer_idx]], dim=1)
+                    LQ_cur_idx = (inner_loop_num-1)*4-3
                     cur_latents = latents[:, :, :6, :, :]
+                    if gpu_manager is not None: # apple 1 block to save Vram
+                        if hasattr(self.denoising_model().LQ_proj_in, 'to'):     
+                            self.denoising_model().LQ_proj_in.to('cpu')
                 else:
                     LQ_latents = None
                     inner_loop_num = 2
+                    if gpu_manager is not None: # apple 1 block to save Vram
+                        if hasattr(self.denoising_model().LQ_proj_in, 'to'):
+                            #print("move LQ_proj_in to gpu for the rest frames")
+                            self.denoising_model().LQ_proj_in.to(gpu_manager.device)
                     for inner_idx in range(inner_loop_num):
                         cur = self.denoising_model().LQ_proj_in.stream_forward(
                             LQ_video[:, :, cur_process_idx*8+17+inner_idx*4:cur_process_idx*8+21+inner_idx*4, :, :]
@@ -456,8 +482,11 @@ class FlashVSRFullPipeline(BasePipeline):
                         else:
                             for layer_idx in range(len(LQ_latents)):
                                 LQ_latents[layer_idx] = torch.cat([LQ_latents[layer_idx], cur[layer_idx]], dim=1)
+                    LQ_cur_idx = cur_process_idx*8+21+(inner_loop_num-2)*4            
                     cur_latents = latents[:, :, 4+cur_process_idx*2:6+cur_process_idx*2, :, :]
-
+                    if gpu_manager is not None: # apple 1 block to save Vram
+                        if hasattr(self.denoising_model().LQ_proj_in, 'to'):     
+                            self.denoising_model().LQ_proj_in.to('cpu')
                 # 推理（无 motion_controller / vace）
                 #print(cur_latents.shape,"cur_latents") #torch.Size([1, 16, 6, 384, 640]) scale 4  cur_latents orch.Size([1, 16, 6, 96, 160]) cur_latents scale 1
                 noise_pred_posi, pre_cache_k, pre_cache_v = model_fn_wan_video(
@@ -478,13 +507,17 @@ class FlashVSRFullPipeline(BasePipeline):
                     t_mod=self.t_mod,
                     t=self.t,
                     local_range = local_range,
+                    gpu_manager=gpu_manager,
                 )
 
                 # 更新 latent
                 cur_latents = cur_latents - noise_pred_posi
                 latents_total.append(cur_latents)
+                LQ_pre_idx = LQ_cur_idx
 
             latents = torch.cat(latents_total, dim=2)  #torch.Size([1, 16, 20, 48, 80])
+            if offload:
+                gpu_manager.unload_all_blocks_to_cpu()
             #self.dit.to("cpu")
             #torch.cuda.empty_cache
             # Decode
@@ -503,7 +536,7 @@ class FlashVSRFullPipeline(BasePipeline):
             # except:
             #     pass
                 
-        return latents #frames[0]
+        return latents,LQ_cur_idx #frames[0]
 
 
 # -----------------------------
@@ -578,6 +611,7 @@ def model_fn_wan_video(
     t_mod : torch.Tensor = None,
     t : torch.Tensor = None,
     local_range: int = 9,
+    gpu_manager=None,
     **kwargs,
 ):
     # patchify
@@ -622,6 +656,17 @@ def model_fn_wan_video(
         x = tea_cache.update(x)
     else:
         for block_id, block in enumerate(dit.blocks):
+            if gpu_manager is not None: # apple 1 block to save Vram
+                if block_id < len(dit.blocks):
+                    module = gpu_manager.managed_modules[block_id]
+                    if hasattr(module, 'to'):
+                        #print(f"move block {block_id} to gpu")
+                        module.to(gpu_manager.device)
+                if block_id > 0 and (block_id - 1) < len(dit.blocks):
+                    prev_module = gpu_manager.managed_modules[block_id - 1]
+                    if hasattr(prev_module, 'to'):
+                        prev_module.to('cpu')
+
             if LQ_latents is not None and block_id < len(LQ_latents):
                 x = x + LQ_latents[block_id]
             x, last_pre_cache_k, last_pre_cache_v = block(

FlashVSR/diffsynth/pipelines/flashvsr_tiny.py

@@ -309,6 +309,7 @@ class FlashVSRTinyPipeline(BasePipeline):
         kv_ratio=3.0,
         local_range = 9,
         color_fix = True,
+        offload=False,
     ):
         # 只接受 cfg=1.0（与原代码一致）
         assert cfg_scale == 1.0, "cfg_scale must be 1.0"
@@ -347,7 +348,8 @@ class FlashVSRTinyPipeline(BasePipeline):
             self.dit.LQ_proj_in.clear_cache()
 
         latents_total = []
-        self.TCDecoder.clean_mem()
+        if self.TCDecoder is not None:
+            self.TCDecoder.clean_mem()
         LQ_pre_idx = 0
         LQ_cur_idx = 0
 

FlashVSR/diffsynth/pipelines/flashvsr_tiny_long.py

@@ -8,7 +8,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 import numpy as np
 from einops import rearrange
-from PIL import Image
+from ..models.wan_video_dit import BlockGPUManager
 from tqdm import tqdm
 # import pyfiglet
 
@@ -157,6 +157,7 @@ class FlashVSRTinyLongPipeline(BasePipeline):
         self.scheduler = FlowMatchScheduler(shift=5, sigma_min=0.0, extra_one_step=True)
         self.dit: WanModel = None
         self.vae: WanVideoVAE = None
+        self.TCDecoder = None
         self.model_names = ['dit', 'vae']
         self.height_division_factor = 16
         self.width_division_factor = 16
@@ -175,10 +176,10 @@ class FlashVSRTinyLongPipeline(BasePipeline):
                          ⚡FlashVSR
 """)
 
-    def enable_vram_management(self, num_persistent_param_in_dit=None):
+    def enable_vram_management(self, num_persistent_param_in_dit=None,):
         # 仅管理 dit / vae
-        dtype = next(iter(self.dit.parameters())).dtype
         from ..vram_management import enable_vram_management, AutoWrappedModule, AutoWrappedLinear
+        dtype = next(iter(self.dit.parameters())).dtype
         enable_vram_management(
             self.dit,
             module_map={
@@ -310,7 +311,14 @@ class FlashVSRTinyLongPipeline(BasePipeline):
         color_fix = True,
         pad_first_frame=False,
         fix_method="wavelet",
+        offload=False,
     ):
+        if not offload:
+            gpu_manager=None
+        else:
+            gpu_manager = BlockGPUManager(device="cuda")
+            gpu_manager.setup_for_inference(self.dit)
+
         # 只接受 cfg=1.0（与原代码一致）
         assert cfg_scale == 1.0, "cfg_scale must be 1.0"
 
@@ -346,11 +354,13 @@ class FlashVSRTinyLongPipeline(BasePipeline):
         # 清理可能存在的 LQ_proj_in cache
         if hasattr(self.dit, "LQ_proj_in"):
             self.dit.LQ_proj_in.clear_cache()
-
-        self.TCDecoder.clean_mem()
+        if self.TCDecoder is not None:
+            self.TCDecoder.clean_mem()
         LQ_pre_idx = 0
         LQ_cur_idx = 0
         frames_total = []
+        lQ_cur_list = []
+
 
         with torch.no_grad():
             for cur_process_idx in tqdm(range(process_total_num)):
@@ -359,6 +369,9 @@ class FlashVSRTinyLongPipeline(BasePipeline):
                     pre_cache_v = [None] * len(self.dit.blocks)
                     LQ_latents = None
                     inner_loop_num = 7
+                    if gpu_manager is not None: # apple 1 block to save Vram
+                        if hasattr(self.denoising_model().LQ_proj_in, 'to'):
+                            self.denoising_model().LQ_proj_in.to(gpu_manager.device)
                     for inner_idx in range(inner_loop_num):
                         cur = self.denoising_model().LQ_proj_in.stream_forward(
                             LQ_video[:, :, max(0, inner_idx*4-3):(inner_idx+1)*4-3, :, :].to(self.device)
@@ -372,9 +385,15 @@ class FlashVSRTinyLongPipeline(BasePipeline):
                                 LQ_latents[layer_idx] = torch.cat([LQ_latents[layer_idx], cur[layer_idx]], dim=1)
                     LQ_cur_idx = (inner_loop_num-1)*4-3
                     cur_latents = latents[:, :, :6, :, :]
+                    if gpu_manager is not None: # apple 1 block to save Vram
+                        if hasattr(self.denoising_model().LQ_proj_in, 'to'):     
+                            self.denoising_model().LQ_proj_in.to('cpu')
                 else:
                     LQ_latents = None
                     inner_loop_num = 2
+                    if gpu_manager is not None: # apple 1 block to save Vram
+                        if hasattr(self.denoising_model().LQ_proj_in, 'to'):
+                            self.denoising_model().LQ_proj_in.to(gpu_manager.device)
                     for inner_idx in range(inner_loop_num):
                         cur = self.denoising_model().LQ_proj_in.stream_forward(
                             LQ_video[:, :, cur_process_idx*8+17+inner_idx*4:cur_process_idx*8+21+inner_idx*4, :, :].to(self.device)
@@ -388,6 +407,9 @@ class FlashVSRTinyLongPipeline(BasePipeline):
                                 LQ_latents[layer_idx] = torch.cat([LQ_latents[layer_idx], cur[layer_idx]], dim=1)
                     LQ_cur_idx = cur_process_idx*8+21+(inner_loop_num-2)*4
                     cur_latents = latents[:, :, 4+cur_process_idx*2:6+cur_process_idx*2, :, :]
+                    if gpu_manager is not None: # apple 1 block to save Vram
+                        if hasattr(self.denoising_model().LQ_proj_in, 'to'):     
+                            self.denoising_model().LQ_proj_in.to('cpu')
 
                 # 推理（无 motion_controller / vace）
                 noise_pred_posi, pre_cache_k, pre_cache_v = model_fn_wan_video(
@@ -408,38 +430,41 @@ class FlashVSRTinyLongPipeline(BasePipeline):
                     t_mod=self.t_mod,
                     t=self.t,
                     local_range = local_range,
+                    gpu_manager=gpu_manager,
                 )
 
                 # 更新 latent
                 cur_latents = cur_latents - noise_pred_posi
                 # Decode
-                cur_LQ_frame = LQ_video[:,:,LQ_pre_idx:LQ_cur_idx,:,:].to(self.device)
-                cur_frames = self.TCDecoder.decode_video(cur_latents.transpose(1, 2),parallel=False, show_progress_bar=False, cond=LQ_video[:,:,LQ_pre_idx:LQ_cur_idx,:,:].to(self.device)).transpose(1, 2).mul_(2).sub_(1)
+                lQ_cur_list.append((LQ_cur_idx,LQ_pre_idx,))
+                #cur_LQ_frame = LQ_video[:,:,LQ_pre_idx:LQ_cur_idx,:,:].to(self.device)
+                #cur_frames = self.TCDecoder.decode_video(cur_latents.transpose(1, 2),parallel=False, show_progress_bar=False, cond=LQ_video[:,:,LQ_pre_idx:LQ_cur_idx,:,:].to(self.device)).transpose(1, 2).mul_(2).sub_(1)
 
-                # 颜色校正（wavelet）
-                try:
-                    if color_fix:
-                        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 = self.ColorCorrector(
-                            cur_frames.to(device=self.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:
-                    pass
+                # # 颜色校正（wavelet）
+                # try:
+                #     if color_fix:
+                #         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 = self.ColorCorrector(
+                #             cur_frames.to(device=self.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:
+                #     pass
 
-                frames_total.append(cur_frames.to('cpu'))
+                #frames_total.append(cur_frames.to('cpu'))
+                frames_total.append(cur_latents)
                 LQ_pre_idx = LQ_cur_idx
 
-            frames = torch.cat(frames_total, dim=2)
+            #frames = torch.cat(frames_total, dim=2)
 
-        return frames[0]
+        return frames_total,lQ_cur_list
 
 def dup_first_frame_1cthw_simple(video_tensor):
     return torch.cat([video_tensor[:, :, :1], video_tensor], dim=2)
@@ -516,6 +541,7 @@ def model_fn_wan_video(
     t_mod : torch.Tensor = None,
     t : torch.Tensor = None,
     local_range: int = 9,
+    gpu_manager=None,
     **kwargs,
 ):
     # patchify
@@ -560,6 +586,16 @@ def model_fn_wan_video(
         x = tea_cache.update(x)
     else:
         for block_id, block in enumerate(dit.blocks):
+            if gpu_manager is not None: # apple 1 block to save Vram
+                if block_id < len(dit.blocks):
+                    module = gpu_manager.managed_modules[block_id]
+                    if hasattr(module, 'to'):
+                        #print(f"move block {block_id} to gpu")
+                        module.to(gpu_manager.device)
+                if block_id > 0 and (block_id - 1) < len(dit.blocks):
+                    prev_module = gpu_manager.managed_modules[block_id - 1]
+                    if hasattr(prev_module, 'to'):
+                        prev_module.to('cpu')
             if LQ_latents is not None and block_id < len(LQ_latents):
                 x = x + LQ_latents[block_id]
             x, last_pre_cache_k, last_pre_cache_v = block(

FlashVSR/examples/WanVSR/infer_flashvsr_full.py

@@ -10,7 +10,7 @@ import torch
 from einops import rearrange
 import folder_paths
 from ...diffsynth import ModelManager, FlashVSRFullPipeline
-from .utils.utils import Buffer_LQ4x_Proj
+from .utils.utils import Buffer_LQ4x_Proj,Causal_LQ4x_Proj
 from comfy.utils import common_upscale
 from safetensors.torch import load_file
 from .utils.utils import calculate_frame_adjustment_simple
@@ -253,151 +253,117 @@ def prepare_input_tensor(path, scale: int = 4,fps=30,dtype=torch.bfloat16, devic
     else:
         raise ValueError(f"Unsupported input: {path}")
 
-def init_pipeline(prompt_path,LQ_proj_in_path="./FlashVSR/LQ_proj_in.ckpt",ckpt_path: str = "./FlashVSR/diffusion_pytorch_model_streaming_dmd.safetensors", vae_path: str = "./FlashVSR/Wan2.1_VAE.pth",decode_vae="none",cur_dir="",device="cuda"):
-    #print(torch.cuda.current_device(), torch.cuda.get_device_name(torch.cuda.current_device()))
+def init_pipeline(prompt_path,LQ_proj_in_path="./FlashVSR/LQ_proj_in.ckpt",ckpt_path: str = "./FlashVSR/diffusion_pytorch_model_streaming_dmd.safetensors",cur_dir="",device="cuda",offload=False,version='1.0'):
     mm = ModelManager(torch_dtype=torch.bfloat16, device="cpu")
-    new_decoder=True if decode_vae!="none"  else False
-    mm.load_models([ckpt_path,vae_path,])
+    mm.load_models([ckpt_path])
     pipe = FlashVSRFullPipeline.from_model_manager(mm, device=device)
-    
-    if new_decoder:
-        pipe.new_decoder = True
-        if "light" in decode_vae.lower() or "tae" in decode_vae.lower():
-            if os.path.basename(decode_vae).split(".")[0]=="lightvaew2_1":
-                from ...vae import WanVAE
-                print("use lightvae decoder")
-                VAE = WanVAE(vae_path=decode_vae,dtype=torch.bfloat16,device=device,use_lightvae=True)
-            elif os.path.basename(decode_vae).split(".")[0]=="taew2_1":
-                from ...vae_tiny import WanVAE_tiny
-                print("use vae_tiny decoder")
-                VAE = WanVAE_tiny(vae_path=decode_vae,dtype=torch.bfloat16,device=device,need_scaled=False)
-            elif os.path.basename(decode_vae).split(".")[0]=="lighttaew2_1":
-                from ...vae_tiny import WanVAE_tiny
-                print("use vae_tiny light decoder")
-                VAE = WanVAE_tiny(vae_path=decode_vae,dtype=torch.bfloat16,device=device,need_scaled=True)
-            else:
-                raise ValueError(f"Unknown vae_name: {decode_vae},only support lightvae,tae,tae_tiny,lighttae_tiny")
-            pipe.VAE=VAE
-        else:    
-            print("use upscale2x decoder")
-            from diffusers import AutoencoderKLWan
-            config=AutoencoderKLWan.load_config(os.path.join(cur_dir,"FlashVSR/examples/config.json"))
-            VAE=AutoencoderKLWan.from_config(config).to(device,dtype=torch.bfloat16)
-            vae_dict=load_file(decode_vae,device="cpu")
-            VAE.load_state_dict(vae_dict,strict=False)
-            pipe.VAE=VAE
-            del vae_dict
-    
-        
-    pipe.denoising_model().LQ_proj_in = Buffer_LQ4x_Proj(in_dim=3, out_dim=1536, layer_num=1).to(device, dtype=torch.bfloat16)
+    if offload:
+        device = "cpu"    
+    if version=="1.0":
+        pipe.denoising_model().LQ_proj_in = Buffer_LQ4x_Proj(in_dim=3, out_dim=1536, layer_num=1).to(device, dtype=torch.bfloat16)
+    else:
+        pipe.denoising_model().LQ_proj_in = Causal_LQ4x_Proj(in_dim=3, out_dim=1536, layer_num=1).to(device, dtype=torch.bfloat16)
     if os.path.exists(LQ_proj_in_path):
         pipe.denoising_model().LQ_proj_in.load_state_dict(torch.load(LQ_proj_in_path, map_location="cpu",weights_only=False), strict=True)
-    pipe.denoising_model().LQ_proj_in.to(device)
-    if pipe.vae is  None:  # safetensors cause error or unknown safetensors or pth
-        from diffusers import AutoencoderKLWan
-        config=AutoencoderKLWan.load_config(os.path.join(cur_dir,"FlashVSR/config.json"))
-        vae=AutoencoderKLWan.from_config(config).to(device,dtype=torch.bfloat16)
-        vae_dict=torch.load(vae_path, map_location="cpu",weights_only=False) if not vae_path.endswith(".safetensors") else load_file(vae_path,device="cpu")
-        vae.load_state_dict(vae_dict,strict=False)
-        del vae_dict
-        pipe.vae=vae
-        pipe.vae.encoder = None
-    else:
-        pipe.vae.model.encoder = None
-        pipe.vae.model.conv1 = None
+    pipe.denoising_model().LQ_proj_in.to(device)  
     pipe.enable_vram_management(num_persistent_param_in_dit=None)
-    pipe.init_cross_kv(prompt_path); pipe.load_models_to_device(["dit","vae"])
+    pipe.init_cross_kv(prompt_path); pipe.load_models_to_device(["dit"])
+    pipe.offload=offload
     return pipe
 
 
 
-def run_inference(pipe,input,seed,scale,kv_ratio=3.0,local_range=9,step=1,cfg_scale=1.0,sparse_ratio=2.0,tiled=True,color_fix=True,fix_method="wavelet",split_num=81,dtype=torch.bfloat16,device="cuda",save_vodeo_=False,):
-    pipe.to('cuda')  #pipe.enable_vram_management(num_persistent_param_in_dit=None)
+def run_inference(pipe,input,seed,scale,kv_ratio=3.0,local_range=9,step=1,cfg_scale=1.0,sparse_ratio=2.0,offload=False,dtype=torch.bfloat16,device="cuda",):
+    if not offload:
+        pipe.to('cuda')  #pipe.enable_vram_management(num_persistent_param_in_dit=None)
     #pipe.init_cross_kv(prompt_path); pipe.load_models_to_device(["dit","vae"])
-    pad_first_frame = True  if "wavelet"== fix_method and color_fix else False
+    #pad_first_frame = True  if "wavelet"== fix_method and color_fix else False
 
     #total,h0,w0,_ = input.shape
     torch.cuda.empty_cache(); torch.cuda.ipc_collect()
 
     LQ, th, tw, F, fps,pad_frames = prepare_input_tensor(input, scale=scale, dtype=dtype, device=device)
 
-    frames = pipe(
-        prompt="", negative_prompt="", cfg_scale=cfg_scale, num_inference_steps=step, seed=seed, tiled=tiled,
+    frames,LQ_cur_idx = pipe(
+        prompt="", negative_prompt="", cfg_scale=cfg_scale, num_inference_steps=step, seed=seed, tiled=True,
         LQ_video=LQ, num_frames=F, height=th, width=tw, is_full_block=False, if_buffer=True,
         topk_ratio=sparse_ratio*768*1280/(th*tw), 
         kv_ratio=kv_ratio,
         local_range=local_range, # Recommended: 9 or 11. local_range=9 → sharper details; 11 → more stable results.
-        color_fix = color_fix,
+        color_fix = True,
+        offload=offload,
     )
-    pipe.dit.to('cpu')
+    
+    if not offload:
+        pipe.dit.to('cpu')
     torch.cuda.empty_cache()
     #torch.Size([1, 16, 20, 48, 80])
-    tiler_kwargs = {"tiled": tiled, "tile_size": (60, 104), "tile_stride": (30, 52)}
-    with torch.no_grad():
-        try:
-            frames = pipe.decode_video(frames, **tiler_kwargs)
-        except:
-            print("vae decode_video OOM.try split latent" )
-            if pipe.new_decoder:
-                if pipe.VAE.__class__.__name__ == "AutoencoderKLWan":
-                    pipe.VAE.to('cpu')
-                else:
-                    if pipe.VAE.__class__.__name__ == "WanVAE":
-                        pipe.VAE.to_cpu()
-                    else: pass
-            else:
-                pipe.vae.to('cpu')
-            torch.cuda.empty_cache()  
-            if pipe.new_decoder:
-                if pipe.VAE.__class__.__name__ == "AutoencoderKLWan":
-                    pipe.VAE.to('cuda') 
-                else:
-                    if pipe.VAE.__class__.__name__ == "WanVAE":
-                        pipe.VAE.to_cuda()
-                    else: pass
-            else:
-                pipe.vae.to('cuda')
-            total_frames = frames.shape[2]
-            segment_size = (split_num-1) * 2 // 4 # 40
-            decoded_frames_list = []
-            for start_idx in range(0, total_frames, segment_size):
-                end_idx = min(start_idx + segment_size, total_frames)
-                frames_segment = frames[:, :, start_idx:end_idx, :, :]          
-                decoded_segment = pipe.decode_video(frames_segment, **tiler_kwargs)
-                decoded_frames_list.append(decoded_segment)
-            frames = torch.cat(decoded_frames_list, dim=2)  
-        try:
-            if color_fix:
-                if pad_first_frame:
-                    frames = dup_first_frame_1cthw_simple(frames)
-                    LQ=dup_first_frame_1cthw_simple(LQ)
-                if pipe.new_decoder and LQ.shape[-1]!=frames.shape[-1]:
-                    scale_=int(frames.shape[-1]/LQ.shape[-1])
-                    LQ=upscale_lq_video_bilinear(LQ,scale_)
-                frames = pipe.ColorCorrector(
-                    frames.to(device=device),
-                    LQ[:, :, :frames.shape[2], :, :],
-                    clip_range=(-1, 1),
-                    chunk_size=16,
-                    method=fix_method
-                )
-                if pad_first_frame:
-                    frames = frames[:, :, 1:, :, :] # remove first  frame torch.Size([1, 3, 78, 1024, 768]) --> torch.Size([1, 3, 77, 1024, 768])
-        except:
-            pass
-        print("Done.")
-        pipe.vae.to('cpu')  
-    del LQ
-    torch.cuda.empty_cache() 
+    # tiler_kwargs = {"tiled": tiled, "tile_size": (60, 104), "tile_stride": (30, 52)}
+    # with torch.no_grad():
+    #     try:
+    #         frames = pipe.decode_video(frames, **tiler_kwargs)
+    #     except:
+    #         print("vae decode_video OOM.try split latent" )
+    #         if pipe.new_decoder:
+    #             if pipe.VAE.__class__.__name__ == "AutoencoderKLWan":
+    #                 pipe.VAE.to('cpu')
+    #             else:
+    #                 if pipe.VAE.__class__.__name__ == "WanVAE":
+    #                     pipe.VAE.to_cpu()
+    #                 else: pass
+    #         else:
+    #             pipe.vae.to('cpu')
+    #         torch.cuda.empty_cache()  
+    #         if pipe.new_decoder:
+    #             if pipe.VAE.__class__.__name__ == "AutoencoderKLWan":
+    #                 pipe.VAE.to('cuda') 
+    #             else:
+    #                 if pipe.VAE.__class__.__name__ == "WanVAE":
+    #                     pipe.VAE.to_cuda()
+    #                 else: pass
+    #         else:
+    #             pipe.vae.to('cuda')
+    #         total_frames = frames.shape[2]
+    #         segment_size = (split_num-1) * 2 // 4 # 40
+    #         decoded_frames_list = []
+    #         for start_idx in range(0, total_frames, segment_size):
+    #             end_idx = min(start_idx + segment_size, total_frames)
+    #             frames_segment = frames[:, :, start_idx:end_idx, :, :]          
+    #             decoded_segment = pipe.decode_video(frames_segment, **tiler_kwargs)
+    #             decoded_frames_list.append(decoded_segment)
+    #         frames = torch.cat(decoded_frames_list, dim=2)  
+    #     try:
+    #         if color_fix:
+    #             if pad_first_frame:
+    #                 frames = dup_first_frame_1cthw_simple(frames)
+    #                 LQ=dup_first_frame_1cthw_simple(LQ)
+    #             if pipe.new_decoder and LQ.shape[-1]!=frames.shape[-1]:
+    #                 scale_=int(frames.shape[-1]/LQ.shape[-1])
+    #                 LQ=upscale_lq_video_bilinear(LQ,scale_)
+    #             frames = pipe.ColorCorrector(
+    #                 frames.to(device=device),
+    #                 LQ[:, :, :frames.shape[2], :, :],
+    #                 clip_range=(-1, 1),
+    #                 chunk_size=16,
+    #                 method=fix_method
+    #             )
+    #             if pad_first_frame:
+    #                 frames = frames[:, :, 1:, :, :] # remove first  frame torch.Size([1, 3, 78, 1024, 768]) --> torch.Size([1, 3, 77, 1024, 768])
+    #     except:
+    #         pass
+    #     print("Done.")
+    #     pipe.vae.to('cpu')  
+    # del LQ
+    # torch.cuda.empty_cache() 
 
-    if pad_frames>0:
-        print("Remove pad frames",pad_frames,frames.shape)
-        frames =  frames = frames[:, :, :-pad_frames, :, :]  #torch.Size([1, 3, 85, 1024, 768]) -> torch.Size([1, 3, 81, 1024, 768]) if 81 output -4 input + 8
+    # if pad_frames>0:
+    #     print("Remove pad frames",pad_frames,frames.shape)
+    #     frames =  frames = frames[:, :, :-pad_frames, :, :]  #torch.Size([1, 3, 85, 1024, 768]) -> torch.Size([1, 3, 81, 1024, 768]) if 81 output -4 input + 8
         
-    frames = tensor2video(frames[0])    
-    if save_vodeo_:
-        save_video(frames, os.path.join(folder_paths.get_output_directory(),f"FlashVSR_Full_seed{seed}.mp4"), fps=fps, quality=6)
-    return frames
+    # frames = tensor2video(frames[0])    
+    # if save_vodeo_:
+    #     save_video(frames, os.path.join(folder_paths.get_output_directory(),f"FlashVSR_Full_seed{seed}.mp4"), fps=fps, quality=6)
+    return frames,LQ_cur_idx,LQ,pad_frames
 
 def upscale_lq_video_bilinear(LQ_video,scale_):
     B, C, T, H, W = LQ_video.shape

FlashVSR/examples/WanVSR/infer_flashvsr_tiny.py

@@ -239,11 +239,13 @@ def prepare_input_tensor(path: str, scale: float = 4,fps=30, dtype=torch.bfloat1
         raise ValueError(f"Unsupported input: {path}")
    
 
-def init_pipeline_tiny(prompt_path,LQ_proj_in_path = "./FlashVSR/LQ_proj_in.ckpt",ckpt_path="./FlashVSR/diffusion_pytorch_model_streaming_dmd.safetensors",TCDecoder_path="./FlashVSR/TCDecoder.ckpt",device="cuda"):
+def init_pipeline_tiny(prompt_path,LQ_proj_in_path = "./FlashVSR/LQ_proj_in.ckpt",ckpt_path="./FlashVSR/diffusion_pytorch_model_streaming_dmd.safetensors",TCDecoder_path="./FlashVSR/TCDecoder.ckpt",device="cuda",offload=False):
 
     mm = ModelManager(torch_dtype=torch.bfloat16, device="cpu")
     mm.load_models([ckpt_path,])
     pipe = FlashVSRTinyPipeline.from_model_manager(mm, device=device)
+    if offload:
+        device = "cpu"
     pipe.denoising_model().LQ_proj_in = Buffer_LQ4x_Proj(in_dim=3, out_dim=1536, layer_num=1).to(device, dtype=torch.bfloat16)
     
     if os.path.exists(LQ_proj_in_path):
@@ -257,10 +259,13 @@ def init_pipeline_tiny(prompt_path,LQ_proj_in_path = "./FlashVSR/LQ_proj_in.ckpt
 
     pipe.enable_vram_management(num_persistent_param_in_dit=None)
     pipe.init_cross_kv(prompt_path); pipe.load_models_to_device(["dit","vae"])
+    pipe.offload=offload
     return pipe
 
-def run_inference_tiny(pipe,input,seed,scale,kv_ratio=3.0,local_range=9,step=1,cfg_scale=1.0,sparse_ratio=2.0,color_fix=True,fix_method="wavelet",split_num=81,dtype=torch.bfloat16,device="cuda", save_vodeo_=False,):
-    pipe.to('cuda')
+def run_inference_tiny(pipe,input,seed,scale,kv_ratio=3.0,local_range=9,step=1,cfg_scale=1.0,sparse_ratio=2.0,color_fix=True,fix_method="wavelet",split_num=81,offload=False,dtype=torch.bfloat16,device="cuda", save_vodeo_=False,):
+    if not offload:
+        pipe.to('cuda')
+    
 
     pad_first_frame = True  if "wavelet"== fix_method and color_fix else False
 
@@ -274,8 +279,10 @@ def run_inference_tiny(pipe,input,seed,scale,kv_ratio=3.0,local_range=9,step=1,c
         kv_ratio=kv_ratio,
         local_range=local_range,  # Recommended: 9 or 11. local_range=9 → sharper details; 11 → more stable results.
         color_fix = color_fix,
+        offload=offload,
     )
-    pipe.dit.to('cpu')  
+    if not offload:
+        pipe.dit.to('cpu')  
     torch.cuda.empty_cache()   
     #print(LQ.shape,frames.shape,LQ_cur_idx)#torch.Size([1, 3, 81, 384, 640]) torch.Size([1, 16, 20, 48, 80]) 77
     with torch.no_grad():

FlashVSR/examples/WanVSR/infer_flashvsr_tiny_long_video.py

@@ -10,7 +10,7 @@ import torch
 from einops import rearrange
 import folder_paths
 from ...diffsynth import ModelManager, FlashVSRTinyLongPipeline
-from .utils.utils import Buffer_LQ4x_Proj
+from .utils.utils import Buffer_LQ4x_Proj,Causal_LQ4x_Proj
 from .utils.TCDecoder import build_tcdecoder
 from .utils.utils import calculate_frame_adjustment_simple
 
@@ -240,27 +240,35 @@ def prepare_input_tensor(path: str, scale: float = 4, fps=30,dtype=torch.bfloat1
 
     raise ValueError(f"Unsupported input: {path}")
 
-def init_pipeline_long(prompt_path,LQ_proj_in_path = "./FlashVSR/LQ_proj_in.ckpt",ckpt_path="./FlashVSR/diffusion_pytorch_model_streaming_dmd.safetensors",TCDecoder_path="./FlashVSR/TCDecoder.ckpt",device="cuda"):
+def init_pipeline_long(prompt_path,LQ_proj_in_path = "./FlashVSR/LQ_proj_in.ckpt",ckpt_path="./FlashVSR/diffusion_pytorch_model_streaming_dmd.safetensors",device="cuda",offload=False,version='1.0'):
     #print(torch.cuda.current_device(), torch.cuda.get_device_name(torch.cuda.current_device()))
     mm = ModelManager(torch_dtype=torch.bfloat16, device="cpu")
     mm.load_models([ckpt_path,])
+    
     pipe = FlashVSRTinyLongPipeline.from_model_manager(mm, device=device)
-    pipe.denoising_model().LQ_proj_in = Buffer_LQ4x_Proj(in_dim=3, out_dim=1536, layer_num=1).to(device, dtype=torch.bfloat16)
+    if offload:
+        device = "cpu"
+    if version=="1.0":
+        pipe.denoising_model().LQ_proj_in = Buffer_LQ4x_Proj(in_dim=3, out_dim=1536, layer_num=1).to(device, dtype=torch.bfloat16)   
+    else:
+        pipe.denoising_model().LQ_proj_in = Causal_LQ4x_Proj(in_dim=3, out_dim=1536, layer_num=1).to(device, dtype=torch.bfloat16)
     if os.path.exists(LQ_proj_in_path):
         pipe.denoising_model().LQ_proj_in.load_state_dict(torch.load(LQ_proj_in_path, map_location="cpu",weights_only=False,), strict=True)
-    pipe.denoising_model().LQ_proj_in.to(device)
+    pipe.denoising_model().LQ_proj_in.to(device, dtype=torch.bfloat16)
 
-    multi_scale_channels = [512, 256, 128, 128]
-    pipe.TCDecoder = build_tcdecoder(new_channels=multi_scale_channels, new_latent_channels=16+768)
-    mis = pipe.TCDecoder.load_state_dict(torch.load(TCDecoder_path,weights_only=False,), strict=False)
-    print(mis)
+    # multi_scale_channels = [512, 256, 128, 128]
+    # pipe.TCDecoder = build_tcdecoder(new_channels=multi_scale_channels, new_latent_channels=16+768)
+    # mis = pipe.TCDecoder.load_state_dict(torch.load(TCDecoder_path,weights_only=False,), strict=False)
+    # print(mis)
 
     pipe.enable_vram_management(num_persistent_param_in_dit=None)
-    pipe.init_cross_kv(prompt_path); pipe.load_models_to_device(["dit","vae"])
+    pipe.init_cross_kv(prompt_path); pipe.load_models_to_device(["dit",])
+    pipe.offload=offload
     return pipe
 
-def run_inference_tiny_long(pipe,input,seed,scale,kv_ratio=3.0,local_range=11,step=1,cfg_scale=1.0,sparse_ratio=2.0,color_fix=True,fix_method="wavelet",split_num=81,dtype=torch.bfloat16,device="cuda", save_vodeo_=False,):
-    pipe.to('cuda')
+def run_inference_tiny_long(pipe,input,seed,scale,kv_ratio=3.0,local_range=11,step=1,cfg_scale=1.0,sparse_ratio=2.0,color_fix=True,fix_method="wavelet",split_num=81,offload=False,dtype=torch.bfloat16,device="cuda", save_vodeo_=False,):
+    if not offload:
+        pipe.to('cuda')
 
     pad_first_frame = True  if "wavelet"== fix_method and color_fix else False
 
@@ -276,14 +284,16 @@ def run_inference_tiny_long(pipe,input,seed,scale,kv_ratio=3.0,local_range=11,st
             color_fix = color_fix,
             pad_first_frame=pad_first_frame,
             fix_method=fix_method,
+            offload=offload,
         )
     print("Done.")
-    pipe.to('cpu')
+    if not offload:
+        pipe.to('cpu')
     del LQ
     torch.cuda.empty_cache()   
     if pad_frames>0:
         print("Remove pad frames",pad_frames,frames.shape)
-        frames =  frames = frames[:, :-pad_frames, :, :]  #torch.Size([ 3, 85, 1024, 768]) -> torch.Size([ 3, 81, 1024, 768]) if 81 output -4 input + 8
+        frames = frames[:, :-pad_frames, :, :]  #torch.Size([ 3, 85, 1024, 768]) -> torch.Size([ 3, 81, 1024, 768]) if 81 output -4 input + 8
         
     frames = tensor2video(frames) 
     if save_vodeo_:

FlashVSR/examples/WanVSR/infer_flashvsr_v11_full.py

@@ -224,65 +224,69 @@ def prepare_input_tensor(path: str, scale: int = 4,fps=30, dtype=torch.bfloat16,
     else:
         raise ValueError(f"Unsupported input: {path}")
 
-def init_pipeline_v11(prompt_path,LQ_proj_in_path="./FlashVSR/LQ_proj_in.ckpt",ckpt_path: str = "./FlashVSR/diffusion_pytorch_model_streaming_dmd.safetensors", vae_path: str = "./FlashVSR/Wan2.1_VAE.pth",decode_vae="none",cur_dir="",device="cuda"):
+def init_pipeline_v11(prompt_path,LQ_proj_in_path="./FlashVSR/LQ_proj_in.ckpt",ckpt_path: str = "./FlashVSR/diffusion_pytorch_model_streaming_dmd.safetensors",cur_dir="",device="cuda",offload=False):
     #print(torch.cuda.current_device(), torch.cuda.get_device_name(torch.cuda.current_device()))
     mm = ModelManager(torch_dtype=torch.bfloat16, device="cpu")
-    mm.load_models([ckpt_path,vae_path,])
-    new_decoder=True if decode_vae!="none"  else False
+    mm.load_models([ckpt_path,])
+    #mm.load_models([ckpt_path,vae_path,])
+    #new_decoder=True if decode_vae!="none"  else False
     pipe = FlashVSRFullPipeline.from_model_manager(mm, device="cuda")
-    if new_decoder:
-        pipe.new_decoder = True
-        if "light" in decode_vae.lower() or "tae" in decode_vae.lower():
-            if os.path.basename(decode_vae).split(".")[0]=="lightvaew2_1":
-                from ...vae import WanVAE
-                print("use lightvae decoder")
-                VAE = WanVAE(vae_path=decode_vae,dtype=torch.bfloat16,device=device,use_lightvae=True)
-            elif os.path.basename(decode_vae).split(".")[0]=="taew2_1":
-                from ...vae_tiny import WanVAE_tiny
-                print("use vae_tiny decoder")
-                VAE = WanVAE_tiny(vae_path=decode_vae,dtype=torch.bfloat16,device=device,need_scaled=False)
-            elif os.path.basename(decode_vae).split(".")[0]=="lighttaew2_1":
-                from ...vae_tiny import WanVAE_tiny
-                print("use vae_tiny light decoder")
-                VAE = WanVAE_tiny(vae_path=decode_vae,dtype=torch.bfloat16,device=device,need_scaled=True)
-            else:
-                raise ValueError(f"Unknown vae_name: {decode_vae},only support lightvae,tae,tae_tiny,lighttae_tiny")
-            pipe.VAE=VAE
-        else:    
-            print("use upscale2x decoder")
-            from diffusers import AutoencoderKLWan
-            config=AutoencoderKLWan.load_config(os.path.join(cur_dir,"FlashVSR/examples/config.json"))
-            VAE=AutoencoderKLWan.from_config(config).to(device,dtype=torch.bfloat16)
-            vae_dict=load_file(decode_vae,device="cpu")
-            VAE.load_state_dict(vae_dict,strict=False)
-            pipe.VAE=VAE
-            del vae_dict
+    # if new_decoder:
+    #     pipe.new_decoder = True
+    #     if "light" in decode_vae.lower() or "tae" in decode_vae.lower():
+    #         if os.path.basename(decode_vae).split(".")[0]=="lightvaew2_1":
+    #             from ...vae import WanVAE
+    #             print("use lightvae decoder")
+    #             VAE = WanVAE(vae_path=decode_vae,dtype=torch.bfloat16,device=device,use_lightvae=True)
+    #         elif os.path.basename(decode_vae).split(".")[0]=="taew2_1":
+    #             from ...vae_tiny import WanVAE_tiny
+    #             print("use vae_tiny decoder")
+    #             VAE = WanVAE_tiny(vae_path=decode_vae,dtype=torch.bfloat16,device=device,need_scaled=False)
+    #         elif os.path.basename(decode_vae).split(".")[0]=="lighttaew2_1":
+    #             from ...vae_tiny import WanVAE_tiny
+    #             print("use vae_tiny light decoder")
+    #             VAE = WanVAE_tiny(vae_path=decode_vae,dtype=torch.bfloat16,device=device,need_scaled=True)
+    #         else:
+    #             raise ValueError(f"Unknown vae_name: {decode_vae},only support lightvae,tae,tae_tiny,lighttae_tiny")
+    #         pipe.VAE=VAE
+    #     else:    
+    #         print("use upscale2x decoder")
+    #         from diffusers import AutoencoderKLWan
+    #         config=AutoencoderKLWan.load_config(os.path.join(cur_dir,"FlashVSR/examples/config.json"))
+    #         VAE=AutoencoderKLWan.from_config(config).to(device,dtype=torch.bfloat16)
+    #         vae_dict=load_file(decode_vae,device="cpu")
+    #         VAE.load_state_dict(vae_dict,strict=False)
+    #         pipe.VAE=VAE
+    #         del vae_dict
     
-    pipe.denoising_model().LQ_proj_in = Causal_LQ4x_Proj(in_dim=3, out_dim=1536, layer_num=1).to("cuda", dtype=torch.bfloat16)
-    #LQ_proj_in_path = "./FlashVSR-v1.1/LQ_proj_in.ckpt"
+    if offload:
+        device = "cpu"
+    pipe.denoising_model().LQ_proj_in = Causal_LQ4x_Proj(in_dim=3, out_dim=1536, layer_num=1).to(device, dtype=torch.bfloat16)
     if os.path.exists(LQ_proj_in_path):
         pipe.denoising_model().LQ_proj_in.load_state_dict(torch.load(LQ_proj_in_path, map_location="cpu",weights_only=False), strict=True)
 
-    pipe.denoising_model().LQ_proj_in.to('cuda')
-    if pipe.vae is  None:  # safetensors cause error or unknown safetensors or pth
-        from diffusers import AutoencoderKLWan
-        config=AutoencoderKLWan.load_config(os.path.join(cur_dir,"FlashVSR/config.json"))
-        vae=AutoencoderKLWan.from_config(config).to(device,dtype=torch.bfloat16)
-        vae_dict=torch.load(vae_path, map_location="cpu",weights_only=False) if not vae_path.endswith(".safetensors") else load_file(vae_path,device="cpu")
-        vae.load_state_dict(vae_dict,strict=False)
-        del vae_dict
-        pipe.vae=vae
-        pipe.vae.encoder = None
-    else:
-        pipe.vae.model.encoder = None
-        pipe.vae.model.conv1 = None
+    pipe.denoising_model().LQ_proj_in.to(device)
+    # if pipe.vae is  None:  # safetensors cause error or unknown safetensors or pth
+    #     from diffusers import AutoencoderKLWan
+    #     config=AutoencoderKLWan.load_config(os.path.join(cur_dir,"FlashVSR/config.json"))
+    #     vae=AutoencoderKLWan.from_config(config).to(device,dtype=torch.bfloat16)
+    #     vae_dict=torch.load(vae_path, map_location="cpu",weights_only=False) if not vae_path.endswith(".safetensors") else load_file(vae_path,device="cpu")
+    #     vae.load_state_dict(vae_dict,strict=False)
+    #     del vae_dict
+    #     pipe.vae=vae
+    #     pipe.vae.encoder = None
+    # else:
+    #     pipe.vae.model.encoder = None
+    #     pipe.vae.model.conv1 = None
     #pipe.to('cuda'); 
     pipe.enable_vram_management(num_persistent_param_in_dit=None)
-    pipe.init_cross_kv(prompt_path); pipe.load_models_to_device(["dit","vae"])
+    pipe.init_cross_kv(prompt_path); pipe.load_models_to_device(["dit",])
+    pipe.offload=offload
     return pipe
 
-def run_inference(pipe,input,seed,scale,kv_ratio=3.0,local_range=9,step=1,cfg_scale=1.0,sparse_ratio=2.0,tiled=True,color_fix=True,fix_method="wavelet",split_num=81,dtype=torch.bfloat16,device="cuda",save_vodeo_=False,):
-    pipe.to('cuda')  #pipe.enable_vram_management(num_persistent_param_in_dit=None)
+def run_inference(pipe,input,seed,scale,kv_ratio=3.0,local_range=9,step=1,cfg_scale=1.0,sparse_ratio=2.0,tiled=True,color_fix=True,fix_method="wavelet",split_num=81,offload=False,dtype=torch.bfloat16,device="cuda",save_vodeo_=False,):
+    if not offload:
+        pipe.to('cuda')  #pipe.enable_vram_management(num_persistent_param_in_dit=None)
     #pipe.init_cross_kv(prompt_path); pipe.load_models_to_device(["dit","vae"])
     pad_first_frame = True  if "wavelet"== fix_method and color_fix else False
 
@@ -299,7 +303,8 @@ def run_inference(pipe,input,seed,scale,kv_ratio=3.0,local_range=9,step=1,cfg_sc
         local_range=local_range, # Recommended: 9 or 11. local_range=9 → sharper details; 11 → more stable results.
         color_fix = color_fix,
     )
-    pipe.dit.to('cpu')
+    if not offload:
+        pipe.dit.to('cpu')
     torch.cuda.empty_cache()
     #torch.Size([1, 16, 20, 48, 80])
     tiler_kwargs = {"tiled": tiled, "tile_size": (60, 104), "tile_stride": (30, 52)}

FlashVSR/examples/WanVSR/infer_flashvsr_v11_tiny.py

@@ -235,16 +235,18 @@ def prepare_input_tensor(path: str, scale: float = 4, fps=30,dtype=torch.bfloat1
 
     raise ValueError(f"Unsupported input: {path}")
 
-def init_pipeline_v11_tiny(prompt_path,LQ_proj_in_path = "./FlashVSR/LQ_proj_in.ckpt",ckpt_path="./FlashVSR/diffusion_pytorch_model_streaming_dmd.safetensors",TCDecoder_path="./FlashVSR/TCDecoder.ckpt",device="cuda"):
+def init_pipeline_v11_tiny(prompt_path,LQ_proj_in_path = "./FlashVSR/LQ_proj_in.ckpt",ckpt_path="./FlashVSR/diffusion_pytorch_model_streaming_dmd.safetensors",TCDecoder_path="./FlashVSR/TCDecoder.ckpt",device="cuda",offload=False):
     #print(torch.cuda.current_device(), torch.cuda.get_device_name(torch.cuda.current_device()))
     mm = ModelManager(torch_dtype=torch.bfloat16, device="cpu")
     mm.load_models([ckpt_path,])
     pipe = FlashVSRTinyPipeline.from_model_manager(mm, device="cuda")
-    pipe.denoising_model().LQ_proj_in = Causal_LQ4x_Proj(in_dim=3, out_dim=1536, layer_num=1).to("cuda", dtype=torch.bfloat16)
+    if offload:
+        device = "cpu"
+    pipe.denoising_model().LQ_proj_in = Causal_LQ4x_Proj(in_dim=3, out_dim=1536, layer_num=1).to(device, dtype=torch.bfloat16)
     #LQ_proj_in_path = "./FlashVSR-v1.1/LQ_proj_in.ckpt"
     if os.path.exists(LQ_proj_in_path):
         pipe.denoising_model().LQ_proj_in.load_state_dict(torch.load(LQ_proj_in_path, map_location="cpu"), strict=True)
-    pipe.denoising_model().LQ_proj_in.to('cuda')
+    pipe.denoising_model().LQ_proj_in.to(device)
 
     multi_scale_channels = [512, 256, 128, 128]
     pipe.TCDecoder = build_tcdecoder(new_channels=multi_scale_channels, new_latent_channels=16+768)
@@ -254,6 +256,7 @@ def init_pipeline_v11_tiny(prompt_path,LQ_proj_in_path = "./FlashVSR/LQ_proj_in.
     #pipe.to('cuda'); 
     pipe.enable_vram_management(num_persistent_param_in_dit=None)
     pipe.init_cross_kv(prompt_path); pipe.load_models_to_device(["dit","vae"])
+    pipe.offload=offload
     return pipe
 
 # def main():

FlashVSR/examples/WanVSR/infer_flashvsr_v11_tiny_long_video.py

@@ -238,16 +238,18 @@ def prepare_input_tensor(path: str, scale: float = 4,fps=30, dtype=torch.bfloat1
 
     raise ValueError(f"Unsupported input: {path}")
 
-def init_pipeline_long_v11(prompt_path,LQ_proj_in_path = "./FlashVSR/LQ_proj_in.ckpt",ckpt_path="./FlashVSR/diffusion_pytorch_model_streaming_dmd.safetensors",TCDecoder_path="./FlashVSR/TCDecoder.ckpt",device="cuda"):
+def init_pipeline_long_v11(prompt_path,LQ_proj_in_path = "./FlashVSR/LQ_proj_in.ckpt",ckpt_path="./FlashVSR/diffusion_pytorch_model_streaming_dmd.safetensors",TCDecoder_path="./FlashVSR/TCDecoder.ckpt",device="cuda",offload=False):
     #print(torch.cuda.current_device(), torch.cuda.get_device_name(torch.cuda.current_device()))
     mm = ModelManager(torch_dtype=torch.bfloat16, device="cpu")
     mm.load_models([ckpt_path,])
     pipe = FlashVSRTinyLongPipeline.from_model_manager(mm, device="cuda")
-    pipe.denoising_model().LQ_proj_in = Causal_LQ4x_Proj(in_dim=3, out_dim=1536, layer_num=1).to("cuda", dtype=torch.bfloat16)
+    if offload:
+        device = "cpu"
+    pipe.denoising_model().LQ_proj_in = Causal_LQ4x_Proj(in_dim=3, out_dim=1536, layer_num=1).to(device, dtype=torch.bfloat16)
     #LQ_proj_in_path = "./FlashVSR-v1.1/LQ_proj_in.ckpt"
     if os.path.exists(LQ_proj_in_path):
         pipe.denoising_model().LQ_proj_in.load_state_dict(torch.load(LQ_proj_in_path, map_location="cpu",weights_only=False,), strict=True)
-    pipe.denoising_model().LQ_proj_in.to('cuda')
+    pipe.denoising_model().LQ_proj_in.to(device)
 
     multi_scale_channels = [512, 256, 128, 128]
     pipe.TCDecoder = build_tcdecoder(new_channels=multi_scale_channels, new_latent_channels=16+768)
@@ -257,6 +259,7 @@ def init_pipeline_long_v11(prompt_path,LQ_proj_in_path = "./FlashVSR/LQ_proj_in.
     #pipe.to('cuda'); 
     pipe.enable_vram_management(num_persistent_param_in_dit=None)
     pipe.init_cross_kv(prompt_path); pipe.load_models_to_device(["dit","vae"])
+    pipe.offload=offload
     return pipe
 
 # def main():

FlashVSR/video_split.bat

@@ -41,7 +41,6 @@ for /f "delims=" %%F in ('ffprobe -v error -select_streams v:0 -count_frames -sh
 :: 获取帧率
 for /f "delims=" %%F in ('ffprobe -v error -select_streams v:0 -show_entries stream^=r_frame_rate -of default^=nokey^=1:noprint_wrappers^=1 "%video_path%"') do set "fps=%%F"
 
-
 :: 检查是否成功获取信息
 if not defined total_frames (
     echo 错误: 无法获取视频信息
@@ -56,11 +55,6 @@ if not exist "%output_dir%" mkdir "%output_dir%"
 echo 视频总帧数: %total_frames%
 echo 视频帧率: %fps%
 
-:: 计算每段时长（秒）
-set /a "segment_time=frame_count*1000/fps"
-set "segment_time=!segment_time:~0,-3!.!segment_time:~-3!"
-
-
 :: 计算按指定帧数分割的段数
 set /a "segments=total_frames/frame_count"
 set /a "remainder=total_frames%%frame_count"
@@ -70,17 +64,6 @@ if %remainder% gtr 0 echo 最后一段将包含 %remainder% 帧
 
 pause
 
-:: 创建临时文件
-set "temp_video=%file_dir%%file_name%_temp%file_ext%"
-ffmpeg -i "%video_path%" -c:v libx264 -x264opts keyint=1:min-keyint=1 -c:a copy "%temp_video%"
-
-:: 检查临时文件是否创建成功
-if not exist "%temp_video%" (
-    echo 错误: 创建临时文件失败
-    pause
-    exit /b 1
-)
-
 :: 使用ffmpeg分割视频
 echo 开始分割视频...
 set /a "segment_num=0"
@@ -97,16 +80,13 @@ if %end_frame% geq %total_frames% set "end_frame=%total_frames%"
 set "formatted_num=00000!segment_num!"
 set "formatted_num=!formatted_num:~-5!"
 
-ffmpeg -i "%temp_video%" -vf "trim=start_frame=%start_frame%:end_frame=%end_frame%" -c:v libx264 -c:a copy "%output_dir%\v%formatted_num%%file_ext%"
+:: 使用更精确的分割方法
+ffmpeg -i "%video_path%" -vf "select='between(n\,%start_frame%,%end_frame%)',setpts=N/FRAME_RATE/TB" -af "aselect='between(n\,%start_frame%,%end_frame%)',asetpts=N/SR/TB" "%output_dir%\v%formatted_num%%file_ext%"
 
 set /a "segment_num+=1"
 if %start_frame% lss %total_frames% goto loop
 :end
 
-
-:: 删除临时文件
-del "%temp_video%"
-
 echo.
 echo 视频分割完成！
 echo 输出目录: "%output_dir%"

FlashVSR_node.py

@@ -4,13 +4,9 @@
 import numpy as np
 import torch
 import os
-from .model_loader_utils import  tensor_upscale,load_images_list,get_video_files
+from .model_loader_utils import  tensor_upscale,decode_latents,get_video_files,clear_comfyui_cache,load_vae_model
 from .FlashVSR.examples.WanVSR.infer_flashvsr_full import init_pipeline,run_inference
-from .FlashVSR.examples.WanVSR.infer_flashvsr_tiny import   init_pipeline_tiny,run_inference_tiny
-from .FlashVSR.examples.WanVSR.infer_flashvsr_tiny_long_video import init_pipeline_long,run_inference_tiny_long
-from .FlashVSR.examples.WanVSR.infer_flashvsr_v11_full import init_pipeline_v11
-from .FlashVSR.examples.WanVSR.infer_flashvsr_v11_tiny import init_pipeline_v11_tiny
-from .FlashVSR.examples.WanVSR.infer_flashvsr_v11_tiny_long_video import init_pipeline_long_v11
+from .FlashVSR.examples.WanVSR.infer_flashvsr_tiny_long_video import init_pipeline_long
 import folder_paths
 from typing_extensions import override
 from comfy_api.latest import ComfyExtension, io
@@ -45,46 +41,24 @@ class FlashVSR_SM_Model(io.ComfyNode):
                 io.Combo.Input("dit",options= ["none"] + [i for i in folder_paths.get_filename_list("FlashVSR") if "dmd" in i.lower()]),
                 io.Combo.Input("proj_pt",options= ["none"] + [i for i in folder_paths.get_filename_list("FlashVSR") if "proj" in i.lower()]),
                 io.Combo.Input("emb_pt",options= ["none"] + [i for i in folder_paths.get_filename_list("FlashVSR") if "prompt" in i.lower()]),
-                io.Combo.Input("vae",options= ["none"] + folder_paths.get_filename_list("vae") ),
-                io.Combo.Input("tcd_encoder",options= ["none"] + [i for i in folder_paths.get_filename_list("FlashVSR") if "tcd" in i.lower()] ),
-                io.Boolean.Input("tiny_long", default=False),
-                io.Combo.Input("decode_vae",options= ["none"] + folder_paths.get_filename_list("vae") ),
                 io.Combo.Input("version",options= ["1.1","1.0"] ),
+                io.Boolean.Input("tiny_long", default=False),
+                io.Boolean.Input("offload", default=False),
             ],
             outputs=[
                 io.Custom("FlashVSR_SM_Model").Output(),
                 ],
             )
     @classmethod
-    def execute(cls, dit,proj_pt,emb_pt,vae,tcd_encoder,tiny_long,decode_vae,version) -> io.NodeOutput:
+    def execute(cls, dit,proj_pt,emb_pt,version,tiny_long,offload) -> io.NodeOutput:
         dit_path=folder_paths.get_full_path("FlashVSR", dit) if dit != "none" else None
         proj_pt_path=folder_paths.get_full_path("FlashVSR", proj_pt) if proj_pt != "none" else None
-        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
         prompt_path=folder_paths.get_full_path("FlashVSR", emb_pt) if emb_pt != "none" else None
-        assert prompt_path is not None  , "Please select the emb"
-        assert dit_path is not None and proj_pt is not None , "Please select the Sdit,proj_pt,checkpoint file"
-        assert vae_path is not None or tcd_encoder_path is not None , "Please select the Sdit,proj_pt,checkpoint file"
-        if tcd_encoder_path is not None:
-            if tiny_long:
-                if "1.0"==version:
-                    model=init_pipeline_long(prompt_path,proj_pt_path,dit_path, tcd_encoder_path, device="cuda")
-                else:
-                    model=init_pipeline_long_v11(prompt_path,proj_pt_path,dit_path, tcd_encoder_path, device="cuda")
-            else:
-                if "1.0"==version:
-                    model=init_pipeline_tiny(prompt_path,proj_pt_path,dit_path, tcd_encoder_path, device="cuda")
-                else:
-                    model=init_pipeline_v11_tiny(prompt_path,proj_pt_path,dit_path, tcd_encoder_path, device="cuda")
-        elif vae_path is not None :
-            decode_vae=folder_paths.get_full_path("vae", decode_vae) if decode_vae != "none" else "none"
-            if "1.0"==version:
-                model=init_pipeline(prompt_path,proj_pt_path,dit_path, vae_path,decode_vae,node_cr_path ,device="cuda")
-            else:
-                model=init_pipeline_v11(prompt_path,proj_pt_path,dit_path, vae_path,decode_vae,node_cr_path ,device="cuda")
+        assert prompt_path is not None and dit_path is not None and proj_pt is not None  , "Please select the emb pt,dit,proj_pt file"
+        if tiny_long:
+            model=init_pipeline_long(prompt_path,proj_pt_path,dit_path, device="cuda",offload=offload,version=version)
         else:
-            raise Exception("Please select the vae or tcd_encoder")
-        model.version = version
+            model=init_pipeline(prompt_path,proj_pt_path,dit_path,node_cr_path ,device="cuda",offload=offload,version=version)
         return io.NodeOutput(model)
     
 
@@ -107,30 +81,64 @@ class FlashVSR_SM_KSampler(io.ComfyNode):
                 io.Int.Input("steps", default=1, min=1, max=10000),
                 io.Float.Input("cfg", default=1.0, min=0.0, max=100.0, step=0.1, round=0.01,),
                 io.Float.Input("sparse_ratio", default=2.0, min=0.0, max=10.0, step=0.1,), 
+            ],
+            outputs=[
+                io.Latent.Output(display_name="latent"),
+            ],
+        )
+    @classmethod
+    def execute(cls, model,image,width,height,seed,scale,kv_ratio,local_range, steps, cfg,sparse_ratio,) -> io.NodeOutput:
+        image=tensor_upscale(image,width, height)
+        offload=getattr(model,"offload",False)
+        clear_comfyui_cache()
+        print(f"start inference with seed:{seed}")
+        images,LQ_cur_idx,LQ,pad_frames=run_inference(model,image,seed,scale,kv_ratio,local_range,steps,cfg,sparse_ratio,offload)
+        outputs={"samples": images,"LQ_cur_idx":LQ_cur_idx,"LQ":LQ,"pad_frames":pad_frames}
+        return io.NodeOutput(outputs)
+    
+class FlashVSR_SM_Vae(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="FlashVSR_SM_Vae",
+            display_name="FlashVSR_SM_Vae",
+            category="FlashVSR",
+            inputs=[
+                io.Combo.Input("vae",options= ["none"] + folder_paths.get_filename_list("vae") ),
+                io.Combo.Input("tcd_encoder",options= ["none"] + [i for i in folder_paths.get_filename_list("FlashVSR") if "tcd" in i.lower()] ),
+                io.Combo.Input("lite_vae",options= ["none"] + folder_paths.get_filename_list("vae") ),
+            ],
+            outputs=[
+                io.Vae.Output(display_name="vae"),
+            ],
+        )
+    @classmethod
+    def execute(cls, vae,tcd_encoder,lite_vae) -> io.NodeOutput:
+        clear_comfyui_cache()
+        vae=load_vae_model(vae,tcd_encoder,lite_vae, device)   
+        return io.NodeOutput(vae)
+
+class FlashVSR_SM_Decoder(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="FlashVSR_SM_Decoder",
+            display_name="FlashVSR_SM_Decoder",
+            category="FlashVSR",
+            inputs=[
+                io.Vae.Input("vae"),
+                io.Latent.Input("latent"),
                 io.Boolean.Input("full_tiled", default=True),
                 io.Boolean.Input("color_fix", default=True),
                 io.Combo.Input("fix_method",options= ["wavelet","adain"]),
-                io.Int.Input("split_num", default=81, min=41, max=MAX_SEED,step=40,),
             ],
             outputs=[
                 io.Image.Output(display_name="images"),
             ],
         )
     @classmethod
-    def execute(cls, model,image,width,height,seed,scale,kv_ratio,local_range, steps, cfg,sparse_ratio,full_tiled,color_fix,fix_method,split_num) -> io.NodeOutput:
-        image=tensor_upscale(image,width, height)
-       
-        if hasattr(model,"TCDecoder") :
-            if model.long_mode:
-                print("infer tiny long mode")
-                images=run_inference_tiny_long(model,image,seed,scale,kv_ratio,local_range,steps,cfg,sparse_ratio,color_fix,fix_method,split_num )
-            else:
-                print("infer tiny mode")
-                images=run_inference_tiny(model,image,seed,scale,kv_ratio,local_range,steps,cfg,sparse_ratio,color_fix,fix_method,split_num )
-        else:
-            print("infer full mode")
-            images=run_inference(model,image,seed,scale,kv_ratio,local_range,steps,cfg,sparse_ratio,full_tiled,color_fix,fix_method,split_num )
-        images=load_images_list(images)
+    def execute(cls, vae,latent,full_tiled,color_fix,fix_method) -> io.NodeOutput:
+        images=decode_latents( vae,latent,color_fix,fix_method,full_tiled, device)
         return io.NodeOutput(images)
 
 
@@ -147,12 +155,8 @@ class FlashVSR_SM_VideoPathLoop(io.ComfyNode):
             inputs=[
                 io.String.Input("video_dir", multiline=False, default="/video"),
                 io.Int.Input("seed", default=0, min=0, max=MAX_SEED),
-                io.Float.Input("start", default=0.0, min=-18446744073709551615, max=18446744073709551615, step=0.01,),
-                io.Float.Input("stop", default=0.0, min=-18446744073709551615, max=18446744073709551615, step=0.01,),
-                io.Float.Input("step", default=1, min=0,max=99999,step=0.01, ),
-                io.Combo.Input("mode",options= ["increment", "decrement", "increment_to_stop", "decrement_to_stop"],),
                 io.Combo.Input("video_file", options=['none', 'webm', 'mp4', 'mkv', 'gif', 'mov']),
-                io.Custom("NUMBER").Input("reset_bool",optional=True),
+                io.Int.Input("reset_bool",default=0,optional=True),
             ],
             outputs=[
                 io.Video.Output(),
@@ -163,30 +167,23 @@ class FlashVSR_SM_VideoPathLoop(io.ComfyNode):
         )
     
     @classmethod
-    def execute(cls, video_dir,seed, mode, start, stop, step,video_file,reset_bool=0,**kwargs) -> io.NodeOutput:
+    def execute(cls, video_dir,seed,  video_file,reset_bool,) -> io.NodeOutput:
         video_path = PureWindowsPath(video_dir).as_posix() if video_dir else None
         video_file = None if video_file == 'none' else video_file
         assert video_path is not None, "video_dir is not set"
         UNIQUE_ID = os.path.normpath(video_path) 
+        start = 0
+        step = 1 
         counter =start
         if cls.counters.__contains__(UNIQUE_ID):
             counter = cls.counters[UNIQUE_ID]
         if round(reset_bool) >= 1:
             counter = start
 
-        if mode == 'increment':
-            counter += step
-        elif mode == 'decrement':
-            counter -= step
-        elif mode == 'increment_to_stop':
-            counter = counter + step if counter < stop else counter
-        elif mode == 'decrement_to_stop':
-            counter = counter - step if counter > stop else counter
-
+        counter += step
         cls.counters[UNIQUE_ID] = counter
         result = int(counter)
 
-        
         video_list = get_video_files(video_path, video_file)
         rows = len(video_list) if video_list else 0
         if rows == 0:
@@ -207,18 +204,14 @@ class FlashVSR_SM_VideoPathLoop(io.ComfyNode):
        return ""
 
 
-from aiohttp import web
-from server import PromptServer
-@PromptServer.instance.routes.get("/FlashVSR_SM_Extension")
-async def get_hello(request):
-    return web.json_response("FlashVSR_SM_Extension")
-
 class FlashVSR_SM_Extension(ComfyExtension):
     @override
     async def get_node_list(self) -> list[type[io.ComfyNode]]:
         return [
             FlashVSR_SM_Model,
             FlashVSR_SM_KSampler,
+            FlashVSR_SM_Vae,
+            FlashVSR_SM_Decoder,
             FlashVSR_SM_VideoPathLoop,
         ]
 async def comfy_entrypoint() -> FlashVSR_SM_Extension:  # ComfyUI calls this to load your extension and its nodes.

README.md

@@ -1,38 +1,36 @@
 # ComfyUI_FlashVSR
 [FlashVSR](https://github.com/OpenImagingLab/FlashVSR): Towards Real-Time Diffusion-Based Streaming Video Super-Resolution,this node ,you can use it in comfyUI
 
-# Tips
-* Good news, the [mit-han-lab/Block-Sparse-Attention](https://github.com/mit-han-lab/Block-Sparse-Attention) library has been updated and now supports 5090 and updated graphics cards ,need cuda>=12.8; 
-* [Block-Sparse-Attention](https://github.com/mit-han-lab/Block-Sparse-Attention) 源码库已更新，支持5090系显卡及更新的设备。注意，需要cuda>=12.8 
+#  Update
+* 简化代码， 分解推理和解码节点（tiny和full模式只是解码不同，最快是lightx2v的vae，其次是tcd，也可以使用原生comfyUI vae节点），修复视频分割bat文件的bug，新增单块卸载（offload按钮）以降低显存占用；
+* Simplify the code, decompose inference and decoding nodes (the only difference between tiny and full modes is in decoding; the fastest is lightx2v's VAE, then tcd, and the native comfyUI VAE node can also be used), fix the bug in the video segmentation bat file, and add a single block offload button to reduce VRAM usage;
 
-# Previous update
-* 11/22 update 增加帧数pad代码 支持任意帧数输入不被裁切
-* vae 现在支持safetensor格式，也就是diffuser模型的vae  
-* 修复一个bug（full 1.1），上传一个简易的window环境的分割bat，
-* 更新适配1.1版本的新模型和代码,降低闪烁，提高保真度和稳定性
-* 新增lightx2v 加速vae decoder支持和Wan2.1-VAE-upscale2x 放大decoder支持，只是在full 模式下有效，light的加速模型目前只支持（lightvaew2_1.pth  #32.2M,taew2_1.pth,lighttaew2_1.pth） 三个文件
-*  满足部分网友需要超分单张图片的奇怪要求,默认输出25帧1秒的视频，详见示例，； 
-*  新增切片视频路径加载节点，输入保存切片视频的路径，开启自动推理，即可推理完路径所有视频，切片的一键bat在插件目录下，window环境使用，点击即可； 
+# Previous
+*  如果觉得项目有用，请给官方项目[FlashVSR](https://github.com/OpenImagingLab/FlashVSR) 打星；
+* [Block-Sparse-Attention](https://github.com/mit-han-lab/Block-Sparse-Attention) 源码库已更新，支持5090系显卡及更新的设备。注意，需要cuda>=12.8 
+* 支持任意帧数输入不被裁切
+* 新增lightx2v 加速vae decoder支持和Wan2.1-VAE-upscale2x 放大decoder支持，light的加速模型目前只支持（lightvaew2_1.pth  #32.2M,taew2_1.pth,lighttaew2_1.pth） 三个文件
+*  满足部分网友需要超分单张图片的奇怪要求,默认输出25帧1秒的视频，详见示例，
+*  新增切片视频路径加载节点，输入保存切片视频的路径，开启自动推理，即可推理完路径所有视频，切片的一键bat在插件目录下，window环境使用，点击即可；
 *  local_range=7这个是会最清晰，local_range=11会比较稳定，color fix 推荐用小波（没重影）； 
 *  编译Block-Sparse-Attention  window的轮子 可以使用 [ smthemex 强制编译版](https://github.com/smthemex/Block-Sparse-Attention) 或者 [lihaoyun6 要联网](https://github.com/lihaoyun6/Block-Sparse-Attention) 两个fork来，不推荐用官方的  
 *  Block-Sparse-Attention 正确安装且能调用才是方法的完全体，当前的函数实现会更容易OOM,但是Block-Sparse-Attention轮子实在不好找，目前只有[CU128 toch2.7](https://github.com/lihaoyun6/ComfyUI-WanVideoWrapper)的，我提供的（[cu128，torch2.8，py311单体](https://pan.quark.cn/s/c9ba067c89bc)）或者自己编译  
 * tile关闭质量更高，需要VRam更高，corlor fix对于非模糊图片可以试试。
-*  Choice vae infer full mode ，encoder infer tiny mode 选择vae跑full模式 效果最好，tiny则是速度，数据集基于4倍训练，所以1 scale是不推荐的；  
-*  如果觉得项目有用，请给官方项目[FlashVSR](https://github.com/OpenImagingLab/FlashVSR) 打星；
-*  11/22 update Add frame rate pad code to support any frame rate input without cropping； 
-* fix when vae is safetesors model's error
-* fix a bug（full 1.1），upload a simple split bat for the window environment.
-* update to version v1.1 
-* add full mode [lightx2v vae encoder](https://huggingface.co/lightx2v/Autoencoders/tree/main) support（only lightvaew2_1.pth,taew2_1.pth,lighttaew2_1.pth） and [Wan2.1-VAE-upscale2x](https://huggingface.co/spacepxl/Wan2.1-VAE-upscale2x) support    
+*  Choice vae infer full mode ，encoder infer tiny mode 选择vae跑full模式 效果最好，tiny则是速度，数据集基于4倍训练，所以1 scale是不推荐的； 
+
+*  if you Like it ， star the official project： [FlashVSR](https://github.com/OpenImagingLab/FlashVSR)  
+* Good news, the [mit-han-lab/Block-Sparse-Attention](https://github.com/mit-han-lab/Block-Sparse-Attention) library has been updated and now supports 5090 and updated graphics cards ,need cuda>=12.8; 
+* Support any frame rate input without cropping
+* Added support for lightx2v accelerated VAE decoder and Wan2.1-VAE-upscale2x upscaling decoder. The light acceleration model currently only supports three files: (lightvaew2_1.pth #32.2M, taew2_1.pth, lighttaew2_1.pth)
 * Satisfying the strange requirement of some netizens to VSR a ‘single image’, the default output is 25 frames per second of video, you can find it in the example.
 * Add a sliced video path loading node, enter the path to save the sliced video, enable automatic inference, and all videos in the path can be inferred. The one click ‘bat、 for slicing can be used in the node directory and window environment. Click to complete the inference;  
-* Local_range=7 will be the clearest, while local_range=11 will be more stable. For color fix, it is recommended to use wavelets (without ghosting);  
+ * Local_range=7 will be the clearest, while local_range=11 will be more stable. For color fix, it is recommended to use wavelets (without ghosting); 
 * The wheel for compiling the Block Spark Attention window can be used[ smthemex ](https://github.com/smthemex/Block-Sparse-Attention) 或者 [lihaoyun6](https://github.com/lihaoyun6/Block-Sparse-Attention)Two forks, not recommended to use official ones
 * The correct installation and ability to call Block Spark Attention are the complete set of methods. The current function implementation will be easier to OOM, but the Block Spark Attention wheel is really hard to find. Currently, only [CU128 toch2.7]（ https://github.com/lihaoyun6/ComfyUI-WanVideoWrapper ）Yes, I provided [cu128，torch2.8，py311](https://pan.quark.cn/s/c9ba067c89bc)Or compile it yourself
 * Tile shutdown has higher quality and requires higher VRam. For non blurry images, corlor fix can be tried.    
 * Choose VAE in full mode, encoder in tiny mode. Choosing VAE to run in full mode yields the best results, while tiny mode is for speed. The dataset is based on 4x training, so 1 scale is not recommended;   
-* If you find the project useful, please provide the official project [FlashVSR]（ https://github.com/OpenImagingLab/FlashVSR ）Strike stars;
-*  if you Like it ， star the official project： [FlashVSR](https://github.com/OpenImagingLab/FlashVSR)  
+
+
 
   
 1.Installation  
@@ -75,23 +73,19 @@ python setup.py install
 |     ├── diffusion_pytorch_model_streaming_dmd.safetensors #v1.1 or v1.0
 |     ├── posi_prompt.pth
 ├── ComfyUI/models/vae
-|        ├──Wan2.1_VAE.pth # or safetensors
-|        ├──lightvaew2_1.pth  #32.2M  or taew2_1.pth,lighttaew2_1.pth
-|        ├──Wan2.1_VAE_upscale2x_imageonly_real_v1_diff.safetensors  # rename from diffusion_pytorch_model.safetensors
+|        ├──Wan2.1_VAE.pth # or safetensors  optional/可选
+|        ├──lightvaew2_1.pth  #32.2M  or taew2_1.pth,lighttaew2_1.pth optional/可选
+|        ├──Wan2.1_VAE_upscale2x_imageonly_real_v1_diff.safetensors  # rename from diffusion_pytorch_model.safetensors optional/可选
 ```
   
 
 # Example
-* upscale2x and ligth lightvaew2_1.pth
-![](https://github.com/smthemex/ComfyUI_FlashVSR/blob/main/example_workflows/example_decoder.png)
+* normal
+![](https://github.com/smthemex/ComfyUI_FlashVSR/blob/main/example_workflows/example.png)
 * single image VSR 
 ![](https://github.com/smthemex/ComfyUI_FlashVSR/blob/main/example_workflows/example_s.png)
-* full old node 
-![](https://github.com/smthemex/ComfyUI_FlashVSR/blob/main/example_workflows/example18.png)
-* tiny new
-![](https://github.com/smthemex/ComfyUI_FlashVSR/blob/main/example_workflows/example1022.png)
 * video files loop
-![](https://github.com/smthemex/ComfyUI_FlashVSR/blob/main/example_workflows/example.png)
+![](https://github.com/smthemex/ComfyUI_FlashVSR/blob/main/example_workflows/example_l.png)
 
 # Acknowledgements
 [DiffSynth Studio](https://github.com/modelscope/DiffSynth-Studio)  

model_loader_utils.py

@@ -6,13 +6,254 @@ 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)
+        vae.load_state_dict(torch.load(tcd_encoder_path,weights_only=False,), strict=False)
+        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()
@@ -144,7 +385,22 @@ def images_generator(img_list: list, ):
         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)