From 96c8d93fbe8446dd27f153ffcad6659a543762d9 Mon Sep 17 00:00:00 2001 From: lum3on Date: Thu, 8 May 2025 10:17:09 +0200 Subject: [PATCH] Add files via upload --- README.md | 88 +++++++ __init__.py | 39 +++ __pycache__/__init__.cpython-312.pyc | Bin 0 -> 1137 bytes __pycache__/nodes.cpython-312.pyc | Bin 0 -> 27466 bytes examples/workflow_quantize.json | 358 ++++++++++++++++++++++++++ nodes.py | 371 +++++++++++++++++++++++++++ requirements.txt | 3 + 7 files changed, 859 insertions(+) create mode 100644 README.md create mode 100644 __init__.py create mode 100644 __pycache__/__init__.cpython-312.pyc create mode 100644 __pycache__/nodes.cpython-312.pyc create mode 100644 examples/workflow_quantize.json create mode 100644 nodes.py create mode 100644 requirements.txt diff --git a/README.md b/README.md new file mode 100644 index 0000000..6f23b85 --- /dev/null +++ b/README.md @@ -0,0 +1,88 @@ +# ComfyUI Quantization Node Pack + +A custom node pack for ComfyUI that provides tools for quantizing model weights to lower precision formats like FP16, BF16, or true FP8 types. + +## Overview + +This node pack provides tools for quantizing models directly within ComfyUI. It includes: + +1. **Model To State Dict**: Extracts the state dictionary from a model object and attempts to normalize keys. +2. **Quantize Model to FP8 Format**: Converts model weights directly to `float8_e4m3fn` or `float8_e5m2` format (requires CUDA). +3. **Quantize Model Scaled**: Applies simulated FP8 scaling (per-tensor or per-channel) and then casts the model to `float16`, `bfloat16`, or keeps the original format. +4. **Save As SafeTensor**: Saves the processed state dictionary to a `.safetensors` file at a specified path. + +## Installation + +1. Clone or download this repository into your ComfyUI's `custom_nodes` directory. + * Example using git: + ```bash + cd ComfyUI/custom_nodes + git clone [https://github.com/YourUsername/YourRepoName.git](https://github.com/YourUsername/YourRepoName.git) ComfyUI_QuantNodes + # Replace with your actual repo URL and desired folder name + ``` + * Alternatively, download the ZIP and extract it into `ComfyUI/custom_nodes/ComfyUI_QuantNodes`. +2. Install dependencies (optional, check if already installed): + ```bash + pip install -r ComfyUI/custom_nodes/ComfyUI_QuantNodes/requirements.txt + ``` +3. Restart ComfyUI. + +## Usage + +### Model To State Dict +* **Category:** `Model Quantization/Utils` +* **Function:** Extracts state dict from a MODEL object, stripping common prefixes. +* **Inputs**: + * `model`: The input `MODEL` object. +* **Outputs**: + * `model_state_dict`: The extracted state dictionary. + +### Quantize Model to FP8 Format +* **Category:** `Model Quantization/FP8 Direct` +* **Function:** Converts model weights directly to a specific FP8 format. Requires CUDA. +* **Inputs**: + * `model_state_dict`: The state dictionary to quantize. + * `fp8_format`: The target FP8 format (`float8_e5m2` or `float8_e4m3fn`). +* **Outputs**: + * `quantized_model_state_dict`: The state dictionary with FP8 tensors. + +### Quantize Model Scaled +* **Category:** `Model Quantization` +* **Function:** Applies simulated FP8 value scaling and then casts to FP16, BF16, or keeps the original dtype. Useful for size reduction with good compatibility. +* **Inputs**: + * `model_state_dict`: The state dictionary to quantize. + * `scaling_strategy`: How to simulate scaling (`per_tensor` or `per_channel`). + * `processing_device`: Where to perform calculations (`Auto`, `CPU`, `GPU`). + * `output_dtype`: Final data type (`Original`, `float16`, `bfloat16`). Defaults to `float16`. +* **Outputs**: + * `quantized_model_state_dict`: The processed state dictionary. + +### Save As SafeTensor +* **Category:** `Model Quantization/Save` +* **Function:** Saves the processed state dictionary to a `.safetensors` file. +* **Inputs**: + * `quantized_model_state_dict`: The state dictionary to save. + * `absolute_save_path`: The full path (including filename) where the model will be saved. +* **Outputs**: None (Output node). + +## Example Workflow + +1. Load a model using a standard loader (e.g., `Load Checkpoint`). +2. Connect the `MODEL` output to the `Model To State Dict` node. +3. Connect the `model_state_dict` output from `Model To State Dict` to `Quantize Model Scaled`. +4. In `Quantize Model Scaled`, select your desired `scaling_strategy` and set `output_dtype` to `float16` (for size reduction). +5. Connect the `quantized_model_state_dict` output from `Quantize Model Scaled` to the `Save Model as SafeTensor` node. +6. Specify the `absolute_save_path` in the `Save Model as SafeTensor` node. +7. Queue the prompt. +8. Restart ComfyUI or refresh loaders to find the saved model. + +## Requirements + +* PyTorch (usually included with ComfyUI) +* `safetensors` +* `tqdm` +* CUDA-enabled GPU is required for the `Quantize Model to FP8 Format` node. + +## License + +MIT (Or your chosen license) diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..d3ae638 --- /dev/null +++ b/__init__.py @@ -0,0 +1,39 @@ +# __init__.py +# This file is necessary to make Python treat the directory as a package. +# It's also where ComfyUI looks for node mappings. + +# Import node classes from nodes.py +# Renamed QuantizeScaled to QuantizeModel +from .nodes import ModelToStateDict, QuantizeFP8Format, QuantizeModel, SaveAsSafeTensor + +# A dictionary that ComfyUI uses to map node class names to node objects +NODE_CLASS_MAPPINGS = { + # Node Utils + "ModelToStateDict": ModelToStateDict, + # Direct FP8 Conversion + "QuantizeFP8Format": QuantizeFP8Format, + # Scaled Quantization + Casting (FP16/BF16) + "QuantizeModel": QuantizeModel, # Renamed class QuantizeScaled -> QuantizeModel + # Saving Node + "SaveAsSafeTensor": SaveAsSafeTensor, +} + +# A dictionary that ComfyUI uses to map node class names to display names +NODE_DISPLAY_NAME_MAPPINGS = { + "ModelToStateDict": "Model To State Dict", + "QuantizeFP8Format": "Quantize Model to FP8 Format", + "QuantizeModel": "Quantize Model Scaled", # Display name for the renamed node + "SaveAsSafeTensor": "Save Model as SafeTensor", +} + +# Optional: Print a message to the console when the extension is loaded +print("----------------------------------------------------") +print("--- ComfyUI Quantization Node Pack Loaded ---") +print("--- Renamed QuantizeScaled to QuantizeModel ---") +print("--- Removed FP8 output from QuantizeModel ---") +print("--- Using User-Provided Scaling Logic ---") # Added note about scaling +print("--- Developed by [Lum3on] ---") # Remember to change this! +print("--- Version 0.5.0 ---") # Incremented version +print("----------------------------------------------------") + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/__pycache__/__init__.cpython-312.pyc b/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..07cdf27566c76420dbfb97fe129344b49c5c0eb0 GIT binary patch literal 1137 zcmbVLO>fgc5M4WVzDn9uRYg+}TJ_W*8d{*VRHZ0tS`;Y_p>FvIMwXl177N=Ot?fwh zjbFeK4*UVcpWx0V5;+u$LnKbz8uft0iCrgw$cPJjSkIezGjGPb`ze!21Ge7wFTMC0 z2jCkAqa_cG^N-AU0}w!h2`r&5SYlnYV)d9M)g=L70gE_>B_zC)>M{~h3`s~1QCtK# zf#M&8K_vl%I1$0g2u?+CI)W!6_(BAa>7fjoJQ3^3_EbLmop+I!g6tka*gPPLtGl>l zG~8hNp{Luf;p5%f`YxfC?gp8WDaIa-SY?C5n%;N` zD?~>aLB{m2bNW8Eb&EMh9S&N9tpAZ`Fwj_nMOrL@k4HSW>ACQj5-UQ7e-sWK!#;tB z4yH>rN?Hb&=PKM=Aty!y4i;Bfnkro9K}%nN?%+Z!q$+fnz6<#Lr>7Y`Z4 zQA?!ZS=e?RwMmJqADLJk9-7+l9G6&HC^~X3g$FP?4XSFIVH>Wd6`JiJxoHvPnRts% zF-3mOo!4wTiwT18V-g5YgwNpdSsLW>z1-E_>_TsDu|J)~73Q7qXMqs!#5cR(?iV0;;;UV7J2ck1VErG)au=+Gp*vl$5gI(bHvu~fTb=a1 KR|%S7J^cbK=QQsC literal 0 HcmV?d00001 diff --git a/__pycache__/nodes.cpython-312.pyc b/__pycache__/nodes.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..964ea035f0f22854ab7921ee39111b8849a5ed42 GIT binary patch literal 27466 zcmeHw3ve4pdgcsX#ES&@ev2R_l7d7*BK4xECq;@B^&lyU)Po`oLzp26J^*R}N+JZ> zx{a$s?>RA@P1aK4JySl{N;)eseY;ntvs-66@vXRtudX3L1~p+W&c(T8@9I+YN>}#o z)>YmA_h5h_1xve|+*PISro`?>cTZ10|Nj0*xBfgm-N?Z6zy6_k@RKtP^EdcIe)RFc zqgf5Zyv}e8$Br;P=&MH9QB{wM#jturJ*w_e(=a=t>Cvc|R)*8O#BkboR8k2&S{kQ= zxU@u^j>hRB&X9;pqj5%vGbQ5mG|mih>4`W4jmv;IOCrw5n>g$6dWM-+35<&2&0OX? ztn_^renXwIAb&Q`3>C=@O!`A5$*>%HM=gbVEcERhcsuvkc-ugB{7Osw%I33&>~bBFKA$DGrYD!r6|{UV#}q=cm-R3`d3>Ir;R^V?ORNC3g|-*U zxv}khzFa!}W#)?bJgxxW9@>)pf*kq-hAU2_*U7O-{|}h&XjYVNkW23=qTeoAk!t2j zlfU_*p`FQ+Im}guoJ9(<4>^r+wH7q@VhSY&VJ7b4?i@4i`=MNIZC?tk4Hlgc%9*y9+qQz zR2)+#@c(-xl!MTw`M+6E;Y;U6j&Y0gum1Q*8c#E5JKB zk2~NHO*Glb1t!LLcbwYNN@>AG`t|5Ja9(EuaZ1QAfkdtlE7K#tC4l+q!?7V2D1hoa zYPmnZ0Wl%Q8K;H}<6+=ee5EqX1Y=jX2MfEs{_(LfpAg_VsmfBF8f?L8hpls>%O|)8 zZJdwy+q}MjZ4d^;78vySZLUB-@LU)V@D96HOdE9hF+$V^Fu!O#?HU>94+(-#5YnJ_ zV!Ge$8iAp7xgmq!t`aqVex#qK`LSI#+s{umHS~7-dBNX%vazkL-F7N4&Ut+P-e#YB ze3bVF{Jmp>FW|Z`!uK}$M*Ana+j`yO{(x`P>Gg5EUyj_=GFIP9V;o}>qTcEBcs&89 zGnlI!GKU%&pYhNY|$QdU1#JeN)mjgZpONS2qFy^FKk zqINn(S=M9@IiyIh{S?c2ckSUY`7R(huoE z`nw6R1+yTH(+5p3?gqtsDOwRyrWYj3lS7;#qz&m(L8!Y%889X4s4Oc5z=Ap&uuf1j zAU|@yNcDsih?_M2Iz1CIrW{G^kN9etG=)s37`rKWR3^VbiEPpuv;}yt-zOlg4Gi)& zKM+<=zX!+~QsaJ~FzO1}#{C}efX(k29UFlhK$GC7L)5s(IhUw*kBy7E{t=%mP#>&4 zucQ}!wlgll3waxCP2J6nloYxyxjZ9CjvTfwVVt)GpSQVv-b+Bhkl+ja#c_`SR1@<% zY>j~cKROn`k3(WpXLpSaXeZ?IkB9XygxDb440j|&a~ z$-x~~)Y4t!@94PP2_uf&7I*z~Oaq|K~M7u@I^!S~$LGWScn9t)4 z2soSqu2KP~u4slFiIzyTT015{Mo|^;iRw|;6;bWF;1|;dc-Xyw!W&`7^ZA`4o?#$# zQjgQGiCV!o?&a)e0hhXvi4HC&(c%+410Jty#3|KHG`n41=eVDD(q2ywxCAI$PP9s6 z>GXM>Qi7NUb#PvCjfh$CoOFwkb}3Oy_Z6`yzS;7eGyXmd@%RXAx14BQo#r<~wQPz< zz+E7jFJ6Q{SgM=?4!IaPIR{?ykHEaV#(ZYUjajP0mg=a*zL;G&b^Nm;+nbwiY?`us zo?r5M@U>tp-x1DtMDn*xnICT1^M>|K^9}RN>1aXC$CdYLKB`$5j~witO8-K~SW9N~ zx0_=%d&2O~wr43li{!S<2V$*f!mVdW%h_dyImb?@z{?4$9y3hFRn<~y#hbl1dT&=o zO1Dh)eO6rl=B^vNX8IzW?XTJ2IP_-QjW$xXBT}+6l2t#IwrI+lc71ct z=jH3~=FVzjTMmS`9Eg-R67!m+?1EVK+Hm$-vThsM(;m(4Am)x`9g|yfoqw%k>d@yo zg|D~0)^?qLyZct}&EDC{`E=6F$GV5Z-NVt!kx0(y)S<^ZrfTD@^Ec0v+Q#`q;o7zZ za85)jPm=DQaAgnayAZCtK=c)mZ%yUf-ExB?rF)6KXt8*Ethhd0Tu=6&AjS2O;*-!0 zsDRZ*^ko84KGB#&QY)}iVw3$Y49n|)cP3#YecDFb%)F|8RpVBB7`Li#8^iR%3Y%0< zY66L^JETbx_Eg)gOfvoKE9D{0(>7noQ5RB&G~ZIavO~j6YUy`?h)PdL`?Pujwqu5H z+OB#4&#O29xpy@(hcT%O>6)27*qr3Hz?-SIN6io}(N&KO6hdmbT_IISy$5o=vO)6< z&;nnD?~R$of!_Z_uS0dcd8YElQL_6utU5@6LYDH2Nl8ttq%K@i7b)2`_uN8r z>{xI3ST8x+M|M9)_)#b;ntyf5{5QqbFCU$1p051NzA0wk7q;)48~^a?2UmYExxhv2 zr-<1$lXKr(wX9`I)-UUslB!to#&Ge*Pm4E6uPs`#N%7|Umd%Ts>WFpSOxJzu2F&ch z<-d;T*Z%BLF4oA8RN&hcP3gK1tl)fP)Pjo)qe`z?D;~Y(b{I39x>o8o7X{j{8s2@XqkcFd^YtKa{p7wk4qk5%<5^?Qb=a?eKKmJOVVHv;!;1n$`k zzfBYAfJ(ApVxK~Okc+Qg;GrKGaXp-pcWosR;DtDBtKW3N6QHOK01_Q(qzQuWiU*(_ zg02fLKSC!j*Gde#|Zv zC0gSM6zB+4iXWI*v6~^~eNIE7}0VJCA8Nz*j)f0k!5Lq%st`w{&u&61S>aYEJ90!C{GK++WBBY|s zkA^ZjYQ+~*?9i~2#z!gDF3f>wO5SOaIp;K3GL(-jls2RVR!0{+Ca;Q>A|;r2@##~8 z;-}HE_#P;&62avi;A?CyF9%aPzO6t8hfoY1M;VEM-}|)c?@f-(eFBQr2WWrJJLm0s z!G>8VpMY$Hv`1<$5Ap&}_gA3`VnrSOrqiyICR$=vhi zqw6`+KOC`+OlhI0iuJcEZ?2iG443aD`ZY^{g^#{=l-PDea(2dY_Jwoy&0U%wkK`Pm zGA!nm%~Xc-s&9j{dCIhCDFSROt)NV(80C#tEFKtgX}K3cNGbRWu}RTiszoACAi)aL zII&2ewD==~x`b3I#qyc_8rcoUQaQ&rl7KiB3{|A2kg|ffGQI?4pH6kvFenrRRY>D4 zSGF>PEiT&>vKfHA@db*DWhfzwNfsct9YAm@S&IsxEu`T*U-iH;PpH8VRD&cY zrF4taVh>HT!ndK0xRO>qASnVuF8u!U-BFQ)=ckK4IpC%gnUvT z(udM;#sEuU2OaHT={>Ee{~EdPlD;DKsjUkcxU{?a z73-xuS&AHz*UQjaIVS1LY5x>!DS-P0eLyama$O;7E6;^*u#{y~P3nPRGzJgI3?miu zB_LybD_8{uF^_xD=jRc6Nhr%98UgpkSu<%bl&~Mk6&n<=5tKP1>Rdcf@YS}z{DYmmCyt5;^^MoWOD=d4n^hW7a#%Cr=%w!9j zY%|C1o3<2_DJ}cjF+x6C(*V?|{jeGd~+66D+hGN`U zm>Y|3zc_UWxW;#T?)1FVJJlRBuL+yiTn~oL8y3sf#>#ew%XZJT%pZ!B9U&pJv${yt&PYl9v|+K* zF%>A(BgOYhSIr%&p^u`7Q;lV(3*Z8z{v3AG%Meh}TuVh*ivsQ$y6*CQC zt7CC(b!_ea@Y?`yZ9B2rmN&#JQO=avrcFzm>SlMmzx%!2(M|g!Hd#BGWZrD0|@!vRZ zockgF;m8Lg3tPz*H}t)c>8@CAbvU>Bw()-M*3Y4x zM{XQ>d;HeK&57B(NOk?3mz?g2mYknH^x5`?_xs-K`<^ql{b+dm(a84W)6KD*s&G!# zOfZ~N_gRWI8KpKoPGho5N!hNsb+Lwya6<>#)k#WEM6ym2^U0t6w2-OjW~mhRM;hZ1 zz2-+24Y)!rOq&%95aJZd0Nn=s@c)K557oDvn$z6X{<9V6>4JM@`kcV-N_JkN&!{T@ z8kt@uzCMB645eJ0ZGAAqzJm9W#C*b5m>+h70HA(lvrXUT|xsoyTN%v{6>X4 z!O9e*4de?W6EZ_WYyM3vfPo1qOWi*Ezan)bd)pTRu?{ZW!I5^J@P4v8ungFhZAN9l zG1Upw1{9oB{fa_3RMmQ=v81EOrd5QT#rP1oYFsih++u4s8@Rs!^ZSigX-H|&Z z%M1kPx8K|I(HZ2f^ho)*?yez(c7Vcv5w+ zCT{pkJcthPc~l2ZQYy!*Irb$+Moc{#P6bvd>K+}Z=}F_XU{#`fiRm#+q}lbNl}e4e ze5cSlq!~nwZkSw&2k9TBf!GmX%EzL(-@|&qKESSN7xi&HBPVJ>zR8b(3{-o(qxsM= z+J&GiWc_|DPTImnQ(JpycbBv4T<4)v!IIS8K~S#%E^NT^sBPSj4r9o-2u;te+B7BBL>>+ zufYAy3?)09G*AE|0n9_65Ohp}5zbS_r(&78gnq-Hk+iS@Sxmg^S|%IIe~(@Xl0syQ(=6K|nWgK-$vD==7_ zLunx$m2sJAj1GTeiIBEe(Yo}PGr(*kLus}Ftwiz0sz+`mGIg&S;`+AgxVA_?2@Ks3 z3KLg7FK5WL6EZ^jWLuN*98&~#l)xGqK&xg4hEP2Tr73D?QS`t1*X@7KI;n-8WI=nS z^_vu>E9_K3*J$L7V4k3Udi(oMCR?2}kpSb3gm_{oh4F?wv!VaXLnbaO2*TY-QvhXo z^a&YK+QG^*!VHELWpdu6A2M@Bxnw945<J8^4Bp!@npK9w{k=2E9W(?F}La**dkA+zmN{%0qS37FEdxv=a|c^ zJ+D2uj zaW@!l@s6N9#XB8wGX~o*KjF9eLEt*VH`smx!a62sq3q=f zijL$2k;^085r2<83RnW{!I(Z6_DL|UYIkC850)YAOx7pWOmU*Q}!%VP$sD_SG z{RD>^HDIV37Xtpvp1`18Evg(_M0QxzjR}0e=ZYVxrOkFtlCv824De`gBJ2y+t?pyT zN}}!9GvWg#(!bAKa2zyTkT9GmZ1?J@ZHAh0_lXbiLt%n4kF( z#Q#QKIs5x*<*X!RWe{M>{LfP0Yl3s>?N=&qOC^j&HxzVsJMX0@S5`tyV zm{@ON5!(Sn!4z83r<079sJ;~pF{j6`4@50hcL!$q_lMsbp5GI7bS{}QNmfhL+)7TI zBAwj}CCjQTLpQrvQGILU&5f~&`fx>kq@sbC%P?oo-XC(L>2$2=e7Na+q^U2m?>XXh zk>_1x-vzSQO|rPCnO|JH19Z};#&!;nBZG4mGVuZ#y7U)C_G#Tx@fuQA2UaxG`o-+x z>lc4#@6F1kqS6SR=-@_*x~J0?vx}~u_}1>3*4u;Obvx#=@2}f4pSDmRIdBS81Ei#J zzG}W;A%_&TM{_zBhbF&dm|+0ypjJl@0Q)f8f?iFl6xX5(U+5t{uD5e!o|BvQh z5PjFr{m5PX`^JXWT;@-9rGxvYX%*<&Yg#S(k8>Mq;m42C8q>hLXVIaXYiQl2zPHBO zx=nr0!D4uu7Tk7_qBTKTNp=MUWN6Vn$_EC0oG=DyqWR>ZuI`iV5_nGhYHvJ#=#;2$ z>27c8YU=<%-UMj4wd3SD0T;k(M#h72=zItM{2Rb|iCNTIA8HOhQZu^DMa{v*?Df}< zkzMV#YZfhK*AA1KmYIS_syuBsD={GWvK#&qkBm+q;lBN=Fd$&oo<*+FtjDigG4t?ZPnR5iWuwrprbd z+eWwcdxe7-m&Wn^E--xrJ+&Wy8&TCK*=S0A=7m@V=R$`dSo%c$lq|~+@D-$*pt&yJ zNSe-)p$S+fA+`-!>RPsOW~!;9%90 zRQ&{8%UeMkrWYzU=%3sQ`bk^-9prq#)1!dW_#f>*F8N;lI{Ojq(siqr(B1KgelQ8ggx zBDu{JN$VB4Eg}6%BpqTCNLv4bo+4@R6-e6J9?U=g#Omp@b@^;TtHX8_Fv)re(%J#j z;yesiq{=TNW@;GY0AUXfQj3s}FO;Nu;S|Nm0ObtTHi{wuI;GX1nN5SOE4VS$EA3#v zEm?9)_S9$&4LDz99l_F6`QbDOKyf&T5-d)YEQ7KP36#AERTZwGL-7`sIiU%HBzU5- z#1t4zGKY8x@`1n+Cwg)fCk-0}0co(PmGBhn71=QXjS&2}wi5_Rf=9Rk1Z5JzC+Ka` z%9xUp5AkguEHVimSwUt$6CzuiK5iy$=VEQoh1;Hsv^^hbc99DladV`ZCr$k%V<4&@ zT(T6Bq78SOqn6t6uC`CMlN007qnCiPHL9ABz#UK>LT`cvPo+0->{7wH*`jE{o~gFQ zwKbBmf*k0KtUUn`X)&h|)GNo) z>W_<94BNEewm&KGoQ5Lkwk5oVq2EF0Tj=~QI&|Jqh>tU2HBR_8#-a1a@aI1O4#f#y zs2J^tib92D8jPcbvkA0tbf#cgWzur&ZIBM5KcD@W`qE~GMFNRUKJp&$s1xFn=mC$= z4ZxoB7TcrYL2#)MWhz9NVBVM3V}t{8Ce8pSSB>yz;?4bfYU-B`rhXY>x=eHEmKgjk zBQAr~>H%4ix*n9TnA1xudKj6Nav6&KEi>M<*;$=1M=8@uvFl5Z?S-c((o7P zv?1T>2U7JECmG%VyaWj^h~&I4vToTRXc2+$V+Gn5L1) z4{DReOt5nVjV&H*kyXBCsWYTA5+D@F>geu5xoq((@PB?OLp=n>|h{56+(>oxP;FkF0x6tlBd-u~0?!ogkGbKP#_@m2VH1Z=W;F z?;<;o6I=TiCS761vRZF7E;D+AF^RB*KftEst^`s}N&gefdJ<11+b{ePyeYg3j#~+e z5^$;qhYihgjP5^=mXG7BZm`jf&~+9>ylWR z(*hJygJsZqsQQ5<(18=DC~(HQuX>UYNe(4}SSpJwIKm%?W^xTCbxA>R@=V*L*1H>O zqNg=cW4*jv;kSYQb_8%Iq|cC54ki6mB3rq3f*OROTCY<4XOe3KAPK(n^!h7K@1-gF z+Z;+;34B28s)rtXiscu8N)V?|P-%cuX~|JCDq5G9Y?U06`%TCN;VFt?Rz0D_nKTqX zdZ9A6p^F&ZDi%U(tu=XSA#5AJZ->uzzaH zAV-~IhL?s6(DQ4cwN=SpkVBAC1}SFwchX-;#Uy2gqEwsQM=*}&ID1hczYXj~4ye-N z^a4(G7SLWQT+vG~lg=@(ydGGmc25RhU~K(GR(;0{Q23y3}vnuU$1SxC0>m$fH`nYjjND>6r` zp1A&D6X+ip)ISLAkO4r+DhW10fd}9z#`0Gn{I&FJ%T?FH_?^Zn!>MrM=Volyg&*2M4`f7{QN#}#7v2+P!Lwp zO6QLvM!1RDYcRVG&Z)z^7Y+2hAP0v*!Fm`EdBU+l2kUsDHei~&8_34pWW%1hBc$~#X*oyEc|c`FJeNtvm8gCK z3~m)-A1eJg9yPZwt*=9|QCD>R5t6fxWOosB7pPLC0Gw;FJBYq=(Nr>{yKkyms;CNA z`$@x*Si^~M!-;6asYtay8b2?35M?fx*n=pDX;K|lJ;GmY=&e1XLgib{X0oLzQqvqZ z+ZWT<%v9b_uU^Ov+gs!xOO~8hgD(eTma4F&Dr#B3P!-b`Lk(N`}PZ;Tb!hKp-IE#4{} zCItc5^-4)y6}31P>-T=B|3FU~dm{DciFNa==Dv0NV#9utRXe-xe%8+AG^S+RvXLn* zBl^PGI%3^T-KI1W-Vn>}X3L-}NHZ^YI|%%xP%*iJeINxO%5hDg?M z)I0(evX)#A-nZD7thprbT-4gLxOQ6@?0t3COpvnKw|qZ$&K$8HA~{u5X}51KH0ND7 zNIG1k{Q?2~U(rA`XK=B)K32UiT)pqp>H{}S*EQGqrQ*sN{%-fXy?1(}#SL@gvHhpR z`%g#rpIrnq-LvnVCC#U2&qlVNfwQ`|yY83Nf%UE&TsAQ48`;N3Cbx)`?3^>tz5u$n zqRwc}iRC=Twk1}!JzTc^)3TkfnWi<Qvhn zY8KXzGv`Rp^W>Zha$f+;FMNqd31M!(YO%I4R@)k`ZT+;iEnKkay5VL6scl=_d?2>D zCA_)i)6Iv&1skvHXY7l`YhfwfFWyah++^?Zg~9uK&-_m{v!#vw#c~>c?-!4YKnDrc zG*r`}ZklhL-%7Gtqvpd>Qy!SBezB~DsU3{_l$lG5D>U^AL7k{O)pfG1y*0SDsbMo->#|uETXAvgA3uC7jtW`;0N;SJMFZSN z0%ndLFAp(qU`!U=tKu7h8?>DCnzJ!iVDKrr&yv;cP)g@E0rTd^ElyH>9=`w!B5oRd zQgs=ZM~nKX8VGlb>X;gSU1g8sFGnZ-u$@G=@T6KAjW#Hz3+#yS>NU!O1oJ7Fla&>c z$tKo@HbYo|{3|q>i9Y?e(6dj9-jc_9R@!f-5SpPrGAqfrfQ4tANd*y-n$vLF0WgF` zL25{~BA-$phxyWy`JQIEnM&}+C4@14E1WUbQ_DN0g{L7|ClItG1!iXGmQ_zG-KWBC z6LhFT<1=DF6RMT$%N=bj!Rx*HED!{q?O5dITSJ`Z!v{BLo{5cWyF$S`3(^I z#lP~(%PYi9R;p3X3~52ApcEBXCrX98QRLoOF&>FKL*p|&8@0#ZQOWZdgcu+?mL3q_ zRDdmLd}V<6CI_uGm8*=dm0T%SDj{xlX?g)J(*gkr_zIB*SY56;zna>BRKs;i;&XvX zS)U?3!NH;ILY7mAs&Nn8+z6~jFgIleumpVaQ}LXgKoMV>`)Nelyw44 z6Ys?Gm0b@L;Z#Y<%;`>bk54QMla^4fabc>n_QkhF|ZHoMHdqHV`0pmd7 z1sED&2cK*i#=wOX<8IX0^pC^+5sqLQjf7q;Pw@h6Hdvct-$RzPcsBrRg_A$2!-iXK ze3y9v8s_Dw^{1%v`9%$GDqw8mfx9Y3;ou-`wO`bax`yFYzTlVGJ0x0?%^`AtFkE2b zC<&NaF$GyU;WyC16;9bZ`Iax-$LIzi8YLzd&Jzk^23(}3V5USf<#U~GFf$DJ#dK;! zC|`_2Z48Bf3FU|x(w2$ViO^oCEr$mO1c{qMa{43_^<#()7GwV@cX$|VdGFeS8{-bJ zsEI4k9b|JOaU3KWO;LR_RgrJ|J{Ps@4mb5hod8$BBDoWBvkg5qim#P`P9@QcAPJ#E(QvUXKV(*NC&iuTSTo@$7FGA8V+eizMc&76#Rlj$E zw`eJst__P=>@iDS*itvUFKRgmK%VprhMNZ=5FPSC34ZoGR%)2-16{q@aD`neEPgZp zM*f>+H_C{8ccgI7)RBjl!nZ4KRo|?><+$kpMSjJeA71!y_=Dkv4aDhwu#b!O4?{^l zwjE0vfL7wI9__QLnUeX&+2e}^b+Ll_a6$d21r1MF`$e-Ii#g>pC*a`U^zr$|dxt+d zeDC;2$Kf)s`^`ObnIGnUko&`YFs;0e82$>dZN!a=|FAo-RBD@9BPjtRr8}lu7DBNO zXSl;jp6@4vo=C?K8TE!cymRLH)A#F-EpFI2?VUMsKd%}Xmcre@uOR2Jb1r@M#rYv} zqK|YwN3xtz^YcsQ+-WbV0-j7*L=Igh`z}i= z|I4WIKfI8+peB{aqGiX)ne*g1j(En%MS+a@$v}YgkCQ&IRlNM5<1*L*T>d$=aQ&_x zqb`pZE=6tFv}qIAH&cP%3Va7*LAXK7w*DlKgoBgQc!~t=Hh9eqB_$qMVDx$90Pi;e z^vQSePQikLlTdoNkyiyLo;38g8h&f1V5aoIod9-?n10;n9&YjpQt2!3xdzHZs`Ka{ z!1+BE&?NZM^D^9t_FEwEmUn#EKQT9bQVFe7qC2RZy#GG!bO_v=i>Q`9|3|-)CVvHP z=T)`;1SvBd5!1pIy}aUbDpBS2in?RKOSs_5U23=>>gWZuwQ;}hj&uQ#XmC2IC9BgZ zN#7;oe|iQln75+e6E^_0;e^BC{puR?#-SUn;gXtYi383R4}LKCq3;9V!UeFvj_m7- z?du8e>xt~^oz;Bb1cWL!B!q{A*w93HXo3VoWb#E23L)_WtDa#n^#O(hJsR1jBF|lT z@SKOv$`vQSP=!H9{2&DIU_`0|5fF=cRy|2;LmADUVkFgsO87g-EBse*)P>0!f-B%ds<+wSW{0rZs~&k`uehWTfHl>9!S#7~nc&md88{ME zpow3!mpG0?-9>u}S#a!)U#N!LMkHy+Ah1)Mh+pk1-OoZF7 zeB__MxR5z}5%!h*5AuId_}*30e&zm-39@mcgs1P{m@-s7Yev$2CH0X+k1i%xp>Vg<3oDNU#T0SG*k*d=C+JWrHYeNV(E zzDA&9OnZhlcLa(S{vMns$pM@PeWHE|qHr9$2u0c}i9+8GoU@vLCmHgrSjLIAUc(sCTtU~#Gn#usQi)(@LW7wD0jM{ZyI*T+d|Ege6Y)J>3- zH1C(rAk(P=REZIO3J$e4kIxK%NQ{8aF)Y>q0hk#}`eXlDFBM4B+z_33ARm1m;WSF_ ze|QSGM?Il};fHg+@$Ln_29DPVE5<%S_E4On6q+bTSN;YPMR@{BU#0FdXb2p&qqpg7 zh@Ue{Y-q_MQf-m)_n)>s!fiY?KxtHyEDIXp_+2ebfvx(5@pvlE92nU!ly)V?QT-*>(FjU&W85Oa@) z-J>!0m9YCt)E$f*xEgi`$#c&6%zJqs<=tEJ(V9r(iCE*AaO0T=jpqo>@w58}u2L#< zcXMNzrXAGUspd`_R?~`320CBKP6yTm2lFN7FZ5TX9GDnPm-c6ebVo}37_{Qlx2+$p zlZI1i!LoyJi=pjQv#ol?8wvAi;rm!zu~CL@E3p|uv*7MkkWKkrA-R>zGf?Ma2E#5L#ID9l6^GHQTDYNNWLS1fG?54+B>W z|CfOzXcdikbtL6~IeL6ZOvf87oqo<47zf3n}1?w9NO zSho5L2Hh_YsB_sXkbv%T9iuT%1taRhhZ+MhwnR0pk2Pv-kiD*bQ-4GMrsam^F@wJ! z>s8vi>-?L8HwNGI-S9nz+ac<{)UnzS6hvJeQ?NdgZzr0ZhneM(%(d5!Eat3<Tec(He=dHL5whSWp=&*c~p|J=Z$lOHQ2ywN#{lBbwZY8+MbdwKI+Pv#RE^t{r`t zmNot2?Tz95?TeWu*PY?aO;fsus2JTcoA-XndnL0OWFOr23rF!h;a(3pKSVNyqxuoN z+$i*N==%7ZlQ$;sTWW~DhFIW~S?;ToFHgoSo3LNz(*3Lm);J2DtPGDzA$ff^v=SAkWy1dxNK*vY;U-1?|lp08@`uV_CB*Qe~sN<^*~d-+{ZE+={m;$ E4YN&Jxc~qF literal 0 HcmV?d00001 diff --git a/examples/workflow_quantize.json b/examples/workflow_quantize.json new file mode 100644 index 0000000..977cc1f --- /dev/null +++ b/examples/workflow_quantize.json @@ -0,0 +1,358 @@ +{ + "id": "0c605522-962a-4c82-9089-e97d4235261f", + "revision": 0, + "last_node_id": 14, + "last_link_id": 14, + "nodes": [ + { + "id": 5, + "type": "UNETLoader", + "pos": [ + 141.08615112304688, + 169.95050048828125 + ], + "size": [ + 287.13525390625, + 83.07095336914062 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 4 + ] + } + ], + "properties": { + "cnr_id": "comfy-core", + "ver": "0.3.32", + "Node name for S&R": "UNETLoader", + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65 + }, + "widgets_values": [ + "flux\\flux1-fill-dev.safetensors", + "default" + ] + }, + { + "id": 11, + "type": "QuantizeFP8Format", + "pos": [ + 466.7684020996094, + 179.52532958984375 + ], + "size": [ + 323.7391357421875, + 59.94718933105469 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "model_state_dict", + "type": "MODEL_STATE_DICT", + "link": 11 + } + ], + "outputs": [ + { + "name": "quantized_model_state_dict", + "type": "MODEL_STATE_DICT", + "links": [ + 12 + ] + } + ], + "properties": { + "Node name for S&R": "QuantizeFP8Format", + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65 + }, + "widgets_values": [ + "float8_e4m3fn" + ] + }, + { + "id": 2, + "type": "ModelToStateDict", + "pos": [ + 142.7964630126953, + 307.8086853027344 + ], + "size": [ + 291.43902587890625, + 79.07095336914062 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 4 + } + ], + "outputs": [ + { + "name": "model_state_dict", + "type": "MODEL_STATE_DICT", + "slot_index": 0, + "links": [ + 11, + 13 + ] + } + ], + "properties": { + "Node name for S&R": "ModelToStateDict", + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65 + }, + "widgets_values": [] + }, + { + "id": 13, + "type": "QuantizeModel", + "pos": [ + 119.18730163574219, + 500.8162536621094 + ], + "size": [ + 345.158203125, + 106 + ], + "flags": {}, + "order": 4, + "mode": 4, + "inputs": [ + { + "name": "model_state_dict", + "type": "MODEL_STATE_DICT", + "link": 13 + } + ], + "outputs": [ + { + "name": "quantized_model_state_dict", + "type": "MODEL_STATE_DICT", + "links": [ + 14 + ] + } + ], + "properties": { + "Node name for S&R": "QuantizeModel", + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65 + }, + "widgets_values": [ + "per_tensor", + "Auto", + "float16" + ] + }, + { + "id": 4, + "type": "SaveAsSafeTensor", + "pos": [ + 471.4664611816406, + 311.7754211425781 + ], + "size": [ + 318.212890625, + 69.432373046875 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "quantized_model_state_dict", + "type": "MODEL_STATE_DICT", + "link": 12 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "SaveAsSafeTensor", + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65 + }, + "widgets_values": [ + "C:\\Users\\RAIIN Studios\\Documents\\protable\\ComfyUI\\models\\fluxfill-fp8e4m3fn.safetensors" + ] + }, + { + "id": 14, + "type": "SaveAsSafeTensor", + "pos": [ + 490.1119689941406, + 509.20977783203125 + ], + "size": [ + 363.431396484375, + 83.39691162109375 + ], + "flags": {}, + "order": 6, + "mode": 4, + "inputs": [ + { + "name": "quantized_model_state_dict", + "type": "MODEL_STATE_DICT", + "link": 14 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "SaveAsSafeTensor", + "enableTabs": false, + "tabWidth": 65, + "tabXOffset": 10, + "hasSecondTab": false, + "secondTabText": "Send Back", + "secondTabOffset": 80, + "secondTabWidth": 65 + }, + "widgets_values": [ + "C:\\Users\\RAIIN Studios\\Documents\\protable\\ComfyUI\\models\\fluxfill-fp8e4m3fn.safetensors" + ] + }, + { + "id": 9, + "type": "Note", + "pos": [ + 911.549560546875, + 467.0907897949219 + ], + "size": [ + 371.62591552734375, + 149.91725158691406 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": {}, + "widgets_values": [ + "Available Scaling Strategies:\n\n- per_tensor : Uses a single scale factor for the entire tensor (fast, less precise).\n\n- per_channel : Computes a separate scale for each output channel (more accurate)." + ], + "color": "#432", + "bgcolor": "#653" + } + ], + "links": [ + [ + 4, + 5, + 0, + 2, + 0, + "MODEL" + ], + [ + 11, + 2, + 0, + 11, + 0, + "MODEL_STATE_DICT" + ], + [ + 12, + 11, + 0, + 4, + 0, + "MODEL_STATE_DICT" + ], + [ + 13, + 2, + 0, + 13, + 0, + "MODEL_STATE_DICT" + ], + [ + 14, + 13, + 0, + 14, + 0, + "MODEL_STATE_DICT" + ] + ], + "groups": [ + { + "id": 1, + "title": "fp8", + "bounding": [ + 131.08615112304688, + 96.35050201416016, + 680.042724609375, + 311.2386169433594 + ], + "color": "#3f789e", + "font_size": 24, + "flags": {} + }, + { + "id": 2, + "title": "fp/bf16", + "bounding": [ + 109.18731689453125, + 427.21624755859375, + 775.532958984375, + 198.77328491210938 + ], + "color": "#3f789e", + "font_size": 24, + "flags": {} + } + ], + "config": {}, + "extra": { + "frontendVersion": "1.18.9", + "ue_links": [], + "VHS_latentpreview": false, + "VHS_latentpreviewrate": 0, + "VHS_MetadataImage": true, + "VHS_KeepIntermediate": true + }, + "version": 0.4 +} \ No newline at end of file diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..75323c6 --- /dev/null +++ b/nodes.py @@ -0,0 +1,371 @@ +# nodes.py +# This file contains the implementation of your custom nodes. + +import torch +from safetensors.torch import save_file +from tqdm import tqdm +import os + +# --- Helper Classes (Using original simpler scaling logic as provided by user) --- + +class FP8Quantizer: + """A class to apply FP8 quantization to a state_dict.""" + + def __init__(self, quant_dtype: str = "float8_e5m2"): + if not hasattr(torch, quant_dtype): + raise ValueError(f"Unsupported quant_dtype: {quant_dtype}. PyTorch does not have this attribute.") + self.quant_dtype = quant_dtype + self.scale_factors = {} # Not used in current quantize_weights but kept for potential future use + + def quantize_weights(self, weight: torch.Tensor, layer_name: str) -> torch.Tensor: + """Quantizes a weight tensor to the specified FP8 format using simple scaling.""" + if not weight.is_floating_point(): + return weight # Only quantize floating point tensors + + original_device = weight.device + # Ensure FP8 conversion happens on CUDA if possible, as FP8 types require it + can_use_cuda = torch.cuda.is_available() + target_device = torch.device("cuda") if can_use_cuda else torch.device("cpu") + + # Warn if trying FP8 on CPU without CUDA + if not can_use_cuda and "float8" in self.quant_dtype: + print(f"[FP8Quantizer] Warning: CUDA not available. True {self.quant_dtype} conversion requires CUDA. Attempting on CPU, but results may be unexpected or errors may occur.") + target_device = torch.device("cpu") # Try on CPU despite warning + + weight_on_target = weight.to(target_device) + + max_val = torch.max(torch.abs(weight_on_target)) + if max_val == 0: + # For zero tensor, just cast to target dtype on the target device + target_torch_dtype = getattr(torch, self.quant_dtype) + return torch.zeros_like(weight_on_target, dtype=target_torch_dtype) + else: + # Using the simple scaling from user's provided script + scale = max_val / 127.0 + # Clamp scale to avoid division by zero if max_val is extremely small + scale = torch.max(scale, torch.tensor(1e-12, device=target_device, dtype=weight_on_target.dtype)) + + # Quantize: scale, round, unscale (simulated int8 range mapping) + quantized_weight_simulated = torch.round(weight_on_target / scale * 127.0) / 127.0 * scale + + # Final cast to the target FP8 dtype + target_torch_dtype = getattr(torch, self.quant_dtype) + quantized_weight = quantized_weight_simulated.to(dtype=target_torch_dtype) + + # Return on the device where conversion happened (target_device, ideally CUDA) + return quantized_weight + + def apply_quantization(self, state_dict: dict) -> dict: + """Applies direct FP8 quantization to all applicable weights.""" + quantized_state_dict = {} + eligible_tensors = {name: param for name, param in state_dict.items() if isinstance(param, torch.Tensor) and param.is_floating_point()} + progress_bar = tqdm(eligible_tensors.items(), desc=f"Quantizing to {self.quant_dtype}", unit="tensor", leave=False) + + for name, param in progress_bar: + # quantize_weights handles device logic now + quantized_state_dict[name] = self.quantize_weights(param.clone(), name) + + for name, param in state_dict.items(): + if name not in quantized_state_dict: + quantized_state_dict[name] = param + return quantized_state_dict + +class FP8ScaledQuantizer: + """ + Simulated FP8 quantizer using 8-bit scaled float approximation based on user provided script. + Operations are performed on the device of the input tensors. + (Used internally by QuantizeModel node for the value simulation step). + """ + def __init__(self, scaling_strategy: str = "per_tensor"): + self.scaling_strategy = scaling_strategy + self.scale_factors = {} # Stores the calculated scales (Python floats or lists) + + def _quantize_fp8_simulated(self, tensor: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + """Simulate quantization by scaling, clamping to 8-bit range, and dequantizing.""" + # Ensure scale is a tensor on the correct device and dtype + scale = scale.to(device=tensor.device, dtype=tensor.dtype) + # Prevent division by zero + scale = torch.where(scale == 0, torch.tensor(1e-9, device=tensor.device, dtype=tensor.dtype), scale) + + # Perform simulation: scale, round, clamp, unscale + quantized_intermediate = tensor / scale * 127.0 + quantized = torch.round(quantized_intermediate).clamp_(-127.0, 127.0) + dequantized = quantized / 127.0 * scale + return dequantized + + def quantize_weights(self, weight: torch.Tensor, layer_name: str) -> torch.Tensor: + """Applies the simulated quantization based on the chosen strategy.""" + if not isinstance(weight, torch.Tensor) or not weight.is_floating_point(): + return weight # Skip non-float tensors + + current_device = weight.device + + if self.scaling_strategy == "per_tensor": + scale_val = torch.max(torch.abs(weight)) + # Ensure scale_val is a tensor for consistent handling in _quantize_fp8_simulated + scale = scale_val if scale_val != 0 else torch.tensor(1.0, device=current_device, dtype=weight.dtype) + self.scale_factors[layer_name] = scale.item() # Store the scale value + quantized_weight = self._quantize_fp8_simulated(weight, scale) + + elif self.scaling_strategy == "per_channel": + if weight.ndim < 2: # Fallback to per-tensor for 1D tensors + scale_val = torch.max(torch.abs(weight)) + scale = scale_val if scale_val != 0 else torch.tensor(1.0, device=current_device, dtype=weight.dtype) + self.scale_factors[layer_name] = scale.item() + quantized_weight = self._quantize_fp8_simulated(weight, scale) + else: + # Assume channel dimension is 0 for typical Conv layers + # For Linear layers (e.g., [out_features, in_features]), dim 0 is also common. + # If weights are [in, out], dim 1 might be needed. Defaulting to dim 0. + channel_dim = 0 + dims_to_reduce = [d for d in range(weight.ndim) if d != channel_dim] + if not dims_to_reduce: # Handle edge case if channel_dim is the only dim somehow + scale_val = torch.max(torch.abs(weight)) + scale = scale_val if scale_val != 0 else torch.tensor(1.0, device=current_device, dtype=weight.dtype) + self.scale_factors[layer_name] = scale.item() + else: + scale = torch.amax(torch.abs(weight), dim=dims_to_reduce, keepdim=True) + # Store scales as a list of floats + self.scale_factors[layer_name] = scale.squeeze().tolist() + + quantized_weight = self._quantize_fp8_simulated(weight, scale) + else: + raise ValueError(f"Unknown scaling strategy: {self.scaling_strategy}") + + # The output tensor retains the original dtype but has modified values + return quantized_weight + + def apply_quantization(self, state_dict: dict) -> dict: + """Applies simulated FP8 quantization to all applicable weights.""" + quantized_state_dict = {} + # Process only floating point tensors + eligible_tensors = {name: param for name, param in state_dict.items() if isinstance(param, torch.Tensor) and param.is_floating_point()} + progress_bar = tqdm(eligible_tensors.items(), desc=f"Applying scaled ({self.scaling_strategy}) quantization", unit="tensor", leave=False) + + for name, param in progress_bar: + # Pass the clone to avoid modifying the original dict if errors occur mid-way + quantized_state_dict[name] = self.quantize_weights(param.clone(), name) + + # Add back non-floating point tensors and non-tensor data + for name, param in state_dict.items(): + if name not in quantized_state_dict: + quantized_state_dict[name] = param + return quantized_state_dict + +# --- ComfyUI Nodes --- + +class ModelToStateDict: + @classmethod + def INPUT_TYPES(s): return {"required": {"model": ("MODEL",)}} + RETURN_TYPES = ("MODEL_STATE_DICT",); RETURN_NAMES = ("model_state_dict",) + FUNCTION = "get_state_dict"; CATEGORY = "Model Quantization/Utils" + def get_state_dict(self, model): + print("[ModelToStateDict] Attempting to extract state_dict...") + if not hasattr(model, 'model') or not hasattr(model.model, 'state_dict'): + print("[ModelToStateDict] Error: Invalid MODEL structure."); return ({},) + try: + original_state_dict = model.model.state_dict() + print(f"[ModelToStateDict] Original keys sample: {list(original_state_dict.keys())[:5]}") + state_dict_to_return = original_state_dict; prefixes_to_try = ["diffusion_model.", "model."]; prefix_found = False + for prefix in prefixes_to_try: + num_keys = len(original_state_dict); + if num_keys == 0: break + matches = sum(1 for k in original_state_dict if k.startswith(prefix)) + if matches > 0 and (matches / num_keys > 0.5 or matches == num_keys): + print(f"[ModelToStateDict] Stripping prefix '{prefix}'...") + state_dict_to_return = {k[len(prefix):] if k.startswith(prefix) else k: v for k, v in original_state_dict.items()} + print(f"[ModelToStateDict] New keys sample: {list(state_dict_to_return.keys())[:5]}"); prefix_found = True; break + if not prefix_found: print("[ModelToStateDict] No common prefixes stripped.") + dtypes = {}; total = 0 + for k, v in state_dict_to_return.items(): + if isinstance(v, torch.Tensor): total += 1; dt = str(v.dtype); dtypes[dt] = dtypes.get(dt, 0) + 1 + print(f"[ModelToStateDict] DEBUG: Output Tensors: {total}, Dtypes: {dtypes}") + return (state_dict_to_return,) + except Exception as e: print(f"[ModelToStateDict] Error: {e}"); return ({},) + +class QuantizeFP8Format: # Direct FP8 conversion node + @classmethod + def INPUT_TYPES(s): return { "required": { "model_state_dict": ("MODEL_STATE_DICT",), "fp8_format": (["float8_e4m3fn", "float8_e5m2"], {"default": "float8_e5m2"}), } } + RETURN_TYPES = ("MODEL_STATE_DICT",); RETURN_NAMES = ("quantized_model_state_dict",) + FUNCTION = "quantize_model"; CATEGORY = "Model Quantization/FP8 Direct" + def quantize_model(self, model_state_dict: dict, fp8_format: str): + print(f"[QuantizeFP8Format] To {fp8_format}. Keys(sample): {list(model_state_dict.keys())[:3]}") + if not isinstance(model_state_dict, dict) or not model_state_dict: print("[QuantizeFP8Format] Invalid input."); return ({},) + try: + quantizer = FP8Quantizer(quant_dtype=fp8_format) # Uses helper class with simple scaling + quantized_state_dict = quantizer.apply_quantization(model_state_dict) + found = False; + for n, p in quantized_state_dict.items(): + if isinstance(p, torch.Tensor) and "float8" in str(p.dtype): print(f"[QuantizeFP8Format] Sample '{n}' dtype: {p.dtype}, dev: {p.device}"); found=True; break + if not found: print(f"[QuantizeFP8Format] No tensor converted to {fp8_format}.") + print("[QuantizeFP8Format] Complete."); return (quantized_state_dict,) + except Exception as e: print(f"[QuantizeFP8Format] Error: {e}"); return (model_state_dict,) + +class QuantizeModel: # <<< RENAMED CLASS from QuantizeScaled + """ + Applies simulated FP8 scaling (per-tensor/per-channel) and casts + to a specified output dtype (float16, bfloat16, or Original). + """ + # Removed "FP8" from the list as requested + OUTPUT_DTYPES_LIST = ["Original", "float16", "bfloat16"] + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "model_state_dict": ("MODEL_STATE_DICT",), + "scaling_strategy": (["per_tensor", "per_channel"], {"default": "per_tensor"}), + "processing_device": (["Auto", "CPU", "GPU"], {"default": "Auto"}), + # Default to float16 for size reduction, user can choose others + "output_dtype": (s.OUTPUT_DTYPES_LIST, {"default": "float16"}), + } + } + + RETURN_TYPES = ("MODEL_STATE_DICT",) + RETURN_NAMES = ("quantized_model_state_dict",) + FUNCTION = "quantize_model_scaled" # Internal function name can stay + CATEGORY = "Model Quantization" # More general category + + def quantize_model_scaled(self, model_state_dict: dict, scaling_strategy: str, processing_device: str, output_dtype: str): + # Log using the new class name + print(f"[QuantizeModel] Strategy: {scaling_strategy}, Device: {processing_device}, Output Dtype: {output_dtype}") + + if not isinstance(model_state_dict, dict) or not model_state_dict: + print("[QuantizeModel] Error: Input model_state_dict is invalid."); + return (model_state_dict if isinstance(model_state_dict, dict) else {},) + + # Determine processing device + current_processing_device_str = "cpu" + if processing_device == "Auto": + first_tensor_device = next((p.device for p in model_state_dict.values() if isinstance(p, torch.Tensor)), torch.device("cpu")) + current_processing_device_str = str(first_tensor_device) + elif processing_device == "CPU": current_processing_device_str = "cpu" + elif processing_device == "GPU": + if torch.cuda.is_available(): current_processing_device_str = "cuda" + else: print("[QuantizeModel] Warning: GPU selected, CUDA unavailable. Defaulting to CPU."); current_processing_device_str = "cpu" + current_processing_device = torch.device(current_processing_device_str) + print(f"[QuantizeModel] Value scaling simulation target device: {current_processing_device}") + + # Move input state_dict to the processing device + state_dict_on_processing_device = {} + for name, param in model_state_dict.items(): + if isinstance(param, torch.Tensor): + state_dict_on_processing_device[name] = param.to(current_processing_device) + else: state_dict_on_processing_device[name] = param + + scaled_state_dict = {}; final_state_dict = {} + + try: + # Perform FP8 value simulation using FP8ScaledQuantizer helper (simple scaling version) + quantizer = FP8ScaledQuantizer(scaling_strategy=scaling_strategy) + scaled_state_dict = quantizer.apply_quantization(state_dict_on_processing_device) + print(f"[QuantizeModel] FP8 value scaling simulation performed on {current_processing_device}.") + + # Cast to final output_dtype (Original, float16, bfloat16) + if output_dtype == "Original": + print("[QuantizeModel] Output Dtype: Original. No further dtype casting.") + final_state_dict = scaled_state_dict + else: + # output_dtype is guaranteed to be 'float16' or 'bfloat16' + try: + target_torch_dtype = getattr(torch, output_dtype) + print(f"[QuantizeModel] Casting output to {output_dtype} ({target_torch_dtype})...") + for name, param in scaled_state_dict.items(): + if isinstance(param, torch.Tensor) and param.is_floating_point(): + final_state_dict[name] = param.to(dtype=target_torch_dtype) + else: + final_state_dict[name] = param # Pass non-float tensors or non-tensors + print(f"[QuantizeModel] Casting to {output_dtype} complete.") + except AttributeError: # Should not happen with the restricted list + print(f"[QuantizeModel] Error: Invalid torch dtype '{output_dtype}'. Using scaled tensors without final casting.") + final_state_dict = scaled_state_dict + except Exception as e_cast: + print(f"[QuantizeModel] Error during casting loop to {output_dtype}: {e_cast}. Using scaled tensors without final casting for affected tensors.") + for name_done, param_done in final_state_dict.items(): pass + for name_rem, param_rem in scaled_state_dict.items(): + if name_rem not in final_state_dict: final_state_dict[name_rem] = param_rem + + # Verification log + for name, param in final_state_dict.items(): + if isinstance(param, torch.Tensor) and param.is_floating_point(): + print(f"[QuantizeModel] Sample output tensor '{name}' final dtype: {param.dtype}, device: {param.device}") + break + print(f"[QuantizeModel] Processing complete.") + return (final_state_dict,) + except Exception as e: + print(f"[QuantizeModel] Major error during processing: {e}") + return (model_state_dict,) + + +class SaveAsSafeTensor: # No changes needed + @classmethod + def INPUT_TYPES(s): return { "required": { "quantized_model_state_dict": ("MODEL_STATE_DICT",), "absolute_save_path": ("STRING", {"default": "C:/temp/quantized_model.safetensors", "multiline": False}), } } + RETURN_TYPES = () ; OUTPUT_NODE = True ; FUNCTION = "save_model"; CATEGORY = "Model Quantization/Save" + def save_model(self, quantized_model_state_dict: dict, absolute_save_path: str): + print(f"[SaveAsSafeTensor] Saving to: {absolute_save_path}") + if not isinstance(quantized_model_state_dict, dict) or not quantized_model_state_dict: print("[SaveAsSafeTensor] Error: Input invalid."); return {"ui": {"text": ["Error: Input invalid."]}} + if not absolute_save_path: print("[SaveAsSafeTensor] Error: Path empty."); return {"ui": {"text": ["Error: Path empty."]}} + if not absolute_save_path.lower().endswith(".safetensors"): absolute_save_path += ".safetensors"; print(f"[SaveAsSafeTensor] Appended .safetensors") + try: + output_dir = os.path.dirname(absolute_save_path); + if output_dir and not os.path.exists(output_dir): os.makedirs(output_dir, exist_ok=True); print(f"[SaveAsSafeTensor] Created dir.") + cpu_state_dict = {}; dtype_counts = {}; total_tensors = 0 + for k, v in quantized_model_state_dict.items(): + if isinstance(v, torch.Tensor): + total_tensors += 1; tensor_to_save = v.cpu() if v.device.type != 'cpu' else v; cpu_state_dict[k] = tensor_to_save + dt_str = str(tensor_to_save.dtype); dtype_counts[dt_str] = dtype_counts.get(dt_str, 0) + 1 + else: cpu_state_dict[k] = v + print(f"[SaveAsSafeTensor] DEBUG: Tensors: {total_tensors}, Dtypes: {dtype_counts}") + save_file(cpu_state_dict, absolute_save_path) + print(f"[SaveAsSafeTensor] Saved successfully."); return {"ui": {"text": [f"Saved: {absolute_save_path}"]}} + except Exception as e: print(f"[SaveAsSafeTensor] Error saving: {e}"); return {"ui": {"text": [f"Error: {e}"]}} + +# --- Main (for testing outside ComfyUI, not strictly necessary for the plugin) --- +# (Test block remains the same as previous version, it already tests QuantizeModel) +if __name__ == '__main__': + print("--- Testing Quantization Nodes (Renamed QuantizeModel) ---") + + class MockCoreModel(torch.nn.Module): + def __init__(self): super().__init__(); self.layer1 = torch.nn.Linear(10,10).float(); self.layer2=torch.nn.Linear(10,10).float() + def forward(self,x): return self.layer2(self.layer1(x)) + def state_dict(self, *args, **kwargs): return {k:v.clone() for k,v in super().state_dict(*args,**kwargs).items()} + class MockModelWithPrefix(torch.nn.Module): + def __init__(self): super().__init__(); self.model = MockCoreModel() # Use "model." prefix + def forward(self,x): return self.model(x) + class MockModelPatcher: + def __init__(self): self.model = MockModelWithPrefix(); [p.data.normal_().float() for p in self.model.parameters()] + + mock_comfy_model = MockModelPatcher() + node_to_sd = ModelToStateDict() + base_sd_tuple = node_to_sd.get_state_dict(mock_comfy_model) + base_sd = base_sd_tuple[0] if base_sd_tuple else {} + if not base_sd or 'layer1.weight' not in base_sd: print("ModelToStateDict failed."); exit() # Check unprefixed key + print(f"Base SD 'layer1.weight' dtype: {base_sd['layer1.weight'].dtype}, device: {base_sd['layer1.weight'].device}") + + # Test the renamed node + node_quantize = QuantizeModel() + print("\n--- Test QuantizeModel ---") + + # Test Case 1: Output float16 + print("Testing QuantizeModel: Output Dtype = float16, Device = CPU") + result_fp16_tuple = node_quantize.quantize_model_scaled(base_sd, "per_tensor", "CPU", "float16") + result_fp16 = result_fp16_tuple[0] if result_fp16_tuple else {} + if result_fp16 and 'layer1.weight' in result_fp16: + tensor = result_fp16['layer1.weight'] + print(f" Output 'layer1.weight' dtype: {tensor.dtype} (Expected float16), device: {tensor.device}") + assert tensor.dtype == torch.float16 + assert tensor.device.type == 'cpu' + else: print(" Test Case 1 Failed.") + + # Test Case 2: Output Original + print("\nTesting QuantizeModel: Output Dtype = Original, Device = CPU") + result_orig_tuple = node_quantize.quantize_model_scaled(base_sd, "per_tensor", "CPU", "Original") + result_orig = result_orig_tuple[0] if result_orig_tuple else {} + if result_orig and 'layer1.weight' in result_orig: + tensor = result_orig['layer1.weight'] + print(f" Output 'layer1.weight' dtype: {tensor.dtype} (Expected {base_sd['layer1.weight'].dtype}), device: {tensor.device}") + assert tensor.dtype == base_sd['layer1.weight'].dtype # Should match original + assert tensor.device.type == 'cpu' + else: print(" Test Case 2 Failed.") + + print("\n--- Testing Complete ---") \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..8a3cd47 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,3 @@ +torch>=2.0.0 +safetensors>=0.3.1 +tqdm>=4.65.0