From baec4b634f9b81d74e881f8172e7b25399741e12 Mon Sep 17 00:00:00 2001
From: Adrien Toupet <cayus.net@gmail.com>
Date: Wed, 24 Dec 2025 09:50:19 +0100
Subject: [PATCH] fix(vae): restore MPS memory leak workaround removed in
 v2.5.23 cleanup

---
 README.md                                         | 4 ++++
 pyproject.toml                                    | 2 +-
 src/models/video_vae_v3/modules/attn_video_vae.py | 8 ++++++++
 src/utils/constants.py                            | 2 +-
 4 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index e26bf35..197ad2d 100644
--- a/README.md
+++ b/README.md
@@ -36,6 +36,10 @@ We're actively working on improvements and new features. To stay informed:
 
 ## 🚀 Release Notes
 
+**2025.12.24 - Version 2.5.24**
+
+- **🍎 Fix: MPS memory leak regression** - Restored MPS cache clearing after VAE encode/decode operations that was accidentally removed during code cleanup in v2.5.23
+
 **2025.12.24 - Version 2.5.23**
 
 - **🔒 Security: Prevent code execution in model loading** - Added protection against malicious .pth files by restricting deserialization to tensors only
diff --git a/pyproject.toml b/pyproject.toml
index aab7f49..71c12e2 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 name = "seedvr2_videoupscaler"
 description = "SeedVR2 official ComfyUI integration: ByteDance-Seed's one-step diffusion-based video/image upscaling with memory-efficient inference"
-version = "2.5.23"
+version = "2.5.24"
 authors = [
     {name = "numz"},
     {name = "adrientoupet"}
diff --git a/src/models/video_vae_v3/modules/attn_video_vae.py b/src/models/video_vae_v3/modules/attn_video_vae.py
index de83771..2a7d70b 100644
--- a/src/models/video_vae_v3/modules/attn_video_vae.py
+++ b/src/models/video_vae_v3/modules/attn_video_vae.py
@@ -1224,6 +1224,10 @@ class VideoAutoencoderKL(diffusers.AutoencoderKL):
         
         output = causal_conv_gather_outputs(output)
         
+        # MPS memory leak workaround (pytorch/pytorch#155060)
+        if self.device.type == 'mps':
+            torch.mps.empty_cache()
+        
         # Only transfer back if needed
         return output if output.device == x.device else output.to(x.device)
 
@@ -1240,6 +1244,10 @@ class VideoAutoencoderKL(diffusers.AutoencoderKL):
         output = self.decoder(_z, memory_state=memory_state)
         output = causal_conv_gather_outputs(output)
         
+        # MPS memory leak workaround (pytorch/pytorch#155060)
+        if self.device.type == 'mps':
+            torch.mps.empty_cache()
+        
         # Only transfer back if needed
         return output if output.device == z.device else output.to(z.device)
 
diff --git a/src/utils/constants.py b/src/utils/constants.py
index 4e2a287..088d538 100644
--- a/src/utils/constants.py
+++ b/src/utils/constants.py
@@ -4,7 +4,7 @@ Only includes constants actually used in the codebase
 """
 
 # Version information
-__version__ = "2.5.23"
+__version__ = "2.5.24"
 
 import os
 import warnings