From fb34e91f0ef4ec8da8457299de4f822eaecfa1f3 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Wed, 18 Dec 2024 10:29:56 +0200
Subject: [PATCH] Support official scaled fp8 model

---
 hyvideo/modules/fp8_map.safetensors | Bin 0 -> 28320 bytes
 hyvideo/modules/fp8_optimization.py | 106 ++++++++++++++++++++++++++++
 nodes.py                            |   7 +-
 3 files changed, 111 insertions(+), 2 deletions(-)
 create mode 100644 hyvideo/modules/fp8_map.safetensors
 create mode 100644 hyvideo/modules/fp8_optimization.py

diff --git a/hyvideo/modules/fp8_map.safetensors b/hyvideo/modules/fp8_map.safetensors
new file mode 100644
index 0000000000000000000000000000000000000000..602950ccca4057aedd338adb1cd63fedb4a2c448
GIT binary patch
literal 28320
zcmbuHZ-{199mglpFw;WQ7x@RaaVIoy%iKTDeeRtxNC@dc3bBYxWSez&ZF6U4+a044
z-<X6_3Qi^ZqA&tUi|CK+x<nu;Q7naMNHHj48Q6=+r8gBqpL6fIvuB?9j_^B2xIFi{
z_uS7tzvp|;b3gms@jP;9ru*-%*}Z$053DWiUs_w*yLWcsuGxdDN0*n@W*26Utj*q?
z&F){lbKl_ud-ok)z3q-;d*=5XTE2bn(%Ra}-Xlj>zv$ZUU;FBj12n(BEM~LWV_#Tu
z?{3L&<xt05J$Ufgfwg07n$N0R-wybtJHBk<RavtWZ~5?%JqK^gO}Hk@i?`<$CSG1-
zdA%cWb^o5jhgJ?O9kns*te98KwXdv2rc%sjh0J9WY0agmvO*@~h%w0&O;*Zi&H!!E
zltos`XgXkAG-aJtBAPs5z*LMF5lw)!MpIQ;C8FtxG0{{_HWJZ<KwC7UA{&WlJTNYr
zQJvK?n$1hx{9RVhXSIxG6KRd6uCiK2bLJwqeV^4$HkQ$x0otM&7ui@w(*fh68P{2(
zqDi^WH1k=bvI&^ha2oD4Dx4k}GjN(FYgIZS(H2izWUY$F6XWt}>nxww0c~FC=D`CI
zs*Y%rYK;gWP#w~lE8aGQ;AOe8f_#3@8Kx~JcrbNP9WyQ{kob#AD!5(Plfn*ir3ZOH
zH7A82<jN2{G-f!#%W@?NTbRbA(1cuh!oZA73RB3XD&*Zo6(`pap{g+zP|Zo<3%L}A
z9vYPt%8*NE2${yDu!dY}gJ(u2g*fEW9Et?1g9ue66;RDdp%1zAhaMW06b6w?g$S9(
zq>zYQT7+jtCWS}jQY4C$F@;IwQY8YaIVoHsmom{qqmsfVa;Xy`)0h-KkxQZQ%*doL
zib5JinJ}gxLRDibpjwjx2viwG4~<C*URFq>2${B|;K5Wz;hAwsfy@eN6lKDgf(TV7
z743cif$F5rc$&@q;AMq0ioX0lePdD>MInvCGb5A2C<<v5RmzycC<<v50o6RFFp5GN
zMGuV{Qy4`djUr?klfozpX%wCrnG{A*NTaAy#uP?TNTUd-=A<x+LK;O6jY<lmD5Oz@
zOk+|QMInvCGb5A2C<<v5qm(g)Q54cB0;)MFjG~Z6(L<w>!YB%96d}`?6h=`<qwvhg
zq%ew78pSAKOhJUI?o$EPniN2w>OR#&W0HcGmC`6erY$LWFqKhwW?WJr`5|BZMo}k>
zDTq*2QUTSR6h=`>qv)YgNxA13(kMcvF)8;bLmGu=Mka+(l+q~blre=-l+q{ysyQi)
zqLfC_L!*+yC`xG*A=8)?Mo~(m@XW}hFp5$d#W-b5VHBk_ihycP3Zp2cQS{KLq%ew7
z8b!!7CWTRy(kMJLGAWFrltwX58B-WVDUBkanv=pPN@)~5G%6{KqLfAvGL1=L6s0r@
z&x}k8qo|}&Gznu0B2;ys3aHkk00LD;(L-aBf|pg&C_<(!DR?lIQFvxtQXsQR8by;Z
zrXWJqNk#ug0T8H8D)u)Dcv&TlqA$Ns-<T9eQAwll%*dqN-$zKJXj8^i_a_o6seo!8
zQ|_-Mq*3(Hs4<06RMIFyrZFjuqLN18nUP6h6qPiJHf2m<6qPiJfND+(qo|}&^w6lJ
zFp5eVMaVQJg;7+}C_FPVDU700MiD$nGC%SK2viKI0BX(&pQw~igb$KzKhJ{8Dy0)4
z(wGu9Q7N17$jFe8iII|tB4IQE1S&&{zCpo*DQ-}GM`OkkxNM|MB1GCk0tZrvgh$4O
z1d+dE6TeIZkB`hF2_R4rQUJA%BzQ1INa5on+eiYJjg&-$NMlGi#7H@WM@EK(LX4C`
zq&$xT2vmj?KavoLkrIgT$&oQ6{9&a0Aw(KOLLWv-ACjI$LF5lO#YhSs8req@@-R~J
z5J1f%33nJNcL*OE8AC!HMoJw*q%kDSVWiB#BO^x=;xJO;P^H|H5QmWxhX86032_)H
zap;**_aww&q{Jaa8bd-HMoJt!GBPB@p;qD$JSsAeB!EE0*IfX$h6E3$xF>~=ifkhZ
zTvjV_2$8msz=0Iv;E{15LCk6;4#86*dq{{wt;8XKnnOYyY9$WgQzBzXh(oQ!Aw(KO
zLL6!(4jvg965>!RaR?p|*+W7cY9$T<)EpAxP%Ci=9}pQsLL6!(4k6MQ65>!Raq!5<
zkPwGji9?-mPXY*3eBA|5b4Z9ot;C^c#@v&@WwjEA5NQkvaj2Czcw}Tqh(oQ!A$Tlg
zA4!Nqt;8XKnnOYyY9$WgV<F>6LL6!(4k6MQ65>!Raq!5<kPwHl5{KZ4kU1oPK*dN3
zpw^J!!4%>UJ`u8o1TGsZaR`yNkidZy;^2{SAwlGS&l3Mm3LXa8LqZ(JN*n^HIV8kk
zti&OF7-S3yaTqIc2$9B+5Qni62ak*l32_)Jap)g~#QmQd0D+2;6hO@(Ar50D4m~sG
zo&+u%D{%;s#*h$)u@VQ5j0_2J7%OoI9s$`$65=pc;t)X1At4T9B@W>uAmd0v9L7o<
zLZmSy#9^$&!6PF>LLA0Q9D=7m_K*;Vu@Z*>Y7Pl;7%OoIpZ*v_LLA0Q973cqB*bB?
z#K9vYLqZ&yUL1}cTDd)(^`W1uz*aXXaNQT3#@T8H1+KG&Y25Pa{yPsZtsPojiJ4RZ
zXx8-3uoawpnz3wrq3L}=<BVk^3QaEx)40U4k%gw0g(>iST-}I5(~H6mJY(TlLepEq
zHeO=o$U@W0!W20&aWtXnHDL#yv2qNd=?!5UFR^kwq3Jzgid@8HjwCd_B<#R5R*oSw
zy&-JlC033lG`%HEk;}Ns5rn1}gdKRs%8`SnmxFD*#L97mrgwxXauruOg3$DWumjIn
zIcm`KYOsx$SUG;s^nNf!9>rCT8#KKe?7%Ztju<q(7;NJuR*oMuy&p`G>$u7>gI1YA
zKJaXn!-6N&0A5<<P$3IT$j4EY!^&DI2LsPoId;%WI~b={jwrNJ6q>lo@q<?S!N4<C
zjvutr51Q1<F@;u|LK{~(e$Yxk7<k6Y@q<?SL7Q4RrqC)=h`XJG0w)H~fHT&O9kj|0
ziqx{vgjQ)n5vQ)Om^-KYF>dSE<tEQqIC9V`Ie?e=^@b+2N)w7Wb%m9+N)9%8#>$a{
zR>=Xp#LCfxR%t>pT%5Llron<IZb6$oW97&}tK<M)V&!N;t2CjAQ&(8<x*xQsuJ(av
ztQ<LDN+&gZ%EZzg*0EJgarWw9j;Oxm2BNKe2ZN{_p{SDUM;BnLu;Lun!5}J6C<dan
zenbJH%4n2aKhgkOr8StRZP+^)L}d%bK(yA6E<jWnrP{2+!dC834E|}YA7y~3(i(`?
z`mqLxDx*}Vby(QaAqESz){i?tRA~)FYyH?mJ}>nl)odNsu~k}wi(2c)ARwx=2BNio
z93r2W5|Qe+4(r$=t>Io}Egy|wT0e2K8i3a7@d$t-q4b*{N&#9#RmOR)gF#d-u}QR+
zk4!*R0j0XG!!lK(5{oj<b{)(SRagVjSU)gFR2iijufxJtVo}C9uY)<FN^2lm>vuZg
zC!XADlk_podmYxXRa(PfHI9B@j;PWah}Qa@PWYK8r|M50|6m<kr8TIkwSK1)e)7qw
z-VQ`-{Z1$R#FN{Isrfsx@bgbju?JPP*6(z}&px>W(OSRL2|xGbHezc1PAvQclvC_M
zRju_qo$!-S?m)EG?{vb?Jh_dSTE7zuKLh0ydr(zt{Z1$R{F6Hnt@S&f@KaB2Bc|5x
z)WT0dImI4S)mpzZ3P1bg4n%AH&M5rcliP@?^*gulb5KsP2UWG!?~KAvKe+?ZSihTj
z5~8|CCCjgyfD*R2TMb{<#`4|llewRMa%=em&|1Co2|xGbG%zuIH~nPp=bzk)j{NR-
zGfx5(New`2@lGcEw3E|4OfBA7g`a(LD?0LvAC?0Ol^B55;+;wONhha!m|DD33P1DY
zR&?aIJ}d_m<7fa{i+38~r<|PbVQTTtC;ZHlThWnU`LF=0Und5jwRq<ce!9u&9;O!W
zWWrB6xfLDxeGdzu`e|YST8noA;isFN?qO>2&Lp}iC-2_1YuDUl-@?qy(+iXJds)t3
zoSU4r=-#jIKDn^6@i_ZlW&3lyKC-ZJ{AIR3!S?SjPS$rXo>=}N>&MytU6!}HwksQd
z?AkldWPJ_ZH@bFnlYcCnSUyGjJ#7C8=Xjp?zvTVjIre+BJG!`W{2>@_`ze+y7uR<`
zz;be7ZgL*`-s1Hpj=i_*pX=Ak#!uPy3zi4F{VN+Uv3@E0Kfk!Lagc3y)8D_qe2hNr
zX8SX2I|t^QobwXa-LeexJG6VF>&uDdXK8;QeRh2~#riyLF5uW3yLOZHzj5r#?0YwT
zo@Kd+eKWkC=G?C4Pj~gY&up?j!?_-%?KOV?dF-EIxt9K2-_^FV@z1V*8^`bG_*u5S
zgY&)2`e}~44Aw2QeGt}DY<tq5`%|3rdHVc&`utn=eT-wS;Mkw@{$aLXO}huW+TAfH
zmfs8K*{*+clMi*ru57%*{zv_J{>pLpbZoMv{gw1FgL#0yJ<q<ESX^#rbL@}U|3A+8
zZ;rp2_YKQ^w0nm455f3Dcg%_9Z^EACyc=wr<Mn;C`xg7or=Km$b-e#L%VQjOk~Xii
zTtU0rIPY@TSFQ`r`wHs?=YFp1tGhPcbuc%%hvjo!JNMc2pLYh=T`#A*_a~Nb;Qjry
z`&!3$*V<(La`>O<@|c<VGkllXwQPHUZ8y=se^KABSYCv)kI()k`uF~>_PNPx{`hm~
z=eODZ1KPa6F~5WVYt}D?@f+G)%yAF1{z7+N_gOm5WWC}z*Qb*l>*~LX^(Sca5nfyR
z@H(tl+1GOHf7w2weaW$pbL?~U=^EPQEZ4E$jnkiT&WqUp!R~W(*J}Tnf1UPcvwj-h
z4EsLHdg0DZ`zP79i}e?~{<v)VYw2d%TuPl!aUOS0H;(^AA1+|~Yix7l;TF#QpYAhr
z<9f1wF71xP`z-6{vtGk5IPP59xf<VO-_soTV_sd(?)?$=U+h@jI{oGKMz*bSt`B#w
zTvzPBf%ANlV=iQQqn(?)lh+ThT+Q~+I0omuoO8a1{-4XTi+-L(U6-+bRkyAG-Z+QP
W^<6BNu<sPdxvVdwtsA%Vi~k4A#QOXI

literal 0
HcmV?d00001

diff --git a/hyvideo/modules/fp8_optimization.py b/hyvideo/modules/fp8_optimization.py
new file mode 100644
index 0000000..3665a53
--- /dev/null
+++ b/hyvideo/modules/fp8_optimization.py
@@ -0,0 +1,106 @@
+import os
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from comfy.utils import load_torch_file
+
+def get_fp_maxval(bits=8, mantissa_bit=3, sign_bits=1):
+    _bits = torch.tensor(bits)
+    _mantissa_bit = torch.tensor(mantissa_bit)
+    _sign_bits = torch.tensor(sign_bits)
+    M = torch.clamp(torch.round(_mantissa_bit), 1, _bits - _sign_bits)
+    E = _bits - _sign_bits - M
+    bias = 2 ** (E - 1) - 1
+    mantissa = 1
+    for i in range(mantissa_bit - 1):
+        mantissa += 1 / (2 ** (i+1))
+    maxval = mantissa * 2 ** (2**E - 1 - bias)
+    return maxval
+
+def quantize_to_fp8(x, bits=8, mantissa_bit=3, sign_bits=1):
+    """
+    Default is E4M3.
+    """
+    bits = torch.tensor(bits)
+    mantissa_bit = torch.tensor(mantissa_bit)
+    sign_bits = torch.tensor(sign_bits)
+    M = torch.clamp(torch.round(mantissa_bit), 1, bits - sign_bits)
+    E = bits - sign_bits - M
+    bias = 2 ** (E - 1) - 1
+    mantissa = 1
+    for i in range(mantissa_bit - 1):
+        mantissa += 1 / (2 ** (i+1))
+    maxval = mantissa * 2 ** (2**E - 1 - bias)
+    minval = - maxval
+    minval = - maxval if sign_bits == 1 else torch.zeros_like(maxval)
+    input_clamp = torch.min(torch.max(x, minval), maxval)
+    log_scales = torch.clamp((torch.floor(torch.log2(torch.abs(input_clamp)) + bias)).detach(), 1.0)
+    log_scales = 2.0 ** (log_scales - M - bias.type(x.dtype))
+    # dequant
+    qdq_out = torch.round(input_clamp / log_scales) * log_scales
+    return qdq_out, log_scales
+
+def fp8_tensor_quant(x, scale, bits=8, mantissa_bit=3, sign_bits=1):
+    for i in range(len(x.shape) - 1):
+        scale = scale.unsqueeze(-1)
+    new_x = x / scale
+    quant_dequant_x, log_scales = quantize_to_fp8(new_x, bits=bits, mantissa_bit=mantissa_bit, sign_bits=sign_bits)
+    return quant_dequant_x, scale, log_scales
+
+def fp8_activation_dequant(qdq_out, scale, dtype):
+    qdq_out = qdq_out.type(dtype)
+    quant_dequant_x = qdq_out * scale.to(dtype)
+    return quant_dequant_x
+
+def fp8_linear_forward(cls, original_dtype, input):
+    weight_dtype = cls.weight.dtype
+    #####
+    if cls.weight.dtype != torch.float8_e4m3fn:
+        maxval = get_fp_maxval()
+        scale = torch.max(torch.abs(cls.weight.flatten())) / maxval
+        linear_weight, scale, log_scales = fp8_tensor_quant(cls.weight, scale)
+        linear_weight = linear_weight.to(torch.float8_e4m3fn)
+        weight_dtype = linear_weight.dtype
+    else:
+        scale = cls.fp8_scale.to(cls.weight.device)
+        linear_weight = cls.weight
+    #####
+
+    if weight_dtype == torch.float8_e4m3fn and cls.weight.sum() != 0:
+        if True or len(input.shape) == 3:
+            cls_dequant = fp8_activation_dequant(linear_weight, scale, original_dtype)
+            if cls.bias != None:
+                output = F.linear(input, cls_dequant, cls.bias)
+            else:
+                output = F.linear(input, cls_dequant)
+            return output
+        else:
+            return cls.original_forward(input.to(original_dtype))
+    else:
+        return cls.original_forward(input)
+
+def convert_fp8_linear(module, original_dtype):
+    setattr(module, "fp8_matmul_enabled", True)
+    script_directory = os.path.dirname(os.path.abspath(__file__))
+
+    # loading fp8 mapping file
+    #fp8_map_path = dit_weight_path.replace('.pt', '_map.pt')
+    fp8_map_path = os.path.join(script_directory,"fp8_map.safetensors")
+    if os.path.exists(fp8_map_path):
+        #fp8_map = torch.load(fp8_map_path, map_location=lambda storage, loc: storage)
+        fp8_map = load_torch_file(fp8_map_path)
+    else:
+        raise ValueError(f"Invalid fp8_map path: {fp8_map_path}.")
+
+    fp8_layers = []
+    for key, layer in module.named_modules():
+        if isinstance(layer, nn.Linear) and ('double_blocks' in key or 'single_blocks' in key):
+            fp8_layers.append(key)
+            original_forward = layer.forward
+            layer.weight = torch.nn.Parameter(layer.weight.to(torch.float8_e4m3fn))
+            setattr(layer, "fp8_scale", fp8_map[key].to(dtype=original_dtype))
+            setattr(layer, "original_forward", original_forward)
+            setattr(layer, "forward", lambda input, m=layer: fp8_linear_forward(m, original_dtype, input))
+    
+
diff --git a/nodes.py b/nodes.py
index b51e6e7..ea8d1a0 100644
--- a/nodes.py
+++ b/nodes.py
@@ -239,7 +239,7 @@ def INPUT_TYPES(s):
                 "model": (folder_paths.get_filename_list("diffusion_models"), {"tooltip": "These models are loaded from the 'ComfyUI/models/diffusion_models' -folder",}),
 
             "base_precision": (["fp32", "bf16"], {"default": "bf16"}),
-            "quantization": (['disabled', 'fp8_e4m3fn', 'fp8_e4m3fn_fast', 'torchao_fp8dq', "torchao_fp8dqrow", "torchao_int8dq", "torchao_fp6", "torchao_int4", "torchao_int8"], {"default": 'disabled', "tooltip": "optional quantization method"}),
+            "quantization": (['disabled', 'fp8_e4m3fn', 'fp8_e4m3fn_fast', 'fp8_scaled', 'torchao_fp8dq', "torchao_fp8dqrow", "torchao_int8dq", "torchao_fp6", "torchao_int4", "torchao_int8"], {"default": 'disabled', "tooltip": "optional quantization method"}),
             "load_device": (["main_device", "offload_device"], {"default": "main_device"}),
             },
             "optional": {
@@ -324,7 +324,7 @@ def loadmodel(self, model, base_precision, load_device,  quantization,
 
         if not "torchao" in quantization:
             log.info("Using accelerate to load and assign model weights to device...")
-            if quantization == "fp8_e4m3fn" or quantization == "fp8_e4m3fn_fast":
+            if quantization == "fp8_e4m3fn" or quantization == "fp8_e4m3fn_fast" or quantization == "fp8_scaled":
                 dtype = torch.float8_e4m3fn
             else:
                 dtype = base_dtype
@@ -363,6 +363,9 @@ def loadmodel(self, model, base_precision, load_device,  quantization,
             if quantization == "fp8_e4m3fn_fast":
                 from .fp8_optimization import convert_fp8_linear
                 convert_fp8_linear(patcher.model.diffusion_model, base_dtype, params_to_keep=params_to_keep)
+            elif quantization == "fp8_scaled":
+                from .hyvideo.modules.fp8_optimization import convert_fp8_linear
+                convert_fp8_linear(patcher.model.diffusion_model, base_dtype)
 
             #compile
             if compile_args is not None: