feat: expose seedvr vae tuning

Author: GACLove

lightx2v/models/runners/seedvr/seedvr_runner.py

@@ -305,6 +305,9 @@ def load_image_encoder(self):
         return None
 
     def load_vae_encoder(self):
+        vae_causal_slice_size = int(self.config.get("vae_causal_slice_size", 4))
+        vae_memory_limit_gb = float(self.config.get("vae_memory_limit_gb", 0.5))
+        vae_memory_limit = None if vae_memory_limit_gb <= 0 else vae_memory_limit_gb
         vae = attn_video_vae_v3_s8_c16_t4_inflation_sd3_init(
             device=AI_DEVICE,
             dtype=GET_DTYPE(),
@@ -314,10 +317,18 @@ def load_vae_encoder(self):
             strict=False,
             cpu_offload=self.config.get("cpu_offload", False),
             use_tiling=self.config.get("use_tiling_vae", False),
+            tile_size=int(self.config.get("vae_tile_size", 512)),
+            tile_overlap=int(self.config.get("vae_tile_overlap", 64)),
         )
         vae.requires_grad_(False).eval()
-        vae.set_causal_slicing(split_size=4, memory_device="same")
-        vae.set_memory_limit(conv_max_mem=0.5, norm_max_mem=0.5)
+        vae.set_causal_slicing(split_size=vae_causal_slice_size if vae_causal_slice_size > 0 else None, memory_device="same")
+        vae.set_memory_limit(conv_max_mem=vae_memory_limit, norm_max_mem=vae_memory_limit)
+        logger.info(
+            f"[SeedVRRunner] VAE config: tiling={self.config.get('use_tiling_vae', False)}, "
+            f"tile={self.config.get('vae_tile_size', 512)}, overlap={self.config.get('vae_tile_overlap', 64)}, "
+            f"causal_slice={vae_causal_slice_size if vae_causal_slice_size > 0 else 'off'}, "
+            f"memory_limit={vae_memory_limit_gb if vae_memory_limit_gb > 0 else 'off'}GiB"
+        )
         return vae
 
     def load_vae_decoder(self):
@@ -340,9 +351,14 @@ def run_vae_decoder(self, latents):
         if self._ori_length < sample.shape[0]:
             sample = sample[: self._ori_length]
 
-        # color fix
-        input = rearrange(self._input[:, None], "c t h w -> t c h w") if self._input.ndim == 3 else rearrange(self._input, "c t h w -> t c h w")
-        sample = wavelet_reconstruction(sample.to("cpu"), input[: sample.size(0)].to("cpu"))
+        color_fix = str(self.config.get("color_fix", "cpu")).lower()
+        if color_fix not in ("cpu", "gpu", "off"):
+            logger.warning(f"[SeedVRRunner] Unknown color_fix={color_fix}; fallback to cpu")
+            color_fix = "cpu"
+        if color_fix != "off":
+            input = rearrange(self._input[:, None], "c t h w -> t c h w") if self._input.ndim == 3 else rearrange(self._input, "c t h w -> t c h w")
+            fix_device = torch.device("cpu") if color_fix == "cpu" else sample.device
+            sample = wavelet_reconstruction(sample.to(fix_device), input[: sample.size(0)].to(fix_device))
         sample = rearrange(sample[:, None], "t c h w -> c t h w") if sample.ndim == 3 else rearrange(sample, "t c h w -> c t h w")
         sample = sample[None, :]
 

lightx2v/models/video_encoders/hf/seedvr/attn_video_vae.py

@@ -1491,6 +1491,8 @@ def __init__(
         freeze_encoder: bool,
         cpu_offload: bool = False,
         use_tiling: bool = False,
+        tile_size: Union[int, Tuple[int, int]] = 512,
+        tile_overlap: Union[int, Tuple[int, int]] = 64,
         **kwargs,
     ):
         self.spatial_downsample_factor = spatial_downsample_factor
@@ -1499,6 +1501,14 @@ def __init__(
         self.cpu_offload = cpu_offload
         super().__init__(*args, **kwargs)
         self.use_tiling = use_tiling
+        self.tile_size = self._as_pair(tile_size)
+        self.tile_overlap = self._as_pair(tile_overlap)
+
+    @staticmethod
+    def _as_pair(value: Union[int, Tuple[int, int]]) -> Tuple[int, int]:
+        if isinstance(value, tuple):
+            return int(value[0]), int(value[1])
+        return int(value), int(value)
 
     def forward(self, x: torch.FloatTensor) -> CausalAutoencoderOutput:
         with torch.no_grad() if self.freeze_encoder else nullcontext():
@@ -1572,10 +1582,10 @@ def vae_encode(self, samples: List[Tensor]) -> List[Tensor]:
                 if hasattr(self, "preprocess"):
                     sample = self.preprocess(sample)
                 if use_sample:
-                    latent = self.encode(sample, tiled=self.use_tiling).latent
+                    latent = self.encode(sample, tiled=self.use_tiling, tile_size=self.tile_size, tile_overlap=self.tile_overlap).latent
                 else:
                     # Deterministic vae encode, only used for i2v inference (optionally)
-                    latent = self.encode(sample, tiled=self.use_tiling).posterior.mode().squeeze(2)
+                    latent = self.encode(sample, tiled=self.use_tiling, tile_size=self.tile_size, tile_overlap=self.tile_overlap).posterior.mode().squeeze(2)
                 latent = latent.unsqueeze(2) if latent.ndim == 4 else latent
                 latent = rearrange(latent, "b c ... -> b ... c")
                 latent = (latent - shift) * scale
@@ -1614,7 +1624,7 @@ def vae_decode(self, latents: List[Tensor]) -> List[Tensor]:
                 latent = latent / scale + shift
                 latent = rearrange(latent, "b ... c -> b c ...")
                 latent = latent.squeeze(2)
-                sample = self.decode(latent, tiled=self.use_tiling).sample
+                sample = self.decode(latent, tiled=self.use_tiling, tile_size=self.tile_size, tile_overlap=self.tile_overlap).sample
                 if hasattr(self, "postprocess"):
                     sample = self.postprocess(sample)
                 samples.append(sample)
@@ -1642,6 +1652,8 @@ def attn_video_vae_v3_s8_c16_t4_inflation_sd3_init(
     strict: bool = True,
     cpu_offload: bool = False,
     use_tiling: bool = False,
+    tile_size: Union[int, Tuple[int, int]] = 512,
+    tile_overlap: Union[int, Tuple[int, int]] = 64,
 ) -> VideoAutoencoderKLWrapper:
     """Example: initialize VideoAutoencoderKLWrapper with SD3 inflation config params."""
     model = VideoAutoencoderKLWrapper(
@@ -1673,6 +1685,8 @@ def attn_video_vae_v3_s8_c16_t4_inflation_sd3_init(
         freeze_encoder=False,
         cpu_offload=cpu_offload,
         use_tiling=use_tiling,
+        tile_size=tile_size,
+        tile_overlap=tile_overlap,
     )
 
     if weights_path is not None: