From de5f6e9dcc7aac4af6ef8e06ddacbecc6549672a Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Fri, 16 Aug 2024 04:29:14 +0200
Subject: [PATCH] mmap loading, lowvram support

This should fix both types of OOM issues!
---
 nodes.py |  3 ---
 ops.py   | 40 +++++++++++++++++++++++++++++++++-------
 2 files changed, 33 insertions(+), 10 deletions(-)

diff --git a/nodes.py b/nodes.py
index 88b3a8d..3ce670c 100644
--- a/nodes.py
+++ b/nodes.py
@@ -22,9 +22,6 @@ def gguf_sd_loader(path):
     for tensor in reader.tensors:
         sd[str(tensor.name)] = GGMLTensor(tensor)
         dt[str(tensor.tensor_type)] = dt.get(str(tensor.tensor_type), 0) + 1
-    
-    # TODO: .to doesn't get passed later on for some reason, this is a hotfix
-    #sd = {k:v.to(comfy.model_management.get_torch_device()) for k,v in sd.items()}
 
     # sanity check debug print
     print("\nggml_sd_loader:")
diff --git a/ops.py b/ops.py
index 20ded45..f1999cd 100644
--- a/ops.py
+++ b/ops.py
@@ -18,7 +18,7 @@ class GGMLTensor(torch.Tensor):
         )
 
     def __new__(cls, tensor, *args, **kwargs):
-        data = torch.tensor(tensor.data)
+        data = torch.from_numpy(tensor.data)
         return super().__new__(cls, data, *args, **kwargs)
 
     def to(self, *args, **kwargs):
@@ -49,6 +49,23 @@ class GGMLLayer(torch.nn.Module):
             else:
                 missing_keys.append(k)
 
+    def _save_to_state_dict(self, destination, prefix, keep_vars):
+        # This is a fake state dict for vram estimation
+        if self.weight is not None:
+            weight = torch.zeros_like(self.weight, device=torch.device("meta"))
+            destination[f"{prefix}.weight"] = weight
+        if self.bias is not None:
+            bias = torch.zeros_like(self.bias, device=torch.device("meta"))
+            destination[f"{prefix}.bias"] = bias
+        return
+
+        # This would return the actual state dict
+        weight, bias = self.get_weights()
+        if weight is not None:
+            destination[f"{prefix}.weight"] = weight
+        if bias is not None:
+            destination[f"{prefix}.bias"] = weight
+
     def _apply(self, fn):
         if self.weight is not None:
             self.weight = fn(self.weight)
@@ -67,13 +84,22 @@ class GGMLOps(comfy.ops.disable_weight_init):
     Dequantize weights on the fly before doing the compute
     """
     class Linear(GGMLLayer):
+        comfy_cast_weights = True
+
         def __init__(self, *args, device=None, dtype=None, **kwargs):
             super().__init__(device=device, dtype=dtype)
-            self.parameters_manual_cast = torch.float32
 
         def forward(self, x):
-            with torch.no_grad():
-                weight, bias = self.get_weights(x.dtype)
-                x = torch.nn.functional.linear(x, weight, bias)
-                del weight, bias
-            return x
+            # lowvram hack
+            device = None
+            if self.weight.device != x.device:
+                device = self.weight.device
+                self.to(x.device)
+
+            weight, bias = self.get_weights(x.dtype)
+            x = torch.nn.functional.linear(x, weight, bias)
+            del weight, bias
+
+            if device:
+                self.to(device)
+            return x
\ No newline at end of file