From 48ca1a4910069883c0490c8f34a7f2e4454d999f Mon Sep 17 00:00:00 2001
From: doctorpangloss <@hiddenswitch.com>
Date: Sun, 25 Aug 2024 22:41:10 -0700
Subject: [PATCH] Include Kijai fp8 nodes. LoRAs are not supported by nf4

---
 comfy/model_downloader.py       | 2 ++
 comfy/model_patcher.py          | 2 +-
 comfy_extras/nodes/nodes_nf4.py | 4 +++-
 requirements.txt                | 3 ++-
 4 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/comfy/model_downloader.py b/comfy/model_downloader.py
index ae1c914b4..56c4ae708 100644
--- a/comfy/model_downloader.py
+++ b/comfy/model_downloader.py
@@ -424,6 +424,8 @@ KNOWN_UNET_MODELS: Final[KnownDownloadables] = KnownDownloadables([
     HuggingFile("ByteDance/Hyper-SD", "Hyper-SDXL-1step-Unet-Comfyui.fp16.safetensors"),
     HuggingFile("black-forest-labs/FLUX.1-schnell", "flux1-schnell.safetensors"),
     HuggingFile("black-forest-labs/FLUX.1-dev", "flux1-dev.safetensors"),
+    HuggingFile("Kijai/flux-fp8", "flux1-dev-fp8.safetensors"),
+    HuggingFile("Kijai/flux-fp8", "flux1-schnell-fp8.safetensors"),
 ], folder_name="diffusion_models")
 
 KNOWN_CLIP_MODELS: Final[KnownDownloadables] = KnownDownloadables([
diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py
index 7d30fe24d..3cb013acf 100644
--- a/comfy/model_patcher.py
+++ b/comfy/model_patcher.py
@@ -386,7 +386,7 @@ class ModelPatcher(ModelManageable):
             logging.debug("loaded partially {} {} {}".format(lowvram_model_memory / (1024 * 1024), mem_counter / (1024 * 1024), patch_counter))
             self._memory_measurements.model_lowvram = True
         else:
-            logging.info("loaded completely {} {} {}".format(lowvram_model_memory / (1024 * 1024), mem_counter / (1024 * 1024), full_load))
+            logging.debug("loaded completely {} {} {}".format(lowvram_model_memory / (1024 * 1024), mem_counter / (1024 * 1024), full_load))
             self._memory_measurements.model_lowvram = False
             if full_load:
                 self.model.to(device_to)
diff --git a/comfy_extras/nodes/nodes_nf4.py b/comfy_extras/nodes/nodes_nf4.py
index cc5414ea3..cae80f71c 100644
--- a/comfy_extras/nodes/nodes_nf4.py
+++ b/comfy_extras/nodes/nodes_nf4.py
@@ -61,6 +61,8 @@ def copy_quant_state(state: QuantState, device: torch.device = None) -> QuantSta
 
 class ForgeParams4bit(Params4bit):
     def to(self, *args, **kwargs):
+        if 'copy' in kwargs:
+            kwargs.pop('copy')
         device, dtype, non_blocking, convert_to_format = torch._C._nn._parse_to(*args, **kwargs)
         if device is not None and device.type == "cuda" and not self.bnb_quantized:
             return self._quantize(device)
@@ -179,7 +181,7 @@ class CheckpointLoaderNF4:
 
     def load_checkpoint(self, ckpt_name):
         if not has_bitsandbytes:
-            raise BitsAndBytesNotFoundError(f"Because your platform is {platform.platform()}, bitsandbytes is not installed, so this cannot be executed")
+            raise BitsAndBytesNotFoundError(f"bitsandbytes is not installed, so {CheckpointLoaderNF4.__name__} cannot be executed")
         ckpt_path = get_or_download("checkpoints", ckpt_name)
         out = comfy.sd.load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, embedding_directory=get_folder_paths("embeddings"), model_options={"custom_operations": OPS})
         return out[:3]
diff --git a/requirements.txt b/requirements.txt
index 16ea5913f..71eb76464 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -10,7 +10,8 @@ sentencepiece
 peft>=0.10.0
 torchinfo
 safetensors>=0.4.2
-bitsandbytes>=0.43.0 ;platform_system == 'Linux' or platform_system == 'Windows'
+bitsandbytes>=0.43.0 ;platform_system != 'Darwin'
+bitsandbytes ;platform_system == 'Darwin'
 aiohttp>=3.8.4
 accelerate>=0.25.0
 pyyaml>=6.0