From 4cbefda3e71cf66b2e48d666eb8d995ecbd295b2 Mon Sep 17 00:00:00 2001 From: flybirdxx <1119577418@qq.com> Date: Fri, 15 Aug 2025 19:14:12 +0800 Subject: [PATCH] Merge model loading nodes, change model paths, add automatic downloads, and add display memory. --- __pycache__/__init__.cpython-311.pyc | Bin 0 -> 339 bytes __pycache__/sdmatte_nodes.cpython-311.pyc | Bin 0 -> 15668 bytes sdmatte_nodes.py | 407 ++++++++++++---------- 3 files changed, 230 insertions(+), 177 deletions(-) create mode 100644 __pycache__/__init__.cpython-311.pyc create mode 100644 __pycache__/sdmatte_nodes.cpython-311.pyc diff --git a/__pycache__/__init__.cpython-311.pyc b/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d79b773b0b62ad3d383031598375a86ffc058e18 GIT binary patch literal 339 zcmZ3^%ge<81hH9jGG+tm#~=<2FhLog<$#Ro3@HpLj5!Rsj8TlaOi@gX3@J=0%;`)~ z%qc9v44SMj8G(v4nQjUD`MbErJNq~W2gmz51_XHexd-2p0t>o$1_$^!M#lR&`ntm9 zikN}=idcYzpC-#K-r|(p#FCQK_`Lj-)MA9(E%x~M#GIV?_>~Nwfl3*E`M6reIOpf4 zRfc-TmuKdslIY~ z;;_lhPbtkwwJQRJ2FNGH20-EiGb1D84F<0Z*w6#}1E>iA DMu1!j literal 0 HcmV?d00001 diff --git a/__pycache__/sdmatte_nodes.cpython-311.pyc b/__pycache__/sdmatte_nodes.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bf8a57f70dd1e348c61a6b7c5777ae0e192dd3cb GIT binary patch literal 15668 zcmcJ0Yj6`+mT9*R)6Y=JB}Es2L$+Bg|9 z$SeuYY%(~+3`#a(oP;EvWF~PKNSIBsVX9`Ty1hMxs;j7?YLT)3?3g06Ma`e@JNLF^ zJxmhz+izP}N4M|e+;h);oqNvFuXQ>#0gv?$zv{WagdqL{y(AuaJn~=$BIgN;ph!2- zCH|6KBz~9lN*xlC=#_QJNiip@O943&x3X8&rHZGiyVSU}rb`3AQn$8O*QJ9z(yo^f zTM0^bilF3|B?R#Se8s0Ln=(=gnml4iq)^J#H>ZqJosxGMp_Do`eNHJ2l*+-aXrUEd zMk^*dn>JDU5lL4rl}#CNERQ1h5Z0V&)NqHD6x4pk?(sRjOs@~#vfW-gW%GL(M~}_v za?_3kedsqq{(}|pa-N_G#HKC@Ev3j)#AQhP0F}Z|myD8h$th`979|6g$xjhoN-7H) zRtPzJHnrRRe!AUD(e5^Aj%FNWq8|LBfsw-Z!7>0aFeD}Mr{o`*jZ@Og5dUEI-ATIO z18vAU1SLa17jQ8&C8!)n`}{VKy_Xg=jvgA`ye^O5XO#$Ar`HXQ+xqPO9-m7J?^vhc^SX`u$xX~9B>4Des8bM z10C}v5}%0o_JEJ-MZFet>-$a$n%0h8dpm8N&+OW~r{6r6DfKD&Mj*mx1u#I&WD&Z8 zsG;Pdn>Vi;KggMP@a7$CS;vf2Qly`e$_(lmLZ(tbfWia}s16V}cXn>JweQ@txlK^x z@$E&d|9dtX2A8LMQM_$OVti3m6@2k1{Fj0tPMgjXev}%2l4B4_t&_;kmX!A&?NXBPYm+%gpe6B%C=;>NUw=;WEq&xe(jUNsmw&RB7z;`%DJ8oMz54+5 z5I<*3-LOdCCPT8Y$w^XKr<6h2rp$iGgL0BgpbVh;q9UEP z7fF9nN(~Jvdr^{J2KX@BcY+u+QCdnz>EWxRv{$k(L*56o`JpU-NfO@?eIia7f@J#g z4=K`GHU<@sYMCNJN)^N~k`V4WL#CY=DxxINIO&pKr{|%4-iDWxRX89jfQ2CVZN~lEoJ@_7}vNAE8A16O@$9p`->VS)A+{aIgqw!VmWKN@>RU zDpR@r3>;7)?LTTgO7ZtNnHUGp>N!2G%jj_#?y-4pkNr!N^aV3!b_m-4p*CSPmBVV% zLB8IngJ8-4`0)xS20DQHN? zLs^c>P%e`nGyuJ+>$M`7{isz#5+3sGQ$hW}-ZTn+talsIC{X?Y1yA@JlXOy*rwHIh zc`_v#OFwTj_o{U4j3lTG>493uV^dqK9=^?v`kkQ|g#6ndCjUXV+dr6XH zFs!rcG`<+qu|R&!MR|H!#*R8awSF4a`p4m1I_~thM~!?iD^i*4Q!*CR!(LsRKBEu7 z=@^oSa?)m0`b^B|8}xTUYJY+`b10Ifr;r3tRF$M78B3obxZJ#CZZZ~}a~cJ+oJfvZ z*fEWyACMvxaS=g~`WOd;c3`Vq(E8~UejDv^fHc7f(qnd7kOe%nUr_l#(4qovNLBg0 zM`@3%KSRbt=er5M4_*QAI(eSxCW6Gl<`DVckurXu?ChBisAr!+N#^O8&c(5&Ao;4~ z?8ZIBdmujw60c8?A-PbHyL^<35oAX|YP3oPN&RAhq=f7~h6cIW<)JAXL-%=os8Ne$ zfXRm6zpK|F@zLS!?OWg9@Cd;!DTCFob9)_jw{JZ}5>;veg>OB8hXVwg;}pZN@pp%R z+x>a(C%xZp{ATNSZU5TFKKU$r*uib0_)QeIg63DyP;#hYwg!V_#y~NE$&o9uk+ImN ztFf!E-@A4-HumbhH$RG9e|?Q5fUY0)mdTMzlOu0Veg1Lm%Tti3gT%jHJ3Box5Ig%) zZ0G|q#S;5)NUZlC126WM?2q^Bpk=QQWJt?lyN|XYR;;mX=@hd3jH}n)2hx673a7Bx ziA8anR}{Mi4d%sh+ST3T7jyw1ZLsh-x>>p1X-Br22|B& zcha_QxA!m%4Wy~tU~F2v9;d4t-i#3QqoHToX}_Ja`|X0%(HB4$3h1&84Je{S(8h7w zydL+-{sl=A@a?J|b-XsUNMk1xpsg0~(f)bKvU3VpoOZaDQEM#y#cP4ze&%GN54fGi zEkH8ezc7j6hg!o0QFB(3K{u`bvPacPHfLE2aku-h-Eq_cQ`>nG_aNn}Xw?X6JSH~E zW}aaKo3U>}tBE9ym>`{nhAU3OmwgN|gYghBhMK%VfS$z!pIt6#OQ zoA%Hr`k3|oiywFKCu+C=k8crxDfqzhU+9Ez^SRO1apPzQTep3@`=)W+EheWzaj`c* z#mV!L$xmOu`}(c>f_JvUjOReYooD`h9`gd#nh<_Xdn6bTI}@IyT5;Fa`dg(ul`AF z2>+g7Sq_TtFyrQa-xq$L)~bt7UXy$CZ<}i-DB^jdtoCG4Cpd1YB4DKn9Bio|M*a} zco|wS!zose=M&(qkLT3)!g99RLGytRIK~2uU@8Cz3OE6<0zqlR8UmY*K?km2+yYCn z!|UJT4S1-{4C7@4c^?DX1xACiby%YiaJ#XFK`@Cc5%wgP-RibaaK^ffN#J~ z5AkkGu?e1XG|0P8(IfJ_ta5!T1xvo*e5FkY`8GXi&6>=#K3 zdWYRZxnO^#ZHP<;-5Uh*D4fV>jSEV@w-09s>g1$g(5m+30_=B- z8N#vx)&v~Ew~hI{asDmistMyN&bWp*t_f%TEsM|=%n+n%CmGdtu;~FZbMs$YdA{iv zP4BMaOm)1eF08y=TskuUQr(5R$ceGMaTQm*kuTmD&W@IJvSLscMiXZ}I>OEM0J$6| zpW_TGqMDo;LaNf=&Z(M_5XMcC(Whq!^ee&}-qbgi@~6Iv@%vJVN&h_oV1|(B^nX(l zDm|-j9Nh^30;gQhE7!Bi^)W+9biwLqVOg};9L+aJYZgSytD^ZOGucX``aS{Rdj!Jr zzvUBo6?|^}XlUHUVfBopNVP1gUB{+} z8Hr4n7d5P9(?ir?VGWkZf=DT6sOJs!EDjHEn@WM{y1XQa8ghqZGg3%<`0(M33?o~| z8q$J@jICVCRxINT%X!0c7Kfiim!;C}Aa9pdj&zNxZphireQen?T-h^x*)v1x+lD+g zey6FqzJqMV zA+F*OUvX$?`|aE!ws@<^ST=VXgon2aEYNSLk_wa6E4FhL+xd#^L-+$Mfs(+J*=I(B zh{CUb(-Vdwmw177)0F1`s4}D|H-|N59Kfdw(=W#-`q9)FtZ8x4|i?>1G za$EV_RyMaanp-rq`C(K8GeoMYeRsiV*Y#Ykay?(U9w)@IbB5LoK7aQ4NX|qy5JKu| zKa|Fg(auA+Ivo?84z82tJ84Md7CL!=wNAd)iHkEM@i zumw-O?jPQMY3GHVTubI&*JXjS&$ z<=Z3q%f9{W0gBs4^ZV%3bTR(umAmtIrr&jrFZgQF)kWNz9sHUd-10Vlc^lWz&NsBP zdRW#wIpr>1xrNDy=?OAPD8-*9;8Rbwur!3@^g{-ph zjv+66BAo5}2=>pf8?v@qh_6@WY|;?ls7Z)_qfu`vlzx*ZLGc0-#S5h!XG>8$M-H{VE09C2?@AH2?=0$V`O@zeRBkg$ZyMzg@35)_INh=C z-zKPSHc{%d*#xc4_FTa3j^}6`ZoAJ1l0uJ{VhqqtI5#(U?%mxX>hT5Tmc1PGY8HrpDD+<7UDUoG71-k`G z&=a{e18vdll7Vf}Y|Fs58QFePUN}QM8j8)KM&6T56Ce>V#~}_tRkPJp!<;)3`qd@64ks-K<%mYQ3jgaTVqA`*t@W=@T zaN)cWW^lUsyl(!$mZ-)UUNL0lG^MFn zUDMpbVNF>Phw{xQ`fJVqUcY<&=hFk1VplK2>*Pys!yOZ9O}#wSuz1O` z+R2fD+2Yp8@Gn=cSXCSQ;?t?O-kKcv_0-R*7$R6=)O(N2E^7uN66l2JVHgVlTZM`{I+y!7*!RpcBXt zsTW(idu=#&_Lopjy#I@;yn9zKPhPw<_0{Rg-@kKr{AF<_!R!&o84AZvp99UJ*r147 zkZ)<*+1%NE;SKamH{BB>Uo^${>zSs26PixN9|QH|)yuVu>l>bEs9%iL5O^xYTmS6L z-SKxPBiEq0#r3Ni#WBV&s~g1}kwBcZqA~W-^~sB0-u>$0z1QD>ogV3$1!=Df4uxL( z3An}j={`IitV%(bFwbMH0NgRr8Ob13te|&!oHPT9&{a`G6bC_I9dQA&=&?e5Zj@@;m<@qx9r+0t^}BB zpkaxPjBbUL;V)w!4g97kE;)PNRP$>+|rWvUm3M{uECkhd$)c9-y*1imk zVa@`8-6X!ov9cMozMy@kL{UJ{*wG;aYG_DfkRoD>Bh?9E)~#X*#7@gj*ZfZ$E6&se zRZMQG>%b8iw1^)n!(b0($p{i<<~)Nn=`sOH4AvY$QX_sS$x(b1Gyej{@7Yv>_$gfHD@GG>8n-zB;IWVR@vvx=#v zcN`?_OlMf#29GXkPaGqk01*Z>i6`1SJ@xzAkpd{(@ zbpS>yqzmao*&#zvAIuKrgp5H+1Q~ATJkTpOrJyvIq0R@hqZ~=$XtGK=?;mf}LniL;2yFuna_~>{H560gPioFvCrQ(jaVs4!R<%%==Uh)P>0vvh#qW`xU0 zoQ7cGk^Drzfu5Z7`3IOxb1==iqoi|vIn4P-uA@z9bQC6911GoC)irA&zml8I6NDFf zO$PJRXKSbkOzZGTFAeKXQ$;ioKG|`J=#+z8bV#%wHUB@bNQhL2f+MqzNkL;U2WA4C zB1!kWU{U&es5n^sXg3ck!_g#y^s#duSP_c!-bT!G3PD2pxIB(a08pl2js>NURkkEp z5>$%rBcak@36&RAqM3ITlM*TmmIllE73iKHGDCkYKTTn%oXQWD2g~Nnv{^1dZt3g< zA*cd8-w`)JwobaVp;8D3RT#2B+#*sGH21@(geq_w6(~Pgjz^;6QIt`cV^$f=qKe}u z!u+(ET$!j3J6OsNnMlnXhm5&U<@tBOToT8;D-Cm10>|TF4pt(51S=!+67^DlsUrMQ zr0=8M5~^s9dsS7iiYkLOrG5!6dr5wWsyrc})f}hw*~gzJk;milyx=^lT%1wC%7kkY z9<}Pn!cjKhRa`MyB z$y4WIXWkSqLQ|iA3bzU{l)~yFrjalN7IOz80_jCApolVYB@h%2JC?oEv<$HP31oG8 z@FxJTOVX7yK_UVRa_*BUz(ueamRT-+NF#{EInQ8*DCkE*X#5$>*rb0fFY{=J;gEDN zn4NO0@DWr2kMFqvO@oe24VNCjt2^Kg_!v7hCP?A>1N_$drd>`_B-}%3JCn`Ic)$yQmEp`7NsTI)HBY;-`BlejYT%Bmdqs z{eR%WY60d|C;}5yO1XOL8Ki`b-=P24#Jq;GYM3qX!)!%h8v^J8FY**~T80}%GegkO zy?y?ZHisP?vY6c{MFAFISS*a7=&}27WFF;mI>CvC@rfpJUp*wpy=jswY818cYXs1} zRgn97?0vML00#61;HL8%lyeG!Pf&IiAVBL7^;}SRytZxz?9)nnzz^OrU~E=6!C>ug zSdmmQf(kZ1bL606@9jgDWj~NDE_X-+R{&-Vp;Lo1nb!v%`bTMSys?9a2)G1*saTL> zq}KV&VMNCp2%y!8M-k69W(2)nMBrBlTtXm%01{jn+(j3}1ahN5ZWG9-1+qoR@-nV& zTTlG=IIQ*F?zr+C#atfdLsUhE9D$fmQB(?(K~RI=6J^8d96<^i0`xx%P%8AHtG|!~ zS^&D6L2d(}I_(B;A@E~Do;4Y1*#XZ{kN3D|wvbNLEF=(Ti|QhO52&Pj;Li+#-t9e( z)D2D-cu0ZOUoRFA8AMaezaaXuuaAo634SulThCk!milmd~pVZ;2}NZYc{Vl!e1!XR`20 z3#+uie;~9uXV(tbOlZI;s#i75NTj+vu(iIXzGbMHFw}5{THa90;_#DpUZq0Ubi2Sj z92#|wQCz`VzF_T;ENTJ+ztp&s{7^9_V^?hFYIpFpJH(`D?mjj>+{rDxmAh~vcOhHX z1mR`_Ti4FjwetXTJNVoVHn-zpf{rF+O%wsP4mXZhu=$M;IKz70u%4Y0u&EVn+)c>T znnmQ7ytaQ(9VWvKVf&CVT2wsTJd_pAf}Gr<^Nqh~j2I)$k>^G$MAw5`rj-+>m7Hlc zZ(4oJv~I$*Zu}|Ew23!u0;dFH?P%feTE_L=nRfA}U18;&qVmywHUXrGF8tX) zZhY_HjjC^~Y-bm1+0R+_^OpTk%*H9pV!E8O2M29uZD2Q3HDD0U&O7@L!_JB9dGMbY zT?3j~hsLt5)7N_0yf!Ycjn8WX>Ws@Vh*s1`%0~TM#WKEP+0b@mLKib|?tN75``N2_~7U=u52-1ws=TAqe)PSNxhGV{(hG7WFInO7=Kz#v4DZdUh?u0 zUSA%`AI)bii#dG*uWuMz!08)#eIu|_Z#;MO%+cYBLC;ywEq&#LzB1Afv2*%bUSE4l zzi2|gNUU;8zkEW!oYSx5^((_tw9;T96sj!IqVh|{7mD8~9VrcOpHUN)b!>gpL`4&; zHAgEJj{4Z;TewAACMvd~1j`b3`Bu)dl{Ie*?*trWmP@S{THn|)vLoCMuMNw>?QC%! zr+osPWg1reRx{qtZa=`UJIFrk;MO@<+R65~xWz~K#YYkap;mcicxTkSaMU|tZoE;> zmF)_5MC%rX+lT8pt@Tb$dAJ2MgvDi-$}W^e7_MjmU$h{szoW??uA9)*pgydMm00+a zg`=B4ZhgP?y&YF}a3!nwl2zfha9dPYINUy=v(6BCs#QQ#X~oFus3GsX>KCe4HLq!g zJ)EH~YB1k2R8JVHBZo&TM-PuxaE4{PVHv0pYxF=~ZvKoyT3T?Q0PsBm;Wl7XY3-_Bp;fJdip7>BW25HUTjnJb<|SiGIP(hLyaL##BrJ=@OcUkH z#?3dYCRVp|<_=&)iRIGj3#;E)JF<4P{FZh7gmpb<-N0Kna3vd1BY?NGVzz{tD_M!t zQnpQ*D(`=^ZDtY1+@S+(b|tH<^!Z?S|3jgAy-z1eMLo79;VsQQfjq z`gKW8%M$6=OXLv$Zc$xpzVv2^cH2_<&G}jsUn<+SQgw5s6vfw&D84ou;(yGSwH7P> zSVjWmYRtRk#DrYEJ4ZTUBmu%77`Ir}j2BvG`Ve>y0R{o2GlE|$F!@Q{hS?5b79)EH=(v)~6e-k*akc$PGyUcU~TFB6Gy2RVjn)rMUZ;}z3>M_iKkD{VE@5m;3(+zh{Bi#p+f>D{9O{e`3C$hw+}lMGOZ{N zQ7L3~iUQs#lyVvY)Ca-T0sk_zwX`+w*<)*O-nFZ>W9uHF7$$)<-pgo#ar z6m0WZA}zog$~FyXX0lu)30@YDhWBy$_o=!2q* zzzvHmi{!ln-fa-MDgP7HNhGJl)OHC8)9KOhK2HBWHFvw5D0w29Yk>`_0DO%AK&NN4 WM9HFPZWSb&QKA_Vi^YKF;{OFX!)izX literal 0 HcmV?d00001 diff --git a/sdmatte_nodes.py b/sdmatte_nodes.py index 64bfe70..4b2e66e 100644 --- a/sdmatte_nodes.py +++ b/sdmatte_nodes.py @@ -3,144 +3,228 @@ import torch import torch.nn.functional as F import numpy as np from torchvision import transforms - import folder_paths import comfy from comfy.utils import load_torch_file - -# 延迟导入核心模型，避免在依赖未安装时阻断节点注册 +def download_with_progress(url, filepath, description="File"): + try: + import requests + from tqdm import tqdm + + print(f"[SDMatte] Starting download: {description}") + print(f"[SDMatte] URL: {url}") + print(f"[SDMatte] Save path: {filepath}") + + response = requests.get(url, stream=True, timeout=30) + response.raise_for_status() + + total_size = int(response.headers.get('content-length', 0)) + + with open(filepath, 'wb') as f: + with tqdm( + total=total_size, + unit='iB', + unit_scale=True, + unit_divisor=1024, + desc=f"Downloading {description}" + ) as pbar: + for chunk in response.iter_content(chunk_size=8192): + size = f.write(chunk) + pbar.update(size) + + print(f"[SDMatte] Download completed: {filepath}") + return True + + except Exception as e: + print(f"[SDMatte] Download failed: {e}") + return False +def verify_file(filepath): + if not os.path.exists(filepath): + return False + + actual_size = os.path.getsize(filepath) + if actual_size > 0: + print(f"[SDMatte] File verification successful, size: {actual_size} bytes") + return True + + return False SDMatteCore = None - - -class SDMatteModelLoader: - - @classmethod - def INPUT_TYPES(s): - return { - "required": { - "ckpt_name": (folder_paths.get_filename_list("checkpoints"), ), - } - } - - RETURN_TYPES = ("SDMATTE_MODEL",) - FUNCTION = "load_model" - CATEGORY = "Matting/SDMatte" - - def load_model(self, ckpt_name): - device = comfy.model_management.get_torch_device() - - global SDMatteCore - if SDMatteCore is None: - from .src.modeling.SDMatte.meta_arch import SDMatte as SDMatteCore # type: ignore - - # 使用本地固定路径作为 diffusers 基底 - base_dir = os.path.dirname(__file__) - pretrained_repo = os.path.join(base_dir, "src", "SDMatte") - required_subdirs = ["text_encoder", "vae", "unet", "scheduler", "tokenizer"] - missing = [d for d in required_subdirs if not os.path.isdir(os.path.join(pretrained_repo, d))] - if missing: - raise FileNotFoundError( - f"本地基底缺少目录: {missing}. 期望路径: {pretrained_repo}，需包含 {required_subdirs} 子目录。" - ) - - print(f"[SDMatte] Using base repo: {pretrained_repo}") - # 实例化模型（使用默认优化配置） - model = SDMatteCore( - pretrained_model_name_or_path=pretrained_repo, - load_weight=False, - use_aux_input=True, - aux_input="trimap", - use_encoder_hidden_states=True, - use_attention_mask=True, - add_noise=False, +SDMATTE_MODEL_LIST = [ + "SDMatte.pth", + "SDMatte_plus.pth" +] +SDMATTE_MODEL_MAPPING = { + "SDMatte.pth": { + "url": "https://huggingface.co/LongfeiHuang/SDMatte/resolve/main/SDMatte.pth", + "filename": "SDMatte.pth", + "description": "SDMatte Standard Version (12.2 GB)" + }, + "SDMatte_plus.pth": { + "url": "https://huggingface.co/LongfeiHuang/SDMatte/resolve/main/SDMatte_plus.pth", + "filename": "SDMatte_plus.pth", + "description": "SDMatte Plus Version (12.1 GB)" + } +} +_sdmatte_model_cache = { + 'model': None, + 'ckpt_name': None, +} +def clear_memory(): + if torch.cuda.is_available(): + torch.cuda.empty_cache() + torch.cuda.ipc_collect() +def get_model_path(ckpt_name): + sdmatte_dir = os.path.join(folder_paths.models_dir, "SDMatte") + os.makedirs(sdmatte_dir, exist_ok=True) + + local_model_path = os.path.join(sdmatte_dir, ckpt_name) + + if os.path.exists(local_model_path): + print(f"[SDMatte] Using local model from models/SDMatte/: {local_model_path}") + return local_model_path + + if ckpt_name not in SDMATTE_MODEL_MAPPING: + print(f"[SDMatte] User custom model path: {local_model_path}") + return local_model_path + + model_info = SDMATTE_MODEL_MAPPING[ckpt_name] + print(f"[SDMatte] Local model {ckpt_name} not found, preparing to download from HuggingFace...") + print(f"[SDMatte] Model description: {model_info['description']}") + + try: + success = download_with_progress( + url=model_info["url"], + filepath=local_model_path, + description=model_info["description"] ) - - # 加载合并后的权重 - ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name) - print(f"[SDMatte] Loading merged weights: {ckpt_path}") - # 更健壮的加载与 state_dict 提取 - state_root = None + + if success and verify_file(local_model_path): + print(f"[SDMatte] Model download and verification successful: {local_model_path}") + return local_model_path + else: + raise FileNotFoundError(f"Model download or verification failed: {local_model_path}") + + except Exception as e: + print(f"[SDMatte] Auto download failed: {e}") + if os.path.exists(local_model_path): + try: + os.remove(local_model_path) + print(f"[SDMatte] Cleaned up incomplete download file: {local_model_path}") + except: + pass + raise FileNotFoundError(f"Unable to get model {ckpt_name}: {e}") +def load_sdmatte_model(ckpt_name, force_cpu=False): + global _sdmatte_model_cache, SDMatteCore + + if (_sdmatte_model_cache['model'] is not None and + _sdmatte_model_cache['ckpt_name'] == ckpt_name): + print(f"[SDMatte] Reusing cached model: {ckpt_name}") + return _sdmatte_model_cache['model'] + + if _sdmatte_model_cache['model'] is not None: + print("[SDMatte] Cleaning up old model cache") + del _sdmatte_model_cache['model'] + clear_memory() + + device = comfy.model_management.get_torch_device() + if force_cpu: + device = torch.device('cpu') + if SDMatteCore is None: + from .src.modeling.SDMatte.meta_arch import SDMatte as SDMatteCore # type: ignore + base_dir = os.path.dirname(__file__) + pretrained_repo = os.path.join(base_dir, "src", "SDMatte") + required_subdirs = ["text_encoder", "vae", "unet", "scheduler", "tokenizer"] + missing = [d for d in required_subdirs if not os.path.isdir(os.path.join(pretrained_repo, d))] + if missing: + raise FileNotFoundError( + f"Local base missing directories: {missing}. Expected path: {pretrained_repo}, must contain {required_subdirs} subdirectories." + ) + print(f"[SDMatte] Using base repo: {pretrained_repo}") + model = SDMatteCore( + pretrained_model_name_or_path=pretrained_repo, + load_weight=False, + use_aux_input=True, + aux_input="trimap", + use_encoder_hidden_states=True, + use_attention_mask=True, + add_noise=False, + ) + ckpt_path = get_model_path(ckpt_name) + print(f"[SDMatte] Loading merged weights: {ckpt_path}") + state_root = None + try: + from torch.serialization import add_safe_globals try: - from torch.serialization import add_safe_globals - try: - from omegaconf.listconfig import ListConfig - add_safe_globals([ListConfig]) - except Exception: - pass - try: - from omegaconf.base import ContainerMetadata - add_safe_globals([ContainerMetadata]) - except Exception: - pass - - try: - state_root = torch.load(ckpt_path, map_location='cpu', weights_only=True) - print("[SDMatte] torch.load(weights_only=True) ok") - except Exception as e: - print(f"[SDMatte] weights_only=True failed: {e}; retry weights_only=False") - try: - state_root = torch.load(ckpt_path, map_location='cpu', weights_only=False) - print("[SDMatte] torch.load(weights_only=False) ok") - except TypeError: - state_root = torch.load(ckpt_path, map_location='cpu') + from omegaconf.listconfig import ListConfig + add_safe_globals([ListConfig]) + except Exception: + pass + try: + from omegaconf.base import ContainerMetadata + add_safe_globals([ContainerMetadata]) + except Exception: + pass + try: + state_root = torch.load(ckpt_path, map_location='cpu', weights_only=True) + print("[SDMatte] torch.load(weights_only=True) ok") except Exception as e: - print(f"[SDMatte] torch.load failed: {e}; fallback comfy.load_torch_file") - state_root = load_torch_file(ckpt_path) - - # 尝试从常见键中提取真正权重 - candidate_keys = [ - 'state_dict','model_state_dict','params','weights', - 'ema','model_ema','ema_state_dict','net','module','model','unet' - ] - state_dict = None - if isinstance(state_root, dict): - for k in candidate_keys: - inner = state_root.get(k) - if isinstance(inner, dict) and len(inner) > 50: - state_dict = inner - print(f"[SDMatte] extracted inner dict via key: {k}") - break - if state_dict is None and isinstance(state_root, dict): - # 如果顶层就是大量权重键 - if len(state_root) > 50 and any('.weight' in x or '.bias' in x for x in state_root.keys()): - state_dict = state_root - if state_dict is None: - top_keys = list(state_root.keys()) if isinstance(state_root, dict) else [] - raise RuntimeError( - f"无法从权重中提取 state_dict。顶层键: {top_keys[:10]} (共{len(top_keys)})。请使用官方脚本导出仅模型权重的 sdmatte.pth。" - ) - - print(f"[SDMatte] Weights loaded, keys: {len(list(state_dict.keys()))}") - missing, unexpected = model.load_state_dict(state_dict, strict=False) - print(f"[SDMatte] load_state_dict done. missing={len(missing)}, unexpected={len(unexpected)}") - - model.eval() - model.to(device) - - return (model,) - - + print(f"[SDMatte] weights_only=True failed: {e}; retry weights_only=False") + try: + state_root = torch.load(ckpt_path, map_location='cpu', weights_only=False) + print("[SDMatte] torch.load(weights_only=False) ok") + except TypeError: + state_root = torch.load(ckpt_path, map_location='cpu') + except Exception as e: + print(f"[SDMatte] torch.load failed: {e}; fallback comfy.load_torch_file") + state_root = load_torch_file(ckpt_path) + candidate_keys = [ + 'state_dict','model_state_dict','params','weights', + 'ema','model_ema','ema_state_dict','net','module','model','unet' + ] + state_dict = None + if isinstance(state_root, dict): + for k in candidate_keys: + inner = state_root.get(k) + if isinstance(inner, dict) and len(inner) > 50: + state_dict = inner + print(f"[SDMatte] extracted inner dict via key: {k}") + break + if state_dict is None and isinstance(state_root, dict): + if len(state_root) > 50 and any('.weight' in x or '.bias' in x for x in state_root.keys()): + state_dict = state_root + if state_dict is None: + top_keys = list(state_root.keys()) if isinstance(state_root, dict) else [] + raise RuntimeError( + f"Unable to extract state_dict from weights. Top-level keys: {top_keys[:10]} (total {len(top_keys)}). Please use official script to export model-only weights to sdmatte.pth." + ) + print(f"[SDMatte] Weights loaded, keys: {len(list(state_dict.keys()))}") + missing, unexpected = model.load_state_dict(state_dict, strict=False) + print(f"[SDMatte] load_state_dict done. missing={len(missing)}, unexpected={len(unexpected)}") + model.eval() + model.to(device) + + _sdmatte_model_cache['model'] = model + _sdmatte_model_cache['ckpt_name'] = ckpt_name + + return model def _resize_norm_image_bchw(image_bchw: torch.Tensor, size_hw=(1024, 1024)) -> torch.Tensor: - # image_bchw in [0,1], float32 resize = transforms.Resize(size_hw, antialias=True) norm = transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) x = resize(image_bchw) x = norm(x) return x - - def _resize_mask_b1hw(mask_b1hw: torch.Tensor, size_hw=(1024, 1024)) -> torch.Tensor: resize = transforms.Resize(size_hw) return resize(mask_b1hw) - - class SDMatteApply: - @classmethod def INPUT_TYPES(s): return { "required": { - "sdmatte_model": ("SDMATTE_MODEL", {"tooltip": "已加载的SDMatte模型，来自SDMatte Model Loader节点"}), + "ckpt_name": (SDMATTE_MODEL_LIST, { + "default": "SDMatte.pth", + "tooltip": "选择SDMatte模型：SDMatte.pth(标准版) 或 SDMatte_plus.pth(增强版)。模型将存储在models/SDMatte/目录，如果本地不存在将自动下载" + }), "image": ("IMAGE", {"tooltip": "需要进行抠图的输入图像"}), "trimap": ("MASK", {"tooltip": "三值图掩码：白色=前景，黑色=背景，灰色=未知区域"}), "inference_size": ([512, 640, 768, 896, 1024], { @@ -152,7 +236,7 @@ class SDMatteApply: "tooltip": "输入图像是否包含透明通道。如果原图有透明背景请启用" }), "output_mode": (["alpha_only", "matted_rgba", "matted_rgb"], { - "default": "alpha_only", + "default": "matted_rgba", "tooltip": "输出模式：alpha_only=只输出遮罩；matted_rgba=透明背景抠图；matted_rgb=黑色背景抠图(推荐，避免干扰)" }), "mask_refine": ("BOOLEAN", { @@ -163,6 +247,10 @@ class SDMatteApply: "default": 0.8, "min": 0.1, "max": 1.0, "step": 0.1, "tooltip": "trimap约束强度(0.1-1.0)。越高约束越严格，0.8=平衡，0.9=严格过滤，0.6=宽松保留" }), + "unload_model_after": ("BOOLEAN", { + "default": False, + "tooltip": "推理完成后卸载模型释放显存。适用于显存紧张的情况，但会影响后续推理速度" + }), }, "optional": { "force_cpu": ("BOOLEAN", { @@ -171,141 +259,106 @@ class SDMatteApply: }), }, } - RETURN_TYPES = ("MASK", "IMAGE") RETURN_NAMES = ("alpha_mask", "matted_image") FUNCTION = "apply_matte" CATEGORY = "Matting/SDMatte" - - def apply_matte(self, sdmatte_model, image, trimap, inference_size, is_transparent, output_mode, mask_refine, trimap_constraint, force_cpu=False): + def apply_matte(self, ckpt_name, image, trimap, inference_size, is_transparent, output_mode, mask_refine, trimap_constraint, unload_model_after=False, force_cpu=False): + sdmatte_model = load_sdmatte_model(ckpt_name, force_cpu) + device = comfy.model_management.get_torch_device() if force_cpu: device = torch.device('cpu') sdmatte_model.to(device) - - # 自动显存优化（内置默认设置） if device.type == 'cuda': try: - torch.cuda.empty_cache() # 预清理显存 + torch.cuda.empty_cache() except Exception: pass - # 启用注意力切片以节省显存 try: unet = getattr(sdmatte_model, 'unet', None) if unet is not None and hasattr(unet, 'set_attn_processor'): from diffusers.models.attention_processor import SlicedAttnProcessor - # 使用适中的切片大小平衡性能和显存 unet.set_attn_processor(SlicedAttnProcessor(slice_size=1)) except Exception as e: - print(f'[SDMatte] 注意力优化跳过: {e}') - + print(f'[SDMatte] Attention optimization skipped: {e}') B, H, W, C = image.shape orig_h, orig_w = H, W - - # IMAGE: (B,H,W,C)->(B,C,H,W) in [0,1] img_bchw = image.permute(0, 3, 1, 2).contiguous().to(device) img_in = _resize_norm_image_bchw(img_bchw, (int(inference_size), int(inference_size))) - - # 构造必要的数据字典字段 is_trans = torch.tensor([1 if is_transparent else 0] * B, device=device) - data = {"image": img_in, "is_trans": is_trans, "caption": [""] * B} - def to_b1hw(x): return _resize_mask_b1hw(x.unsqueeze(1).contiguous().to(device), (int(inference_size), int(inference_size))) - - # 处理 trimap 输入 tri = to_b1hw(trimap) * 2 - 1 data["trimap"] = tri data["trimap_coords"] = torch.tensor([[0,0,1,1]]*B, dtype=tri.dtype, device=device) - with torch.no_grad(): if device.type == 'cuda': with torch.autocast(device_type='cuda', dtype=torch.float16): pred_alpha = sdmatte_model(data) else: pred_alpha = sdmatte_model(data) - - # 还原到原尺寸 out = transforms.Resize((orig_h, orig_w))(pred_alpha) out = out.squeeze(1).clamp(0, 1).detach().cpu() - # 遮罩优化：使用trimap约束来过滤不需要的区域 if mask_refine: - trimap_cpu = trimap.cpu() # (B, H, W) + trimap_cpu = trimap.cpu() - # 创建严格的前景约束 - # 只在trimap明确标记为前景(白色区域)的地方保留高alpha值 - # 在未知区域(灰色)允许模型输出，但在背景区域(黑色)强制为0 - foreground_regions = trimap_cpu > trimap_constraint # 前景区域 - background_regions = trimap_cpu < (1.0 - trimap_constraint) # 背景区域 - unknown_regions = ~(foreground_regions | background_regions) # 未知区域 + foreground_regions = trimap_cpu > trimap_constraint + background_regions = trimap_cpu < (1.0 - trimap_constraint) + unknown_regions = ~(foreground_regions | background_regions) - # 优化alpha遮罩 refined_alpha = out.clone() - # 背景区域强制为0 refined_alpha[background_regions] = 0.0 - # 前景区域保持原值或增强 refined_alpha[foreground_regions] = torch.clamp(refined_alpha[foreground_regions] * 1.2, 0, 1) - # 未知区域保持模型输出，但应用阈值过滤 - alpha_threshold = 0.3 # 过滤掉低置信度的区域 + alpha_threshold = 0.3 low_confidence = (refined_alpha < alpha_threshold) & unknown_regions refined_alpha[low_confidence] = 0.0 out = refined_alpha - # 根据输出模式创建不同的抠图结果 - # out: (B, H, W), image: (B, H, W, C) - alpha_expanded = out.unsqueeze(-1) # (B, H, W, 1) + alpha_expanded = out.unsqueeze(-1) if output_mode == "alpha_only": - # 只输出alpha遮罩，图像输出为黑色 matted_image = torch.zeros_like(image.cpu()) elif output_mode == "matted_rgba": - # 创建RGBA图像：RGB通道保持原图，A通道为alpha遮罩 - # 背景变为透明，只保留前景对象 matted_image = torch.cat([ - image.cpu() * alpha_expanded, # RGB通道应用遮罩 - alpha_expanded.expand(-1, -1, -1, 1) # A通道为遮罩本身 + image.cpu() * alpha_expanded, + alpha_expanded.expand(-1, -1, -1, 1) ], dim=-1) - # 确保输出仍为3通道（ComfyUI IMAGE格式） matted_image = matted_image[..., :3] elif output_mode == "matted_rgb": - # 只保留前景对象，背景变为黑色 - # 使用trimap的前景信息来更精确地提取对象 - trimap_cpu = trimap.cpu() # (B, H, W) - trimap_expanded = trimap_cpu.unsqueeze(-1) # (B, H, W, 1) + trimap_cpu = trimap.cpu() + trimap_expanded = trimap_cpu.unsqueeze(-1) - # 只在trimap标记为前景(>0.8)或未知区域(0.2-0.8)的地方保留原图 - # 结合alpha遮罩进一步细化 foreground_mask = (trimap_expanded > 0.2) & (alpha_expanded > 0.1) matted_image = image.cpu() * foreground_mask.float() else: - # 默认：直接应用alpha遮罩 matted_image = image.cpu() * alpha_expanded - # 推理后清理显存 if device.type == 'cuda': try: torch.cuda.empty_cache() except Exception: pass - + + if unload_model_after: + global _sdmatte_model_cache + print("[SDMatte] Unloading model to free GPU memory") + if _sdmatte_model_cache['model'] is not None: + del _sdmatte_model_cache['model'] + _sdmatte_model_cache['model'] = None + _sdmatte_model_cache['ckpt_name'] = None + clear_memory() return (out, matted_image) - - NODE_CLASS_MAPPINGS = { - "SDMatteModelLoader": SDMatteModelLoader, "SDMatteApply": SDMatteApply, } - NODE_DISPLAY_NAME_MAPPINGS = { - "SDMatteModelLoader": "Load SDMatte Model", "SDMatteApply": "Apply SDMatte", } - -