diff --git a/onediff_comfy_nodes/__init__.py b/onediff_comfy_nodes/__init__.py
index d32e96ffe..64f16b288 100644
--- a/onediff_comfy_nodes/__init__.py
+++ b/onediff_comfy_nodes/__init__.py
@@ -47,9 +47,10 @@ def lazy_load_extra_nodes():
update_node_mappings(nodes_compare)
if is_nexfort_available():
- from .extras_nodes import nodes_nexfort_booster
+ from .extras_nodes import nodes_nexfort_booster, nodes_flux_speedup
update_node_mappings(nodes_nexfort_booster)
+ update_node_mappings(nodes_flux_speedup)
from .extras_nodes import nodes_prompt_styler
diff --git a/onediff_comfy_nodes/benchmarks/resources/workflows/nexfort/flux_speedup.json b/onediff_comfy_nodes/benchmarks/resources/workflows/nexfort/flux_speedup.json
new file mode 100644
index 000000000..a60163838
--- /dev/null
+++ b/onediff_comfy_nodes/benchmarks/resources/workflows/nexfort/flux_speedup.json
@@ -0,0 +1,1071 @@
+{
+ "last_node_id": 38,
+ "last_link_id": 117,
+ "nodes": [
+ {
+ "id": 11,
+ "type": "DualCLIPLoader",
+ "pos": {
+ "0": 48,
+ "1": 288
+ },
+ "size": {
+ "0": 315,
+ "1": 106
+ },
+ "flags": {},
+ "order": 0,
+ "mode": 0,
+ "inputs": [],
+ "outputs": [
+ {
+ "name": "CLIP",
+ "type": "CLIP",
+ "links": [
+ 10
+ ],
+ "slot_index": 0,
+ "shape": 3
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "DualCLIPLoader"
+ },
+ "widgets_values": [
+ "t5xxl_fp16.safetensors",
+ "clip_l.safetensors",
+ "flux"
+ ]
+ },
+ {
+ "id": 17,
+ "type": "BasicScheduler",
+ "pos": {
+ "0": 480,
+ "1": 1008
+ },
+ "size": {
+ "0": 315,
+ "1": 106
+ },
+ "flags": {},
+ "order": 13,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "model",
+ "type": "MODEL",
+ "link": 55,
+ "slot_index": 0
+ }
+ ],
+ "outputs": [
+ {
+ "name": "SIGMAS",
+ "type": "SIGMAS",
+ "links": [
+ 20
+ ],
+ "shape": 3
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "BasicScheduler"
+ },
+ "widgets_values": [
+ "simple",
+ 20,
+ 1
+ ]
+ },
+ {
+ "id": 16,
+ "type": "KSamplerSelect",
+ "pos": {
+ "0": 480,
+ "1": 912
+ },
+ "size": {
+ "0": 315,
+ "1": 58
+ },
+ "flags": {},
+ "order": 1,
+ "mode": 0,
+ "inputs": [],
+ "outputs": [
+ {
+ "name": "SAMPLER",
+ "type": "SAMPLER",
+ "links": [
+ 19
+ ],
+ "shape": 3
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "KSamplerSelect"
+ },
+ "widgets_values": [
+ "euler"
+ ]
+ },
+ {
+ "id": 26,
+ "type": "FluxGuidance",
+ "pos": {
+ "0": 480,
+ "1": 144
+ },
+ "size": {
+ "0": 317.4000244140625,
+ "1": 58
+ },
+ "flags": {},
+ "order": 12,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "conditioning",
+ "type": "CONDITIONING",
+ "link": 41
+ }
+ ],
+ "outputs": [
+ {
+ "name": "CONDITIONING",
+ "type": "CONDITIONING",
+ "links": [
+ 42
+ ],
+ "slot_index": 0,
+ "shape": 3
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "FluxGuidance"
+ },
+ "widgets_values": [
+ 3.5
+ ],
+ "color": "#233",
+ "bgcolor": "#355"
+ },
+ {
+ "id": 22,
+ "type": "BasicGuider",
+ "pos": {
+ "0": 576,
+ "1": 48
+ },
+ "size": {
+ "0": 222.3482666015625,
+ "1": 46
+ },
+ "flags": {},
+ "order": 14,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "model",
+ "type": "MODEL",
+ "link": 54,
+ "slot_index": 0
+ },
+ {
+ "name": "conditioning",
+ "type": "CONDITIONING",
+ "link": 42,
+ "slot_index": 1
+ }
+ ],
+ "outputs": [
+ {
+ "name": "GUIDER",
+ "type": "GUIDER",
+ "links": [
+ 30
+ ],
+ "slot_index": 0,
+ "shape": 3
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "BasicGuider"
+ },
+ "widgets_values": []
+ },
+ {
+ "id": 13,
+ "type": "SamplerCustomAdvanced",
+ "pos": {
+ "0": 864,
+ "1": 192
+ },
+ "size": {
+ "0": 272.3617858886719,
+ "1": 124.53733825683594
+ },
+ "flags": {},
+ "order": 15,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "noise",
+ "type": "NOISE",
+ "link": 37,
+ "slot_index": 0
+ },
+ {
+ "name": "guider",
+ "type": "GUIDER",
+ "link": 30,
+ "slot_index": 1
+ },
+ {
+ "name": "sampler",
+ "type": "SAMPLER",
+ "link": 19,
+ "slot_index": 2
+ },
+ {
+ "name": "sigmas",
+ "type": "SIGMAS",
+ "link": 20,
+ "slot_index": 3
+ },
+ {
+ "name": "latent_image",
+ "type": "LATENT",
+ "link": 116,
+ "slot_index": 4
+ }
+ ],
+ "outputs": [
+ {
+ "name": "output",
+ "type": "LATENT",
+ "links": [
+ 24
+ ],
+ "slot_index": 0,
+ "shape": 3
+ },
+ {
+ "name": "denoised_output",
+ "type": "LATENT",
+ "links": null,
+ "shape": 3
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "SamplerCustomAdvanced"
+ },
+ "widgets_values": []
+ },
+ {
+ "id": 25,
+ "type": "RandomNoise",
+ "pos": {
+ "0": 480,
+ "1": 768
+ },
+ "size": {
+ "0": 315,
+ "1": 82
+ },
+ "flags": {},
+ "order": 2,
+ "mode": 0,
+ "inputs": [],
+ "outputs": [
+ {
+ "name": "NOISE",
+ "type": "NOISE",
+ "links": [
+ 37
+ ],
+ "shape": 3
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "RandomNoise"
+ },
+ "widgets_values": [
+ 708172115877408,
+ "randomize"
+ ],
+ "color": "#2a363b",
+ "bgcolor": "#3f5159"
+ },
+ {
+ "id": 8,
+ "type": "VAEDecode",
+ "pos": {
+ "0": 866,
+ "1": 367
+ },
+ "size": {
+ "0": 210,
+ "1": 46
+ },
+ "flags": {},
+ "order": 16,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "samples",
+ "type": "LATENT",
+ "link": 24
+ },
+ {
+ "name": "vae",
+ "type": "VAE",
+ "link": 12
+ }
+ ],
+ "outputs": [
+ {
+ "name": "IMAGE",
+ "type": "IMAGE",
+ "links": [
+ 9
+ ],
+ "slot_index": 0
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "VAEDecode"
+ },
+ "widgets_values": []
+ },
+ {
+ "id": 6,
+ "type": "CLIPTextEncode",
+ "pos": {
+ "0": 384,
+ "1": 240
+ },
+ "size": {
+ "0": 422.84503173828125,
+ "1": 164.31304931640625
+ },
+ "flags": {},
+ "order": 9,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "clip",
+ "type": "CLIP",
+ "link": 10
+ }
+ ],
+ "outputs": [
+ {
+ "name": "CONDITIONING",
+ "type": "CONDITIONING",
+ "links": [
+ 41
+ ],
+ "slot_index": 0
+ }
+ ],
+ "title": "CLIP Text Encode (Positive Prompt)",
+ "properties": {
+ "Node name for S&R": "CLIPTextEncode"
+ },
+ "widgets_values": [
+ "cute anime girl with massive fluffy fennec ears and a big fluffy tail blonde messy long hair blue eyes wearing a maid outfit with a long black gold leaf pattern dress and a white apron mouth open holding a fancy black forest cake with candles on top in the kitchen of an old dark Victorian mansion lit by candlelight with a bright window to the foggy forest and very expensive stuff everywhere"
+ ],
+ "color": "#232",
+ "bgcolor": "#353"
+ },
+ {
+ "id": 30,
+ "type": "ModelSamplingFlux",
+ "pos": {
+ "0": 480,
+ "1": 1152
+ },
+ "size": {
+ "0": 315,
+ "1": 130
+ },
+ "flags": {},
+ "order": 11,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "model",
+ "type": "MODEL",
+ "link": 117,
+ "slot_index": 0
+ },
+ {
+ "name": "width",
+ "type": "INT",
+ "link": 115,
+ "slot_index": 1,
+ "widget": {
+ "name": "width"
+ }
+ },
+ {
+ "name": "height",
+ "type": "INT",
+ "link": 114,
+ "slot_index": 2,
+ "widget": {
+ "name": "height"
+ }
+ }
+ ],
+ "outputs": [
+ {
+ "name": "MODEL",
+ "type": "MODEL",
+ "links": [
+ 54,
+ 55
+ ],
+ "slot_index": 0,
+ "shape": 3
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "ModelSamplingFlux"
+ },
+ "widgets_values": [
+ 1.15,
+ 0.5,
+ 1024,
+ 1024
+ ]
+ },
+ {
+ "id": 27,
+ "type": "EmptySD3LatentImage",
+ "pos": {
+ "0": 480,
+ "1": 624
+ },
+ "size": {
+ "0": 315,
+ "1": 106
+ },
+ "flags": {},
+ "order": 10,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "width",
+ "type": "INT",
+ "link": 112,
+ "widget": {
+ "name": "width"
+ }
+ },
+ {
+ "name": "height",
+ "type": "INT",
+ "link": 113,
+ "widget": {
+ "name": "height"
+ }
+ }
+ ],
+ "outputs": [
+ {
+ "name": "LATENT",
+ "type": "LATENT",
+ "links": [
+ 116
+ ],
+ "slot_index": 0,
+ "shape": 3
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "EmptySD3LatentImage"
+ },
+ "widgets_values": [
+ 1024,
+ 1024,
+ 1
+ ]
+ },
+ {
+ "id": 34,
+ "type": "PrimitiveNode",
+ "pos": {
+ "0": 432,
+ "1": 480
+ },
+ "size": {
+ "0": 210,
+ "1": 82
+ },
+ "flags": {},
+ "order": 3,
+ "mode": 0,
+ "inputs": [],
+ "outputs": [
+ {
+ "name": "INT",
+ "type": "INT",
+ "links": [
+ 112,
+ 115
+ ],
+ "slot_index": 0,
+ "widget": {
+ "name": "width"
+ }
+ }
+ ],
+ "title": "width",
+ "properties": {
+ "Run widget replace on values": false
+ },
+ "widgets_values": [
+ 1024,
+ "fixed"
+ ],
+ "color": "#323",
+ "bgcolor": "#535"
+ },
+ {
+ "id": 35,
+ "type": "PrimitiveNode",
+ "pos": {
+ "0": 672,
+ "1": 480
+ },
+ "size": {
+ "0": 210,
+ "1": 82
+ },
+ "flags": {},
+ "order": 4,
+ "mode": 0,
+ "inputs": [],
+ "outputs": [
+ {
+ "name": "INT",
+ "type": "INT",
+ "links": [
+ 113,
+ 114
+ ],
+ "slot_index": 0,
+ "widget": {
+ "name": "height"
+ }
+ }
+ ],
+ "title": "height",
+ "properties": {
+ "Run widget replace on values": false
+ },
+ "widgets_values": [
+ 1024,
+ "fixed"
+ ],
+ "color": "#323",
+ "bgcolor": "#535"
+ },
+ {
+ "id": 9,
+ "type": "SaveImage",
+ "pos": {
+ "0": 1155,
+ "1": 196
+ },
+ "size": {
+ "0": 985.3012084960938,
+ "1": 1060.3828125
+ },
+ "flags": {},
+ "order": 17,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "images",
+ "type": "IMAGE",
+ "link": 9
+ }
+ ],
+ "outputs": [],
+ "properties": {},
+ "widgets_values": [
+ "ComfyUI"
+ ]
+ },
+ {
+ "id": 37,
+ "type": "Note",
+ "pos": {
+ "0": 480,
+ "1": 1344
+ },
+ "size": {
+ "0": 314.99755859375,
+ "1": 117.98363494873047
+ },
+ "flags": {},
+ "order": 5,
+ "mode": 0,
+ "inputs": [],
+ "outputs": [],
+ "properties": {
+ "text": ""
+ },
+ "widgets_values": [
+ "The reference sampling implementation auto adjusts the shift value based on the resolution, if you don't want this you can just bypass (CTRL-B) this ModelSamplingFlux node.\n"
+ ],
+ "color": "#432",
+ "bgcolor": "#653"
+ },
+ {
+ "id": 28,
+ "type": "Note",
+ "pos": {
+ "0": 48,
+ "1": 576
+ },
+ "size": {
+ "0": 336,
+ "1": 288
+ },
+ "flags": {},
+ "order": 6,
+ "mode": 0,
+ "inputs": [],
+ "outputs": [],
+ "properties": {
+ "text": ""
+ },
+ "widgets_values": [
+ "If you get an error in any of the nodes above make sure the files are in the correct directories.\n\nSee the top of the examples page for the links : https://comfyanonymous.github.io/ComfyUI_examples/flux/\n\nflux1-dev.safetensors goes in: ComfyUI/models/unet/\n\nt5xxl_fp16.safetensors and clip_l.safetensors go in: ComfyUI/models/clip/\n\nae.safetensors goes in: ComfyUI/models/vae/\n\n\nTip: You can set the weight_dtype above to one of the fp8 types if you have memory issues."
+ ],
+ "color": "#432",
+ "bgcolor": "#653"
+ },
+ {
+ "id": 38,
+ "type": "OneDiffFluxLoader",
+ "pos": {
+ "0": 20,
+ "1": 61
+ },
+ "size": {
+ "0": 394.79998779296875,
+ "1": 106
+ },
+ "flags": {},
+ "order": 7,
+ "mode": 0,
+ "inputs": [],
+ "outputs": [
+ {
+ "name": "MODEL",
+ "type": "MODEL",
+ "links": [
+ 117
+ ],
+ "slot_index": 0
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "OneDiffFluxLoader"
+ },
+ "widgets_values": [
+ "flux1-dev.sft",
+ true,
+ true
+ ]
+ },
+ {
+ "id": 10,
+ "type": "VAELoader",
+ "pos": {
+ "0": 48,
+ "1": 432
+ },
+ "size": {
+ "0": 311.81634521484375,
+ "1": 60.429901123046875
+ },
+ "flags": {},
+ "order": 8,
+ "mode": 0,
+ "inputs": [],
+ "outputs": [
+ {
+ "name": "VAE",
+ "type": "VAE",
+ "links": [
+ 12
+ ],
+ "slot_index": 0,
+ "shape": 3
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "VAELoader"
+ },
+ "widgets_values": [
+ "ae.sft"
+ ]
+ }
+ ],
+ "links": [
+ [
+ 9,
+ 8,
+ 0,
+ 9,
+ 0,
+ "IMAGE"
+ ],
+ [
+ 10,
+ 11,
+ 0,
+ 6,
+ 0,
+ "CLIP"
+ ],
+ [
+ 12,
+ 10,
+ 0,
+ 8,
+ 1,
+ "VAE"
+ ],
+ [
+ 19,
+ 16,
+ 0,
+ 13,
+ 2,
+ "SAMPLER"
+ ],
+ [
+ 20,
+ 17,
+ 0,
+ 13,
+ 3,
+ "SIGMAS"
+ ],
+ [
+ 24,
+ 13,
+ 0,
+ 8,
+ 0,
+ "LATENT"
+ ],
+ [
+ 30,
+ 22,
+ 0,
+ 13,
+ 1,
+ "GUIDER"
+ ],
+ [
+ 37,
+ 25,
+ 0,
+ 13,
+ 0,
+ "NOISE"
+ ],
+ [
+ 41,
+ 6,
+ 0,
+ 26,
+ 0,
+ "CONDITIONING"
+ ],
+ [
+ 42,
+ 26,
+ 0,
+ 22,
+ 1,
+ "CONDITIONING"
+ ],
+ [
+ 54,
+ 30,
+ 0,
+ 22,
+ 0,
+ "MODEL"
+ ],
+ [
+ 55,
+ 30,
+ 0,
+ 17,
+ 0,
+ "MODEL"
+ ],
+ [
+ 112,
+ 34,
+ 0,
+ 27,
+ 0,
+ "INT"
+ ],
+ [
+ 113,
+ 35,
+ 0,
+ 27,
+ 1,
+ "INT"
+ ],
+ [
+ 114,
+ 35,
+ 0,
+ 30,
+ 2,
+ "INT"
+ ],
+ [
+ 115,
+ 34,
+ 0,
+ 30,
+ 1,
+ "INT"
+ ],
+ [
+ 116,
+ 27,
+ 0,
+ 13,
+ 4,
+ "LATENT"
+ ],
+ [
+ 117,
+ 38,
+ 0,
+ 30,
+ 0,
+ "MODEL"
+ ]
+ ],
+ "groups": [],
+ "config": {},
+ "extra": {
+ "ds": {
+ "scale": 1.2100000000000006,
+ "offset": [
+ -40.02256171358903,
+ -95.08194516611923
+ ]
+ },
+ "groupNodes": {
+ "EmptyLatentImage": {
+ "nodes": [
+ {
+ "type": "PrimitiveNode",
+ "pos": [
+ 432,
+ 480
+ ],
+ "size": {
+ "0": 210,
+ "1": 82
+ },
+ "flags": {},
+ "order": 6,
+ "mode": 0,
+ "outputs": [
+ {
+ "name": "INT",
+ "type": "INT",
+ "links": [],
+ "widget": {
+ "name": "height"
+ },
+ "slot_index": 0
+ }
+ ],
+ "title": "height",
+ "properties": {
+ "Run widget replace on values": false
+ },
+ "color": "#323",
+ "bgcolor": "#535",
+ "index": 0
+ },
+ {
+ "type": "PrimitiveNode",
+ "pos": [
+ 672,
+ 480
+ ],
+ "size": {
+ "0": 210,
+ "1": 82
+ },
+ "flags": {},
+ "order": 7,
+ "mode": 0,
+ "outputs": [
+ {
+ "name": "INT",
+ "type": "INT",
+ "links": [],
+ "slot_index": 0,
+ "widget": {
+ "name": "width"
+ }
+ }
+ ],
+ "title": "width",
+ "properties": {
+ "Run widget replace on values": false
+ },
+ "color": "#323",
+ "bgcolor": "#535",
+ "index": 1
+ },
+ {
+ "type": "EmptySD3LatentImage",
+ "pos": [
+ 480,
+ 624
+ ],
+ "size": {
+ "0": 315,
+ "1": 106
+ },
+ "flags": {},
+ "order": 10,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "width",
+ "type": "INT",
+ "link": null,
+ "widget": {
+ "name": "width"
+ }
+ },
+ {
+ "name": "height",
+ "type": "INT",
+ "link": null,
+ "widget": {
+ "name": "height"
+ }
+ }
+ ],
+ "outputs": [
+ {
+ "name": "LATENT",
+ "type": "LATENT",
+ "links": [],
+ "shape": 3,
+ "slot_index": 0
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "EmptySD3LatentImage"
+ },
+ "widgets_values": [
+ 1024,
+ 1024,
+ 1
+ ],
+ "index": 2
+ }
+ ],
+ "links": [
+ [
+ 1,
+ 0,
+ 2,
+ 0,
+ 34,
+ "INT"
+ ],
+ [
+ 0,
+ 0,
+ 2,
+ 1,
+ 35,
+ "INT"
+ ]
+ ],
+ "external": [
+ [
+ 0,
+ 0,
+ "INT"
+ ],
+ [
+ 1,
+ 0,
+ "INT"
+ ],
+ [
+ 2,
+ 0,
+ "LATENT"
+ ]
+ ],
+ "config": {
+ "0": {
+ "output": {
+ "0": {
+ "name": "height"
+ }
+ },
+ "input": {
+ "value": {
+ "visible": true
+ }
+ }
+ },
+ "1": {
+ "output": {
+ "0": {
+ "name": "width"
+ }
+ },
+ "input": {
+ "value": {
+ "visible": true
+ }
+ }
+ },
+ "2": {
+ "input": {
+ "width": {
+ "visible": false
+ },
+ "height": {
+ "visible": false
+ }
+ }
+ }
+ }
+ }
+ }
+ },
+ "version": 0.4
+}
\ No newline at end of file
diff --git a/onediff_comfy_nodes/benchmarks/resources/workflows/nexfort/flux_speedup.png b/onediff_comfy_nodes/benchmarks/resources/workflows/nexfort/flux_speedup.png
new file mode 100644
index 000000000..084af7e05
Binary files /dev/null and b/onediff_comfy_nodes/benchmarks/resources/workflows/nexfort/flux_speedup.png differ
diff --git a/onediff_comfy_nodes/docs/FLUX/README.md b/onediff_comfy_nodes/docs/FLUX/README.md
new file mode 100644
index 000000000..2930fdb11
--- /dev/null
+++ b/onediff_comfy_nodes/docs/FLUX/README.md
@@ -0,0 +1,76 @@
+## Accelerate FLUX by using onediff
+huggingface: https://huggingface.co/black-forest-labs/FLUX.1-dev
+
+## Environment setup
+### Set UP requirements
+```shell
+# python 3.10
+COMFYUI_DIR=$pwd/ComfyUI
+# install ComfyUI
+git clone https://github.com/comfyanonymous/ComfyUI.git
+
+# install onediff & onediff_comfy_nodes
+git clone https://github.com/siliconflow/onediff.git
+cd onediff && pip install -r onediff_comfy_nodes/FLUX/requirements.txt && pip install -e .
+ln -s $pwd/onediff/onediff_comfy_nodes $COMFYUI_DIR/custom_nodes
+```
+
+
+### Download relevant models
+
+- step1: Get User Access Tokens here https://huggingface.co/settings/tokens
+
+- step2: Download relevant models
+```shell
+export ACCESS_TOKEN="User Access Tokens"
+wget --header="Authorization: Bearer $ACCESS_TOKEN" \
+https://huggingface.co/black-forest-labs/FLUX.1-dev/resolve/main/flux1-dev.safetensors -O models/checkpoints/flux1-dev.safetensors
+
+wget --header="Authorization: Bearer $ACCESS_TOKEN" \
+https://huggingface.co/stabilityai/stable-diffusion-3-medium/resolve/main/text_encoders/clip_l.safetensors -O models/clip/clip_l.safetensors
+
+# wget --header="Authorization: Bearer $ACCESS_TOKEN" \
+# https://huggingface.co/stabilityai/stable-diffusion-3-medium/resolve/main/text_encoders/t5xxl_fp16.safetensors -O models/clip/t5xxl_fp16.safetensors
+
+wget --header="Authorization: Bearer $ACCESS_TOKEN" \
+https://huggingface.co/stabilityai/stable-diffusion-3-medium/resolve/main/text_encoders/t5xxl_fp8_e4m3fn.safetensors -O models/clip/t5xxl_fp8_e4m3fn.safetensors
+```
+
+
+## Usage Example
+### Run ComfyUI
+```shell
+cd $COMFYUI_DIR && python main.py --gpu-only --disable-cuda-malloc
+```
+
+### WorkFlow
+Here is a very basic example how to use it:
+[workflow_flux_speedup.json](../../benchmarks/resources/workflows/nexfort/flux_speedup.json)
+
+
+
+## Performance Comparison
+
+- Testing on NVIDIA GeForce RTX 4090, with image size of 1024*1024, iterating 20 steps.
+
+
+| Metric | NVIDIA GeForce RTX 4090 (1024 * 1024) |
+| ------------------------------------------------ | ------------------------------------- |
+| Data update date(yyyy-mm-dd) | 2024-11-19 |
+| PyTorch E2E time | 12.44 s |
+| OneDiff E2E time | 7.12 s (-42%) |
+| PyTorch Max Mem Used | 20.510 GiB |
+| OneDiff Max Mem Used | 23.094 GiB |
+| PyTorch Warmup with Run time | 49.09 s |
+| OneDiff Warmup with Compilation time1 | 113.73 s |
+| OneDiff Warmup with Cache time | 20.27 s |
+
+ 1 OneDiff Warmup with Compilation time is tested on AMD EPYC 7543 32-Core Processor CPU. Note this is just for reference, and it varies a lot on different CPU.
+
+
+
+
+## Quality
+
+The following table shows the comparison of the plot, seed=0, Baseline (non optimized) on the left, and OneDiff (optimized) on the right
+
\ No newline at end of file
diff --git a/onediff_comfy_nodes/docs/FLUX/flux_quality_compare.png b/onediff_comfy_nodes/docs/FLUX/flux_quality_compare.png
new file mode 100644
index 000000000..5953209aa
Binary files /dev/null and b/onediff_comfy_nodes/docs/FLUX/flux_quality_compare.png differ
diff --git a/onediff_comfy_nodes/extras_nodes/nodes_flux_speedup.py b/onediff_comfy_nodes/extras_nodes/nodes_flux_speedup.py
new file mode 100644
index 000000000..07298ac9b
--- /dev/null
+++ b/onediff_comfy_nodes/extras_nodes/nodes_flux_speedup.py
@@ -0,0 +1,139 @@
+import comfy
+from comfy.model_patcher import ModelPatcher
+import folder_paths
+
+
+from comfy.ldm.flux.layers import RMSNorm as SrcRMSNorm
+from nexfort.nn.modules import RMSNorm4Comfy4D
+
+from comfy.ldm.flux.layers import MLPEmbedder as SrcMLPEmbedder
+from nexfort.nn.modules import MLPEmbedder
+
+from comfy.ldm.flux.layers import Modulation as SrcModulation
+from nexfort.nn.modules import Modulation
+
+from comfy.ldm.flux.layers import SingleStreamBlock as SrcSingleStreamBlock
+from nexfort.nn.modules import SingleStreamBlock
+
+from comfy.ldm.flux.layers import DoubleStreamBlock as SrcDoubleStreamBlock
+from nexfort.nn.modules import DoubleStreamBlock
+
+from comfy.ldm.flux.layers import LastLayer as SrcLastLayer
+from nexfort.nn.modules import LastLayer
+
+
+
+from contextlib import contextmanager
+from comfy.model_management import VRAMState
+
+@contextmanager
+def comfy_cpu_context():
+ orig_state = comfy.model_management.vram_state
+ comfy.model_management.vram_state = VRAMState.NORMAL_VRAM
+ try:
+ yield
+ finally:
+ comfy.model_management.vram_state = orig_state
+
+comfy2nexfort_class_map = {
+ SrcRMSNorm: RMSNorm4Comfy4D,
+ SrcMLPEmbedder: MLPEmbedder,
+ SrcModulation: Modulation,
+ SrcSingleStreamBlock: SingleStreamBlock,
+ SrcDoubleStreamBlock: DoubleStreamBlock,
+ SrcLastLayer: LastLayer,
+}
+
+from nexfort.compilers.custom_transform import transform_mgr
+
+transform_mgr.update_class_proxies(comfy2nexfort_class_map)
+
+
+from nexfort.compilers.transform_model import transform_model
+from nexfort.quantization import quantize
+
+
+def speedup(model: ModelPatcher, enable_quantize: bool, enable_transform: bool):
+ import os
+
+ os.environ["CUBLASLT_WORKSPACE_SIZE"] = str(1024)
+
+ if enable_quantize:
+ quantize(
+ model.model.diffusion_model,
+ quant_type="fp8_e4m3_e4m3_dynamic_per_tensor",
+ )
+ if enable_transform:
+ transform_model(model.model.diffusion_model)
+
+
+class OneDiffFluxSpeedup:
+ @classmethod
+ def INPUT_TYPES(s):
+ return {
+ "required": {
+ "model": ("MODEL",),
+ "quantize_model": (
+ "BOOLEAN",
+ {"default": True, "label_on": "yes", "label_off": "no"},
+ ),
+ "transform": (
+ "BOOLEAN",
+ {"default": True, "label_on": "yes", "label_off": "no"},
+ ),
+ },
+ }
+
+ RETURN_TYPES = ("MODEL",)
+ FUNCTION = "speedup"
+ CATEGORY = "OneDiffEnterprise"
+
+ def speedup(self, model, quantize_model=True, transform=True):
+ speedup(model, quantize_model, transform)
+ return (model,)
+
+
+class OneDiffFluxLoader:
+ @classmethod
+ def INPUT_TYPES(s):
+ return {
+ "required": {
+ "unet_name": (folder_paths.get_filename_list("diffusion_models"),),
+ "enable_quant": (
+ "BOOLEAN",
+ {"default": True, "label_on": "yes", "label_off": "no"},
+ ),
+ "enable_transform": (
+ "BOOLEAN",
+ {"default": True, "label_on": "yes", "label_off": "no"},
+ ),
+ }
+ }
+
+ RETURN_TYPES = ("MODEL",)
+ FUNCTION = "load_unet"
+
+ CATEGORY = "OneDiffEnterprise"
+
+ def load_unet(self, unet_name, enable_quant, enable_transform):
+ model_options = {}
+ unet_path = folder_paths.get_full_path("diffusion_models", unet_name)
+ with comfy_cpu_context():
+ model = comfy.sd.load_diffusion_model(
+ unet_path, model_options=model_options
+ )
+ model.offload_device = comfy.model_management.unet_offload_device()
+ speedup(model, enable_quant, enable_transform)
+
+ return (model,)
+
+
+NODE_CLASS_MAPPINGS = {
+ "OneDiffFluxSpeedup": OneDiffFluxSpeedup,
+ "OneDiffFluxLoader": OneDiffFluxLoader,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+ "OneDiffFluxSpeedup": "OneDiff Flux Speedup - OneDiffEnterprise",
+ "OneDiffFluxLoader": "OneDiff Flux Speedup Loader - OneDiffEnterprise",
+}