Update README.md

Support for schnell-compat / svdquant

.github/workflows/publish_action.yml

@@ -0,0 +1,24 @@
+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,5 +1,6 @@
-# Chroma Manual [WIP]
+# DEPRECATED! 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,7 +22,10 @@ def ensure_gguf():
     if gguf:
         return
     import sys
-    gguf = sys.modules.get("ComfyUI-GGUF")
+    gguf = next(
+        (mod for path, mod in sys.modules.items() if path.endswith("ComfyUI-GGUF")),
+        None
+    )
     if gguf is None:
         raise RuntimeError("Could not find ComfyUI-GGUF node: GGUF support requires ComfyUI-GGUF")
 
@@ -30,7 +33,7 @@ class ExternalFlux(comfy.supported_models_base.BASE):
     unet_config = {}
     unet_extra_config = {}
     latent_format = comfy.latent_formats.Flux
-    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.
+    memory_usage_factor = 2.8
       
     def __init__(self,):
         self.unet_config = {}
@@ -174,11 +177,10 @@ 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(
@@ -197,20 +199,11 @@ 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=load_device
+        device=model_management.get_torch_device()
     )
     unet_config = model_conf.unet_config
     if is_gguf:
@@ -245,10 +238,9 @@ 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:
-        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"})
+        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_management.force_channels_last():
             model.diffusion_model.to(memory_format=torch.channels_last)
 

flux_mod/nodes.py

@@ -124,6 +124,8 @@ 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 = ""
-version = "1.0.0"
+name = "ComfyUI_FluxMod"
+description = "Nodes for use of Chroma model and other prototype models"
+version = "1.0.1"
 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 = "horizon"
+PublisherId = "lodestone-rock"
 DisplayName = "ComfyUI_FluxMod"
 Icon = ""