From 62db11683b7290e8fef40802c769f4b097d2b1ed Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Thu, 29 Jun 2023 11:19:58 -0400
Subject: [PATCH 1/2] Move unet to device right after loading on highvram mode.

---
 comfy/model_management.py | 12 ++++++++----
 comfy/sd.py               |  1 +
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/comfy/model_management.py b/comfy/model_management.py
index 4e0e6a0ae..4f3f28571 100644
--- a/comfy/model_management.py
+++ b/comfy/model_management.py
@@ -216,6 +216,11 @@ current_gpu_controlnets = []
 
 model_accelerated = False
 
+def unet_offload_device():
+    if vram_state == VRAMState.HIGH_VRAM or vram_state == VRAMState.SHARED:
+        return get_torch_device()
+    else:
+        return torch.device("cpu")
 
 def unload_model():
     global current_loaded_model
@@ -228,10 +233,9 @@ def unload_model():
             accelerate.hooks.remove_hook_from_submodules(current_loaded_model.model)
             model_accelerated = False
 
-        #never unload models from GPU on high vram
-        if vram_state != VRAMState.HIGH_VRAM:
-            current_loaded_model.model.cpu()
-            current_loaded_model.model_patches_to("cpu")
+
+        current_loaded_model.model.to(unet_offload_device())
+        current_loaded_model.model_patches_to(unet_offload_device())
         current_loaded_model.unpatch_model()
         current_loaded_model = None
 
diff --git a/comfy/sd.py b/comfy/sd.py
index 52d016b10..542f704a6 100644
--- a/comfy/sd.py
+++ b/comfy/sd.py
@@ -1142,6 +1142,7 @@ def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, o
             clipvision = clip_vision.load_clipvision_from_sd(sd, model_config.clip_vision_prefix, True)
 
     model = model_config.get_model(sd)
+    model = model.to(model_management.unet_offload_device())
     model.load_model_weights(sd, "model.diffusion_model.")
 
     if output_vae:

From 9920367d3c65db363a8fcf1d2acf59f3a7d9f18d Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Thu, 29 Jun 2023 20:42:19 -0400
Subject: [PATCH 2/2] Fix embeddings not working with --gpu-only

---
 comfy/sd1_clip.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/comfy/sd1_clip.py b/comfy/sd1_clip.py
index 0ee314ad5..02a998e5b 100644
--- a/comfy/sd1_clip.py
+++ b/comfy/sd1_clip.py
@@ -95,7 +95,7 @@ class SD1ClipModel(torch.nn.Module, ClipTokenWeightEncoder):
             out_tokens += [tokens_temp]
 
         if len(embedding_weights) > 0:
-            new_embedding = torch.nn.Embedding(next_new_token, current_embeds.weight.shape[1])
+            new_embedding = torch.nn.Embedding(next_new_token, current_embeds.weight.shape[1], device=self.device)
             new_embedding.weight[:token_dict_size] = current_embeds.weight[:]
             n = token_dict_size
             for x in embedding_weights: