From 7b5663a35ade8b0ef3c94e30f9fa853125930fcd Mon Sep 17 00:00:00 2001 From: Justin John <34035011+justinjohn0306@users.noreply.github.com> Date: Fri, 13 Oct 2023 08:30:38 +0530 Subject: [PATCH] Add files via upload --- .../__pycache__/animate.cpython-38.pyc | Bin 0 -> 6919 bytes animation/SadTalker/src/facerender/animate.py | 257 ++++++++ .../__pycache__/dense_motion.cpython-38.pyc | Bin 0 -> 3933 bytes .../__pycache__/generator.cpython-38.pyc | Bin 0 -> 6601 bytes .../keypoint_detector.cpython-38.pyc | Bin 0 -> 4843 bytes .../__pycache__/make_animation.cpython-38.pyc | Bin 0 -> 4766 bytes .../__pycache__/mapping.cpython-38.pyc | Bin 0 -> 1700 bytes .../modules/__pycache__/util.cpython-38.pyc | Bin 0 -> 17173 bytes .../src/facerender/modules/dense_motion.py | 121 ++++ .../src/facerender/modules/discriminator.py | 90 +++ .../src/facerender/modules/generator.py | 255 ++++++++ .../facerender/modules/keypoint_detector.py | 179 ++++++ .../src/facerender/modules/make_animation.py | 170 ++++++ .../src/facerender/modules/mapping.py | 47 ++ .../SadTalker/src/facerender/modules/util.py | 564 ++++++++++++++++++ .../src/facerender/sync_batchnorm/__init__.py | 12 + .../__pycache__/__init__.cpython-38.pyc | Bin 0 -> 413 bytes .../__pycache__/batchnorm.cpython-38.pyc | Bin 0 -> 12945 bytes .../__pycache__/comm.cpython-38.pyc | Bin 0 -> 4849 bytes .../__pycache__/replicate.cpython-38.pyc | Bin 0 -> 3498 bytes .../facerender/sync_batchnorm/batchnorm.py | 315 ++++++++++ .../src/facerender/sync_batchnorm/comm.py | 137 +++++ .../facerender/sync_batchnorm/replicate.py | 94 +++ .../src/facerender/sync_batchnorm/unittest.py | 29 + 24 files changed, 2270 insertions(+) create mode 100644 animation/SadTalker/src/facerender/__pycache__/animate.cpython-38.pyc create mode 100644 animation/SadTalker/src/facerender/animate.py create mode 100644 animation/SadTalker/src/facerender/modules/__pycache__/dense_motion.cpython-38.pyc create mode 100644 animation/SadTalker/src/facerender/modules/__pycache__/generator.cpython-38.pyc create mode 100644 animation/SadTalker/src/facerender/modules/__pycache__/keypoint_detector.cpython-38.pyc create mode 100644 animation/SadTalker/src/facerender/modules/__pycache__/make_animation.cpython-38.pyc create mode 100644 animation/SadTalker/src/facerender/modules/__pycache__/mapping.cpython-38.pyc create mode 100644 animation/SadTalker/src/facerender/modules/__pycache__/util.cpython-38.pyc create mode 100644 animation/SadTalker/src/facerender/modules/dense_motion.py create mode 100644 animation/SadTalker/src/facerender/modules/discriminator.py create mode 100644 animation/SadTalker/src/facerender/modules/generator.py create mode 100644 animation/SadTalker/src/facerender/modules/keypoint_detector.py create mode 100644 animation/SadTalker/src/facerender/modules/make_animation.py create mode 100644 animation/SadTalker/src/facerender/modules/mapping.py create mode 100644 animation/SadTalker/src/facerender/modules/util.py create mode 100644 animation/SadTalker/src/facerender/sync_batchnorm/__init__.py create mode 100644 animation/SadTalker/src/facerender/sync_batchnorm/__pycache__/__init__.cpython-38.pyc create mode 100644 animation/SadTalker/src/facerender/sync_batchnorm/__pycache__/batchnorm.cpython-38.pyc create mode 100644 animation/SadTalker/src/facerender/sync_batchnorm/__pycache__/comm.cpython-38.pyc create mode 100644 animation/SadTalker/src/facerender/sync_batchnorm/__pycache__/replicate.cpython-38.pyc create mode 100644 animation/SadTalker/src/facerender/sync_batchnorm/batchnorm.py create mode 100644 animation/SadTalker/src/facerender/sync_batchnorm/comm.py create mode 100644 animation/SadTalker/src/facerender/sync_batchnorm/replicate.py create mode 100644 animation/SadTalker/src/facerender/sync_batchnorm/unittest.py diff --git a/animation/SadTalker/src/facerender/__pycache__/animate.cpython-38.pyc b/animation/SadTalker/src/facerender/__pycache__/animate.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4efa6fc968612c2e0f6525c9606bff4b54eab485 GIT binary patch literal 6919 zcma)B$&cLFc_&#c7HjS5?&?k5l1fW!R1>H(nMn}DktdS1jN_%p8Obvfq79m=$LeM` zm(n9kEj3t15KUqQi0ttx2gg7uw+ZqG<%05Jdk$FhD%{eUC*}X+y#S zAK&tQ-`l>uzTa%t6g6*Yr3EmMwFEz>PeYOM53osQJ4x*Ta(4f$*teb79^HGCxWe%E!-%^*P1L`6>sFk2fs;gB)V}?_n{MLvQ3p!xyUuXvh5Ng1 z0>HHo`+a{Ddr^4(*x~L6Zs>A}wDsb^o$p=0d8-%&nO5K&xwaG19KfcvTpuwny5sH# zZb+mz-SEH(`!2WlbE@{SmmJu>8;WKz=zDPj?CQ{o6W1PkeKa;==YeZK@R%EAD929X z^1$IocW?J;@AFR;jca)0t0-KBxB?uZGVL?PDzODtX2xeKc%#iK%tWos7Fm_mP&3#P zYqAz<6;{HzE{pbcxgqy9kAfSKI~ep!AeukKsWjUSyzxJS9}^{2nYy7oR)5+amlAcN zrQocw{)zH?%AY77DW57yX`rM!(>4@V+Q@7AhVnPsPt^gH*{M;~Dn-pKYUPRfOu4E2pf;{DGp*Xy`~L)tm-~L5R$l4*&(wUS z?@OG4a#a0X<=Md#nk+qo457AA)Y?UDaaz*= zz04Hc`@*iZFmS-BZpgr{kDS1_8bQQd-yS;L31X`~W8thvH2YByL}3oFbTqWtRLJta z$%Evqx&F3E>`Y|pSyzJ8;i z)Qu0KkhiyKcC&2k|>9y_f+UTBZVbCs`Br6%D4iJQj%Jx zyQhs!Xc{#&j||>TtB_W60-XpetX);^cT;UrPAgJxl_dUOzZPGY_&+T0%TxRxrlkTO z9suV5~r8Y2k<%OTr3hCPs3h!d8C(U`@+L7FF(wvvfot`1; zGbWBj8BEwe5LL;h1J8Fw*-P9Y78ULeeQ>^5kVA0-3uiO0pY&>+E?X?jBs=3u$;wL@ z%D9l2lrIuw{o%~Q3iU{+4qLI~H1Q>7k&Ln~jpZTKd!m3@YR_OcpK9a!lwsZm3zX72nMuh8*)?9n#ZWH{ zr|g9p!R3AoGe^BJqH`?smnTbL1L}t%rGB_loy?U=bY?t8Pv%aghNSYh#GZhL>6BC| zviILXr6T*yg38;vl9Z2X{NGe%+FBO}ky_oK$AIc+$ZSmA$zy@_^hO1nzE zFpBN8D*Ivl(#lWqg>^F4ho6$qUo36!`YM@mrmB&r6snyThn|&@=cBm`muyl2T#wcd1<- zau`z4_-GVjPnnVEHKj!H_W+Qc>>p6AI9~oWYW^V=7479i$HAU^ z^TFoMf!^oQ@Hu=#rB9B0{*YQEeO1C6{|Y6XPPp|GB~cIIP>qba&2*vMY}E~ zTRGGvYzw=r$!IE@f@Gj6c5g#HJNLZcPX zZ$DEei_Z|;f>L82=6AA`G!K_SZ8=?>bW+q*<#YA^HFG(A5kxobV*IQ;uXE|S!*l5-JH<9Bs7cOajhpEj z+hV6FOiC`WGpY7-g`G`Tey;pNCu_sHuuK=9DG0ALP9YcV`W$9BU##pr_H!Y*;e{{@@bYn;tC_HpSoXLB)M zb+S6SoNg2E$l7XA+oW2$4&8NWS2W^i#QUy|jK#hy+S69-1`bR`Kekp9hwmfyDSB*5 zO^E6v2T>M!?~D4-OZo@XMvX_lKdn~>h!R~p90hlKNG6P81bz>PJ@xrlc%Ey}DjX_D zc}_*2iu|D=a+1imlh$M(QR@#bkI7Aff`6hU+7smo;_D5iLtiMRO|bnyd#3$NyRTzo z8s7!3B#-d1)PL}`!A?r@6TXKLg;zI#HqEHxxlcqgQNihP#xfy|* z97N1&+}pQD!`KOizRP<(p(l@qZsvW|Z^Qja?#h6b)1ERzN2cYqK_i_0_`n$=2$+80 zkAWApa44#XSi_iNa-lN^toU0XBFbPXgn)%g6#Ep$IOH zc*0w>p=$(O_M>B$b265kynWuFUT|CxJ#c%Mv(QC-XtlDNmw7zqabY5{C$Csk>9i5m z4kBUZNkFlD&*>lS^Jo;Z{2GwAMBqH+Q6JGWze{u7rs4w>R{LdH<=>>SKc?baRD7GD z7D0fVVdN*{ESo6Ga$OJ&_7IVn)C43hkPZF`;e~`(aUNuAt#daSaZ2(aD_~(LXIA1a zs}&;uGGks*gZs@BoBU3z4D}K+@GL2Yn0Csh1IH0(zaw<|SPL`xknD*x6qt0H_@qg# zS?%d^i<7gq3IMr9q+Hyv$h~HAeus8T!J#m7>G=g(@!z3Pv^HGU zHnIrxGf`X8^q0TLCm?cEZ+!f4D>HiThmakOLzJd!KpRMTf8o3k?*3X@9yqO>FHHb zlTX$xQ$4E!hTp_S`7IPZlkzNdrM4|cvK8f z2oqWQ5ZQYq5aJ_wxl#x4D>S$`mwDL95)UoNHPUZ679SEN89CJ@zyptR|JG%spmt^{ zsGaNx;+-S+k<3gMxy>xWR|Xa9%o}j>NPx0NR`(4+)7wIZj^Z4dK(w~s0Q=%U4ZKQp z%fm-(bWbWUYmM+miRZ_13xycW6qL0*11bJzkV?#2oQ+6P^CGG3METggH}WjggWyn2 x-pt%B%S+GX=W2FZ@6(OO6y%uT{_IE6W2!Ys0e&TtD7@fwDO=!e>Nmj8{|h9MMTYf5$@)DIIq=ethHoWSOhVmao`OiAumUX97~WKAX*raeIXVgxNOcYXNTmt z$=Q`Q93X%c7ziLB=Ny8Wqb~U|I`%EdX>-dpNT7U0a#3VU%v}KHWT+`Z>-ej5dQ*P)*4A5K~sGJqTOX&T2tr3x6H| z{I3uND|pGCX$x%$rxkk1OGEGtR@x*7tkg?QXfKSC`9@)mt&)w6v0e9}QxRsRb1I!l zqlEQBNAFhQqK0#Np`q3;9rGA{HlL`Vq^FTZFQ&t8Ieq*iJt=s%AN)S zu-O}O*@yZoK*b71qd1L*eP!meX(S~LR*vt-Xv6Y)YGq+KnN?e|$nJybu%(PVnv5iQN_I%J zNqk7l^_xBW;eL=}GpdUP_Je8Ui)dQx%lE0%l_)Rbgd*b2f|L6s3F7p$x|7RrXB32yj8YNF zouq3bg|W1#4XCDG zI^k_>ke<-EwBBcPSD2-CV98hLV_pSy9kk@D!dzIG8FiHGTC?O!ZQd^13rAQ>epXhk ztuw1)_N6^({SO@P7?p=ps&r9RQ!OBHOD4M^)MTbek1dtodY&f6)pu~|B$jrc`If%vtd zvxrNm1YbDYRn+eNVet2d|2Ter7uQ$^hgHNubB`h2@ec!ucm@^IWm);|gzhH!Lx`0F zO=|cNgffpKndMb@%b!py`F7Yk(5Ethw#=`>C`H#@cM6B>a0WA`-mxAG4dge zx)53_kULx-1 zh_{7SmkJ&JsoDzj8AC)cus-=+U+Q3J2Il9bTd_V{IAy!hbc)v4Ejwi!p5)c7114Uy z;q&I!L5y}Wx>I(sUJs>i*+Z#c=;L15M;1jlRln?)En(haWq;8pMu5JxSS$M|t)X-Y zJ*@*)Qj3n>?qVJJl_4Cgi98fH0^K3~0UKYgX5{lr<)x}mr56UQKnvDrV+}SQ{DsTy zve(?B>-9Za_8Pl!moANY59_`Mdk17M?6_QB#`q4#6IXP%8QB%hTN_+o( z`R`YccE7w^!SbX6AjRs50ioJ;AXa&->YmF}hF5CNpy=VaRAQQwAl=@-u3s zD3!k@LAb85OvEI~@sdoyMU;(z%O@4am0Ka-pTq}|GLQC=x*!dVWHKuvrRTFm(glDq z@m`X}VmNqfNW6>GltURHS?6`n-7VhQMb435&^lI}PBB8ct1o~ot&H{>Gbx9ZzoVW! zl#lSf7*{sP@+gqvOXzZ9U9h0$B5+*ILLhaJY^~$ga&;Fq2dyq&*9Kt9bbck8Gvs)@YXE11d{{ ziE>ZNHEKRUp1O6Kr*1WQ>ej4?CsjsYCp&~r_2GrSIX)fzUHWFElwHzSle|hDt@;D5 fvgyo5sx8q*lyj$E=hwu@IVAucP@2OBT1Wpsz-JYn literal 0 HcmV?d00001 diff --git a/animation/SadTalker/src/facerender/modules/__pycache__/generator.cpython-38.pyc b/animation/SadTalker/src/facerender/modules/__pycache__/generator.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..074112649e9ef51a9f36223aa1adeb27aec41931 GIT binary patch literal 6601 zcmeHL&2t<_74M#zot>Ti&`Mr^#z_>%32aD&V-pC$Byp_RBo!;E*bcBLHH^1M(n_z-N55wmbmD09G6ZWI)U6`aI>k}D^wuM`LU2NVaw@Ab@n*jNdvq6!X-G;h0K_jJF0 z{oe1r?wi%B%kXr!p9%i@JY#>M%IssI@;cu5mmq=(9?#ulVSmho9q`TCuGyLcrzjk7 zoSah!rzBi(+?+En_;uDSA3!3u(&WmBqNc6f_k;M2u-D!?FI4%u-wno<(#`(3SYGaJ zN8`%mg9}2{ufKO``Av)v1&MCQunWsU6bILONv9XB1j%+!ZnbI3*@x6aPkbIEU@b0K z%MiR}3ilpcB*gQ?joK^eQ~W9hP~~0api-zmfyOx=*#v-CkfigKn6>?>y?#uRST1m1^kront0>C zamEss@RaQuBV!MHm6iqDD#>$+xyDj6HBwjb-(M_2>n59g*FVIOBXx1Lr>TAH6|=^be4bZTWXq_q=coyz87)&iew#%pi*Q51yg zSP&^PAf5h(M5ka~t2`=)M@4QeigG6)AP^SqGF zGpY~|`hk>$S<3OePSi;}Pm-Cdh9GzKy4qho_%S)GBnMT6wT_GvRgiwPj)}URNZEeB zA4EdhXqP0qa8nq}h~-)k{(|KHKO>%(B)SNQ^S`36@Uk zhv5EQS)fvB&>u^OwMBj?Ys9W%AJ->xH5Zf>2jQCh7TP31gRD@mIW377-0k}j;f|_c zTrW3^Sj~~4%&cENM_+!8zN}KyWGUI4M;=gn--Hf_9EjEH zzPRCsu*w&@ekYn%S7X^;UGv+4#8wA#HPd;#x;_TkrT#~95wntV;_rblqh>hV;l^i{ zVf@3jY-7Q&c!e+UWBdfK7!AYbwow7+3}}Nlbos3AbBv$DANp{c?<^hy>q}#!ZPS*H zAF^jMg3Ea0pMs>uCR;5G&7n0c4DF$r^1N;ii$iBv0+tlF%=l&Cn4^E~w6tYPnHnQ5 zax@Tyy~5B)j7>uKBJi`cW&+)c_syY^_XBDZa%BBRz>(D0Gp@6JXGM}HACUKC&U}Y(Hi{Zqs_lPaTEHo{XT5l#N)*_xR9Ccw!*pn7drZpM@bJ z2((-{$x>?$FVLR)rq*U*)xsX;`!=#u8@&<02DgTG_JYAHg$EA||I*@Xd=XvINyC!U7kOaa6_# zwLK%3h3u&~oXf08OH)ke8cfS#p$|he&Ln4}YF?^Nc?BE`arCcH?Y@oYU0kV+`{Es zXizS2Ms(U_oXv(jLw%nnz8MA?uE{3xt`X0=(+RexkSUj_iD7QDD$9E z;|f|5w#!DG@^Zn47;y+x9>V8=HH6&q}~=%EFPC=RDnl4Bpz++4&*-;vszJ z5PsYjAg0lvjeazm#vz`BE%fZ@3VE(z?;J{$+=JYYF+26rIiCHf*rgL z?X-@@NHBxAkCTq{5qILmj}Xu&&Z@&xlj~#=hqg-G_QSwK&_@>5oR?2xjLbeL$&+6N zQN~(Z(&j7os@LrZf&S-a$}i58&(D-E%#^<|Q+`QVNXO1;`l51>Bi@0tt}EweKMQrr z81$Qk%$to)TN%AJ(s9H)@2YY_{}xi|s4cZ|YFpJ7x(ES3!W&b8!F{ak&p3ggHvMbB zAsc+>=&Y$ta|~MZ__0tRo4brR{yj*_kq59mrkN;DO)q;!c5to zDHl@{v5ZD>9?31ua9w?xfu-*m^bKm&{A&|FWUQ?J_BZr=`dYIl69_E_MD7sTCUTbu z=}GPo`5BR)5J`y)i4f8%Hw^r(kC42E11Z30T6#a^kEqd%gLTF4&LD-}0s<4d_K2SV zVNkyd#GC$64KFf8*Vt)1umNK*v*i}Plp&EWkwqdjuIv%f+pUS-eLZU*4YaD5HVk-r zL=xv;0Et@Zi7CR20ZEiJBoP#q07?341-B3SLZwyE&}2ict~&i73~S*hR}BwcBv95j zxkz|hgVq1cMM8cp@dXzN8Pt54iv(wKw5N-N(idDLkQA;cKjI=GoB5%_lNgyGQIrBD6hOmO*AQ zO;O}osy;(xa?q8}Qk9Nl8V$4Db z9f~%91(N{s|3lG#{7{5r(&TU>Ti8fA*BRR4r-|GmqFrTj`jBnnO`S`7@KeMiQEYuS)DRGP$aprv`xD-nJoIx?e~Ye59!mYKJoPtS11bMQNACw?&C)nj>b$@B$atxR zCwUGJn)lGUhc~W(IGM|SR(lrNJhe0zGFUTEZXMCrb4yXQv^I$Jp$->x?zEic3y*b1{?PRruWa)I!%%i@ahk^cR}{1WGg>8UO$Q literal 0 HcmV?d00001 diff --git a/animation/SadTalker/src/facerender/modules/__pycache__/keypoint_detector.cpython-38.pyc b/animation/SadTalker/src/facerender/modules/__pycache__/keypoint_detector.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7827245533d03f33811bbdf8c4da19c1157e2f1a GIT binary patch literal 4843 zcmZu#OLN@D5ys%XZ@xs4mLF?HPO?eNM3f(?#O2tMEz7aRQdx2(+0caocd#UuSYS24 z@)htTr<{vY{sB2Gt8(!zzh@4)nTRin4(>10IjjpMBjK~_=T}Sa|gR?K6#A-2PVMcEARWXGXUPiNjq;+^M9+^~YMLat^y%heT4NTdiwxmsjw#ztc>u>)_ELjU< zuOIkuWsu7H&_60$-35sD{83O!WMd}b!6@64)p$JgcA})WpPrM`Lwy5acy>GrI@hI@jz@tIF2vIHyf6+k&l41A zvN93@4=GfHPHp>M5+B^*-5S~IhY>0`o-`3PskC~@=-87MEcEuIodt242x;}B#Lr|c3x=4P@|mos2urxgSF^EcHYQIF z<4!}GX%O{gOF763jz)f*B0Obl>h76Sq)n5=WE_3W`*5 zi)<|hX@=L3zq8L0Z_#_pDg26g&7}10KQLMGydI^$=l!7lckkYBw!co{s_nrTqilbI zY2yQVx_vJh!oz7A@a@Mwf8s}2wYP_U7+*BEQ_a~)L(a+>^jI&M1*hZl1A%o>Dp+E84&!qX=+mQ2V13)!UeF!}0F}N-hHbyShrEGbOTF_KxXx*oTSYV*g5Xfl zr=Qbgu>uAUhtfLmqjAu&5#y9#vNDd-r{f?%2wG5+CZ&L^mQqg|>3G<&#ZRc$nc5Za zD_sI^@g7*QPUM#*`4x57cKvah0tq+wN56m|r2x`Qa`H{Kz+C-gbiD8XI_U;)HTSgt${4*rMl0WlLJ2PnbE#=tx^b7NxVCa0A!Ia}Abab}!anGMh+ zD0FXVv_U-A$@c?j{#^s7O)9yCF$i9(T-&dF&jgKJRUoJP)_746i`-Ai+^F&mTQgC>C1LL?MA?x2`4J!#GkcM84S=9Le$+(wjC zw-n~7)fTtmpkl$_A#RxnExouy`2yT$bz8kR;L(AkW=epj|b>du!#@wk~Dyn-TFK_ zs0m(<9;cs(Ltk*RmQwg=h?}ZTx6M}B3h{a!vQ=hl_6rj&m#windzB(B^}(2~=f*jewvP%Tg$w--Wd&J=*4G7tNa52JUnXcyhTA8JO72fpgOW z0=eQLn;;=1qYd84jDf3qQ8PX-$`52BG%ASWs*+SQ#BupXEi)$dGb$h&S(EF}*eBXI zH<1b2pfPFjTDFk8&)9V>TT~;|QF7=^dwKvnCM1n=)Fq7CC`Vm}egjzreT=+P+PHSf z#x>Yz&e@>xW;Ud<(crDp#wu*IN*mW})3oHM`d$04Q+&&|Jyzx@2 z{Zi}AqQ&87Xx^N)Xvb2LAE^>M;Q@u11QJMDdM`g~v7U&5EZkWRnni6n}o3uYAJ*d=Xn$Nr%-@yt{UsXs+)@l$e3 zhX{d<_&E_O_@t|9>M%|P!L4MWP>7}gPQe=avWq+Y2t-5Ww91-nNnZk21@N|syMn(t zC%>sf13FFpWMvMyt}e6)h^W|7TmOBv^-oYxO}BZbHe)^0ryD+kbz&AWa8}9LSLiuq zuUm4JEA)yDpGK~gdi9cPT=7P;>}{3YLdh*&@y1fwyIl6Jpnzm&I?5S+axK>ft5m?C zSRjcuxvum8t&=sy;dC`wR~(ftuP6?s%jAaQa5|iw9l7}rDlo0fr2`IQ2E~h2gpaAh ziDlBDuz6Md0c#-sMC2n7X^J2ki$7C)o5(jrJR-EO2~H#+(kHS@WRFNlWI$w}NJL~v zBnFYoKKDmico2C0!LCQCx5>rtsGDjOF(UGmh#-;@$v`@;szC(Bh`3MWkjN2{VmEYIo(H%j%%G5yu<- zaonSSAplL)3tp{R_cffEHZIOg8!7&(@pkAx8UhzV+r9WB4R_;`;`DH}G0wtBoq*Lb sLC}l}fU_T}2A+>ZLin6P8Ko;qA($`I1nI_SrRsKI1o$Obs#hjt#()0?r~4| zxVp#dwR$9^-6*FO5F`!}2m2U_BM1J19QYIJ2q%jxC!{3&Ud{M33O%Y3Mp8R>+pVA9LU3;kMevSRj1w;gSFHn?(~~M-a@*MRQ&^_vDJf` zHQ2&3L|uCmF-0fVv$eJIea-hYKQIl3{8!$8WApt=s%o=8s_Ja>?ebBPmx-u0@0Ww2l2s+e z=EFp6B>A3HuMU!|xTtPcD&5>pQmJGiq}m*mVwB5jbCB%GI4QD0QfFncGJLGVRgtB2 ziYf2Sk+>w0ho8lcx<#o5NuC|acyG9}nwM#kSNB#f257&Zq~%tY6ulW`VW4(IQYYF@ zha>G4qXFjk4$6EqklHGS+Nt(ct=+7agG#$~snT8TL6pHL*R2eIQkCq@=wNRcSLH~h zQimuCl^tZoPPA}QjmFj;L<^RQ$gA4MJHM&Yr@wzs3s*;rG((3zsUtlO85jjlj-7lc_3zl-0J zBq2EO0mCM+`HoLIs|?%_t*jcsszWl_=~LD#^Rz_3HG^~b#7e>8x- z*lEvFv@>mI;_#w%vj1>3&#DQcUO=CU_F?e!yrzFy83epR-C*(l}T(~MJcBEQFtDafIwWv2_QI$%!>c>MMywX-(YQCmhRk>Xc zl0$7*qe0Kn4lS)nE*(PlxSkBdyd=r>t=#?wiNA&MRU3qHk1y~g9`ZX}y^MNF(w=>E zMAYj(QuQe;1X;i5Vf_}_7&hy?42y7Hwi@dghMTkqt2Qa|hD}S*ec!4%bPhcucfKJv zH~`-x{tQN)_cSf2n;tuXfCYVK5gvdTyYCAZ*4Y!@32f_hEtimFW|-vOwYq(+ZurwR zCnL-Z=W98QgSDJ1YdPm@0n0{i?+qhWEy|R{zQGvj2 z8>ewk;CUJkcfnsj{TtM@^TJ2}e(&$!4&F4_`;2Q_WCLU76)F8_ZY%4*ven;vWovKn z>b%%O`*DKSVOGO7I;zZb9hgcCJE>b>nu1T9Ns}sJ|C0JTbZ_p+ItI@edeb30e35YY z0EF>2{j3nc4XwM}XFdxJ^}n<%>r2mGoVNNr1n;8t+ml|``p#ME9%j7){tg=G{u=PV zJO%$UnI0H1IA<6!L!D(99gB^fv1@_=2Tt5UFyJD0k$cEJ0l4mY>K7OVfQ}sj7^A+4 zI-K*^g*vI8quxT@t2rqbW9>0G!Z#}g%e0P%)M(T07h1~y?-~JjOEyO=|M28Hq=-@heqJvpFnAvYI_iZ+wX#xN;04D?o@fLVN zgI4g`f;a7q1B`R)imayk5s@Dgc^yQ1q<%ohpCAIImcaoEyw+MlC-8>2-7){x406vdG|HwY4iF z?@<|JJ#Fm`uSuzXLIaKYCe<4kZg41tb}&z^XNkgqSKEj$V2vr{m=NPXZ|}*+!!m<6 zKksBf{1&=ZBqD?B1ZL@H89FVutp)CI-=4L*ybH(b+l%O--o=QG)L3g8jPEv5fMlTg z9~5eK%um?w|HAMtNJ^_rMSZ}irKON7( zejC^{a|Da=JVwrAD!{XnfD{z7*!v1Q_`mF}A5^=5?eZZO} zfGk)ys5qKkL!8)sU;;&{-cUcqfv7i$(22v-({ZUgAUY)PiEt6Iqk4sy{=@~hiZb3& zi7;oaZ0g@7LKxCsG8|%Zn7TSi_jKTAX$XcdUY46HdA~=TqBi}#3N{AG7K=&lY_?ZkkG;^`LO1%ToqX6?y#50Ew1=}ha$QXe`v~*p_*zlqdL#8pLIkTt*fhi)?b~&n} z*-LB0b#)#4vsVN3Z3t0)7lb(>cMRA=>yAyaqt5~ENwPx3k+Z=Z#UQ^$-;@Si{zxJ8 znW-d;xL|lV^2q*}Z3D%^e(Fcg<;-+`EWL-A>;mt^G0`Y8*ufw&v z)o3V{iZP}3WAlcHV`JOoAxsS+O1%%J`ha>k5)nT|jFrr{C+%BjOvD929H)6wRdIY_ zx>SM=pCDDUAdJrgVo*`r&R4|C;g`IM~l*+(|%JYM`IDsWw}G!zMv%QUVNPIHW? zPrz-C;S`NKgq<>OutjfZ6W8E%jZ?u%^=-IDI6_z;|5*q(MDJJJe2M}T1^Q}($ghb+ zM1BLZ(R0-1*PGEw@Va&uy+8@HYJT|tC)M+pp?{DGSzbqf3ov<x|%B1cqd6+$U7PIJVU`elp&GR4ozFnURzu1@wJ}2zP`S()|=OzD$RzESE?br zSCX6ev9XrXO#m(3T`7u{?NMQjDLK`m(BVa0eTNpiL1dl?T?X?^yyUpau^EkM4>uHB xG=f|@R3p0ass8HZ5bkk|sc7yttj7XUoBlXr?5l9e_x-!!ODF@mNMD36{ST;dmvI09 literal 0 HcmV?d00001 diff --git a/animation/SadTalker/src/facerender/modules/__pycache__/mapping.cpython-38.pyc b/animation/SadTalker/src/facerender/modules/__pycache__/mapping.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ee400d03a4dfd648bba44fcb7fb94eb741fd74a7 GIT binary patch literal 1700 zcmY*ZTW=gS6t+Fro$Lm3A)(wQTEPk$R3Q*T2&fXlOBx~4yo?q`)3G<5WX3bGy$Kt6 zpOUIi`~VOS$zvt{0Dr~eji>z!NGKe8vu*0peEyCv$H(?}yw~d`2ugi&husSi`o|hK zix0-P@XIX#iYTs7id{`o;%cwNcBGz*3wJX@=Mmr%%Y(ixjMJBzlvPTYZSpbBs=@R|7Ocy-%KOetiS*d`<_`iWP-CQ zljSKhQJYf=XJ!0HWiF@oX!@S8_bUPb6VtNFRoQS8F`O?53MGEJ7ksk$d3Nl!d(j+7Cy1?!E^W!7>KwLC%K^S zfgx=S2q$on%lJ0b%U%NLK61Kz<65sL^+OY98P975shDKhOg_N7S#~tat67hekoXK- zo8Gl@3szNGb|aQM&?{^}F2Q3U(mo0??ykF*eqsX1i*^jD-ijsYGOo19 z9=o%)1rE$spT2$W?fX3E+;h%7y-$seKj-IhpLp2t3x4r=!!HKr>f}|svH zIh|>LcJiLx#ha+ti`>nupY3|SH*IWs zes0S|E7n0{%iPFya~OTTb1FZF-;=7+43{v_xysyHaHZXYhfW*6crC2vzmMQl%QGpy>zdLZjuc z*MsoFdQ__)Z(oTqgp=vQ6wzA9CJo24#!dd^%_%c)o(2H2M>G z=o%NS4U06o2^!rpH?6J-8gb;=>OE(Claxi8J!s58eu2gu;4)on2$$(*HfUWq-x+SAL&hp<$=qheo*tK>!Y=w!L<&S@dCGYZQafH)gLxI$8x zd+$9b?mqGC+~qXUTAEv~T@0FYgKK?UUREB4sA-WSzn)c-xW#S-M86!U0&~X+CWyAH zAf)_yGkk762s%M|KBs1w$zrj&5}TJ|^KxuniLEQRE?jFN9HI?<(K? z!#bniT`#%k^?ezV4(nd%7O?J_g>8$bI1B0^ze|qr!Rk0Z-oxJmh>XbWVxQQwwxAtE zlXwG49Iq1^PHY~&?nT|axxLb==yOkOZ)exb`! z#jdeh^sO~peX7?Nw&h)OkUKW?#<*`_xVeQJRD6W#uG?{;)+|?7D$Qn4591=FWdErE z=?AokW?AIQ(_!Pmh8%ALL;K}PIn0INZ)(o(v3>Q`>@W%WEwukk0 zpwwM_-z%4E&0170D+qg;ZVT8DI(V(4eCT&~tDVTi_s{3sgK8vq`~748KvAh7qEhjC5e zZ`^W0i8E$r?-02hPlqrh9Y`;{&fp542k07`paw{8tQsUaB=?3DS(~;_0SC#Oes^^M+NnD z^=6_TAb1PGw-a#5)Pn?X1GsLhQ%nmt?_+A1Z{sZVluH%mKZCsR5P*Rh8};V_+qq?U z9+P*kyQGR{r4f|NalTyE8!Xbra{0OSN?qs3LaB$*LOhnNV>PJP)3ua-okp&bRVA%G zf-5BDduDN5sIQ8Lp&u8Fzwl0gRCf4g$rkwmR)zd1Y`+ePeqZeDMx7K zt$ZX63Duy!Q%pZZK;fWRpOQt8<>>KO6L`JNib7<)aNO)n_ownzFxERz6No-x ztcc1{%0xRDYbZBiF-)1TNf*p<-CVft>fxMa$KOsMY$7}%BR)r*45mLK(%>ar;UvH? zy?bV7YG6DmfoBzIQBlmOcLK!Oa@lWHLAx^5tWZ5h&~K?v%3P8fu8>sXnT}J;$uhn6 z#C)HxU^N^6ALOaN0}ZD{NG|9Tq243VH(464%D~0|bdMg@yZPP-^7|n0uSkCP`Txh{ zr=DPU&lBuMWglV3U!Ka)?S#tIldNhK1#vzbrnZWdf0iyAkKlqePeopwgtb{LY^KKoY8)x3#sv_0(qHu4Kx||4CduG9Aw<;=T!F z3b9dQhrX$_r{0Su<`c?VVsrhC^*i|}6a5TRGTZNCY6Q)E68PO|=H=wp|CVN;{>i$D z%`RY(ee?--Ba7c;Z-p$dyKf_lzOcFpvf#}9A}eKAxc`AHh6vC8(x1 z(cwsMaqgju%a6BOb(mjH8+Y|S78*f-p9MY~2hg0Zg%(t|nSp8BU2su?i{QE&iN5ZJ z)XNz~5L&W9LPGWB5ob{gr8-l)k>D4QnMkJ86zNP4_Vwmv6Q}djtagA1iTkDyA*2%( zzE>lizBzQImf8F&KuS*3;0MTQg!t-_6M8SJ^FF>Zf|~vY_>`J5=@vLdO)n?X{&&=* z8tl|=bn{Q_ncmSx(9QoycyZpo#AB#4lC?ayyDh+3{u-`6LFN^4H$jlEWD8{9u6M`IK_SL>wAR%! z-EMOUUa%tEZor{xR5NHk{c#9%?hDCrYKBQ}8Rv76-NjkmS>xGrILF9rI#H$zKlz%a z?(JI5Ntgie&%*oaPUBSMs?~q?ah|7f*ivn5s_=W8#&ZXv>?VBY@PQ9(vFMgD;Phs% zu3LQ{M=BK~7Nb%ZXXN+-p3rt1r$oZyvE8UOwU;yhE;%!LtQJOb!G~s17FSW2&~7{- zk7b;c;p7i)B5hmbaqg8AjEjxR<#Nx?_-;07)<#L!Ji4CRx)$izJPHRU==#OH-I+`l zU^F#5sFYIb3@-79bkxr+qId~`{|*wn5x8&Q&wRn&&)I((q@ElOfzaRVTTQt7 zB1cHw2KielY0uzTx{uf7E=NO@%RGF>!9{<(2%7OB1rAuQQR($4znP`%%|H*4pFF zJ&L&3kUMj|4d)w2wrxNa1SHg=pvqIBLGRNJE=`wHxby+^Zs;lb8@Hgr zXBUg^u4a)VJ}9wcQa`|lP?_cc3Z*z3=)F=VA%@QBU%a<0e9kvWa@qd-k1Fh9iUwE>@@~7IxUV zC@qmmr8&F;g{}vU8`gjvKNq1@>z<7Ddh#^(f!>?l>};a?APi{;>pQ)i9C|*6o_oEX zkE4~`P;Y(#JC0GF=%Q?SGC5q19d7Ij)3aS=QnCxs?2oPW_WT~b2V7%8*#k%^asGUO z;EpCkg0%X;ogvQSlsP#FMDS3g)C%8{{i1Iee1Nzm79YbA8$*!Fe$b9qH2dT=t@hX_ zo=gk%^r#+i5qW4=$5ZXf!GW=0_L>LGV>meO%nVT3 zXa?j?$?FWR@KJz~_Jg%*U2qs=vXOxnzyoB4lo_iIf(-y+d%0T~aRT?wxKh@2^^D4H zFoybdWQ~lPbC;SOkDUz3~^X~Ah-wN zf4(`3K$3$9O7tz7HogtB1)+gHWCwx{p|5eB%aYn__zon=6I^$`*AzNHZUz(qMS?+! zT^n~C6~j#4QlHQ-)Rr3+9SG#= zSfNrb?MJBCD5W*Rc^ZrhRJRb3>(x3+qF&-8-ivbd!AwG$=FP(v6iwTK7X}miqVtmD z?uDei!ydP$kXJNa{NbSe$SY3xj`*COuw6xaU?LDD*lD~AT}OfiKyuK@DO8jD5~JX( z$V8kkHF+v_*vzjX>JMfx473*pAo9EyY+`k4$w8zs%8Aat%n=&K*WZSP{wXA60bsgt zTJ;f}jle4PGkn|xh+UpKp7i^L)z1M9ay`Y|j7T1l+1m(IB4j?wNlitxV;#b@-f4B_ zcAc<9Uhi&HlN3W%N+QJ>j9omch}&I8A46_>g8T}>c$xE$KXmr()j`QguIfonE>>h* zg<>iJcr=?Nwoj~3xN;CGYi(v=PGNTxtq3yj8qBNB?3S@bx#CG4-@;i6^MK2Zz~x2S zWEakR~8YH*Q^4yQuF zFRMPu$vR6A5Qr4}6{dcb;MV|PN1bB&2tdgV1IVLxoaba!>ikNeTpVT zWs_g#r0H|5ym=@&teiAEV|v0;u6PIA>nlR1`6wpxFCp~_)*<9-7K@6Q9X!EY{z6hg z$`gJ;xB@c;w?4d)aPK2THFfVNUyZmDa>m&!#*@;iCxWmUMB?o~CzCd0uy8dg<3!>3 z6D0a19weKkR+BeC3uGwG zPCizB%~cx(In>P&*$V5r&H+61KFuACJR-LP4&?3_KnnKcaGFj>~1G(DKlwA;B zgaqul_G8blA!30JQ_7`tASvd>&LVWZ*bCdR^i`!kpE#dnE1kqA1c-6iKs`c-eTV=q->U_2`9yLLSiAga|B!tm9J-Z zVD8h|{WNNiGvvvp`eUTk9};{PU_ObXAty6vclktYUoMBSFW539`$m0bA<7QYm z^nlqNV@?=%NYIf9Icym+XA@**r~?>`9&9qu4I4VONoQR;7+d{5N~=F2_#D8XXIuRl zv0otQ%Lfq+e+pd8&XOe!SuDpKu{`NK4XG+5LiShs%57XeAHbfsrM>OTANweOpQAEy z7TIwgM7?&UEazN548-={d>*(jvp1|6zL~u{rc4%S-qJD!bo&Tvd6aPN`*~wG*01~Na9HlNK0y2{bNH-CXSrL%@IHeF5 zD(il&b+Xo|EC)M6$vAK}7TygoBo!eM#rU9EEoK%VXc^F@I-s3*+5Wa4=a%u&fT|#R zDlX#drQkAT3}dO{sgC}$Bn*6R||OPnyXMz>3)753a~H7 zYH(3gpTcx(nTa#njxyNVDFjSvRS*g`>{F=5K0cK5U4D4frCF6-*kO}#&c!C1L_%yb zZv*~CZHM*A6omXS6o4rL=HJFjA4kU5IbzMimeG0B;7C{SNm#oD=?w$jFGnr7fR>?r zjWV`w$Xrqi$X0TLjxKW_`K^9GuY3#aE^xHlla68%W7dYh0C&GYbnAxoS0dO zcNUQ|Xjc6--YVtM`EsORLY?fzwz0(fcIoX``zz|=&|wm~$gI_Q-K!X9m9vumC~Qa< zXK;ldL!DHV;bO3GECd19OH15%g2zH{;u}{iE)V#0O3o+*XHpA$203u;C#{XXV~^Z2du>63LR;=o zTP|wJ!QG)ZBgi&m1htE~_gInJS}=IU=H-JsUw=a^cN zj_oUzOL3-M!^gGwETmPhmpsBUq*y5@Xh~*Kj0t>d*@~FhPTJGUK~zSb=wni#cKN?Z z)=!{iZrAwJXGgF{Jy>@*PXhE_v|!}JPD!x%PGkssSeM#XcGBwHHpm>vHRnAJN~S3dpvUhKB5^MOAbrP zUFsPkh8(Kv5d(D`SypgyqWeOIh_(wUm@F{92*)z=yIJbQw0CJe(N9F{2$mCUV*QP* zwE*W~4ucI5^->LFCJ@bc5pnR?G5TA(?DYSRwiP%2g&ChJ`IWJ;-LMUq3fe zKkQSM3{2F;$Jnv{hAtbh=rMnejK8ONpf(RXCgx zJU(RC#Th}W&#*WNP~Abm&c?2e=23SM^Dcs?3Eo4nM6g28Bxn;v1n(#K5WzJ9I?u#g zAbPQkdSIG*>R#+UpudKD5L2c7q(yi^J=%NLtnKadroHK@xvBm5o1VI5YG!J7YA^nd QO-)SQK6PSh-_-nn1OD+U*#H0l literal 0 HcmV?d00001 diff --git a/animation/SadTalker/src/facerender/modules/dense_motion.py b/animation/SadTalker/src/facerender/modules/dense_motion.py new file mode 100644 index 0000000..a286ead --- /dev/null +++ b/animation/SadTalker/src/facerender/modules/dense_motion.py @@ -0,0 +1,121 @@ +from torch import nn +import torch.nn.functional as F +import torch +from src.facerender.modules.util import Hourglass, make_coordinate_grid, kp2gaussian + +from src.facerender.sync_batchnorm import SynchronizedBatchNorm3d as BatchNorm3d + + +class DenseMotionNetwork(nn.Module): + """ + Module that predicting a dense motion from sparse motion representation given by kp_source and kp_driving + """ + + def __init__(self, block_expansion, num_blocks, max_features, num_kp, feature_channel, reshape_depth, compress, + estimate_occlusion_map=False): + super(DenseMotionNetwork, self).__init__() + # self.hourglass = Hourglass(block_expansion=block_expansion, in_features=(num_kp+1)*(feature_channel+1), max_features=max_features, num_blocks=num_blocks) + self.hourglass = Hourglass(block_expansion=block_expansion, in_features=(num_kp+1)*(compress+1), max_features=max_features, num_blocks=num_blocks) + + self.mask = nn.Conv3d(self.hourglass.out_filters, num_kp + 1, kernel_size=7, padding=3) + + self.compress = nn.Conv3d(feature_channel, compress, kernel_size=1) + self.norm = BatchNorm3d(compress, affine=True) + + if estimate_occlusion_map: + # self.occlusion = nn.Conv2d(reshape_channel*reshape_depth, 1, kernel_size=7, padding=3) + self.occlusion = nn.Conv2d(self.hourglass.out_filters*reshape_depth, 1, kernel_size=7, padding=3) + else: + self.occlusion = None + + self.num_kp = num_kp + + + def create_sparse_motions(self, feature, kp_driving, kp_source): + bs, _, d, h, w = feature.shape + identity_grid = make_coordinate_grid((d, h, w), type=kp_source['value'].type()) + identity_grid = identity_grid.view(1, 1, d, h, w, 3) + coordinate_grid = identity_grid - kp_driving['value'].view(bs, self.num_kp, 1, 1, 1, 3) + + # if 'jacobian' in kp_driving: + if 'jacobian' in kp_driving and kp_driving['jacobian'] is not None: + jacobian = torch.matmul(kp_source['jacobian'], torch.inverse(kp_driving['jacobian'])) + jacobian = jacobian.unsqueeze(-3).unsqueeze(-3).unsqueeze(-3) + jacobian = jacobian.repeat(1, 1, d, h, w, 1, 1) + coordinate_grid = torch.matmul(jacobian, coordinate_grid.unsqueeze(-1)) + coordinate_grid = coordinate_grid.squeeze(-1) + + + driving_to_source = coordinate_grid + kp_source['value'].view(bs, self.num_kp, 1, 1, 1, 3) # (bs, num_kp, d, h, w, 3) + + #adding background feature + identity_grid = identity_grid.repeat(bs, 1, 1, 1, 1, 1) + sparse_motions = torch.cat([identity_grid, driving_to_source], dim=1) #bs num_kp+1 d h w 3 + + # sparse_motions = driving_to_source + + return sparse_motions + + def create_deformed_feature(self, feature, sparse_motions): + bs, _, d, h, w = feature.shape + feature_repeat = feature.unsqueeze(1).unsqueeze(1).repeat(1, self.num_kp+1, 1, 1, 1, 1, 1) # (bs, num_kp+1, 1, c, d, h, w) + feature_repeat = feature_repeat.view(bs * (self.num_kp+1), -1, d, h, w) # (bs*(num_kp+1), c, d, h, w) + sparse_motions = sparse_motions.view((bs * (self.num_kp+1), d, h, w, -1)) # (bs*(num_kp+1), d, h, w, 3) !!!! + sparse_deformed = F.grid_sample(feature_repeat, sparse_motions) + sparse_deformed = sparse_deformed.view((bs, self.num_kp+1, -1, d, h, w)) # (bs, num_kp+1, c, d, h, w) + return sparse_deformed + + def create_heatmap_representations(self, feature, kp_driving, kp_source): + spatial_size = feature.shape[3:] + gaussian_driving = kp2gaussian(kp_driving, spatial_size=spatial_size, kp_variance=0.01) + gaussian_source = kp2gaussian(kp_source, spatial_size=spatial_size, kp_variance=0.01) + heatmap = gaussian_driving - gaussian_source + + # adding background feature + zeros = torch.zeros(heatmap.shape[0], 1, spatial_size[0], spatial_size[1], spatial_size[2]).type(heatmap.type()) + heatmap = torch.cat([zeros, heatmap], dim=1) + heatmap = heatmap.unsqueeze(2) # (bs, num_kp+1, 1, d, h, w) + return heatmap + + def forward(self, feature, kp_driving, kp_source): + bs, _, d, h, w = feature.shape + + feature = self.compress(feature) + feature = self.norm(feature) + feature = F.relu(feature) + + out_dict = dict() + sparse_motion = self.create_sparse_motions(feature, kp_driving, kp_source) + deformed_feature = self.create_deformed_feature(feature, sparse_motion) + + heatmap = self.create_heatmap_representations(deformed_feature, kp_driving, kp_source) + + input_ = torch.cat([heatmap, deformed_feature], dim=2) + input_ = input_.view(bs, -1, d, h, w) + + # input = deformed_feature.view(bs, -1, d, h, w) # (bs, num_kp+1 * c, d, h, w) + + prediction = self.hourglass(input_) + + + mask = self.mask(prediction) + mask = F.softmax(mask, dim=1) + out_dict['mask'] = mask + mask = mask.unsqueeze(2) # (bs, num_kp+1, 1, d, h, w) + + zeros_mask = torch.zeros_like(mask) + mask = torch.where(mask < 1e-3, zeros_mask, mask) + + sparse_motion = sparse_motion.permute(0, 1, 5, 2, 3, 4) # (bs, num_kp+1, 3, d, h, w) + deformation = (sparse_motion * mask).sum(dim=1) # (bs, 3, d, h, w) + deformation = deformation.permute(0, 2, 3, 4, 1) # (bs, d, h, w, 3) + + out_dict['deformation'] = deformation + + if self.occlusion: + bs, c, d, h, w = prediction.shape + prediction = prediction.view(bs, -1, h, w) + occlusion_map = torch.sigmoid(self.occlusion(prediction)) + out_dict['occlusion_map'] = occlusion_map + + return out_dict diff --git a/animation/SadTalker/src/facerender/modules/discriminator.py b/animation/SadTalker/src/facerender/modules/discriminator.py new file mode 100644 index 0000000..d4459b0 --- /dev/null +++ b/animation/SadTalker/src/facerender/modules/discriminator.py @@ -0,0 +1,90 @@ +from torch import nn +import torch.nn.functional as F +from facerender.modules.util import kp2gaussian +import torch + + +class DownBlock2d(nn.Module): + """ + Simple block for processing video (encoder). + """ + + def __init__(self, in_features, out_features, norm=False, kernel_size=4, pool=False, sn=False): + super(DownBlock2d, self).__init__() + self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size) + + if sn: + self.conv = nn.utils.spectral_norm(self.conv) + + if norm: + self.norm = nn.InstanceNorm2d(out_features, affine=True) + else: + self.norm = None + self.pool = pool + + def forward(self, x): + out = x + out = self.conv(out) + if self.norm: + out = self.norm(out) + out = F.leaky_relu(out, 0.2) + if self.pool: + out = F.avg_pool2d(out, (2, 2)) + return out + + +class Discriminator(nn.Module): + """ + Discriminator similar to Pix2Pix + """ + + def __init__(self, num_channels=3, block_expansion=64, num_blocks=4, max_features=512, + sn=False, **kwargs): + super(Discriminator, self).__init__() + + down_blocks = [] + for i in range(num_blocks): + down_blocks.append( + DownBlock2d(num_channels if i == 0 else min(max_features, block_expansion * (2 ** i)), + min(max_features, block_expansion * (2 ** (i + 1))), + norm=(i != 0), kernel_size=4, pool=(i != num_blocks - 1), sn=sn)) + + self.down_blocks = nn.ModuleList(down_blocks) + self.conv = nn.Conv2d(self.down_blocks[-1].conv.out_channels, out_channels=1, kernel_size=1) + if sn: + self.conv = nn.utils.spectral_norm(self.conv) + + def forward(self, x): + feature_maps = [] + out = x + + for down_block in self.down_blocks: + feature_maps.append(down_block(out)) + out = feature_maps[-1] + prediction_map = self.conv(out) + + return feature_maps, prediction_map + + +class MultiScaleDiscriminator(nn.Module): + """ + Multi-scale (scale) discriminator + """ + + def __init__(self, scales=(), **kwargs): + super(MultiScaleDiscriminator, self).__init__() + self.scales = scales + discs = {} + for scale in scales: + discs[str(scale).replace('.', '-')] = Discriminator(**kwargs) + self.discs = nn.ModuleDict(discs) + + def forward(self, x): + out_dict = {} + for scale, disc in self.discs.items(): + scale = str(scale).replace('-', '.') + key = 'prediction_' + scale + feature_maps, prediction_map = disc(x[key]) + out_dict['feature_maps_' + scale] = feature_maps + out_dict['prediction_map_' + scale] = prediction_map + return out_dict diff --git a/animation/SadTalker/src/facerender/modules/generator.py b/animation/SadTalker/src/facerender/modules/generator.py new file mode 100644 index 0000000..5a9edcb --- /dev/null +++ b/animation/SadTalker/src/facerender/modules/generator.py @@ -0,0 +1,255 @@ +import torch +from torch import nn +import torch.nn.functional as F +from src.facerender.modules.util import ResBlock2d, SameBlock2d, UpBlock2d, DownBlock2d, ResBlock3d, SPADEResnetBlock +from src.facerender.modules.dense_motion import DenseMotionNetwork + + +class OcclusionAwareGenerator(nn.Module): + """ + Generator follows NVIDIA architecture. + """ + + def __init__(self, image_channel, feature_channel, num_kp, block_expansion, max_features, num_down_blocks, reshape_channel, reshape_depth, + num_resblocks, estimate_occlusion_map=False, dense_motion_params=None, estimate_jacobian=False): + super(OcclusionAwareGenerator, self).__init__() + + if dense_motion_params is not None: + self.dense_motion_network = DenseMotionNetwork(num_kp=num_kp, feature_channel=feature_channel, + estimate_occlusion_map=estimate_occlusion_map, + **dense_motion_params) + else: + self.dense_motion_network = None + + self.first = SameBlock2d(image_channel, block_expansion, kernel_size=(7, 7), padding=(3, 3)) + + down_blocks = [] + for i in range(num_down_blocks): + in_features = min(max_features, block_expansion * (2 ** i)) + out_features = min(max_features, block_expansion * (2 ** (i + 1))) + down_blocks.append(DownBlock2d(in_features, out_features, kernel_size=(3, 3), padding=(1, 1))) + self.down_blocks = nn.ModuleList(down_blocks) + + self.second = nn.Conv2d(in_channels=out_features, out_channels=max_features, kernel_size=1, stride=1) + + self.reshape_channel = reshape_channel + self.reshape_depth = reshape_depth + + self.resblocks_3d = torch.nn.Sequential() + for i in range(num_resblocks): + self.resblocks_3d.add_module('3dr' + str(i), ResBlock3d(reshape_channel, kernel_size=3, padding=1)) + + out_features = block_expansion * (2 ** (num_down_blocks)) + self.third = SameBlock2d(max_features, out_features, kernel_size=(3, 3), padding=(1, 1), lrelu=True) + self.fourth = nn.Conv2d(in_channels=out_features, out_channels=out_features, kernel_size=1, stride=1) + + self.resblocks_2d = torch.nn.Sequential() + for i in range(num_resblocks): + self.resblocks_2d.add_module('2dr' + str(i), ResBlock2d(out_features, kernel_size=3, padding=1)) + + up_blocks = [] + for i in range(num_down_blocks): + in_features = max(block_expansion, block_expansion * (2 ** (num_down_blocks - i))) + out_features = max(block_expansion, block_expansion * (2 ** (num_down_blocks - i - 1))) + up_blocks.append(UpBlock2d(in_features, out_features, kernel_size=(3, 3), padding=(1, 1))) + self.up_blocks = nn.ModuleList(up_blocks) + + self.final = nn.Conv2d(block_expansion, image_channel, kernel_size=(7, 7), padding=(3, 3)) + self.estimate_occlusion_map = estimate_occlusion_map + self.image_channel = image_channel + + def deform_input(self, inp, deformation): + _, d_old, h_old, w_old, _ = deformation.shape + _, _, d, h, w = inp.shape + if d_old != d or h_old != h or w_old != w: + deformation = deformation.permute(0, 4, 1, 2, 3) + deformation = F.interpolate(deformation, size=(d, h, w), mode='trilinear') + deformation = deformation.permute(0, 2, 3, 4, 1) + return F.grid_sample(inp, deformation) + + def forward(self, source_image, kp_driving, kp_source): + # Encoding (downsampling) part + out = self.first(source_image) + for i in range(len(self.down_blocks)): + out = self.down_blocks[i](out) + out = self.second(out) + bs, c, h, w = out.shape + # print(out.shape) + feature_3d = out.view(bs, self.reshape_channel, self.reshape_depth, h ,w) + feature_3d = self.resblocks_3d(feature_3d) + + # Transforming feature representation according to deformation and occlusion + output_dict = {} + if self.dense_motion_network is not None: + dense_motion = self.dense_motion_network(feature=feature_3d, kp_driving=kp_driving, + kp_source=kp_source) + output_dict['mask'] = dense_motion['mask'] + + if 'occlusion_map' in dense_motion: + occlusion_map = dense_motion['occlusion_map'] + output_dict['occlusion_map'] = occlusion_map + else: + occlusion_map = None + deformation = dense_motion['deformation'] + out = self.deform_input(feature_3d, deformation) + + bs, c, d, h, w = out.shape + out = out.view(bs, c*d, h, w) + out = self.third(out) + out = self.fourth(out) + + if occlusion_map is not None: + if out.shape[2] != occlusion_map.shape[2] or out.shape[3] != occlusion_map.shape[3]: + occlusion_map = F.interpolate(occlusion_map, size=out.shape[2:], mode='bilinear') + out = out * occlusion_map + + # output_dict["deformed"] = self.deform_input(source_image, deformation) # 3d deformation cannot deform 2d image + + # Decoding part + out = self.resblocks_2d(out) + for i in range(len(self.up_blocks)): + out = self.up_blocks[i](out) + out = self.final(out) + out = F.sigmoid(out) + + output_dict["prediction"] = out + + return output_dict + + +class SPADEDecoder(nn.Module): + def __init__(self): + super().__init__() + ic = 256 + oc = 64 + norm_G = 'spadespectralinstance' + label_nc = 256 + + self.fc = nn.Conv2d(ic, 2 * ic, 3, padding=1) + self.G_middle_0 = SPADEResnetBlock(2 * ic, 2 * ic, norm_G, label_nc) + self.G_middle_1 = SPADEResnetBlock(2 * ic, 2 * ic, norm_G, label_nc) + self.G_middle_2 = SPADEResnetBlock(2 * ic, 2 * ic, norm_G, label_nc) + self.G_middle_3 = SPADEResnetBlock(2 * ic, 2 * ic, norm_G, label_nc) + self.G_middle_4 = SPADEResnetBlock(2 * ic, 2 * ic, norm_G, label_nc) + self.G_middle_5 = SPADEResnetBlock(2 * ic, 2 * ic, norm_G, label_nc) + self.up_0 = SPADEResnetBlock(2 * ic, ic, norm_G, label_nc) + self.up_1 = SPADEResnetBlock(ic, oc, norm_G, label_nc) + self.conv_img = nn.Conv2d(oc, 3, 3, padding=1) + self.up = nn.Upsample(scale_factor=2) + + def forward(self, feature): + seg = feature + x = self.fc(feature) + x = self.G_middle_0(x, seg) + x = self.G_middle_1(x, seg) + x = self.G_middle_2(x, seg) + x = self.G_middle_3(x, seg) + x = self.G_middle_4(x, seg) + x = self.G_middle_5(x, seg) + x = self.up(x) + x = self.up_0(x, seg) # 256, 128, 128 + x = self.up(x) + x = self.up_1(x, seg) # 64, 256, 256 + + x = self.conv_img(F.leaky_relu(x, 2e-1)) + # x = torch.tanh(x) + x = F.sigmoid(x) + + return x + + +class OcclusionAwareSPADEGenerator(nn.Module): + + def __init__(self, image_channel, feature_channel, num_kp, block_expansion, max_features, num_down_blocks, reshape_channel, reshape_depth, + num_resblocks, estimate_occlusion_map=False, dense_motion_params=None, estimate_jacobian=False): + super(OcclusionAwareSPADEGenerator, self).__init__() + + if dense_motion_params is not None: + self.dense_motion_network = DenseMotionNetwork(num_kp=num_kp, feature_channel=feature_channel, + estimate_occlusion_map=estimate_occlusion_map, + **dense_motion_params) + else: + self.dense_motion_network = None + + self.first = SameBlock2d(image_channel, block_expansion, kernel_size=(3, 3), padding=(1, 1)) + + down_blocks = [] + for i in range(num_down_blocks): + in_features = min(max_features, block_expansion * (2 ** i)) + out_features = min(max_features, block_expansion * (2 ** (i + 1))) + down_blocks.append(DownBlock2d(in_features, out_features, kernel_size=(3, 3), padding=(1, 1))) + self.down_blocks = nn.ModuleList(down_blocks) + + self.second = nn.Conv2d(in_channels=out_features, out_channels=max_features, kernel_size=1, stride=1) + + self.reshape_channel = reshape_channel + self.reshape_depth = reshape_depth + + self.resblocks_3d = torch.nn.Sequential() + for i in range(num_resblocks): + self.resblocks_3d.add_module('3dr' + str(i), ResBlock3d(reshape_channel, kernel_size=3, padding=1)) + + out_features = block_expansion * (2 ** (num_down_blocks)) + self.third = SameBlock2d(max_features, out_features, kernel_size=(3, 3), padding=(1, 1), lrelu=True) + self.fourth = nn.Conv2d(in_channels=out_features, out_channels=out_features, kernel_size=1, stride=1) + + self.estimate_occlusion_map = estimate_occlusion_map + self.image_channel = image_channel + + self.decoder = SPADEDecoder() + + def deform_input(self, inp, deformation): + _, d_old, h_old, w_old, _ = deformation.shape + _, _, d, h, w = inp.shape + if d_old != d or h_old != h or w_old != w: + deformation = deformation.permute(0, 4, 1, 2, 3) + deformation = F.interpolate(deformation, size=(d, h, w), mode='trilinear') + deformation = deformation.permute(0, 2, 3, 4, 1) + return F.grid_sample(inp, deformation) + + def forward(self, source_image, kp_driving, kp_source): + # Encoding (downsampling) part + out = self.first(source_image) + for i in range(len(self.down_blocks)): + out = self.down_blocks[i](out) + out = self.second(out) + bs, c, h, w = out.shape + # print(out.shape) + feature_3d = out.view(bs, self.reshape_channel, self.reshape_depth, h ,w) + feature_3d = self.resblocks_3d(feature_3d) + + # Transforming feature representation according to deformation and occlusion + output_dict = {} + if self.dense_motion_network is not None: + dense_motion = self.dense_motion_network(feature=feature_3d, kp_driving=kp_driving, + kp_source=kp_source) + output_dict['mask'] = dense_motion['mask'] + + # import pdb; pdb.set_trace() + + if 'occlusion_map' in dense_motion: + occlusion_map = dense_motion['occlusion_map'] + output_dict['occlusion_map'] = occlusion_map + else: + occlusion_map = None + deformation = dense_motion['deformation'] + out = self.deform_input(feature_3d, deformation) + + bs, c, d, h, w = out.shape + out = out.view(bs, c*d, h, w) + out = self.third(out) + out = self.fourth(out) + + # occlusion_map = torch.where(occlusion_map < 0.95, 0, occlusion_map) + + if occlusion_map is not None: + if out.shape[2] != occlusion_map.shape[2] or out.shape[3] != occlusion_map.shape[3]: + occlusion_map = F.interpolate(occlusion_map, size=out.shape[2:], mode='bilinear') + out = out * occlusion_map + + # Decoding part + out = self.decoder(out) + + output_dict["prediction"] = out + + return output_dict \ No newline at end of file diff --git a/animation/SadTalker/src/facerender/modules/keypoint_detector.py b/animation/SadTalker/src/facerender/modules/keypoint_detector.py new file mode 100644 index 0000000..62a38a9 --- /dev/null +++ b/animation/SadTalker/src/facerender/modules/keypoint_detector.py @@ -0,0 +1,179 @@ +from torch import nn +import torch +import torch.nn.functional as F + +from src.facerender.sync_batchnorm import SynchronizedBatchNorm2d as BatchNorm2d +from src.facerender.modules.util import KPHourglass, make_coordinate_grid, AntiAliasInterpolation2d, ResBottleneck + + +class KPDetector(nn.Module): + """ + Detecting canonical keypoints. Return keypoint position and jacobian near each keypoint. + """ + + def __init__(self, block_expansion, feature_channel, num_kp, image_channel, max_features, reshape_channel, reshape_depth, + num_blocks, temperature, estimate_jacobian=False, scale_factor=1, single_jacobian_map=False): + super(KPDetector, self).__init__() + + self.predictor = KPHourglass(block_expansion, in_features=image_channel, + max_features=max_features, reshape_features=reshape_channel, reshape_depth=reshape_depth, num_blocks=num_blocks) + + # self.kp = nn.Conv3d(in_channels=self.predictor.out_filters, out_channels=num_kp, kernel_size=7, padding=3) + self.kp = nn.Conv3d(in_channels=self.predictor.out_filters, out_channels=num_kp, kernel_size=3, padding=1) + + if estimate_jacobian: + self.num_jacobian_maps = 1 if single_jacobian_map else num_kp + # self.jacobian = nn.Conv3d(in_channels=self.predictor.out_filters, out_channels=9 * self.num_jacobian_maps, kernel_size=7, padding=3) + self.jacobian = nn.Conv3d(in_channels=self.predictor.out_filters, out_channels=9 * self.num_jacobian_maps, kernel_size=3, padding=1) + ''' + initial as: + [[1 0 0] + [0 1 0] + [0 0 1]] + ''' + self.jacobian.weight.data.zero_() + self.jacobian.bias.data.copy_(torch.tensor([1, 0, 0, 0, 1, 0, 0, 0, 1] * self.num_jacobian_maps, dtype=torch.float)) + else: + self.jacobian = None + + self.temperature = temperature + self.scale_factor = scale_factor + if self.scale_factor != 1: + self.down = AntiAliasInterpolation2d(image_channel, self.scale_factor) + + def gaussian2kp(self, heatmap): + """ + Extract the mean from a heatmap + """ + shape = heatmap.shape + heatmap = heatmap.unsqueeze(-1) + grid = make_coordinate_grid(shape[2:], heatmap.type()).unsqueeze_(0).unsqueeze_(0) + value = (heatmap * grid).sum(dim=(2, 3, 4)) + kp = {'value': value} + + return kp + + def forward(self, x): + if self.scale_factor != 1: + x = self.down(x) + + feature_map = self.predictor(x) + prediction = self.kp(feature_map) + + final_shape = prediction.shape + heatmap = prediction.view(final_shape[0], final_shape[1], -1) + heatmap = F.softmax(heatmap / self.temperature, dim=2) + heatmap = heatmap.view(*final_shape) + + out = self.gaussian2kp(heatmap) + + if self.jacobian is not None: + jacobian_map = self.jacobian(feature_map) + jacobian_map = jacobian_map.reshape(final_shape[0], self.num_jacobian_maps, 9, final_shape[2], + final_shape[3], final_shape[4]) + heatmap = heatmap.unsqueeze(2) + + jacobian = heatmap * jacobian_map + jacobian = jacobian.view(final_shape[0], final_shape[1], 9, -1) + jacobian = jacobian.sum(dim=-1) + jacobian = jacobian.view(jacobian.shape[0], jacobian.shape[1], 3, 3) + out['jacobian'] = jacobian + + return out + + +class HEEstimator(nn.Module): + """ + Estimating head pose and expression. + """ + + def __init__(self, block_expansion, feature_channel, num_kp, image_channel, max_features, num_bins=66, estimate_jacobian=True): + super(HEEstimator, self).__init__() + + self.conv1 = nn.Conv2d(in_channels=image_channel, out_channels=block_expansion, kernel_size=7, padding=3, stride=2) + self.norm1 = BatchNorm2d(block_expansion, affine=True) + self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + + self.conv2 = nn.Conv2d(in_channels=block_expansion, out_channels=256, kernel_size=1) + self.norm2 = BatchNorm2d(256, affine=True) + + self.block1 = nn.Sequential() + for i in range(3): + self.block1.add_module('b1_'+ str(i), ResBottleneck(in_features=256, stride=1)) + + self.conv3 = nn.Conv2d(in_channels=256, out_channels=512, kernel_size=1) + self.norm3 = BatchNorm2d(512, affine=True) + self.block2 = ResBottleneck(in_features=512, stride=2) + + self.block3 = nn.Sequential() + for i in range(3): + self.block3.add_module('b3_'+ str(i), ResBottleneck(in_features=512, stride=1)) + + self.conv4 = nn.Conv2d(in_channels=512, out_channels=1024, kernel_size=1) + self.norm4 = BatchNorm2d(1024, affine=True) + self.block4 = ResBottleneck(in_features=1024, stride=2) + + self.block5 = nn.Sequential() + for i in range(5): + self.block5.add_module('b5_'+ str(i), ResBottleneck(in_features=1024, stride=1)) + + self.conv5 = nn.Conv2d(in_channels=1024, out_channels=2048, kernel_size=1) + self.norm5 = BatchNorm2d(2048, affine=True) + self.block6 = ResBottleneck(in_features=2048, stride=2) + + self.block7 = nn.Sequential() + for i in range(2): + self.block7.add_module('b7_'+ str(i), ResBottleneck(in_features=2048, stride=1)) + + self.fc_roll = nn.Linear(2048, num_bins) + self.fc_pitch = nn.Linear(2048, num_bins) + self.fc_yaw = nn.Linear(2048, num_bins) + + self.fc_t = nn.Linear(2048, 3) + + self.fc_exp = nn.Linear(2048, 3*num_kp) + + def forward(self, x): + out = self.conv1(x) + out = self.norm1(out) + out = F.relu(out) + out = self.maxpool(out) + + out = self.conv2(out) + out = self.norm2(out) + out = F.relu(out) + + out = self.block1(out) + + out = self.conv3(out) + out = self.norm3(out) + out = F.relu(out) + out = self.block2(out) + + out = self.block3(out) + + out = self.conv4(out) + out = self.norm4(out) + out = F.relu(out) + out = self.block4(out) + + out = self.block5(out) + + out = self.conv5(out) + out = self.norm5(out) + out = F.relu(out) + out = self.block6(out) + + out = self.block7(out) + + out = F.adaptive_avg_pool2d(out, 1) + out = out.view(out.shape[0], -1) + + yaw = self.fc_roll(out) + pitch = self.fc_pitch(out) + roll = self.fc_yaw(out) + t = self.fc_t(out) + exp = self.fc_exp(out) + + return {'yaw': yaw, 'pitch': pitch, 'roll': roll, 't': t, 'exp': exp} + diff --git a/animation/SadTalker/src/facerender/modules/make_animation.py b/animation/SadTalker/src/facerender/modules/make_animation.py new file mode 100644 index 0000000..3360c53 --- /dev/null +++ b/animation/SadTalker/src/facerender/modules/make_animation.py @@ -0,0 +1,170 @@ +from scipy.spatial import ConvexHull +import torch +import torch.nn.functional as F +import numpy as np +from tqdm import tqdm + +def normalize_kp(kp_source, kp_driving, kp_driving_initial, adapt_movement_scale=False, + use_relative_movement=False, use_relative_jacobian=False): + if adapt_movement_scale: + source_area = ConvexHull(kp_source['value'][0].data.cpu().numpy()).volume + driving_area = ConvexHull(kp_driving_initial['value'][0].data.cpu().numpy()).volume + adapt_movement_scale = np.sqrt(source_area) / np.sqrt(driving_area) + else: + adapt_movement_scale = 1 + + kp_new = {k: v for k, v in kp_driving.items()} + + if use_relative_movement: + kp_value_diff = (kp_driving['value'] - kp_driving_initial['value']) + kp_value_diff *= adapt_movement_scale + kp_new['value'] = kp_value_diff + kp_source['value'] + + if use_relative_jacobian: + jacobian_diff = torch.matmul(kp_driving['jacobian'], torch.inverse(kp_driving_initial['jacobian'])) + kp_new['jacobian'] = torch.matmul(jacobian_diff, kp_source['jacobian']) + + return kp_new + +def headpose_pred_to_degree(pred): + device = pred.device + idx_tensor = [idx for idx in range(66)] + idx_tensor = torch.FloatTensor(idx_tensor).type_as(pred).to(device) + pred = F.softmax(pred) + degree = torch.sum(pred*idx_tensor, 1) * 3 - 99 + return degree + +def get_rotation_matrix(yaw, pitch, roll): + yaw = yaw / 180 * 3.14 + pitch = pitch / 180 * 3.14 + roll = roll / 180 * 3.14 + + roll = roll.unsqueeze(1) + pitch = pitch.unsqueeze(1) + yaw = yaw.unsqueeze(1) + + pitch_mat = torch.cat([torch.ones_like(pitch), torch.zeros_like(pitch), torch.zeros_like(pitch), + torch.zeros_like(pitch), torch.cos(pitch), -torch.sin(pitch), + torch.zeros_like(pitch), torch.sin(pitch), torch.cos(pitch)], dim=1) + pitch_mat = pitch_mat.view(pitch_mat.shape[0], 3, 3) + + yaw_mat = torch.cat([torch.cos(yaw), torch.zeros_like(yaw), torch.sin(yaw), + torch.zeros_like(yaw), torch.ones_like(yaw), torch.zeros_like(yaw), + -torch.sin(yaw), torch.zeros_like(yaw), torch.cos(yaw)], dim=1) + yaw_mat = yaw_mat.view(yaw_mat.shape[0], 3, 3) + + roll_mat = torch.cat([torch.cos(roll), -torch.sin(roll), torch.zeros_like(roll), + torch.sin(roll), torch.cos(roll), torch.zeros_like(roll), + torch.zeros_like(roll), torch.zeros_like(roll), torch.ones_like(roll)], dim=1) + roll_mat = roll_mat.view(roll_mat.shape[0], 3, 3) + + rot_mat = torch.einsum('bij,bjk,bkm->bim', pitch_mat, yaw_mat, roll_mat) + + return rot_mat + +def keypoint_transformation(kp_canonical, he, wo_exp=False): + kp = kp_canonical['value'] # (bs, k, 3) + yaw, pitch, roll= he['yaw'], he['pitch'], he['roll'] + yaw = headpose_pred_to_degree(yaw) + pitch = headpose_pred_to_degree(pitch) + roll = headpose_pred_to_degree(roll) + + if 'yaw_in' in he: + yaw = he['yaw_in'] + if 'pitch_in' in he: + pitch = he['pitch_in'] + if 'roll_in' in he: + roll = he['roll_in'] + + rot_mat = get_rotation_matrix(yaw, pitch, roll) # (bs, 3, 3) + + t, exp = he['t'], he['exp'] + if wo_exp: + exp = exp*0 + + # keypoint rotation + kp_rotated = torch.einsum('bmp,bkp->bkm', rot_mat, kp) + + # keypoint translation + t[:, 0] = t[:, 0]*0 + t[:, 2] = t[:, 2]*0 + t = t.unsqueeze(1).repeat(1, kp.shape[1], 1) + kp_t = kp_rotated + t + + # add expression deviation + exp = exp.view(exp.shape[0], -1, 3) + kp_transformed = kp_t + exp + + return {'value': kp_transformed} + + + +def make_animation(source_image, source_semantics, target_semantics, + generator, kp_detector, he_estimator, mapping, + yaw_c_seq=None, pitch_c_seq=None, roll_c_seq=None, + use_exp=True, use_half=False): + with torch.no_grad(): + predictions = [] + + kp_canonical = kp_detector(source_image) + he_source = mapping(source_semantics) + kp_source = keypoint_transformation(kp_canonical, he_source) + + for frame_idx in tqdm(range(target_semantics.shape[1]), 'Face Renderer:'): + # still check the dimension + # print(target_semantics.shape, source_semantics.shape) + target_semantics_frame = target_semantics[:, frame_idx] + he_driving = mapping(target_semantics_frame) + if yaw_c_seq is not None: + he_driving['yaw_in'] = yaw_c_seq[:, frame_idx] + if pitch_c_seq is not None: + he_driving['pitch_in'] = pitch_c_seq[:, frame_idx] + if roll_c_seq is not None: + he_driving['roll_in'] = roll_c_seq[:, frame_idx] + + kp_driving = keypoint_transformation(kp_canonical, he_driving) + + kp_norm = kp_driving + out = generator(source_image, kp_source=kp_source, kp_driving=kp_norm) + ''' + source_image_new = out['prediction'].squeeze(1) + kp_canonical_new = kp_detector(source_image_new) + he_source_new = he_estimator(source_image_new) + kp_source_new = keypoint_transformation(kp_canonical_new, he_source_new, wo_exp=True) + kp_driving_new = keypoint_transformation(kp_canonical_new, he_driving, wo_exp=True) + out = generator(source_image_new, kp_source=kp_source_new, kp_driving=kp_driving_new) + ''' + predictions.append(out['prediction']) + predictions_ts = torch.stack(predictions, dim=1) + return predictions_ts + +class AnimateModel(torch.nn.Module): + """ + Merge all generator related updates into single model for better multi-gpu usage + """ + + def __init__(self, generator, kp_extractor, mapping): + super(AnimateModel, self).__init__() + self.kp_extractor = kp_extractor + self.generator = generator + self.mapping = mapping + + self.kp_extractor.eval() + self.generator.eval() + self.mapping.eval() + + def forward(self, x): + + source_image = x['source_image'] + source_semantics = x['source_semantics'] + target_semantics = x['target_semantics'] + yaw_c_seq = x['yaw_c_seq'] + pitch_c_seq = x['pitch_c_seq'] + roll_c_seq = x['roll_c_seq'] + + predictions_video = make_animation(source_image, source_semantics, target_semantics, + self.generator, self.kp_extractor, + self.mapping, use_exp = True, + yaw_c_seq=yaw_c_seq, pitch_c_seq=pitch_c_seq, roll_c_seq=roll_c_seq) + + return predictions_video \ No newline at end of file diff --git a/animation/SadTalker/src/facerender/modules/mapping.py b/animation/SadTalker/src/facerender/modules/mapping.py new file mode 100644 index 0000000..0e3a1c2 --- /dev/null +++ b/animation/SadTalker/src/facerender/modules/mapping.py @@ -0,0 +1,47 @@ +import numpy as np + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class MappingNet(nn.Module): + def __init__(self, coeff_nc, descriptor_nc, layer, num_kp, num_bins): + super( MappingNet, self).__init__() + + self.layer = layer + nonlinearity = nn.LeakyReLU(0.1) + + self.first = nn.Sequential( + torch.nn.Conv1d(coeff_nc, descriptor_nc, kernel_size=7, padding=0, bias=True)) + + for i in range(layer): + net = nn.Sequential(nonlinearity, + torch.nn.Conv1d(descriptor_nc, descriptor_nc, kernel_size=3, padding=0, dilation=3)) + setattr(self, 'encoder' + str(i), net) + + self.pooling = nn.AdaptiveAvgPool1d(1) + self.output_nc = descriptor_nc + + self.fc_roll = nn.Linear(descriptor_nc, num_bins) + self.fc_pitch = nn.Linear(descriptor_nc, num_bins) + self.fc_yaw = nn.Linear(descriptor_nc, num_bins) + self.fc_t = nn.Linear(descriptor_nc, 3) + self.fc_exp = nn.Linear(descriptor_nc, 3*num_kp) + + def forward(self, input_3dmm): + out = self.first(input_3dmm) + for i in range(self.layer): + model = getattr(self, 'encoder' + str(i)) + out = model(out) + out[:,:,3:-3] + out = self.pooling(out) + out = out.view(out.shape[0], -1) + #print('out:', out.shape) + + yaw = self.fc_yaw(out) + pitch = self.fc_pitch(out) + roll = self.fc_roll(out) + t = self.fc_t(out) + exp = self.fc_exp(out) + + return {'yaw': yaw, 'pitch': pitch, 'roll': roll, 't': t, 'exp': exp} \ No newline at end of file diff --git a/animation/SadTalker/src/facerender/modules/util.py b/animation/SadTalker/src/facerender/modules/util.py new file mode 100644 index 0000000..b916dee --- /dev/null +++ b/animation/SadTalker/src/facerender/modules/util.py @@ -0,0 +1,564 @@ +from torch import nn + +import torch.nn.functional as F +import torch + +from src.facerender.sync_batchnorm import SynchronizedBatchNorm2d as BatchNorm2d +from src.facerender.sync_batchnorm import SynchronizedBatchNorm3d as BatchNorm3d + +import torch.nn.utils.spectral_norm as spectral_norm + + +def kp2gaussian(kp, spatial_size, kp_variance): + """ + Transform a keypoint into gaussian like representation + """ + mean = kp['value'] + + coordinate_grid = make_coordinate_grid(spatial_size, mean.type()) + number_of_leading_dimensions = len(mean.shape) - 1 + shape = (1,) * number_of_leading_dimensions + coordinate_grid.shape + coordinate_grid = coordinate_grid.view(*shape) + repeats = mean.shape[:number_of_leading_dimensions] + (1, 1, 1, 1) + coordinate_grid = coordinate_grid.repeat(*repeats) + + # Preprocess kp shape + shape = mean.shape[:number_of_leading_dimensions] + (1, 1, 1, 3) + mean = mean.view(*shape) + + mean_sub = (coordinate_grid - mean) + + out = torch.exp(-0.5 * (mean_sub ** 2).sum(-1) / kp_variance) + + return out + +def make_coordinate_grid_2d(spatial_size, type): + """ + Create a meshgrid [-1,1] x [-1,1] of given spatial_size. + """ + h, w = spatial_size + x = torch.arange(w).type(type) + y = torch.arange(h).type(type) + + x = (2 * (x / (w - 1)) - 1) + y = (2 * (y / (h - 1)) - 1) + + yy = y.view(-1, 1).repeat(1, w) + xx = x.view(1, -1).repeat(h, 1) + + meshed = torch.cat([xx.unsqueeze_(2), yy.unsqueeze_(2)], 2) + + return meshed + + +def make_coordinate_grid(spatial_size, type): + d, h, w = spatial_size + x = torch.arange(w).type(type) + y = torch.arange(h).type(type) + z = torch.arange(d).type(type) + + x = (2 * (x / (w - 1)) - 1) + y = (2 * (y / (h - 1)) - 1) + z = (2 * (z / (d - 1)) - 1) + + yy = y.view(1, -1, 1).repeat(d, 1, w) + xx = x.view(1, 1, -1).repeat(d, h, 1) + zz = z.view(-1, 1, 1).repeat(1, h, w) + + meshed = torch.cat([xx.unsqueeze_(3), yy.unsqueeze_(3), zz.unsqueeze_(3)], 3) + + return meshed + + +class ResBottleneck(nn.Module): + def __init__(self, in_features, stride): + super(ResBottleneck, self).__init__() + self.conv1 = nn.Conv2d(in_channels=in_features, out_channels=in_features//4, kernel_size=1) + self.conv2 = nn.Conv2d(in_channels=in_features//4, out_channels=in_features//4, kernel_size=3, padding=1, stride=stride) + self.conv3 = nn.Conv2d(in_channels=in_features//4, out_channels=in_features, kernel_size=1) + self.norm1 = BatchNorm2d(in_features//4, affine=True) + self.norm2 = BatchNorm2d(in_features//4, affine=True) + self.norm3 = BatchNorm2d(in_features, affine=True) + + self.stride = stride + if self.stride != 1: + self.skip = nn.Conv2d(in_channels=in_features, out_channels=in_features, kernel_size=1, stride=stride) + self.norm4 = BatchNorm2d(in_features, affine=True) + + def forward(self, x): + out = self.conv1(x) + out = self.norm1(out) + out = F.relu(out) + out = self.conv2(out) + out = self.norm2(out) + out = F.relu(out) + out = self.conv3(out) + out = self.norm3(out) + if self.stride != 1: + x = self.skip(x) + x = self.norm4(x) + out += x + out = F.relu(out) + return out + + +class ResBlock2d(nn.Module): + """ + Res block, preserve spatial resolution. + """ + + def __init__(self, in_features, kernel_size, padding): + super(ResBlock2d, self).__init__() + self.conv1 = nn.Conv2d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size, + padding=padding) + self.conv2 = nn.Conv2d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size, + padding=padding) + self.norm1 = BatchNorm2d(in_features, affine=True) + self.norm2 = BatchNorm2d(in_features, affine=True) + + def forward(self, x): + out = self.norm1(x) + out = F.relu(out) + out = self.conv1(out) + out = self.norm2(out) + out = F.relu(out) + out = self.conv2(out) + out += x + return out + + +class ResBlock3d(nn.Module): + """ + Res block, preserve spatial resolution. + """ + + def __init__(self, in_features, kernel_size, padding): + super(ResBlock3d, self).__init__() + self.conv1 = nn.Conv3d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size, + padding=padding) + self.conv2 = nn.Conv3d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size, + padding=padding) + self.norm1 = BatchNorm3d(in_features, affine=True) + self.norm2 = BatchNorm3d(in_features, affine=True) + + def forward(self, x): + out = self.norm1(x) + out = F.relu(out) + out = self.conv1(out) + out = self.norm2(out) + out = F.relu(out) + out = self.conv2(out) + out += x + return out + + +class UpBlock2d(nn.Module): + """ + Upsampling block for use in decoder. + """ + + def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1): + super(UpBlock2d, self).__init__() + + self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, + padding=padding, groups=groups) + self.norm = BatchNorm2d(out_features, affine=True) + + def forward(self, x): + out = F.interpolate(x, scale_factor=2) + out = self.conv(out) + out = self.norm(out) + out = F.relu(out) + return out + +class UpBlock3d(nn.Module): + """ + Upsampling block for use in decoder. + """ + + def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1): + super(UpBlock3d, self).__init__() + + self.conv = nn.Conv3d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, + padding=padding, groups=groups) + self.norm = BatchNorm3d(out_features, affine=True) + + def forward(self, x): + # out = F.interpolate(x, scale_factor=(1, 2, 2), mode='trilinear') + out = F.interpolate(x, scale_factor=(1, 2, 2)) + out = self.conv(out) + out = self.norm(out) + out = F.relu(out) + return out + + +class DownBlock2d(nn.Module): + """ + Downsampling block for use in encoder. + """ + + def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1): + super(DownBlock2d, self).__init__() + self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, + padding=padding, groups=groups) + self.norm = BatchNorm2d(out_features, affine=True) + self.pool = nn.AvgPool2d(kernel_size=(2, 2)) + + def forward(self, x): + out = self.conv(x) + out = self.norm(out) + out = F.relu(out) + out = self.pool(out) + return out + + +class DownBlock3d(nn.Module): + """ + Downsampling block for use in encoder. + """ + + def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1): + super(DownBlock3d, self).__init__() + ''' + self.conv = nn.Conv3d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, + padding=padding, groups=groups, stride=(1, 2, 2)) + ''' + self.conv = nn.Conv3d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, + padding=padding, groups=groups) + self.norm = BatchNorm3d(out_features, affine=True) + self.pool = nn.AvgPool3d(kernel_size=(1, 2, 2)) + + def forward(self, x): + out = self.conv(x) + out = self.norm(out) + out = F.relu(out) + out = self.pool(out) + return out + + +class SameBlock2d(nn.Module): + """ + Simple block, preserve spatial resolution. + """ + + def __init__(self, in_features, out_features, groups=1, kernel_size=3, padding=1, lrelu=False): + super(SameBlock2d, self).__init__() + self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, + kernel_size=kernel_size, padding=padding, groups=groups) + self.norm = BatchNorm2d(out_features, affine=True) + if lrelu: + self.ac = nn.LeakyReLU() + else: + self.ac = nn.ReLU() + + def forward(self, x): + out = self.conv(x) + out = self.norm(out) + out = self.ac(out) + return out + + +class Encoder(nn.Module): + """ + Hourglass Encoder + """ + + def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256): + super(Encoder, self).__init__() + + down_blocks = [] + for i in range(num_blocks): + down_blocks.append(DownBlock3d(in_features if i == 0 else min(max_features, block_expansion * (2 ** i)), + min(max_features, block_expansion * (2 ** (i + 1))), + kernel_size=3, padding=1)) + self.down_blocks = nn.ModuleList(down_blocks) + + def forward(self, x): + outs = [x] + for down_block in self.down_blocks: + outs.append(down_block(outs[-1])) + return outs + + +class Decoder(nn.Module): + """ + Hourglass Decoder + """ + + def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256): + super(Decoder, self).__init__() + + up_blocks = [] + + for i in range(num_blocks)[::-1]: + in_filters = (1 if i == num_blocks - 1 else 2) * min(max_features, block_expansion * (2 ** (i + 1))) + out_filters = min(max_features, block_expansion * (2 ** i)) + up_blocks.append(UpBlock3d(in_filters, out_filters, kernel_size=3, padding=1)) + + self.up_blocks = nn.ModuleList(up_blocks) + # self.out_filters = block_expansion + self.out_filters = block_expansion + in_features + + self.conv = nn.Conv3d(in_channels=self.out_filters, out_channels=self.out_filters, kernel_size=3, padding=1) + self.norm = BatchNorm3d(self.out_filters, affine=True) + + def forward(self, x): + out = x.pop() + # for up_block in self.up_blocks[:-1]: + for up_block in self.up_blocks: + out = up_block(out) + skip = x.pop() + out = torch.cat([out, skip], dim=1) + # out = self.up_blocks[-1](out) + out = self.conv(out) + out = self.norm(out) + out = F.relu(out) + return out + + +class Hourglass(nn.Module): + """ + Hourglass architecture. + """ + + def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256): + super(Hourglass, self).__init__() + self.encoder = Encoder(block_expansion, in_features, num_blocks, max_features) + self.decoder = Decoder(block_expansion, in_features, num_blocks, max_features) + self.out_filters = self.decoder.out_filters + + def forward(self, x): + return self.decoder(self.encoder(x)) + + +class KPHourglass(nn.Module): + """ + Hourglass architecture. + """ + + def __init__(self, block_expansion, in_features, reshape_features, reshape_depth, num_blocks=3, max_features=256): + super(KPHourglass, self).__init__() + + self.down_blocks = nn.Sequential() + for i in range(num_blocks): + self.down_blocks.add_module('down'+ str(i), DownBlock2d(in_features if i == 0 else min(max_features, block_expansion * (2 ** i)), + min(max_features, block_expansion * (2 ** (i + 1))), + kernel_size=3, padding=1)) + + in_filters = min(max_features, block_expansion * (2 ** num_blocks)) + self.conv = nn.Conv2d(in_channels=in_filters, out_channels=reshape_features, kernel_size=1) + + self.up_blocks = nn.Sequential() + for i in range(num_blocks): + in_filters = min(max_features, block_expansion * (2 ** (num_blocks - i))) + out_filters = min(max_features, block_expansion * (2 ** (num_blocks - i - 1))) + self.up_blocks.add_module('up'+ str(i), UpBlock3d(in_filters, out_filters, kernel_size=3, padding=1)) + + self.reshape_depth = reshape_depth + self.out_filters = out_filters + + def forward(self, x): + out = self.down_blocks(x) + out = self.conv(out) + bs, c, h, w = out.shape + out = out.view(bs, c//self.reshape_depth, self.reshape_depth, h, w) + out = self.up_blocks(out) + + return out + + + +class AntiAliasInterpolation2d(nn.Module): + """ + Band-limited downsampling, for better preservation of the input signal. + """ + def __init__(self, channels, scale): + super(AntiAliasInterpolation2d, self).__init__() + sigma = (1 / scale - 1) / 2 + kernel_size = 2 * round(sigma * 4) + 1 + self.ka = kernel_size // 2 + self.kb = self.ka - 1 if kernel_size % 2 == 0 else self.ka + + kernel_size = [kernel_size, kernel_size] + sigma = [sigma, sigma] + # The gaussian kernel is the product of the + # gaussian function of each dimension. + kernel = 1 + meshgrids = torch.meshgrid( + [ + torch.arange(size, dtype=torch.float32) + for size in kernel_size + ] + ) + for size, std, mgrid in zip(kernel_size, sigma, meshgrids): + mean = (size - 1) / 2 + kernel *= torch.exp(-(mgrid - mean) ** 2 / (2 * std ** 2)) + + # Make sure sum of values in gaussian kernel equals 1. + kernel = kernel / torch.sum(kernel) + # Reshape to depthwise convolutional weight + kernel = kernel.view(1, 1, *kernel.size()) + kernel = kernel.repeat(channels, *[1] * (kernel.dim() - 1)) + + self.register_buffer('weight', kernel) + self.groups = channels + self.scale = scale + inv_scale = 1 / scale + self.int_inv_scale = int(inv_scale) + + def forward(self, input): + if self.scale == 1.0: + return input + + out = F.pad(input, (self.ka, self.kb, self.ka, self.kb)) + out = F.conv2d(out, weight=self.weight, groups=self.groups) + out = out[:, :, ::self.int_inv_scale, ::self.int_inv_scale] + + return out + + +class SPADE(nn.Module): + def __init__(self, norm_nc, label_nc): + super().__init__() + + self.param_free_norm = nn.InstanceNorm2d(norm_nc, affine=False) + nhidden = 128 + + self.mlp_shared = nn.Sequential( + nn.Conv2d(label_nc, nhidden, kernel_size=3, padding=1), + nn.ReLU()) + self.mlp_gamma = nn.Conv2d(nhidden, norm_nc, kernel_size=3, padding=1) + self.mlp_beta = nn.Conv2d(nhidden, norm_nc, kernel_size=3, padding=1) + + def forward(self, x, segmap): + normalized = self.param_free_norm(x) + segmap = F.interpolate(segmap, size=x.size()[2:], mode='nearest') + actv = self.mlp_shared(segmap) + gamma = self.mlp_gamma(actv) + beta = self.mlp_beta(actv) + out = normalized * (1 + gamma) + beta + return out + + +class SPADEResnetBlock(nn.Module): + def __init__(self, fin, fout, norm_G, label_nc, use_se=False, dilation=1): + super().__init__() + # Attributes + self.learned_shortcut = (fin != fout) + fmiddle = min(fin, fout) + self.use_se = use_se + # create conv layers + self.conv_0 = nn.Conv2d(fin, fmiddle, kernel_size=3, padding=dilation, dilation=dilation) + self.conv_1 = nn.Conv2d(fmiddle, fout, kernel_size=3, padding=dilation, dilation=dilation) + if self.learned_shortcut: + self.conv_s = nn.Conv2d(fin, fout, kernel_size=1, bias=False) + # apply spectral norm if specified + if 'spectral' in norm_G: + self.conv_0 = spectral_norm(self.conv_0) + self.conv_1 = spectral_norm(self.conv_1) + if self.learned_shortcut: + self.conv_s = spectral_norm(self.conv_s) + # define normalization layers + self.norm_0 = SPADE(fin, label_nc) + self.norm_1 = SPADE(fmiddle, label_nc) + if self.learned_shortcut: + self.norm_s = SPADE(fin, label_nc) + + def forward(self, x, seg1): + x_s = self.shortcut(x, seg1) + dx = self.conv_0(self.actvn(self.norm_0(x, seg1))) + dx = self.conv_1(self.actvn(self.norm_1(dx, seg1))) + out = x_s + dx + return out + + def shortcut(self, x, seg1): + if self.learned_shortcut: + x_s = self.conv_s(self.norm_s(x, seg1)) + else: + x_s = x + return x_s + + def actvn(self, x): + return F.leaky_relu(x, 2e-1) + +class audio2image(nn.Module): + def __init__(self, generator, kp_extractor, he_estimator_video, he_estimator_audio, train_params): + super().__init__() + # Attributes + self.generator = generator + self.kp_extractor = kp_extractor + self.he_estimator_video = he_estimator_video + self.he_estimator_audio = he_estimator_audio + self.train_params = train_params + + def headpose_pred_to_degree(self, pred): + device = pred.device + idx_tensor = [idx for idx in range(66)] + idx_tensor = torch.FloatTensor(idx_tensor).to(device) + pred = F.softmax(pred) + degree = torch.sum(pred*idx_tensor, 1) * 3 - 99 + + return degree + + def get_rotation_matrix(self, yaw, pitch, roll): + yaw = yaw / 180 * 3.14 + pitch = pitch / 180 * 3.14 + roll = roll / 180 * 3.14 + + roll = roll.unsqueeze(1) + pitch = pitch.unsqueeze(1) + yaw = yaw.unsqueeze(1) + + roll_mat = torch.cat([torch.ones_like(roll), torch.zeros_like(roll), torch.zeros_like(roll), + torch.zeros_like(roll), torch.cos(roll), -torch.sin(roll), + torch.zeros_like(roll), torch.sin(roll), torch.cos(roll)], dim=1) + roll_mat = roll_mat.view(roll_mat.shape[0], 3, 3) + + pitch_mat = torch.cat([torch.cos(pitch), torch.zeros_like(pitch), torch.sin(pitch), + torch.zeros_like(pitch), torch.ones_like(pitch), torch.zeros_like(pitch), + -torch.sin(pitch), torch.zeros_like(pitch), torch.cos(pitch)], dim=1) + pitch_mat = pitch_mat.view(pitch_mat.shape[0], 3, 3) + + yaw_mat = torch.cat([torch.cos(yaw), -torch.sin(yaw), torch.zeros_like(yaw), + torch.sin(yaw), torch.cos(yaw), torch.zeros_like(yaw), + torch.zeros_like(yaw), torch.zeros_like(yaw), torch.ones_like(yaw)], dim=1) + yaw_mat = yaw_mat.view(yaw_mat.shape[0], 3, 3) + + rot_mat = torch.einsum('bij,bjk,bkm->bim', roll_mat, pitch_mat, yaw_mat) + + return rot_mat + + def keypoint_transformation(self, kp_canonical, he): + kp = kp_canonical['value'] # (bs, k, 3) + yaw, pitch, roll = he['yaw'], he['pitch'], he['roll'] + t, exp = he['t'], he['exp'] + + yaw = self.headpose_pred_to_degree(yaw) + pitch = self.headpose_pred_to_degree(pitch) + roll = self.headpose_pred_to_degree(roll) + + rot_mat = self.get_rotation_matrix(yaw, pitch, roll) # (bs, 3, 3) + + # keypoint rotation + kp_rotated = torch.einsum('bmp,bkp->bkm', rot_mat, kp) + + + + # keypoint translation + t = t.unsqueeze_(1).repeat(1, kp.shape[1], 1) + kp_t = kp_rotated + t + + # add expression deviation + exp = exp.view(exp.shape[0], -1, 3) + kp_transformed = kp_t + exp + + return {'value': kp_transformed} + + def forward(self, source_image, target_audio): + pose_source = self.he_estimator_video(source_image) + pose_generated = self.he_estimator_audio(target_audio) + kp_canonical = self.kp_extractor(source_image) + kp_source = self.keypoint_transformation(kp_canonical, pose_source) + kp_transformed_generated = self.keypoint_transformation(kp_canonical, pose_generated) + generated = self.generator(source_image, kp_source=kp_source, kp_driving=kp_transformed_generated) + return generated \ No newline at end of file diff --git a/animation/SadTalker/src/facerender/sync_batchnorm/__init__.py b/animation/SadTalker/src/facerender/sync_batchnorm/__init__.py new file mode 100644 index 0000000..bc8709d --- /dev/null +++ b/animation/SadTalker/src/facerender/sync_batchnorm/__init__.py @@ -0,0 +1,12 @@ +# -*- coding: utf-8 -*- +# File : __init__.py +# Author : Jiayuan Mao +# Email : maojiayuan@gmail.com +# Date : 27/01/2018 +# +# This file is part of Synchronized-BatchNorm-PyTorch. +# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch +# Distributed under MIT License. + +from .batchnorm import SynchronizedBatchNorm1d, SynchronizedBatchNorm2d, SynchronizedBatchNorm3d +from .replicate import DataParallelWithCallback, patch_replication_callback diff --git a/animation/SadTalker/src/facerender/sync_batchnorm/__pycache__/__init__.cpython-38.pyc b/animation/SadTalker/src/facerender/sync_batchnorm/__pycache__/__init__.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..52145db84d2318c5f974a0a517d1eed5a1696d0e GIT binary patch literal 413 zcmZ{g%}N6?5XZB-?dmGM35s{YJt(L*5fR&S5o!^ZKuD7*jm^h0Nh$S>d<ubzAb zPbT2kr33jf|M`=dNv4Cr9*JSrC-9aL@}9*1$fLN6$Dc+>A~Q`))?_Aca#J)#mTI}~ znSRqZWm867Oi11Ph`mT%bcd6T6$^A$y@0&uK`h1&&4t`zTy8V2WSw_MBOdqzN3J#K zClwaMn9aCYcE@dM%@A6x1P{ttCceFKz0*~*WZO6y(|fQ;=tXqZ&lSLaDwYwi1gn%M ze9nj0bm{>;eO~#X=*X?Dc3gTobf!h{9weP``N;JW@XT;!|0MM&=$s2cuo95^_!;bX c7s{BjDlm55Zo1tY7pV4pw*UYD literal 0 HcmV?d00001 diff --git a/animation/SadTalker/src/facerender/sync_batchnorm/__pycache__/batchnorm.cpython-38.pyc b/animation/SadTalker/src/facerender/sync_batchnorm/__pycache__/batchnorm.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e7f0ba10795c7d1c0baa6fafb669935fffb65894 GIT binary patch literal 12945 zcmeHN%WoXXdGE)(Iix6wdhD(vm*cFLvl)&hYK2W0E!VrGbyZjQS5@EhSKsuV$;pa=qqFs<_{F?o{F)l$M*)rR;Eewb1vj`E7)^aOo2F`8 z9osTYqmygqP3p0NLZ{d)s=6JNI^|}$Q)yPv&hdOO(Wy48s26yVm)dr7l9zeqkB&7t2?K0=SXrJQ+v&^OjM&_yp@QR%&KiaXxetZ(}8rotl-N`76PDp7?WmqHo*-`?>XniK|UlEF-t`X58zrXj4d*YzWr!Wt=cKT3&~;RpN&V2RrnvP$=IF&34U zXk{_v@M>4Xyp?xt|Mepl4;kckLdX*DEm6K>Lj4$H4$AW0IHmr%e8 z%8ui~WMWK$%s`R?NQUAqoU!te=0&S&70iOU^Yv%5w=r6Wr_>oQ(ox1q?<+(8Eh=Ql zvWn7-=M!ts25Z4goOzrDoJE{;uY|L_4~h28y~@5ZfDra3s7?lezA8+>9+=r3>aq9D z8Dn54lYMj79OUTA?Xw2zM16{z_bvHPsGZuIPEPl8{e0iv#XGCkpfD(g7twQO@9fh( zxIZI*)i;xK{lcz!(ZH2vS(pLk8ERqHyXFOBPzq-bKSlk6Qon?mYlF?L9N6F!VpQ$C z{3b4GDUom{e%MaU8)-%P3WwZAT9v(!?mHdfhUvtpwdu;VuqFKVMv~^%eK*zuCDB6W zWPQ0!Uu;Q|pz8%Pr&$?R$p_?UvXY!zT8MfH-aYDGimga)xsqSTGjRijVO7nXX~B%u z&s61abNP~a-kvc_X4S5kuUi-Gg7qzH#@@Mfj1~^bdSu0;gX(lH;fxXI;X`4s25j1J zNR<#*pco7_D4Ao2Q+@NEy_X}A;fegvqHpUkQc&_5a4$1RpBa&ipa!e-4>LhDk1;j8pN1 z;7tNUqX-MAV1Ho^wByCSWilTN?l=Dk92>TZs-^qv#Dr6(Ue#8HXWg^=HZQ+s@Jh?- zTR*pcYQV_y0~k2?n(r6-xqTaVii1*;hZC@2P$igixor$8NrB9F*|=AD}N4mxyfTSnhH&bWkv=1k|IQHgEGj7n_NGS^IdY$~S7>v%Mx^t9<6h9kRni|u6zculJ$uHufVyLi^@4`#Eo8kpK$7B~0z&%= zxU?*oWNzmtwD}mhD~hbzpC#DoAcD@3i-8M^Lbl~68|-O+*I>6d{Frq_vJvqL8y@f+ z5u+Upb_QXrTwy2L1kzxDj}oAsfg(9b4LY&GAWzm%y1w&Kw$~f^(wOx^!Xbi_M?EsL zI@{Xty$vSCgPsov$_SybMRujp2CH%W>X)@*I`skeSCUrpT`41}5e;F3G)MkVew(U! zAXRL|$dlDdN*Cn{>RP=hxGNoN7QhBzq-I6W%AtCR=0~-o%Jnf`dw?@0;A_~^1UK+2 zS(nU9*3Rs4N_Et5;B?48$v05khRp%mhi(S}9NTQC6Y*X^^{V3_5(T`M#Wu?iqHr)^Z?m@B>A1MLE)rK$(yFhRkBmU8h0Q%{ z$*7}Ld%foe$($duo7=Y|2?fp6`U4oSNp-s?7i@px2Z)!7HAOVyrO;hUOA8 z-jifgBX~bt?hUX=JcJ?#f@}kc0(ToU?+e&5bJ=i%%LE`fJOZLn3qU;1B&e3InTBb?nfq>tu!Il?Gt^sU z6#_Ad`=XoZnCL=+Ef>;cqjU_Bz1+BRcw%DH4SI^z!>Nr361%@W8jIEkOXW)NBg+~W z(1*fv@ql7N=1}M@PDY-kE7z9RbdFMMu2xa<6$kpOIUKUJ*_Ar`er=7#*b@qDRQv=+ zjyF0Kqk+TxFbbl!vM-1pnoSmRTXeb+@<<0>O}oSoWIJ9^q{+>Wr-l{|wX~lE+z=1D zkT_PXMLHCDMuDiTDQB+oe(d34DxW=s5O}h2cP)DZqF#WkiQR%Bg9*~`45Tz%gb`Nb zLccVB34~%F91vI5{lHJQwH4B!bS!JZF0d@FUgz=)QX=q<{~9Efy|eN6e91-77Ks?we%Xb*Tgen+WPHsT!nOO>g4 zp|{v9`7Wk1K@=#JV$}3xzN_Qfn2IefF0ya3PYIoVKXWG(-;_JieuNfG&vg*>$Du}M zKv($emCKiF*?bR(V)O8O)5YFjrDxx}fnX&LkD|5yJ%|txq2czDs4ZRI00QyJG0y7y z5^gh`oxfaXG|j^l^!5fT>&y;0Maw;ukf9AH?b8XTsXRC`o5nu&=l5;-_ONbI9qAT( zkFufY-AAJ2hh}0Y$e`|8*g7CLivO1=8G%eAZT+4Qy`8r*3!YzA=xljt%Cqe#VV@6% z&04J>D|kMwd_tC~PMmvv%N$kZq`c?SPUM zD$d?L!NmXnn7Gcg9%Vv!HV7U)_9dX+`B#K`HQhfj>aDY@*Xrz@YriM@r7Z3k`n?3i zYrO#CwO5`{27l)p8Q6X2`7l?~9}4m_DipXSc&9PtU8+3;$o?1Z4}om;`9OA_KLupv z-w@SjAypdqmqeQ8{|7;;z^H~hcrm1U?PZYa2-bc@kt%@n2~s^lswZjNETeRgJ~~O; zekp0&082=Na86V8_gxJEWexKp%((Q9C`K);Mo&kOMm2j zC6c{$mZ@DP%j2+Z_T-lUZ?C>0c&q8YB6v#~-Y*T}QjT~G;$C|>h)cNe`4IPAd@TQB z5clfKKwKJ_zW*IT+&`;bNGH4q-#B~p!AG3tRTe!xtCQ*Y54_b2RhJv2Rv4z0gSM=L zQ1kWdN&W9?+)$ft`a@1qGQIh=?19nO)D34uBv#)*=`Wpf^go*Z`Wl&K{CkGXGJT7$ zzJN_D^q~e%diWbye=K`}y6Dpjl?<0^)1cUTfb3O`KG8@^*Y)HU>7)CYK3XYRC48|| ZHK*}Cb;T~_OZhYAv_0LNet-J?{{`wLmk9s> literal 0 HcmV?d00001 diff --git a/animation/SadTalker/src/facerender/sync_batchnorm/__pycache__/comm.cpython-38.pyc b/animation/SadTalker/src/facerender/sync_batchnorm/__pycache__/comm.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..93ecaae35de6b64a4c88d78c6b5f989663ec9683 GIT binary patch literal 4849 zcmai2TW{OQ8RbP6TD~Nk%Wj-*i`i`!8@0C3F0eomEQ;Kk0!8C)67L0q3yS89B-#|I z%ux1PkYAEu``Eu=qfhBiBKp>+{R@5S`DRFovW-^842Q$He78C08-KID-m>r;9e*Z* z4=n57^iusb@bVBv9iieDw_~eoo_5#P?@rg@4tIHN;C9_VTD;DkcNTXv|y{6~4xc$m%um6O-t(|=B*(94t z@lvQsoaHO8;^0Vp9gRiadUcY7&jXbS8Pd?|&%w(>6m=h!ut;bKY;P}HZWxHYqUE3?m>b85M0)n?X#dv4>c&aGG0dwZwd$eX?_ zv^MfY$sn)&ISmiL}k(5 z)ATrr(}1hJC+TR6)uF`RtAOtY@u860qaaEauX{>{d;K63QY2i+Jq3IF`#}~ClT?oO z!gMt1j8AgU_oF1reE&A-N1 zArjUX5L@t0yDn)zNn++3kD$LK^Zh}}RLWKAY$RzG^-tQ4T-UE=Ep1LyU88PyJjw2& zp*&QU?b$b=^Yr?X#u%~`2c#^WJFVZWqsKXONj);aDNLZ1w!gESLfwBr z*@fD&j%C`lg-%9?+HnwNr3NlExFr5SWDHZ)8Y&XMWw)HhWfd@Nr`;qMC&5Vge%|u^ zQOYMVwby+A^(2UkFHPU)Y3Td%I#!gosQQR1io;4p*3-xA1{&IWo?Uaiy4UD#n8W;1 z3?h}u6L}LO@(q7};vWYJr#|91gHb;cGOxd$h=~X>Z0SdKeu$#(pt}0-BVg*%#pmfg z!V=a`B^zld*;WKFjDmZTM4>tCOEF8+YAD)`TWF9Q&nPmTnca8JZS6cArKgVh#%sjOa9no5a`}p!B)%5$x0q8z-wn8L&7R(zTdFf3 zwS&#b404kyidgwGsy?CWD%X&99TY_>l50Xdmwq(jA&Pp6>T2^|#8$onjBSDdR`K-2 z3gCYIo%Po4w)h5Ya!t}^0fXZnB`ny7_k%EFy#>JcSh{~8!mQJ3G5ogK(}@Hc0;dHI z`#3HfwMAG=eLTnlHV$ME$0BCOQ5>^OMuUNnEJ#?0?)@MrCFrCC)Z%VURPwMCt@_xkIJuC(e*v|l)v~F+ifF8@zN(>2zkPWqDCdxXO)j$Xy z2a-c-rM8M-WJU}YB}eHYr0J1g%?dpu?FbMa4DhKh)6oUlwc7g>%dkS23BILwC7Edx z&=-UZU_Ux<0C_$)Oex-@KKd~r6wy&>?sCViiaMp7##(F;zz0%)j}quL(x*rpiwmqc ziolcDyI87_MhmCif=}8lORg|7EuBO<=v1DdsC%e1M^uTmc2++p2yJj;k0z%R=uY35 zD56MuBnRq&-kIL$XUyxDG4cQwD%&J8>}kBgS)i>m476$HMlQ9w$|Xhz+U40!+sm&P zzMOBSaUrk2AXv*AKv#OM7E6PsRKEHb2`U!+X%r$ELVjSd61a>L6CDNca*6xbE@)CY z?iz#%r?urYg*Aw>LE1MLhi*aH7qgkU(TUE;vFbd>XoLb>;0|} zMa@Ndf#Fw+vKX@@ii-(zQT~J%vVy#iTIY4@PUf1#g@Xiq`l^G`wz9@s>(t4bGv_T? zt~Rs3vNDfC_rOM%c2^`{*j;Uof*Ef&T?F9xI# zED=9sIAKU?gE*QR`hAAXoDz9SXFRWq(KtJixVKpHbKOA;NS##V*Zk+2#*NjGM39>b zR+YrHdBf!M7lUD=gh9iuyBMm7AFWnH=XNb5M;rEZvl6AtGnRsb(o0Isbe#SZ4UIpz zAW?=eu;G4kHv-_P_19ac?hHv^86XtLukRwR8#Ck*r?pv)JDZqE<`~jK zOu8O$9Lah7JtuNQ$|X;9mPG7ofTaX|`b;ERC9nn1?b7&?q>ceytRG1Qrdk5Zd)jB6Me$DCPh z!MTWp040+i!z|hsWV_x40_b(vGvv9T;2?^z+JrKPEo>}U&!@N^$~1|lfoQ0NB40_QIkzfr1miJGsj*KSc1l!tZI4xks4b*p>HC(P7`C5+4N=4$=z5a8u_jwSotO9{0+sBnKQjr*>i5C9oX-xYk|JI{Tj37 zHdS9x^#E1wjb%EUonHeL+E^h$dIn^iWb4mq2;Z5F8*$KpEx%5(|4ZF@J%*R z;h!5mcnH_XhT#>1-z~aL@@#@&lfR*{A5%pjuI~s}ZJc{wnf(3tbbl{)_iWE;x!3Q& J4es^Ue*t?m+N1yg literal 0 HcmV?d00001 diff --git a/animation/SadTalker/src/facerender/sync_batchnorm/__pycache__/replicate.cpython-38.pyc b/animation/SadTalker/src/facerender/sync_batchnorm/__pycache__/replicate.cpython-38.pyc new file mode 100644 index 0000000000000000000000000000000000000000..798aeb00d914cca2bb52eb6dd7f6f0deb2850178 GIT binary patch literal 3498 zcmc&%&u`l{6ecCvR@^uZx?)}1Z6(?bi-yE#iVYZwCTN#50|qP@kQN<07YZZNwh~!# zNIE|XZFIgkT_! zT9~1hW%g+a8@kp#J9B+MrhV@FS>5;h2^&WGecJbThcwE+yh1|9)l`3USMIhSOD<%4 zdni@dUQPDnD4|TYSCalfa49)!Z&3D#MmtKpXf}97Z{C8+#esIFd|Y^`LRo<#2T-MOm{-veR!HF#A6qAAgh%+uVwOHa zWvA8%uR`CEt(+8oSJBh2+SZYCRDpIS9}j&up|5&WQ zklnT2JFEf+U8Uc`zUl;?hxtho8k5oiI80I=CfW3T23YV11sG1x4g$SNU?itOkK~~z zf#6heMz#(~BEoJM(`eFavu>7r)Z=6O`(YH3El$FCFWG^q`v4b0MR(Zev6AmxelAc4 z^CU>a+*_>jnm?vMJ1(ZLnYmAT3FV?QJku%c6hAN(wSmUp zXN&KYpys_O0XO6_yp;2^>+gW_x>w8Wh{u_u$7i)12r{ccgu6VO;xPbH0JyVSk4mbP z5Eo%9o-4F-tqXaeeGSG|Jr;!)$(92cl^G7RYHl&Z+Zpb^=B3bHU56rPph687&*Dqg zEWV8E_zhXqQPqzmP>p4=WxqfjY!t(7IWLfdTAJyOsK794mvq<4OiE!d$TmsQQ(_h{ zd+^U>p>swSN}!b-ryUK-Ev`W}mW&)?rEF}LYQ#55XBY=M#gq)1!1?Q@|E+vHJfMBx zwY+O~qXlxA4@3Brm#8ehigK zKWJnfTPaSh6ARLfo!TReQEG3XS9slXvWgrIxG=?IzVd;u>C1We(PC1;fN6Hd{#{WDU| zapfgC!G?N3zUM9^X-MHI@@)n&^BmeJnfF~1@9@I~a{tC2)(d@;+ji%Fl2nD382`~XvOo>_Ixv&Hj+JWilz zEUbIgtRBiRmWsx35!SwBGY{a}EVMvB6e1Bh0jG>;CrJc3bYIYcgcRRklQ&86iT1#d z@|J!xq|_>~21V+AG&5sWEf>#JyeaV+v|Be+B7$Bsj+^Cuw0YL1*$_wpTw&lJfoe`B tbHy;hGr}eGA`>atY{7)fy_aX{W!;jOp>nNh+rsAWPsg2g+$*z}{{pt*-&g 1, 'BatchNorm computes unbiased standard-deviation, which requires size > 1.' + mean = sum_ / size + sumvar = ssum - sum_ * mean + unbias_var = sumvar / (size - 1) + bias_var = sumvar / size + + self.running_mean = (1 - self.momentum) * self.running_mean + self.momentum * mean.data + self.running_var = (1 - self.momentum) * self.running_var + self.momentum * unbias_var.data + + return mean, bias_var.clamp(self.eps) ** -0.5 + + +class SynchronizedBatchNorm1d(_SynchronizedBatchNorm): + r"""Applies Synchronized Batch Normalization over a 2d or 3d input that is seen as a + mini-batch. + + .. math:: + + y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta + + This module differs from the built-in PyTorch BatchNorm1d as the mean and + standard-deviation are reduced across all devices during training. + + For example, when one uses `nn.DataParallel` to wrap the network during + training, PyTorch's implementation normalize the tensor on each device using + the statistics only on that device, which accelerated the computation and + is also easy to implement, but the statistics might be inaccurate. + Instead, in this synchronized version, the statistics will be computed + over all training samples distributed on multiple devices. + + Note that, for one-GPU or CPU-only case, this module behaves exactly same + as the built-in PyTorch implementation. + + The mean and standard-deviation are calculated per-dimension over + the mini-batches and gamma and beta are learnable parameter vectors + of size C (where C is the input size). + + During training, this layer keeps a running estimate of its computed mean + and variance. The running sum is kept with a default momentum of 0.1. + + During evaluation, this running mean/variance is used for normalization. + + Because the BatchNorm is done over the `C` dimension, computing statistics + on `(N, L)` slices, it's common terminology to call this Temporal BatchNorm + + Args: + num_features: num_features from an expected input of size + `batch_size x num_features [x width]` + eps: a value added to the denominator for numerical stability. + Default: 1e-5 + momentum: the value used for the running_mean and running_var + computation. Default: 0.1 + affine: a boolean value that when set to ``True``, gives the layer learnable + affine parameters. Default: ``True`` + + Shape: + - Input: :math:`(N, C)` or :math:`(N, C, L)` + - Output: :math:`(N, C)` or :math:`(N, C, L)` (same shape as input) + + Examples: + >>> # With Learnable Parameters + >>> m = SynchronizedBatchNorm1d(100) + >>> # Without Learnable Parameters + >>> m = SynchronizedBatchNorm1d(100, affine=False) + >>> input = torch.autograd.Variable(torch.randn(20, 100)) + >>> output = m(input) + """ + + def _check_input_dim(self, input): + if input.dim() != 2 and input.dim() != 3: + raise ValueError('expected 2D or 3D input (got {}D input)' + .format(input.dim())) + super(SynchronizedBatchNorm1d, self)._check_input_dim(input) + + +class SynchronizedBatchNorm2d(_SynchronizedBatchNorm): + r"""Applies Batch Normalization over a 4d input that is seen as a mini-batch + of 3d inputs + + .. math:: + + y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta + + This module differs from the built-in PyTorch BatchNorm2d as the mean and + standard-deviation are reduced across all devices during training. + + For example, when one uses `nn.DataParallel` to wrap the network during + training, PyTorch's implementation normalize the tensor on each device using + the statistics only on that device, which accelerated the computation and + is also easy to implement, but the statistics might be inaccurate. + Instead, in this synchronized version, the statistics will be computed + over all training samples distributed on multiple devices. + + Note that, for one-GPU or CPU-only case, this module behaves exactly same + as the built-in PyTorch implementation. + + The mean and standard-deviation are calculated per-dimension over + the mini-batches and gamma and beta are learnable parameter vectors + of size C (where C is the input size). + + During training, this layer keeps a running estimate of its computed mean + and variance. The running sum is kept with a default momentum of 0.1. + + During evaluation, this running mean/variance is used for normalization. + + Because the BatchNorm is done over the `C` dimension, computing statistics + on `(N, H, W)` slices, it's common terminology to call this Spatial BatchNorm + + Args: + num_features: num_features from an expected input of + size batch_size x num_features x height x width + eps: a value added to the denominator for numerical stability. + Default: 1e-5 + momentum: the value used for the running_mean and running_var + computation. Default: 0.1 + affine: a boolean value that when set to ``True``, gives the layer learnable + affine parameters. Default: ``True`` + + Shape: + - Input: :math:`(N, C, H, W)` + - Output: :math:`(N, C, H, W)` (same shape as input) + + Examples: + >>> # With Learnable Parameters + >>> m = SynchronizedBatchNorm2d(100) + >>> # Without Learnable Parameters + >>> m = SynchronizedBatchNorm2d(100, affine=False) + >>> input = torch.autograd.Variable(torch.randn(20, 100, 35, 45)) + >>> output = m(input) + """ + + def _check_input_dim(self, input): + if input.dim() != 4: + raise ValueError('expected 4D input (got {}D input)' + .format(input.dim())) + super(SynchronizedBatchNorm2d, self)._check_input_dim(input) + + +class SynchronizedBatchNorm3d(_SynchronizedBatchNorm): + r"""Applies Batch Normalization over a 5d input that is seen as a mini-batch + of 4d inputs + + .. math:: + + y = \frac{x - mean[x]}{ \sqrt{Var[x] + \epsilon}} * gamma + beta + + This module differs from the built-in PyTorch BatchNorm3d as the mean and + standard-deviation are reduced across all devices during training. + + For example, when one uses `nn.DataParallel` to wrap the network during + training, PyTorch's implementation normalize the tensor on each device using + the statistics only on that device, which accelerated the computation and + is also easy to implement, but the statistics might be inaccurate. + Instead, in this synchronized version, the statistics will be computed + over all training samples distributed on multiple devices. + + Note that, for one-GPU or CPU-only case, this module behaves exactly same + as the built-in PyTorch implementation. + + The mean and standard-deviation are calculated per-dimension over + the mini-batches and gamma and beta are learnable parameter vectors + of size C (where C is the input size). + + During training, this layer keeps a running estimate of its computed mean + and variance. The running sum is kept with a default momentum of 0.1. + + During evaluation, this running mean/variance is used for normalization. + + Because the BatchNorm is done over the `C` dimension, computing statistics + on `(N, D, H, W)` slices, it's common terminology to call this Volumetric BatchNorm + or Spatio-temporal BatchNorm + + Args: + num_features: num_features from an expected input of + size batch_size x num_features x depth x height x width + eps: a value added to the denominator for numerical stability. + Default: 1e-5 + momentum: the value used for the running_mean and running_var + computation. Default: 0.1 + affine: a boolean value that when set to ``True``, gives the layer learnable + affine parameters. Default: ``True`` + + Shape: + - Input: :math:`(N, C, D, H, W)` + - Output: :math:`(N, C, D, H, W)` (same shape as input) + + Examples: + >>> # With Learnable Parameters + >>> m = SynchronizedBatchNorm3d(100) + >>> # Without Learnable Parameters + >>> m = SynchronizedBatchNorm3d(100, affine=False) + >>> input = torch.autograd.Variable(torch.randn(20, 100, 35, 45, 10)) + >>> output = m(input) + """ + + def _check_input_dim(self, input): + if input.dim() != 5: + raise ValueError('expected 5D input (got {}D input)' + .format(input.dim())) + super(SynchronizedBatchNorm3d, self)._check_input_dim(input) diff --git a/animation/SadTalker/src/facerender/sync_batchnorm/comm.py b/animation/SadTalker/src/facerender/sync_batchnorm/comm.py new file mode 100644 index 0000000..922f8c4 --- /dev/null +++ b/animation/SadTalker/src/facerender/sync_batchnorm/comm.py @@ -0,0 +1,137 @@ +# -*- coding: utf-8 -*- +# File : comm.py +# Author : Jiayuan Mao +# Email : maojiayuan@gmail.com +# Date : 27/01/2018 +# +# This file is part of Synchronized-BatchNorm-PyTorch. +# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch +# Distributed under MIT License. + +import queue +import collections +import threading + +__all__ = ['FutureResult', 'SlavePipe', 'SyncMaster'] + + +class FutureResult(object): + """A thread-safe future implementation. Used only as one-to-one pipe.""" + + def __init__(self): + self._result = None + self._lock = threading.Lock() + self._cond = threading.Condition(self._lock) + + def put(self, result): + with self._lock: + assert self._result is None, 'Previous result has\'t been fetched.' + self._result = result + self._cond.notify() + + def get(self): + with self._lock: + if self._result is None: + self._cond.wait() + + res = self._result + self._result = None + return res + + +_MasterRegistry = collections.namedtuple('MasterRegistry', ['result']) +_SlavePipeBase = collections.namedtuple('_SlavePipeBase', ['identifier', 'queue', 'result']) + + +class SlavePipe(_SlavePipeBase): + """Pipe for master-slave communication.""" + + def run_slave(self, msg): + self.queue.put((self.identifier, msg)) + ret = self.result.get() + self.queue.put(True) + return ret + + +class SyncMaster(object): + """An abstract `SyncMaster` object. + + - During the replication, as the data parallel will trigger an callback of each module, all slave devices should + call `register(id)` and obtain an `SlavePipe` to communicate with the master. + - During the forward pass, master device invokes `run_master`, all messages from slave devices will be collected, + and passed to a registered callback. + - After receiving the messages, the master device should gather the information and determine to message passed + back to each slave devices. + """ + + def __init__(self, master_callback): + """ + + Args: + master_callback: a callback to be invoked after having collected messages from slave devices. + """ + self._master_callback = master_callback + self._queue = queue.Queue() + self._registry = collections.OrderedDict() + self._activated = False + + def __getstate__(self): + return {'master_callback': self._master_callback} + + def __setstate__(self, state): + self.__init__(state['master_callback']) + + def register_slave(self, identifier): + """ + Register an slave device. + + Args: + identifier: an identifier, usually is the device id. + + Returns: a `SlavePipe` object which can be used to communicate with the master device. + + """ + if self._activated: + assert self._queue.empty(), 'Queue is not clean before next initialization.' + self._activated = False + self._registry.clear() + future = FutureResult() + self._registry[identifier] = _MasterRegistry(future) + return SlavePipe(identifier, self._queue, future) + + def run_master(self, master_msg): + """ + Main entry for the master device in each forward pass. + The messages were first collected from each devices (including the master device), and then + an callback will be invoked to compute the message to be sent back to each devices + (including the master device). + + Args: + master_msg: the message that the master want to send to itself. This will be placed as the first + message when calling `master_callback`. For detailed usage, see `_SynchronizedBatchNorm` for an example. + + Returns: the message to be sent back to the master device. + + """ + self._activated = True + + intermediates = [(0, master_msg)] + for i in range(self.nr_slaves): + intermediates.append(self._queue.get()) + + results = self._master_callback(intermediates) + assert results[0][0] == 0, 'The first result should belongs to the master.' + + for i, res in results: + if i == 0: + continue + self._registry[i].result.put(res) + + for i in range(self.nr_slaves): + assert self._queue.get() is True + + return results[0][1] + + @property + def nr_slaves(self): + return len(self._registry) diff --git a/animation/SadTalker/src/facerender/sync_batchnorm/replicate.py b/animation/SadTalker/src/facerender/sync_batchnorm/replicate.py new file mode 100644 index 0000000..b71c7b8 --- /dev/null +++ b/animation/SadTalker/src/facerender/sync_batchnorm/replicate.py @@ -0,0 +1,94 @@ +# -*- coding: utf-8 -*- +# File : replicate.py +# Author : Jiayuan Mao +# Email : maojiayuan@gmail.com +# Date : 27/01/2018 +# +# This file is part of Synchronized-BatchNorm-PyTorch. +# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch +# Distributed under MIT License. + +import functools + +from torch.nn.parallel.data_parallel import DataParallel + +__all__ = [ + 'CallbackContext', + 'execute_replication_callbacks', + 'DataParallelWithCallback', + 'patch_replication_callback' +] + + +class CallbackContext(object): + pass + + +def execute_replication_callbacks(modules): + """ + Execute an replication callback `__data_parallel_replicate__` on each module created by original replication. + + The callback will be invoked with arguments `__data_parallel_replicate__(ctx, copy_id)` + + Note that, as all modules are isomorphism, we assign each sub-module with a context + (shared among multiple copies of this module on different devices). + Through this context, different copies can share some information. + + We guarantee that the callback on the master copy (the first copy) will be called ahead of calling the callback + of any slave copies. + """ + master_copy = modules[0] + nr_modules = len(list(master_copy.modules())) + ctxs = [CallbackContext() for _ in range(nr_modules)] + + for i, module in enumerate(modules): + for j, m in enumerate(module.modules()): + if hasattr(m, '__data_parallel_replicate__'): + m.__data_parallel_replicate__(ctxs[j], i) + + +class DataParallelWithCallback(DataParallel): + """ + Data Parallel with a replication callback. + + An replication callback `__data_parallel_replicate__` of each module will be invoked after being created by + original `replicate` function. + The callback will be invoked with arguments `__data_parallel_replicate__(ctx, copy_id)` + + Examples: + > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False) + > sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1]) + # sync_bn.__data_parallel_replicate__ will be invoked. + """ + + def replicate(self, module, device_ids): + modules = super(DataParallelWithCallback, self).replicate(module, device_ids) + execute_replication_callbacks(modules) + return modules + + +def patch_replication_callback(data_parallel): + """ + Monkey-patch an existing `DataParallel` object. Add the replication callback. + Useful when you have customized `DataParallel` implementation. + + Examples: + > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False) + > sync_bn = DataParallel(sync_bn, device_ids=[0, 1]) + > patch_replication_callback(sync_bn) + # this is equivalent to + > sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False) + > sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1]) + """ + + assert isinstance(data_parallel, DataParallel) + + old_replicate = data_parallel.replicate + + @functools.wraps(old_replicate) + def new_replicate(module, device_ids): + modules = old_replicate(module, device_ids) + execute_replication_callbacks(modules) + return modules + + data_parallel.replicate = new_replicate diff --git a/animation/SadTalker/src/facerender/sync_batchnorm/unittest.py b/animation/SadTalker/src/facerender/sync_batchnorm/unittest.py new file mode 100644 index 0000000..0675c02 --- /dev/null +++ b/animation/SadTalker/src/facerender/sync_batchnorm/unittest.py @@ -0,0 +1,29 @@ +# -*- coding: utf-8 -*- +# File : unittest.py +# Author : Jiayuan Mao +# Email : maojiayuan@gmail.com +# Date : 27/01/2018 +# +# This file is part of Synchronized-BatchNorm-PyTorch. +# https://github.com/vacancy/Synchronized-BatchNorm-PyTorch +# Distributed under MIT License. + +import unittest + +import numpy as np +from torch.autograd import Variable + + +def as_numpy(v): + if isinstance(v, Variable): + v = v.data + return v.cpu().numpy() + + +class TorchTestCase(unittest.TestCase): + def assertTensorClose(self, a, b, atol=1e-3, rtol=1e-3): + npa, npb = as_numpy(a), as_numpy(b) + self.assertTrue( + np.allclose(npa, npb, atol=atol), + 'Tensor close check failed\n{}\n{}\nadiff={}, rdiff={}'.format(a, b, np.abs(npa - npb).max(), np.abs((npa - npb) / np.fmax(npa, 1e-5)).max()) + )