From 76cdd93680c536a6f5683f60539eeabf27b39775 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Sun, 4 Aug 2024 16:37:16 -0400
Subject: [PATCH] Flux needs to be bf16.

---
 pyproject.toml      |  2 +-
 tensorrt_convert.py | 10 ++++++++--
 tensorrt_loader.py  |  1 +
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 7dcc0bb..42d1e9c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 name = "comfyui_tensorrt"
 description = "TensorRT Node for ComfyUI\nThis node enables the best performance on NVIDIA RTX™ Graphics Cards  (GPUs) for Stable Diffusion by leveraging NVIDIA TensorRT."
-version = "0.1.5"
+version = "0.1.6"
 license = "LICENSE"
 dependencies = [
   "tensorrt>=10.0.1",
diff --git a/tensorrt_convert.py b/tensorrt_convert.py
index 513b348..5774aa3 100644
--- a/tensorrt_convert.py
+++ b/tensorrt_convert.py
@@ -163,6 +163,7 @@ class TRT_MODEL_CONVERSION_BASE:
         context_len_min = context_len
         y_dim = model.model.adm_channels
         extra_input = {}
+        dtype = torch.float16
 
         if isinstance(model.model, comfy.model_base.SD3): #SD3
             context_embedder_config = model.model.model_config.unet_config.get("context_embedder_config", None)
@@ -179,6 +180,7 @@ class TRT_MODEL_CONVERSION_BASE:
             context_len = 256
             y_dim = model.model.model_config.unet_config.get("vec_in_dim", None)
             extra_input = {"guidance": ()}
+            dtype = torch.bfloat16
 
         if context_dim is not None:
             input_names = ["x", "timesteps", "context"]
@@ -268,7 +270,7 @@ class TRT_MODEL_CONVERSION_BASE:
                     torch.zeros(
                         shape,
                         device=comfy.model_management.get_torch_device(),
-                        dtype=torch.float16,
+                        dtype=dtype,
                     ),
                 )
 
@@ -325,7 +327,11 @@ class TRT_MODEL_CONVERSION_BASE:
                 input_names[k], encode(min_shape), encode(opt_shape), encode(max_shape)
             )
 
-        config.set_flag(trt.BuilderFlag.FP16)
+        if dtype == torch.float16:
+            config.set_flag(trt.BuilderFlag.FP16)
+        if dtype == torch.bfloat16:
+            config.set_flag(trt.BuilderFlag.BF16)
+
         config.add_optimization_profile(profile)
 
         if is_static:
diff --git a/tensorrt_loader.py b/tensorrt_loader.py
index dce734e..faffdd8 100644
--- a/tensorrt_loader.py
+++ b/tensorrt_loader.py
@@ -158,6 +158,7 @@ class TensorRTLoader:
             conf = comfy.supported_models.Flux({})
             conf.unet_config["disable_unet_model_creation"] = True
             model = conf.get_model({})
+            unet.dtype = torch.bfloat16 #TODO: autodetect
         model.diffusion_model = unet
         model.memory_required = lambda *args, **kwargs: 0 #always pass inputs batched up as much as possible, our TRT code will handle batch splitting