From 1b103e0cb2d7aeb05fc8b7e006d4438e7bceca20 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Sat, 30 Dec 2023 05:38:21 -0500
Subject: [PATCH] Add argument to run the VAE on the CPU.

---
 comfy/cli_args.py         | 2 ++
 comfy/model_management.py | 5 +++++
 2 files changed, 7 insertions(+)

diff --git a/comfy/cli_args.py b/comfy/cli_args.py
index 8de0adb..50d7b62 100644
--- a/comfy/cli_args.py
+++ b/comfy/cli_args.py
@@ -66,6 +66,8 @@ fpvae_group.add_argument("--fp16-vae", action="store_true", help="Run the VAE in
 fpvae_group.add_argument("--fp32-vae", action="store_true", help="Run the VAE in full precision fp32.")
 fpvae_group.add_argument("--bf16-vae", action="store_true", help="Run the VAE in bf16.")
 
+parser.add_argument("--cpu-vae", action="store_true", help="Run the VAE on the CPU.")
+
 fpte_group = parser.add_mutually_exclusive_group()
 fpte_group.add_argument("--fp8_e4m3fn-text-enc", action="store_true", help="Store text encoder weights in fp8 (e4m3fn variant).")
 fpte_group.add_argument("--fp8_e5m2-text-enc", action="store_true", help="Store text encoder weights in fp8 (e5m2 variant).")
diff --git a/comfy/model_management.py b/comfy/model_management.py
index c0cb413..fefd3c8 100644
--- a/comfy/model_management.py
+++ b/comfy/model_management.py
@@ -186,6 +186,9 @@ except:
 if is_intel_xpu():
     VAE_DTYPE = torch.bfloat16
 
+if args.cpu_vae:
+    VAE_DTYPE = torch.float32
+
 if args.fp16_vae:
     VAE_DTYPE = torch.float16
 elif args.bf16_vae:
@@ -555,6 +558,8 @@ def intermediate_device():
         return torch.device("cpu")
 
 def vae_device():
+    if args.cpu_vae:
+        return torch.device("cpu")
     return get_torch_device()
 
 def vae_offload_device():