feat: initial Scaled FP8 support

.github/workflows/publish_action.yml

@@ -1,24 +0,0 @@
-name: Publish to Comfy registry
-on:
-  workflow_dispatch:
-  push:
-    branches:
-      - main
-    paths:
-      - "pyproject.toml"
-
-permissions:
-  issues: write
-
-jobs:
-  publish-node:
-    name: Publish Custom Node to registry
-    runs-on: ubuntu-latest
-    if: ${{ github.repository_owner == 'lodestone-rock' }}
-    steps:
-      - name: Check out code
-        uses: actions/checkout@v4
-      - name: Publish Custom Node
-        uses: Comfy-Org/publish-node-action@v1
-        with:
-          personal_access_token: ${{ secrets.POTATO }} ## Add your own personal access token to your Github Repository secrets and reference it 

README.md

@@ -1,6 +1,5 @@
-# DEPRECATED! Chroma Manual [WIP]
+# Chroma Manual [WIP]
 Chroma is a 8.9B parameter model based on FLUX.1-schnell.
-use comfyui default workflow here instead [https://comfyanonymous.github.io/ComfyUI_examples/chroma/](https://comfyanonymous.github.io/ComfyUI_examples/chroma/)
 
 ## Requirements
 

flux_mod/loader.py

@@ -22,10 +22,7 @@ def ensure_gguf():
     if gguf:
         return
     import sys
-    gguf = next(
-        (mod for path, mod in sys.modules.items() if path.endswith("ComfyUI-GGUF")),
-        None
-    )
+    gguf = sys.modules.get("ComfyUI-GGUF")
     if gguf is None:
         raise RuntimeError("Could not find ComfyUI-GGUF node: GGUF support requires ComfyUI-GGUF")
 
@@ -33,7 +30,7 @@ class ExternalFlux(comfy.supported_models_base.BASE):
     unet_config = {}
     unet_extra_config = {}
     latent_format = comfy.latent_formats.Flux
-    memory_usage_factor = 2.8
+    memory_usage_factor = 2.0 # Lower memory usage factor than FLUX, if we get OOMs this may need to be raised by ~0.2 or so.
       
     def __init__(self,):
         self.unet_config = {}
@@ -177,10 +174,11 @@ def load_flux_mod(model_path, timestep_guidance_path=None, linear_dtypes=torch.b
             n_layers = 4
 
     param_count = sum(x.numel() for x in state_dict.values())
-    unet_dtype = torch.bfloat16
 
     load_device = model_management.get_torch_device()
     offload_device = model_management.unet_offload_device()
+    unet_dtype = model_management.unet_dtype(model_params=param_count)
+    manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device)
 
 
     params=FluxParams(
@@ -199,11 +197,20 @@ def load_flux_mod(model_path, timestep_guidance_path=None, linear_dtypes=torch.b
     )
 
     model_conf = ExternalFlux()
+    if state_dict["scaled_fp8"] is not None:
+        scaled_fp8_weight = state_dict.pop("scaled_fp8")
+        model_conf.scaled_fp8 = scaled_fp8_weight.dtype
+        if model_conf.scaled_fp8 == torch.float32:
+            model_conf.scaled_fp8 = torch.float8_e4m3fn
+        if scaled_fp8_weight.nelement() == 2:
+            model_conf.optimizations["fp8"] = False
+        else:
+            model_conf.optimizations["fp8"] = True
     model_class = ChromaFluxModel if timestep_guidance_path is None else ExternalFluxModel
     model = model_class(
         model_conf,
         model_type=comfy.model_base.ModelType.FLUX,
-        device=model_management.get_torch_device()
+        device=load_device
     )
     unet_config = model_conf.unet_config
     if is_gguf:
@@ -238,9 +245,10 @@ def load_flux_mod(model_path, timestep_guidance_path=None, linear_dtypes=torch.b
     model.diffusion_model.dtype = unet_dtype
     model.diffusion_model.eval()
     if not is_gguf:
-        model.diffusion_model.to(unet_dtype)
-        # we cast to fp8 for mixed matmul ops but omit the picky and sensitive layers
-        cast_layers(model.diffusion_model, nn.Linear, dtype=linear_dtypes, exclude_keywords={"img_in", "final_layer", "scale"})
+        if model_conf.scaled_fp8 is None:
+            #model.diffusion_model.to(linear_dtypes)
+            # we cast to fp8 for mixed matmul ops but omit the picky and sensitive layers
+            cast_layers(model.diffusion_model, nn.Linear, dtype=linear_dtypes, exclude_keywords={"img_in", "txt_in", "final_layer", "in_proj", "in_layer", "out_layer", "out_proj"})
         if model_management.force_channels_last():
             model.diffusion_model.to(memory_format=torch.channels_last)
 

flux_mod/nodes.py

@@ -124,8 +124,6 @@ class ChromaPromptTruncation:
     def append(self, conditioning):
         pruning_idx = conditioning[0][1]["attention_mask"].sum() + 1
         conditioning[0][0] = conditioning[0][0][:, :pruning_idx]
-        del conditioning[0][1]["attention_mask"]
-        conditioning[0][1]["pooled_output"] = torch.zeros(1, 768, dtype=torch.float32)
         c = node_helpers.conditioning_set_values(conditioning, {"guidance": 0})
         return (c,)
 

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
-name = "ComfyUI_FluxMod"
-description = "Nodes for use of Chroma model and other prototype models"
-version = "1.0.1"
+name = "comfyui_fluxmod"
+description = ""
+version = "1.0.0"
 license = {file = "LICENSE"}
 
 [project.urls]
@@ -9,6 +9,6 @@ Repository = "https://github.com/lodestone-rock/ComfyUI_FluxMod"
 #  Used by Comfy Registry https://comfyregistry.org
 
 [tool.comfy]
-PublisherId = "lodestone-rock"
+PublisherId = "horizon"
 DisplayName = "ComfyUI_FluxMod"
 Icon = ""