From 8e1239fcc3f7998ddbe83dd41f846073fb4e1d66 Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Thu, 27 Jun 2024 14:47:22 +0300 Subject: [PATCH] Faster model loading, autodownload node, optimizations --- empty_text_embed.pt | Bin 8966 -> 0 bytes ldm/modules/attention.py | 51 +++++---- ldm/modules/diffusionmodules/model.py | 63 +++++------ ldm/modules/diffusionmodules/util.py | 11 +- ldm/modules/encoders/modules.py | 2 + ldm/modules/midas/midas/blocks.py | 21 ++-- model/q_sampler.py | 13 ++- nodes.py | 151 ++++++++++++++++++-------- 8 files changed, 195 insertions(+), 117 deletions(-) delete mode 100644 empty_text_embed.pt diff --git a/empty_text_embed.pt b/empty_text_embed.pt deleted file mode 100644 index dcedcdecbf8ec649a9f088f3a9e50996dce8f89b..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 8966 zcmZ{K2{=~K*1xIDb21j12$d2k?>U=F5h9gR6d?)8SPCJE3?Va>D3VMiiFdE{Hb_dT zWL7H8X^;k`;nV+q|L?i?`|e%m*=O&w&)UCduQRQ+*ZJ952@3J?iHPw1A5rEL<@4~} z85FuX$YXcVW)I&j9`1A8U4vX_@AURrws;)hhQAUi;1=W;;I?h{=HMVNpFpwA0Uld| zy?orq2JHy+3)mboPt2cBTYJ6m7=GZGrmLriKmR|5h%NK+a}8Skk4`avf%W`jI{N;C z>jlSZx1GWMLN+aKB7Wj++X}}}{}#LdXt5%n^1ruuj{dU6 z3jeV8{~7&HyTzdea9`^=RE-zr7KV61Rj(2I?YSQ;)pmi|(sEoEUIBQsmaIa%JAE_$ z8VKHOr2Ro3P1H`Ng6rpO=4MhAy|{2a$nFhhjtKNJ%O2YE=OKiltElKUVAtpe~9C?HXme+RG{R0FUNdH74x}y+{tMx z853(UuIcq-&{oHRbs?%SIvCG<5L979%lql@NDaO2#*af~O~9Kp#C03!pq`XKw82Bz zFL@eNZXLv@hZb|9s?8W-!K-wBX*PW%l?m?Q{WL)<5@h(U!E`}==vgFAXY85`FUnZh zml_NzEuBnth&z1gHHYdbQ^>UMgZH?Rwnz>`VqYv+Tg*hBS0I#KK{(YtU}rLI<;pAWtxf%I)t5b4w0*uWg2(3l$ooa^ur?$c9zKr-5;w@ei)fxzuMRkSCqtIL9E#tS zggE0pkidBWSzkio&y7R0x>^Pp?{_rTRURUrWpQ74sPmFDC&Mf)ABg^b6~dN%q%YJ; zX_A^goX9GMtz8)){H>UdQ)>ge#N#l*S%mAiu)u`F?*S?`Dv+S_1{5sBpv^=J!mCQC z*Ki!8=Tyci?Hs0b#{;@CJ_ky)zrdB8LfUj>H~0IYT==9_1<(9qK#+Y7wG*0QkLNwE z@2(us`S6fY)kuJO#+QI!Y(Gpdvxh`EH~PM97jxVA4Y>FFg2KzW&|1|BvvszE^eRVs z-R~8*cFkfS!v=Lu7qxNeyE&jx`307{NQ3mPyKq%xG8<%QWs+apWzy%j1~#Am31;KO zO(RTR!H7#L*yiNIDvSj4`Dz%DSDQ zx&uNW;q(e}S`(QWMZKJ-p)+A~LMNP+Qirdi-c|ZxgfBS)uSgThPBG^6A@lju#Md14lK2$Ui2 zfdICdzXj%{IxuCs7(eZKjPG^{co&gG3$~LDDoH?Jr1z`QPS1|DHLZ8a@D}@|)qE;vtxQr4er2J^{bab;01O!*HY{ z8YGfu!xrC2X29DNe1{)49N6s-_0wzUjzV|1I#nNryT5_U=F{-Tk}#J?bZN|~0s3n< z<%G_5Y)HMLjP8Dyn8jjd5MOZsCa>nf#RKEf_VNd6lXw{vuIz@lBAsAL;+XlVl~Cdk z$2?wN4K}6XaP6Koj1Rg7{tYiUTIB|e(&R+$V{3?yJoyhhHTnB+K zt=z^jTRrRBfq%ypNIt#-Mgl}I@Io25ElXxXRY9${mcmVqkD(HFIHlXhWY|1z3K`qpzBxAg$ZHUhMcT zI^bCYksfkPsKaJhzs(oMJ(&$2H?ldxhmo0c(Gq$UMVZ{yGBobMe1@F0hf7&nGq zd05UEfumsw*a}~O56T)4Hd-`xZ6r9k7uu=Cj0UiB$b+{RD&fMK`yghnO3&VXO}8dz z!la07Mta|7@VbAUbEDN0K;P{_h=vWa<_+y^Wxz5aRJmlxC#AlEFov|9ISgej&`k_Ouw0@(s|?WP`7$@P>j0< z4~vW7tj_@Bc&{05Unydo+xW0hSQA`54w|T*9i<(scEW9g40!L@Mn#XO!7A}wj^IUA z5Vt=G{nAG{37^OAn~VgG?wi6se!d@;eU5+?JCDQWZ%tGsd+s+nk+}mqpKHVSnWdaM-2(8`YNHD~_%OB50G7$vz`*UfaJ_LAB+fF0D{H>f zsL$c_LzNV@O}2rRqE4{GGz?_hIvak-UZ=AUya)BYEkHJIfuq~5(Q1|-s@u8r_aX(r zs!b58nhsc#0yo4LK+lR{&bt@o(qDqV`DGqHf8~-0B9ExjVsS zt2{N-Gyp+f032)o0tG*VL1^_=+NvA@C!Cw0cVr&?c}1Zjb~bm8wIal>+Sbr>vlzD7 zBE?#1pu@jbSF3ap2z6we)y+ z7(8ryP6xv7(B8{4z-~o0%#ykV(Ro8mX#azTkh)^XS3L(-MjYzk*+@&us$iWG;VL;- z!?)lxsODq9AU_H2jW}>$ti1`UUILgn{~5hkz(RtgGt}oTgj7*CroSi?UQCT;PT!Y@ zD_f%J*&GCk+;0#!=SPM?7X>hUA00cD4D67#v7DHU(RGn;xRQda? ze(OH)mWbq}ofW00XH2KhRtup;z#6KfaDi6LwuN)S2$Ftj&|y0tIL|U5P=rE;Lm6!@ zt_3SzBV3z#fC(GYVOxhBAiZ$7VXB}IJ9|%T!-tcbpt3O?T4Lt;dF_{cUo-w366C3(FRX~#Rr6el`Mv9x2$Q=sWg!7 zYh&^!&w;&4P0$!Az%JlF3mf0&!sn$y5W%kircqO%B7HY#hRD)o>7S{eWdQ^iykTm5 zj=-NK;#e1c6db-*(L>oGFgP<6raSh6S!gBYchiF<$D&!0HTg_-O(PiD+d`9*2t?&Q zq~Y`r*pAjSz5bTa@HLN{iz*=KQw1KA3gJSvI$fzD&GJ%2*n93u^m;-z*JP_0WKBw; zb3R;w+r|}e$~}p?E_Y@3EOz8E%~ zP$-)ZqY9SwdP?0;k{VAr3g_|3*n+?4Xe<2R`9qItN+e&c#%Rvx8 zD~NE>hk9iVc3-dq=jvK5dJattQ*6v697HT0L-1<>|Zs_)I44jn24 zj4W2dTDlOL$JQf(!`I-l`FJ!?w&YBU*Z>poC?t)W4~nsQa5hI6ufJD^fFoH@^<+C! zFH6BGu#ei?mxJ|3d+1TA>`q$zaWX4@ z?FFN%EeDJ4h`^ma19b6P9=s1oh68seQq}cOsgi6sn;9Mjk1A!5bR*>Il)&4*4Q%a^ z%k<>&vk)nC7G{611l?;bbTNMuY_giknRr;4oh!NuybW$b?YSm6+Rg`eA5361{30Ra zc|6_Q)=hV33Q@_ri_nsElNPk-L+-Ct%#fP}G%V!6otqT4xvI0OBbBhbCJp8fCDVmt z{$8i=FbICv0(wmt1+^uhxtAaM&ey=v_Y=XN^BM*X>)`IpL;l7vOVod%7uzb94SRw#xQW4;P;GOMX0|#(ynH-Zm93+! zjU8JyNfui7FzmpCpUjWBO|-{X44>Hb!J^BTXb%;`j^OPuAQDO0ff4F@_ZRq{TMg!| zZQP=6b*5qMD{vMof^TX)5E4t_-IPsGWn9kKk2j%OngdM5!lFg@OINdVA21Z1{Gek^ z76{7vGI95{V0Q3E`Zy>StX-5~#;2)ZuO^NL+CHH5LINv{7eM7MeI_A4uZRyIacFHa6r(EkJkikRPb>R?Kaf>U64QGLM?;Gy0;RZPMaSijp!+|!x zi-L=9)0y1%eK5CH8K$}4qI^#)AwVmd9p5QLbyKq-v80B)vK%q7P3;B6n2U_z$$Cy( zs5GT#5|Mjt8SC#ZYSgvp9)^*VU7-w#SCnS}UOJv_6i(UKO-7Qo)##3_N5V z$r{CXK)aD7b}q;Ri&k@@v2O<%%#nrJC~NBWLIxk*iXxn<5d12xi^(5t@nm-;nb&j> z2k3rm8|lCcpEL;aezN)F^@!cpePn@sGzMjwQms;i)M@^#ObQaV=mx4idCsPOSi^>& zKZxT*^x??G$+$CZ8Vs#8BT}y?;&FLpFPv{lz6@P;+{J@#&;UKwJ&Kmkwa?;}G|&8%c{KDnIB!=OppxcqJ!DGYpw`EDPu zqxvfGzC8;o{?uXQk|v^Zq6D7H6Vr9|3NZJIA#aj<9=o|>0sczAPSeL?IXRPA#3WOK zH@C|Pk9`ou?8je|7|W1fcP_II2d&A}tI=qicn{i3Wle>|^1$fS0Cu*=5^-BfPQ@6KOim+a zI!>`uP%BK9_aXf-@}^az(mVlFdlrKwr~UBKdIgLZ zi^tZaBz(r5!7FHwV9vE%Wp5PRhpurHMjQ9y>xM;`(zp=h?hJw7QCa*l<0;jdqRUg6 zm(0pUZ$yVLzQiE?2TAIXWX}hF!PKU!C^{%g)`bg@ZLXd;!heM}IgI0-U+6<7FloT| zbtTrsAPxVmjN-Qo@V<=}aqB2yTl4eik)7M1iZ2(0&kNzI&9_l|;1>C|;sO~K_r`fg zlHlRw#aNJ&fodaf_49f9quI&Et?+et?+0KaZ1N zA4bC*L3D{bh-<|j;^nU?$m2d_pFf_)3lkOP{Z_QX-;R95IBt~8TNOrDygJKXPFap& zFE)^s4`$%ipWncB>nu~#E-4~#bs|hUFNyN9X4ql-2%To7kc|rykuM+yc2=;a)3VCg zM>8+7Rar>mHf6D^lLJBSWHYfFe!%=Z)=gelEW(C&L2YhD6}_%gf0|XE}HwHb^zsh>%XL5cJr$pMGpSN7xH(BvI%#=1y3K z0y~n}0D~?lGIS=}4v3KJcDL9?L3!vS-pmdL@!>}+UwD0yOANi*iNu@tWX9HF46*WH zy*CJAz0zawiqAl~DsxnE>m*NCsT0K+5~j_`$~ftz7OuUYhzXZpf^=yHC@-3B>S$ku zul#qish4^%Gpz=NKk=je(hYcc^)~bjbRz!CrFhL%Ram=l5!t7_k$fm!juYAiOuuFI z!J$4&(sLjWIiKoj)+IsmcGp*o`y>SqCZw}r4^>FJ+fMRpJr0sn_J7O0}!qa%Hb&xK% zUwI9hLxgGOTV%PSO3O%jY+1*nN@= zNsK4);v1OEtE-`Ly%oxOWuR2=3dWst7^d`H;cR-spfkqd9XluTy(g0lDtVD~$5$}k z@-=o_`rsltKJ0&Wk#!2RA$yOM;}qpn;J-Q<6Zm4$p=Ugf%s++aj*hJ@zvSZY@Ui&I zPSkY9Eq&S;Mq&430p7SdeXu@T8aIs0gu%%7FwiW=yYe|6kKaweTxg?*ils5o%$r&n z{KncaL7b&MpXf`e^G-c4#%7szlnDOHPH=F)ehxZ@ z93!8}buc`fM|TC^=WLVD!^@BB7@egeyc1vUp@60uj1C8|&jKn*<2`B9T8#!WHLQT$ zo6SdPxg<=>Da7c-9OCAoP5d7BL3>LXcIDn-96qgOm*1Ys^%Yr5LYAE-Mt)r|Gwv21 zzp8?%{sC<0?%Vj*Qw_ItH=#h}UGl_24_oqPqj>8$ykol^mE=Fd1lM@BOYI1&%iQB8 zO&MToVp@sY{sfF&EQQM=VuVq`iAzyF@l2n} z&UY%omacsyJw$_dVoDy0C>yZr`m*ri@KF@JJw!(>9%HSqFte`KkG=k38E^lav2Sm& z8tGDV!RW{ntm%>@?0XZ-p1z%eJw^m)GLPWqrz*03X&XptKEYqVmykW*kD16fh_RXi z=BQjc9`b**uoinRk|>=t#?B&!;kTZMOIs(PbwoOdzbnS9enHl*E^#bJtw?OvDVVyQ z636R%rSR#{Gxq%9m1yPJz(jS=$G+ZL@M=lK>sO_4c)?Sc`>BOp9;AaOR2y*8!9`?w zN(+`(tFhz~KWx5Z$jklt7B}^W!L_(qg#T(kJ8jNt(}N!Aq}6Z%8XK<0B=MK5k(v^( zZSq76SJ;GRjVGAkRbp&LQ85ULjj%EX8APIO8YqvbLc#}ed=|D6-6yB9+B0<6f~|Zw zPw@bT6gHDRp?plLOE7D+TpP7I`JlexAU1iF5EbWKc+_Z5P<2?A`2= z!5a`)FqteVnTXuvUfkX103%;lqo+eD8^sQ4-h_wkHAvD*F6w2fgV>+9EZ_7p;?n8MUXj&eN4mAibEiXN=i3?6eoYd~ zR}3`=W3gH20Dd`Li8j`bctt^o8$;VA9N3`>B!!SuT6yi>ayo>QtPjc>~kG zZzZ>KjiBPJFH)2`H!8LCCj5#H;hp#p^9*NZh_vxW3c)@;2Y;% zWS-4q_RB_fqTXG>(OB%yM)sb*awEohTxT0Fge?hMC|`;C*mg;6E91K ztU294r7Rir=j|qI3U7kcM^l_gSS=GO+QQZvi18l$ODyO zc3mT7b|tML@=NC8+-3`sv%wvAzM0I6l#*tQRnMa4wqf{_ng)tdCb-XV341Di6r|hD zc?Bnm@WD@Q{OP3v7C-bDA>;ix8X~|d3j9Wi>}E*+dK`cMNX4uJ17zCqy==Cj4W5T( zWPinVc(OdyhMI zpF!7DVN)mPa+JHKjc^K4RP-(@FEWm|U{?uxKeAzVJXXdou_$u-yeD(_)h;|cuA8%< zYAUO!u#PC|9Kk;^bLoJ+9jrN&NLHAqVfXhG_%7QEAA1r=X{R1$q)kE-B~8=sTh|eZ z5AO*_$bjv8dl#&i4RJ*mCqZRsF#cJ4n_~5G(y^kDDld2qs?|%${pn34OZ_9-)Q^Xz z_Hx|x!;%y~UP*o^#o&o2!k9Fq2>cGua8~JYe6nv2d%|Wp?a#Bpmib6UM)!{SJOi@n zwkVmWWy?OhpM|;l;$%a63VBp!O-@>su>8>{khkLvJ7||qwi;A{lAAWks&!^t{IiKm zY$Dt-E8&LgIHRc64$`B0kR)`?BZ<3JAl#}HcHD6!D^!vgJ^ng2;miSO?f%0|sY~GS zyN6@d=L&K+Qk>immL?moRk3atXP|=GeQ5r;61bH!VM|yhBXxH@p>8kmac&ew zWe9_)eK4uqsf)yO1$H<~@G_dZ!A<)i_5{S@+oT0V?ch_0pxG#0If*EAT)~;+G>PD$ z`52l!i`d(27|SJfV32JB-ncsh9VevXNv#Q{m$Pet-*)Vu%}Rb=i(es5QaQ#x%{>gE z=rx|JdMn~)BFujwq3d{4!#$hG0K>NE% zrvx%Gmjf|)aSQiwZUD~Tm;>ULGaBwhT_7*)l4xGE(5tS78pC`Jm$avP@x=b%mc3QTc4gl8T^`UslADqAM6Gd-S041|BCI=UczwKzo)b1y_rt4ajcVDD>WIA#@(WnsQLEe#rU)2 z`S5E9i`q<1e(<2OHg4z@5QcNyZV&_WKWy?3DdM8+3Tb*8ru#IYn~*D|aX;^D=~?{LqG)95R+35{Ns!i0xYSi3U;JwLGv+e3(XVl}ZE@WS-67!u>Ihw1V>@|RsueW7j^pT-!q zV66OSc7YP#qW{Y-2=NFA9HTERTWs|2=>`A9H2C*T1F^p|xOdERt-#m>{uVnc{=e1# za60~(`=I>y#J|XYWO)2@n0Z5jUD}4m^u3Y8|0tosq`-(cAEbNjQh(I`78gwzW 0: - self.temb_proj = torch.nn.Linear(temb_channels, + self.temb_proj = ops.Linear(temb_channels, out_channels) self.norm2 = Normalize(out_channels) self.dropout = torch.nn.Dropout(dropout) - self.conv2 = torch.nn.Conv2d(out_channels, + self.conv2 = ops.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1) if self.in_channels != self.out_channels: if self.use_conv_shortcut: - self.conv_shortcut = torch.nn.Conv2d(in_channels, + self.conv_shortcut = ops.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) else: - self.nin_shortcut = torch.nn.Conv2d(in_channels, + self.nin_shortcut = ops.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, @@ -157,22 +158,22 @@ class AttnBlock(nn.Module): self.in_channels = in_channels self.norm = Normalize(in_channels) - self.q = torch.nn.Conv2d(in_channels, + self.q = ops.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) - self.k = torch.nn.Conv2d(in_channels, + self.k = ops.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) - self.v = torch.nn.Conv2d(in_channels, + self.v = ops.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) - self.proj_out = torch.nn.Conv2d(in_channels, + self.proj_out = ops.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, @@ -216,22 +217,22 @@ class MemoryEfficientAttnBlock(nn.Module): self.in_channels = in_channels self.norm = Normalize(in_channels) - self.q = torch.nn.Conv2d(in_channels, + self.q = ops.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) - self.k = torch.nn.Conv2d(in_channels, + self.k = ops.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) - self.v = torch.nn.Conv2d(in_channels, + self.v = ops.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) - self.proj_out = torch.nn.Conv2d(in_channels, + self.proj_out = ops.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, @@ -318,14 +319,14 @@ class Model(nn.Module): # timestep embedding self.temb = nn.Module() self.temb.dense = nn.ModuleList([ - torch.nn.Linear(self.ch, + ops.Linear(self.ch, self.temb_ch), - torch.nn.Linear(self.temb_ch, + ops.Linear(self.temb_ch, self.temb_ch), ]) # downsampling - self.conv_in = torch.nn.Conv2d(in_channels, + self.conv_in = ops.Conv2d(in_channels, self.ch, kernel_size=3, stride=1, @@ -394,7 +395,7 @@ class Model(nn.Module): # end self.norm_out = Normalize(block_in) - self.conv_out = torch.nn.Conv2d(block_in, + self.conv_out = ops.Conv2d(block_in, out_ch, kernel_size=3, stride=1, @@ -467,7 +468,7 @@ class Encoder(nn.Module): self.in_channels = in_channels # downsampling - self.conv_in = torch.nn.Conv2d(in_channels, + self.conv_in = ops.Conv2d(in_channels, self.ch, kernel_size=3, stride=1, @@ -512,7 +513,7 @@ class Encoder(nn.Module): # end self.norm_out = Normalize(block_in) - self.conv_out = torch.nn.Conv2d(block_in, + self.conv_out = ops.Conv2d(block_in, 2*z_channels if double_z else z_channels, kernel_size=3, stride=1, @@ -571,7 +572,7 @@ class Decoder(nn.Module): self.z_shape, np.prod(self.z_shape))) # z to block_in - self.conv_in = torch.nn.Conv2d(z_channels, + self.conv_in = ops.Conv2d(z_channels, block_in, kernel_size=3, stride=1, @@ -613,7 +614,7 @@ class Decoder(nn.Module): # end self.norm_out = Normalize(block_in) - self.conv_out = torch.nn.Conv2d(block_in, + self.conv_out = ops.Conv2d(block_in, out_ch, kernel_size=3, stride=1, @@ -658,7 +659,7 @@ class Decoder(nn.Module): class SimpleDecoder(nn.Module): def __init__(self, in_channels, out_channels, *args, **kwargs): super().__init__() - self.model = nn.ModuleList([nn.Conv2d(in_channels, in_channels, 1), + self.model = nn.ModuleList([ops.Conv2d(in_channels, in_channels, 1), ResnetBlock(in_channels=in_channels, out_channels=2 * in_channels, temb_channels=0, dropout=0.0), @@ -668,11 +669,11 @@ class SimpleDecoder(nn.Module): ResnetBlock(in_channels=4 * in_channels, out_channels=2 * in_channels, temb_channels=0, dropout=0.0), - nn.Conv2d(2*in_channels, in_channels, 1), + ops.Conv2d(2*in_channels, in_channels, 1), Upsample(in_channels, with_conv=True)]) # end self.norm_out = Normalize(in_channels) - self.conv_out = torch.nn.Conv2d(in_channels, + self.conv_out = ops.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, @@ -719,7 +720,7 @@ class UpsampleDecoder(nn.Module): # end self.norm_out = Normalize(block_in) - self.conv_out = torch.nn.Conv2d(block_in, + self.conv_out = ops.Conv2d(block_in, out_channels, kernel_size=3, stride=1, @@ -744,7 +745,7 @@ class LatentRescaler(nn.Module): super().__init__() # residual block, interpolate, residual block self.factor = factor - self.conv_in = nn.Conv2d(in_channels, + self.conv_in = ops.Conv2d(in_channels, mid_channels, kernel_size=3, stride=1, @@ -759,7 +760,7 @@ class LatentRescaler(nn.Module): temb_channels=0, dropout=0.0) for _ in range(depth)]) - self.conv_out = nn.Conv2d(mid_channels, + self.conv_out = ops.Conv2d(mid_channels, out_channels, kernel_size=1, ) @@ -841,7 +842,7 @@ class Resize(nn.Module): raise NotImplementedError() assert in_channels is not None # no asymmetric padding in torch conv, must do it ourselves - self.conv = torch.nn.Conv2d(in_channels, + self.conv = ops.Conv2d(in_channels, in_channels, kernel_size=4, stride=2, diff --git a/ldm/modules/diffusionmodules/util.py b/ldm/modules/diffusionmodules/util.py index 637363d..76d221e 100644 --- a/ldm/modules/diffusionmodules/util.py +++ b/ldm/modules/diffusionmodules/util.py @@ -17,6 +17,9 @@ from einops import repeat from ldm.util import instantiate_from_config +import comfy.ops +ops = comfy.ops.manual_cast + def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3): if schedule == "linear": @@ -223,11 +226,11 @@ def conv_nd(dims, *args, **kwargs): Create a 1D, 2D, or 3D convolution module. """ if dims == 1: - return nn.Conv1d(*args, **kwargs) + return ops.Conv1d(*args, **kwargs) elif dims == 2: - return nn.Conv2d(*args, **kwargs) + return ops.Conv2d(*args, **kwargs) elif dims == 3: - return nn.Conv3d(*args, **kwargs) + return ops.Conv3d(*args, **kwargs) raise ValueError(f"unsupported dimensions: {dims}") @@ -235,7 +238,7 @@ def linear(*args, **kwargs): """ Create a linear module. """ - return nn.Linear(*args, **kwargs) + return ops.Linear(*args, **kwargs) def avg_pool_nd(dims, *args, **kwargs): diff --git a/ldm/modules/encoders/modules.py b/ldm/modules/encoders/modules.py index ab8b31d..9358ff3 100644 --- a/ldm/modules/encoders/modules.py +++ b/ldm/modules/encoders/modules.py @@ -143,8 +143,10 @@ class FrozenOpenCLIPEmbedder(AbstractEncoder): # def __init__(self, arch="ViT-H-14", version="laion2b_s32b_b79k", device="cuda", max_length=77, def __init__(self, arch="ViT-H-14", version="laion2b_s32b_b79k", max_length=77, freeze=True, layer="last"): + super().__init__() assert layer in self.LAYERS + return model, _, _ = open_clip.create_model_and_transforms(arch, device=torch.device('cpu'), pretrained=version) del model.visual self.model = model diff --git a/ldm/modules/midas/midas/blocks.py b/ldm/modules/midas/midas/blocks.py index 2145d18..8837f60 100644 --- a/ldm/modules/midas/midas/blocks.py +++ b/ldm/modules/midas/midas/blocks.py @@ -1,6 +1,9 @@ import torch import torch.nn as nn +import comfy.ops +ops = comfy.ops.manual_cast + from .vit import ( _make_pretrained_vitb_rn50_384, _make_pretrained_vitl16_384, @@ -59,16 +62,16 @@ def _make_scratch(in_shape, out_shape, groups=1, expand=False): out_shape3 = out_shape*4 out_shape4 = out_shape*8 - scratch.layer1_rn = nn.Conv2d( + scratch.layer1_rn = ops.Conv2d( in_shape[0], out_shape1, kernel_size=3, stride=1, padding=1, bias=False, groups=groups ) - scratch.layer2_rn = nn.Conv2d( + scratch.layer2_rn = ops.Conv2d( in_shape[1], out_shape2, kernel_size=3, stride=1, padding=1, bias=False, groups=groups ) - scratch.layer3_rn = nn.Conv2d( + scratch.layer3_rn = ops.Conv2d( in_shape[2], out_shape3, kernel_size=3, stride=1, padding=1, bias=False, groups=groups ) - scratch.layer4_rn = nn.Conv2d( + scratch.layer4_rn = ops.Conv2d( in_shape[3], out_shape4, kernel_size=3, stride=1, padding=1, bias=False, groups=groups ) @@ -164,11 +167,11 @@ class ResidualConvUnit(nn.Module): """ super().__init__() - self.conv1 = nn.Conv2d( + self.conv1 = ops.Conv2d( features, features, kernel_size=3, stride=1, padding=1, bias=True ) - self.conv2 = nn.Conv2d( + self.conv2 = ops.Conv2d( features, features, kernel_size=3, stride=1, padding=1, bias=True ) @@ -244,11 +247,11 @@ class ResidualConvUnit_custom(nn.Module): self.groups=1 - self.conv1 = nn.Conv2d( + self.conv1 = ops.Conv2d( features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups ) - self.conv2 = nn.Conv2d( + self.conv2 = ops.Conv2d( features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups ) @@ -310,7 +313,7 @@ class FeatureFusionBlock_custom(nn.Module): if self.expand==True: out_features = features//2 - self.out_conv = nn.Conv2d(features, out_features, kernel_size=1, stride=1, padding=0, bias=True, groups=1) + self.out_conv = ops.Conv2d(features, out_features, kernel_size=1, stride=1, padding=0, bias=True, groups=1) self.resConfUnit1 = ResidualConvUnit_custom(features, activation, bn) self.resConfUnit2 = ResidualConvUnit_custom(features, activation, bn) diff --git a/model/q_sampler.py b/model/q_sampler.py index 8c6450c..883e753 100644 --- a/model/q_sampler.py +++ b/model/q_sampler.py @@ -577,7 +577,8 @@ class SpacedSampler: # fuse by tile_weights on noise (score) noise_buffer /= count pred_x0 = self._predict_xstart_from_eps(x_t=img, t=index, eps=noise_buffer) - tao_index = torch.tensor(torch.round(index * t_max), dtype=torch.int64) + tao_index = torch.round(index * t_max).clone().detach().to(torch.int64) + img = self.q_sample(pred_x0, tao_index) @@ -606,11 +607,11 @@ class SpacedSampler: tile_cond_img = cond_img[:, :, hi * 8:hi_end * 8, wi * 8: wi_end * 8] tile_cond = { "c_latent": [self.model.apply_condition_encoder(tile_cond_img)], - "c_crossattn": [self.model.get_learned_conditioning([positive_prompt] * b)] + "c_crossattn": [empty_text_embed * b] } tile_uncond = { "c_latent": [self.model.apply_condition_encoder(tile_cond_img)], - "c_crossattn": [self.model.get_learned_conditioning([negative_prompt] * b)] + "c_crossattn": [empty_text_embed * b] } # predict noise for this tile tile_noise = self.predict_noise(tile_img, ts, tile_cond, cfg_scale, tile_uncond) @@ -728,11 +729,12 @@ class SpacedSampler: # accumulate noise noise_buffer[:, :, hi:hi_end, wi:wi_end] += tile_noise count[:, :, hi:hi_end, wi:wi_end] += 1 - pbar.update(1) + # average on noise (score) noise_buffer.div_(count) pred_x0 = self._predict_xstart_from_eps(x_t=img, t=index, eps=noise_buffer) - tao_index = torch.tensor(torch.round(index * t_max), dtype=torch.int64) + tao_index = torch.round(index * t_max).clone().detach().to(torch.int64) + img = self.q_sample(pred_x0, tao_index) noise_buffer.zero_() @@ -795,6 +797,7 @@ class SpacedSampler: noise_buffer.zero_() count.zero_() + pbar.update(1) img = pred_x0 # decode samples of each diffusion process diff --git a/nodes.py b/nodes.py index c16ba43..83bee8e 100644 --- a/nodes.py +++ b/nodes.py @@ -9,13 +9,12 @@ from .model.ccsr_stage1 import ControlLDM from .utils.common import instantiate_from_config, load_state_dict -import comfy.model_management +import comfy.model_management as mm import comfy.utils import folder_paths from nodes import ImageScaleBy from nodes import ImageScale - script_directory = os.path.dirname(os.path.abspath(__file__)) class CCSR_Upscale: @@ -64,29 +63,24 @@ class CCSR_Upscale: CATEGORY = "CCSR" @torch.no_grad() - def process(self, ccsr_model, image, resize_method, scale_by, steps, t_max, t_min, tile_size, tile_stride, color_fix_type, keep_model_loaded, vae_tile_size_encode, vae_tile_size_decode, sampling_method, seed): + def process(self, ccsr_model, image, resize_method, scale_by, steps, t_max, t_min, tile_size, tile_stride, + color_fix_type, keep_model_loaded, vae_tile_size_encode, vae_tile_size_decode, sampling_method, seed): torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) - comfy.model_management.unload_all_models() - device = comfy.model_management.get_torch_device() - config_path = os.path.join(script_directory, "configs/model/ccsr_stage2.yaml") - empty_text_embed = torch.load(os.path.join(script_directory, "empty_text_embed.pt"), map_location=device) - dtype = torch.float16 if comfy.model_management.should_use_fp16() and not comfy.model_management.is_device_mps(device) else torch.float32 - if not hasattr(self, "model") or self.model is None: - config = OmegaConf.load(config_path) - self.model = instantiate_from_config(config) - - load_state_dict(self.model, comfy.utils.load_torch_file(ccsr_model), strict=True) - # reload preprocess model if specified + mm.unload_all_models() + device = mm.get_torch_device() + offload_device = mm.unet_offload_device() + dtype = ccsr_model['dtype'] + model = ccsr_model['model'] + + #empty_text_embed = torch.load(os.path.join(script_directory, "empty_text_embed.pt"), map_location=device) + empty_text_embed_sd = comfy.utils.load_torch_file(os.path.join(script_directory, "empty_text_embed.safetensors")) + empty_text_embed = empty_text_embed_sd['empty_text_embed'].to(dtype).to(device) - self.model.freeze() - self.model.to(device, dtype=dtype) - sampler = SpacedSampler(self.model, var_type="fixed_small") + sampler = SpacedSampler(model, var_type="fixed_small") - batch_size = image.shape[0] image, = ImageScaleBy.upscale(self, image, resize_method, scale_by) - # Assuming 'image' is a PyTorch tensor with shape [B, H, W, C] and you want to resize it. B, H, W, C = image.shape # Calculate the new height and width, rounding down to the nearest multiple of 64. @@ -95,28 +89,26 @@ class CCSR_Upscale: # Reorder to [B, C, H, W] before using interpolate. image = image.permute(0, 3, 1, 2).contiguous() - - # Resize the image tensor. - resized_image = F.interpolate(image, size=(new_height, new_width), mode='bicubic', align_corners=False) - - # Move the tensor to the GPU. - #resized_image = resized_image.to(device) + resized_image = F.interpolate(image, size=(new_height, new_width), mode='bilinear', align_corners=False) + strength = 1.0 - self.model.control_scales = [strength] * 13 + model.control_scales = [strength] * 13 + model.to(device, dtype=dtype).eval() + height, width = resized_image.size(-2), resized_image.size(-1) shape = (1, 4, height // 8, width // 8) - x_T = torch.randn(shape, device=self.model.device, dtype=torch.float32) - autocast_condition = dtype == torch.float16 and not comfy.model_management.is_device_mps(device) - out = [] + x_T = torch.randn(shape, device=model.device, dtype=torch.float32) - pbar = comfy.utils.ProgressBar(batch_size) - - with torch.autocast(comfy.model_management.get_autocast_device(device), dtype=dtype) if autocast_condition else nullcontext(): - for i in range(batch_size): + out = [] + if B > 1: + pbar = comfy.utils.ProgressBar(B) + autocast_condition = dtype == torch.float16 and not mm.is_device_mps(device) + with torch.autocast(mm.get_autocast_device(device), dtype=dtype) if autocast_condition else nullcontext(): + for i in range(B): img = resized_image[i].unsqueeze(0).to(device) if sampling_method == 'ccsr_tiled_mixdiff': - self.model.reset_encoder_decoder() + model.reset_encoder_decoder() print("Using tiled mixdiff") samples = sampler.sample_with_mixdiff_ccsr( empty_text_embed, tile_size=tile_size, tile_stride=tile_stride, @@ -126,7 +118,7 @@ class CCSR_Upscale: color_fix_type=color_fix_type ) elif sampling_method == 'ccsr_tiled_vae_gaussian_weights': - self.model._init_tiled_vae(encoder_tile_size=vae_tile_size_encode // 8, decoder_tile_size=vae_tile_size_decode // 8) + model._init_tiled_vae(encoder_tile_size=vae_tile_size_encode // 8, decoder_tile_size=vae_tile_size_decode // 8) print("Using gaussian weights") samples = sampler.sample_with_tile_ccsr( empty_text_embed, tile_size=tile_size, tile_stride=tile_stride, @@ -136,7 +128,7 @@ class CCSR_Upscale: color_fix_type=color_fix_type ) else: - self.model.reset_encoder_decoder() + model.reset_encoder_decoder() print("no tiling") samples = sampler.sample_ccsr( empty_text_embed, steps=steps, t_max=t_max, t_min=t_min, shape=shape, cond_img=img, @@ -145,9 +137,10 @@ class CCSR_Upscale: color_fix_type=color_fix_type ) out.append(samples.squeeze(0).cpu()) - comfy.model_management.throw_exception_if_processing_interrupted() - pbar.update(1) - print("Sampled image ", i, " out of ", batch_size) + mm.throw_exception_if_processing_interrupted() + if B > 1: + pbar.update(1) + print("Sampled image ", i, " out of ", B) original_height, original_width = H, W processed_height = samples.size(2) @@ -156,8 +149,8 @@ class CCSR_Upscale: resized_back_image, = ImageScale.upscale(self, out_stacked, "lanczos", target_width, processed_height, crop="disabled") if not keep_model_loaded: - self.model = None - comfy.model_management.soft_empty_cache() + model.to(offload_device) + mm.soft_empty_cache() return(resized_back_image,) class CCSR_Model_Select: @@ -173,15 +166,85 @@ class CCSR_Model_Select: CATEGORY = "CCSR" def load_ccsr_checkpoint(self, ckpt_name): + device = mm.get_torch_device() + offload_device = mm.unet_offload_device() ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name) + config_path = os.path.join(script_directory, "configs/model/ccsr_stage2.yaml") + dtype = torch.float16 if mm.should_use_fp16() and not mm.is_device_mps(device) else torch.float32 - return (ckpt_path,) + if not hasattr(self, "model") or self.model is None: + config = OmegaConf.load(config_path) + self.model = instantiate_from_config(config) + + load_state_dict(self.model, comfy.utils.load_torch_file(ckpt_path), strict=True) + # reload preprocess model if specified + + ccsr_model = { + 'model': self.model, + 'dtype': dtype + } + return (ccsr_model,) + +class DownloadAndLoadCCSRModel: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "model": ( + [ + 'real-world_ccsr-fp16.safetensors', + 'real-world_ccsr-fp32.safetensors' + ], + ), + }, + } + + RETURN_TYPES = ("CCSRMODEL",) + RETURN_NAMES = ("ccsr_model",) + FUNCTION = "loadmodel" + CATEGORY = "CCSR" + + def loadmodel(self, model): + device = mm.get_torch_device() + offload_device = mm.unet_offload_device() + dtype = torch.float16 if 'fp16' in model else torch.float32 + + model_path = os.path.join(folder_paths.models_dir, "CCSR") + safetensors_path = os.path.join(model_path, model) + + if not os.path.exists(safetensors_path): + print(f"Downloading CCSR model to: {model_path}") + from huggingface_hub import snapshot_download + snapshot_download(repo_id="Kijai/ccsr-safetensors", + allow_patterns=[f'*{model}*'], + local_dir=model_path, + local_dir_use_symlinks=False) + + config_path = os.path.join(script_directory, "configs/model/ccsr_stage2.yaml") + config = OmegaConf.load(config_path) + + model = instantiate_from_config(config) + + sd = comfy.utils.load_torch_file(safetensors_path) + + model.load_state_dict(sd, strict=False) + del sd + mm.soft_empty_cache() + + ccsr_model = { + 'model': model, + 'dtype': dtype, + } + + return (ccsr_model,) + NODE_CLASS_MAPPINGS = { "CCSR_Upscale": CCSR_Upscale, - "CCSR_Model_Select": CCSR_Model_Select + "CCSR_Model_Select": CCSR_Model_Select, + "DownloadAndLoadCCSRModel": DownloadAndLoadCCSRModel } NODE_DISPLAY_NAME_MAPPINGS = { "CCSR_Upscale": "CCSR_Upscale", - "CCSR_Model_Select": "CCSR_Model_Select" + "CCSR_Model_Select": "CCSR_Model_Select", + "DownloadAndLoadCCSRModel": "DownloadAndLoad CCSRModel" } \ No newline at end of file