From e0a58c6b0406c14147bb4e3e2715fbbc70b0e8cd Mon Sep 17 00:00:00 2001
From: wailovet <wailovet@163.com>
Date: Mon, 12 Aug 2024 20:02:07 +0800
Subject: [PATCH] init

---
 mz_flux_quantize.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/mz_flux_quantize.py b/mz_flux_quantize.py
index 796b708..f88c23a 100644
--- a/mz_flux_quantize.py
+++ b/mz_flux_quantize.py
@@ -31,7 +31,7 @@ def MZ_EmptyModel_call(kwargs):
             "qkv_bias": True,
             "guidance_embed": True
         })
-        model_config.manual_cast_dtype = torch.float32
+        model_config.manual_cast_dtype = torch.bfloat16
 
         model = comfy.model_base.Flux(
             model_config=model_config,
@@ -72,8 +72,7 @@ def MZ_FluxQuantizeUNETLoader_call(kwargs):
     unet_name = kwargs["unet_name"]
     import folder_paths
     unet_path = folder_paths.get_full_path("unet", unet_name)
-    
-    
+
     model = MZ_EmptyModel_call({"model_type": "Flux"})[0]
 
     model_patcher = MZ_ModelLoadStateDict_call({
@@ -81,4 +80,4 @@ def MZ_FluxQuantizeUNETLoader_call(kwargs):
         "safetensors_file": unet_path,
     })[0]
 
-    return (model_patcher,)
\ No newline at end of file
+    return (model_patcher,)