From 0c3bdf336e4a948328acee4b4a10204c9a12583c Mon Sep 17 00:00:00 2001
From: "Xiangxi Guo (Ryan)" <ryan.guo99@gmail.com>
Date: Wed, 9 Apr 2025 17:46:33 -0700
Subject: [PATCH] Enable more torch.compile optimization by reducing graph
 breaks (#243)

As title, there were 2 fundamental sources of graph breaks:
1. using a pre-existing set in a `torch.compile` region, this is a
   [known issue](https://github.com/pytorch/pytorch/issues/145761) in
   PyTorch and the fix is WIP, but it's a lot easier to change the set
   into a tuple here to fix the graph break.
2. using tensor subclass `GGMLTensor`. The full support of tensor
   subclass has been gated for `torch.compile`, but with [a recent WIP
   stack](https://github.com/pytorch/pytorch/pull/149792),
   we'll be able to remove the gating and have `torch.compile` fully
   support the `GGMLTensor` in this repo. The only required change here
   is avoid changing the `__class__` attribute, and instead using
   `torch.Tensor` to convert the subclass object back to Tensor.
---
 dequant.py |  2 +-
 ops.py     | 34 +++++++++++++++++++++++++++-------
 2 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/dequant.py b/dequant.py
index 8232500..9e545b7 100644
--- a/dequant.py
+++ b/dequant.py
@@ -4,7 +4,7 @@ import torch
 from tqdm import tqdm
 
 
-TORCH_COMPATIBLE_QTYPES = {None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}
+TORCH_COMPATIBLE_QTYPES = (None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16)
 
 def is_torch_compatible(tensor):
     return tensor is None or getattr(tensor, "tensor_type", None) in TORCH_COMPATIBLE_QTYPES
diff --git a/ops.py b/ops.py
index cc82b08..c74296e 100644
--- a/ops.py
+++ b/ops.py
@@ -6,15 +6,35 @@ import comfy.ops
 import comfy.model_management
 from .dequant import dequantize_tensor, is_quantized
 
-# to avoid breaking really old pytorch versions
-if hasattr(torch, "compiler") and hasattr(torch.compiler, "disable"):
-    torch_compiler_disable = torch.compiler.disable
-else:
-    def torch_compiler_disable(*args, **kwargs):
+def chained_hasattr(obj, chained_attr):
+    probe = obj
+    for attr in chained_attr.split('.'):
+        if hasattr(probe, attr):
+            probe = getattr(probe, attr)
+        else:
+            return False
+    return True
+
+# A bakcward and forward compatible way to get `torch.compiler.disable`.
+def get_torch_compiler_disable_decorator():
+    def dummy_decorator(*args, **kwargs):
         def noop(x):
             return x
         return noop
 
+    from packaging import version
+
+    if not chained_hasattr(torch, "compiler.disable"):
+        return dummy_decorator # torch too old
+    elif version.parse(torch.__version__) >= version.parse("2.8"):
+        return dummy_decorator # torch compile works
+    if chained_hasattr(torch, "_dynamo.config.nontraceable_tensor_subclasses"):
+        return dummy_decorator # torch compile works, nightly before 2.8 release
+    else:
+        return torch.compiler.disable
+
+torch_compiler_disable = get_torch_compiler_disable_decorator()
+
 class GGMLTensor(torch.Tensor):
     """
     Main tensor-like class for storing quantized weights
@@ -149,7 +169,7 @@ class GGMLLayer(torch.nn.Module):
 
         # prevent propagating custom tensor class
         if isinstance(weight, GGMLTensor):
-            weight.__class__ = torch.Tensor
+            weight = torch.Tensor(weight)
 
         # apply patches
         if patch_list:
@@ -189,7 +209,7 @@ class GGMLLayer(torch.nn.Module):
 
         # non-ggml forward might still propagate custom tensor class
         if isinstance(out, GGMLTensor):
-            out.__class__ = torch.Tensor
+            out = torch.Tensor(out)
         return out
 
     def forward_ggml_cast_weights(self, input):