From 708f002955c82d46b3d0438e5eab1997b90e1d17 Mon Sep 17 00:00:00 2001 From: Rattus Date: Tue, 21 Oct 2025 18:03:33 +1000 Subject: [PATCH] sd: Implement RAM getter for VAE --- comfy/sd.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/comfy/sd.py b/comfy/sd.py index 28bee248d..ad2f3bd44 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -293,6 +293,7 @@ class VAE: self.working_dtypes = [torch.bfloat16, torch.float32] self.disable_offload = False self.not_video = False + self.size = None self.downscale_index_formula = None self.upscale_index_formula = None @@ -595,6 +596,16 @@ class VAE: self.patcher = comfy.model_patcher.ModelPatcher(self.first_stage_model, load_device=self.device, offload_device=offload_device) logging.info("VAE load device: {}, offload device: {}, dtype: {}".format(self.device, offload_device, self.vae_dtype)) + self.model_size() + + def model_size(self): + if self.size is not None: + return self.size + self.size = comfy.model_management.module_size(self.first_stage_model) + return self.size + + def get_ram_usage(self): + return self.model_size() def throw_exception_if_invalid(self): if self.first_stage_model is None: