From f22b80ef9429474283bbc40023ca1e8ef5328bf6 Mon Sep 17 00:00:00 2001
From: layerdiffusion <19834515+lllyasviel@users.noreply.github.com>
Date: Mon, 26 Aug 2024 06:16:13 -0700
Subject: [PATCH] restrict baking to 16bits

---
 backend/operations_gguf.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/backend/operations_gguf.py b/backend/operations_gguf.py
index d2c43c8c..3d9cec66 100644
--- a/backend/operations_gguf.py
+++ b/backend/operations_gguf.py
@@ -70,6 +70,11 @@ class ParameterGGUF(torch.nn.Parameter):
 
 def bake_gguf_model(model):
     computation_dtype = model.computation_dtype
+
+    if computation_dtype not in [torch.float16, torch.bfloat16]:
+        # Baking only supports 16bits otherwise super slow
+        computation_dtype = torch.float16
+
     backed_layer_counter = 0
 
     for m in model.modules():