From 295dc00c22e2495af873469749dbe5dc234c3f1a Mon Sep 17 00:00:00 2001 From: Jake Austin Date: Fri, 6 Jan 2023 02:01:21 +0000 Subject: [PATCH] Small patches and configs --- HF_TOKEN | 1 + nerf_accumulation.jpg | Bin 0 -> 1705 bytes nerf_output.png | Bin 0 -> 8374 bytes nerfstudio/engine/dreamfusion_trainer.py | 35 +++++++++++++------ nerfstudio/models/dreamfusion.py | 34 ++++++++++++++++-- nerfstudio/pipelines/dreamfusion_pipeline.py | 10 +++--- 6 files changed, 61 insertions(+), 19 deletions(-) create mode 100644 HF_TOKEN create mode 100644 nerf_accumulation.jpg create mode 100644 nerf_output.png diff --git a/HF_TOKEN b/HF_TOKEN new file mode 100644 index 000000000..82d1555a8 --- /dev/null +++ b/HF_TOKEN @@ -0,0 +1 @@ +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 \ No newline at end of file diff --git a/nerf_accumulation.jpg b/nerf_accumulation.jpg new file mode 100644 index 0000000000000000000000000000000000000000..d4f10cb84a6ed1bc75c4ba4fe40d9f8aa26facdf GIT binary patch literal 1705 zcmbW!do!nvqayR9dr1aw!^G zxyI&BkxP*#mpCpVgfd-dQ!b;Z%#PhX?K%6`?mo|X|9YPDobTtn;$iV5u-@L5W(zx|%8mi`CFJ)YZUk)52o6 zn&@r&j$mwTtgcJ7G$mRXl8lLKM<7Tf5+#dLmXlK^YGO5s|GC8X0JJo)1K0sVaDWsV z0-_<}2Y@;NKwy%zYk~hNh!hBg!4c9(8Cl6f&3Zrz0)kRd5C(%nCEF(@a{!8lDQFU` z;2Y?^2wb=#k(-<&tz}(tU&*=mgEonIG!7}Vag*|9jLsI_t$KK4GR4G{YG$)@m+fvl zn*BZ(S2u>cho|4cL;eAQhlALhh+~n*qoU88Jr{rerv%>Rl+?8JD_1izZ{+6X7ZeuV zEUv7ouBp9!r>?%G^+DT1!K3z$zW#wHgG0klN1nf!625%(`pxvr-2B3#XzAmp<&`xT z1OWeJeUtseg_gLapimHsSaU(7A|($*Lt&Z(xPlcO;TygIN90N?S|{gJ+(&AWoIfZr zkM_!J)HZ&uGq*)J+U^_53@}!?beb2l5%G7CHO!7NM zgHev)lxh84@6lY}y^HKsr8j3!O7l`AueG zvOS+QDMWke=iBCZRXKsTloJka6P&57m)#A7K+~b=xDrQ=x;tnAAB-C>?7Q9K z7p_=l)jn^eg>}OuI=DHE$NG~;Z9+i{`j6!n6f=g=t>cHwI2n2Kt#?w>WY`M9d$4Q$Ng@!f=<1u;og+l*Yht5X}pCZgh z-4#p=^v<+`eRoe1KSz`=%(?|2%&%Mu?(Cae?MSM!_*i^f4c8Pm+>&yn{aM8x&lT@h z-HSwuo73mv28zo1!db8PmOR?6)Tn|FHYumEQ`duXizcpD6}O(?dT})tMP=IM1G6lz zBPj%C7(Ci7X(W%gX{77eVtXmxn=EtGnca9dFF2uaze(YZ(3&2{>c-0xd+@BOOOSYf z&89SKxmfvH(=a>my>MM@+P&7B?TV;?ZuJY%N%l8JQXgv)%8}EVd5L(myyg|KjnQ?h!hCw)BsM>hEG(9BH z=E#t|mK@KCH0{ouHw=w1Z(CU;@PnW0ssv5Ugbgw}Et}d^Uen{u)uWR-;_b3``A34y zK?0PE5Q|Ay$9ffGnDGLG@qfsPAum)eJ|%!PumLaLfJUf4Suc2 zf#uW^MO|U0&&-t`QXmQIg!JHMt@u(*zN0n!6FvOp<#$8Jlt`JyLM!Tp3wl%~U-yT( z1!E$z(yJ$_wEOU)!A#nG{x9}AhS?tYXPGH31H;1x=~k)yN0Up7#f(Tv7FKo^qOVAS U)=u40VEds}&$J^rgf4#iJB|+V#sB~S literal 0 HcmV?d00001 diff --git a/nerf_output.png b/nerf_output.png new file mode 100644 index 0000000000000000000000000000000000000000..494994a2f036f896be911bbdad1a8974b7ee4799 GIT binary patch literal 8374 zcmV;nAW7eeP)V>bE;TMQEFfrfbZ~PzFE4FjbZ~5MbZlv2E^l&Y zFU>2t(f|Me32;bRa{vGjVE_ORVF9Q=r)dBHAR|dcK~#90ZM#X6WJz-0`BPio;#+K) zm08u*g>G~c=$Rp9gd&s?itxe%k8_yIT=_}(^|%Omo8b@ufo`Ces>(GscYm9ism{S( zL}nHcAs+7j{KfO9YN~(T{wDwDfB3&mc=`p_XIWo;dTr0oVC#EBOxM>6rU>C}*{=kwa%oM{=tH;d@HN&vy_hwe_ z&I~iw&-io<;ORYp8NtY~3jRraK0CJlbziO{(zf@;4}oS>bW~f!pJop9s7v;vWfL1^lFd&m?*J z90-tB!JY*FTmjFj;?uv_^W#bIpY49Zd|J5w;f(h}+Xv<$(X@eiOx(UX0?K#4zDBK} zxsY;3i*bH5wg*pW16Djpt`7{*^cE&#!*L``Oyu2U;k7gxCVS9)?+?_Z+ z4s7-T6D6lY&KVJC7C{S4W1;DV(0a@i)p|>YU?zZ~1@aR4#m&EFkq7?p?!SL&?)K#> z+p2+2R`I1&s|CnsV#;R%^DOY^tN-Hgiq7|xSxG7J;q6D}Y2mBaH-s1vV7K3)Iq}=S z{v$UpFF9Ow#Kw_xW|=bMFkxzJwjF7P*)nEEDS}fF5llfSSrxE6*%|h%geR5MS7OGi z%byb^;?zL_QzI{l+jkGV{oy0aGGPYS*B5AU{NZ%)kXhM43lLi=#}h`14nn{PxX% z`*P(!7x*(pJiQKPX%+gD1;}T<1MuDTU(|4BU`a6rD~5P@`@=iV4 zd7xMUjjaqxfqyWjA4fBk!yg|uYyqD;ekE-v%&bNVKzX)#Y^J)Q4R**3=-);W|(|KVyjWt$dl#&q#B8KyjQ{nr6c+Zc2 zxaIBlx4iuN0_WZH<#Wmm=aF|m-ZBmos)eU`EE4e2<7qK-oR2fplt@WQb3t{D>!P?| z#MU7mv>3>m{;D8OD8*K#!8~S~jYquUti@SF7xvs7UavUt^!rH*0m>ID`NieGWH~BY z3eE{7E8+w*BSc4$LP?pn3!J@h_x_P(P7LRXzHd1l2LN&|e0+CL-#3(O>@Pd|{f^i; z#A5wQbu4Gk!|i}LxO%az0E^fa8WcjVc?fjEOKAfE)+qNG$9P7AT|XlgJD5*#MV zl9w<1|MP?(VCDJE&#rz=KEZUF$xFruAON&vLkdI&MQ`gyZ00K9|pQjL)+!5Ag{#eF|-05 zLQb`K5)UGdoD?n9Ri+`)be`CFN-;`S5byyIVbkte(!?~Mf9eSYp*$b)^MPe4Xfc{D zFrO{>?`_q@IPdxrV+xh;IMfL~nwE7Hj*xzP3wLD1wFPKv3K=1f`&P?*Pz zbHaAt(r+XC%a(rI@Njn|=R%C3KDU81XF?2k@0iD#ACDuy{j2X$W6k%CrirxOE5r$H zA1m}(1zXLIQVQOC)RbjXa>=BrAnr-{HDgpk1;*1%Yy)lQX|@f)N1U6v+P^|gnU>)* z1)#jX{Bz3SSjLQI!!;FPn$FYpk$Ftyl+gn7ysYmF^E7kw)d3eAUwyl0vkxqD=JEbS zN}0ZIIiDxyX+eu}Iu1BjL)qi~z%gE?j-J#|XsE;oF97bB!Dve|rD7pSge;&HR$neP$jLAyyL? zY68-2V$GD*YABoofB4&X7!?$!Dd7VwlhCxD7q1W8e6>ThMvI5Lf#3h#TU3pCUdXB7 zoUcl;dhq98a7A!ns& z8^$3sjhWbZ9zIMQZhC@`s2Ia3p(@N%VwnLxW7FUqQ}&Dyssd)+4ba97gMuQ zq#%o68qi{>x=$7mB;pz)L?K3x7R7rHra12iQIM{h;Zsb$xZZL;U2r~*EKBBms^&MI zX4H)FJR?pRhk&Zkv<)FRCNo3?=N*zWh)`5%+JLGu4m0z-us?J-5AhTbS3rf+rO-~Jo^=7;}CFck5L zwu?-yC(VTr9bkAb&*4P`l!DJNofk~40^76@qv!N6;k;XGdSw|hkIO(Q#&DXLr-XB^ zHcHf3mRu7N?`hh|JTBao`>GrOlND?|4o@xx)y%se?@-NLT<@8tg?`f!+rZuXN8Ea? z5Ioy`RifsnzM~^eM$Se~g>>wm!l2<0oD-TR(6*j_=huQn!NB{q_^}#WLv+OWWKzmJ zCB!?1(@bnVP^}>)WtlSlrseo}Mpf~?y4G&LVVM&Bwj)H(I85ArnCLehXFp>Z#0kq( zz?5dw5*tsyX*pb7)tWH zJS}B$)${{6=g6s6TFwdcT8f4!h=;t?P}y%gR7Xlt%oHb%qRRQ>#Q9iPToB4)JWOY_ zz%nP6IgxXz#rGNtLlAD?pIGJ$>Y#JHe|w9XqFLDNI@A?JT0-;{jL!>g-_Y+`9zLF! z&I|Ffmf&WFZ)%+oW8ic*b3RR&!gk+craXR}`1a<1B{*>uW6ikW@phyXW3!EjbC^N1 zmZHm9DMcxzu*`|gu4OpQjOPV4Bh872yEAD?bbU)+49KJ40s>d9RjKa=K?+i`o6`7 zKuHxl<}t6>1l~JLjk_Obf-{eb;6w;f*lZheE|g--r=^ymeIPXTpO!h(Z6a+SYnfVV zd682gUi3B0`7N=DY&IR!n0Wu=9ghzM@c|5)Ht_gxrt2D-uA%EYni#oxagFz$$NNXp zQV5~taMiKd^k^}Z7VM&=>& z@%;n8`d_{wpdK+nMDQLqdykrD9t$CQN-}cJ3=a#T33wMkAeYRtWV{!4hYiR3fo|K- zZ#ukpIPYlNmQo74{f^)N_6HvCkCbe97uak%+OFm5`hX8G%(YrXgw1x#IjpRiM(KAgG(>Mb?+g%4{Jlq+_^l1|)DcA0gI+i67VqLxAm>8#oDKxP`)i_*kaKSN;iBf7wd;k89-~QG2 zqzQ1<^h4`emV_C^HH|eh=Bajzoi`2_J5o9md?e?>o9{o6m%{#HN2i&bYQC||3u#ua zZZ1hlX_}5DEj&J)S>}b^VM|_=H-CJ`JT7Y+t&mfuX*zZnJ;Aw}cc)o7-cQ8Q*Z3(l zH6Dk=G8b|#IB}%KR-}jVyzudKq~A7lo0dL$+CH$%g@@aLWhp=nr59H_PRA4TGP5iT zF-CU#EzSu!RSYsU@CqnmN=^&Ys1!?xKx`{;IuWLE;ltbenzNcCFAK+qf$Nu-#OT>{ zTc(H1hmUuBe1FH)^?}`?XBuZNF1K7OTjseoH4V@rbX|*(`0iJ)NU6~Ek)nmH>AA;;6H2C5cuS$( zM51rdBIJ_s;_>3}As`Mmo1S(PdAvVjX5^*zC6{?2hKTcy@jNn=>M}Pk4@|?v<<*X+ z^9<*a*aSjT^8_;^SaH>lN=+K_0;k8JhPEIzTPua*m5y7i=Udd&!nZ6>nA&N_2PoCR*Xdp7dJg~T6pp5g2UAYQ{!|TnWu$$T=@9z zf#YK&eE4zV&7Y2>+1Tz|#`D6%!@&7CvE6lS_Z>OaS>5%`p000*O-(wS3mhNMoX->Q ze!N3OxV+xczS@&g=6F9+a$%Y#Qd+pXd*IK0`87GeM2m8KI8#dE>gJODVPC_ePK@UP z=Y-e_u?c+h=a+1EE$7pSnbD>WD^M~Fr;$<$AKu-eC9~absyl8PL>xXCT^~6;PHP>Y z48wx+Qo(R{Xz)(BKi3-he4IJnP2Al*a({c``sJ3*wq>{P5fO5#@zlFOE``%^L_~P; z>Vn;V%QR2ieK@kr8PgiF%#3A8OvB9Gha)itx;r_(y zZ(p$6cH}aXi{s;)N4EQ(zHgzx-TM>HdoHgJH6;3)_=ppVWp+*_xCd& z-`(@!-97t5Pw0B#+h1N1Vn9np`Q5(d;&Q|7$0Nged9J-5@9PTec02a_mSvo2+m@z@ z_~Xb0}{q2dS4K!`FnA_WXmSthH-Oz2f zYl-XCe0+OP$%)5D<9wV5 z&J+BCIN^Mph+&U!9GM!wRh4ejaCyC@luQ&yE(I3^r4T~xDSGF4yd628X1@RZdoHea zwSG{e-!ueqD@nmTmp2{rn7R47!3V)qxV&kYC-~#vd;Et}ZEBQSDmP7Jo@T~jr0;9( zpR)mu-M&Sfu)FX~!^HVCvfXy%lDW9(xwzS|eIZ~AFJ5oJG9pHd4LK(^n~u%4r`tA! z=r}!0q?CxQ$1jEZ?~i=<>sM?K9m~A1jER(M4T4eU1q8OchP)I$o{t1d#Wc({XV`n337QHy z71}1Szw8)J6EA+YV}H>xJf;;>7N&7wf7o*KqBbq&dEw$_%Q7X7w-fuT9_JkQAI`k{ z;gO~f+}|F#xY^)iP1B{w%1Zz(@bT>78(voQqEk{1rwmmqNY@|ssaKk)tUKQaz8 zn9;T#&4ufiTeK*TcOyW)TT;qgUhSCYnQqe%((KVvCqA)pxKI_M zYAl2B|NNIXwFGTyclP?lj*<-L9Yu{09BtRs?C`RJ)w!! z=goNhI5C`NF0S^x{ox+x>)>-77GmpBt7UV)3;gOYzh-+~&3YM)oD#3Uy+XCn^?_wc zHL+1xrox~8;R9DUwKv(cfn{0fHja5#wp~L?3q>8N1SzschS!qN-6v5$w@%49y zn%1Vo<@KH>M22C;N$nJ=R`!-it+hoC+*&6KGSLS3Oh{PqL=Hqvi9F0Nag3uK+x9QM5a z_5cRcxN!Gj;NfoI@~ftX%f^9IT&&Gk4En8Scj?KCJ-6n>95}1ftdx2Evr8WD=Gvq^ zFC6a&V(0Oprt43XtFZ>XMhQ4qV(s;wB?LgPbvQ`2D;#9$K)rVul-8G(ixWU5695 z;tMpLtK>=nq_V|)W0a}moA367AUN}gFPMPJh~S7_BuFJviq(8Yf?{H*C4BE_x(&@3 zQEhMtr6}Vm@#6Ka;>4V5U7<=aHzG8+;0Wq*Zq;^ixB#&&IA3Tr0G`XQGKCX!e#DBQ z7~}i^2yJt~yM^dME#ss?#Sn{_n=v8~E55{p$^xbsO45nOUt$Hiosnj$fCZe-6k8~k zQOl?nyac>!F)OI0T3lhM71~f}cRQ{_!+0|8KNMcP+Tk3`OQPF2QZl;U5!5n4TEx{s znG05FGNBNfg`_7^$#_@W$ZDB!z9;1w5l=1)Su@^6w9Et-$$G}&5hC6;K%r0&>WWro zs8Q@p(I<+I;V|MRA#uW|4Vn$#WeP`%juc%;WyVbL;)$WhiC>jcA0m#TBOx?2U0@j< z*DqQ`U>*v&DBeMA9GjgdWF6NUU-7X)TtRKFEDp8G-PM&%YKoIeZmP%Y6xfMJFuboE zL5Bx%^?YkkKUeM_E2r3L6~0X9${+T=#bice#u1Qy#$-WlCgn4^Oe@zo6I@dpD3(zQ z4u?@2DNY<}GiEc;(zFh6V8)V^AP^#K_JL-T2)Rsn2{>s$1f#4MvsD;}trvLuapG1c znd0P$TV|jclSHwEb1hyX_=E#&RT6*|hq7RnG0lj0gt7_{Fq=>^xs0f0K=ICFCb)H6 z=ftnyA1k&XRhbk7N754^HgsEGhb)t^+0|*Y3l)e}9tdWYk7}hwRR}R6lJTysayhBy zZJAtVmU5(M!MlKIex}#C)jdQj$2|Bgo)%n26;$d37j+OD9KIOkDYRs~mAYkZ3rhd9HuR9C^+9^A@68TiwO|Ci8& zO+`|tEBy_`8kChyVTs^&Ea}L!JYrf8(}Jq6v|J z6G7cfumyc>qkwhoC>(JKTK_Z6KHQaQ1d4*+;yIlH;c4AU(sk2rC79|$hu z#6RPNcY)|zoH&AyMBf66>hy_)o_j1u#AQ5j#gP@b2w)lKGpbXaOasw3G$GWw3o$f6wYGW>%d-+v9$C_aWJf6@zV|p<(iF(WIOhL> zTWOI2;N9Kd@bdCk6}8=ps#Yg8P3SRKBSiGIFX$y=nkmfdx%#Z2Y8|rTMCIzKI-C=n zGmuQtp#mI+$&7amsxu(f%%xgYa2^5L*t4WFC69O?@y=J;jEfjr!?7+n7x7L|SNj84 zJw7cT*>xSU+0sVkcz5E7A`zZKThWCO+I72%u6`GB67X@e#`=m0tbioqaAqK$CUl56 zf^RCpii9U(2;<4)Jv1GND5?vl3q==-$g?{cWM{Nc?S?YJ2f{7u7=eUm) zlu|2GO?8$KYvP7fyEdDNZOf0h|LJ-4pA{g$-Qyp5aq-L5w^`$?v7{s3MVhchR=jWvHTB=n#k4Q_-<8wk5L$^k8zutFYE`NL${&~NEd?wDvhkxL5|4scT-%-j) zDXE4c%`{09uf%@ASa7%+x3$J?=PIyPD~LGf zR~#xBGrsb339e4D)e4wV@=&u!TZmz^Hbi?;840es2|(~$Fk@QocsTvIV#X(aeE1}O z{TZe^`eu(4Ps&fB&&Z~jX8P_5C!tz|K`sk|L$ww*!SA75V!$%O{oTYj-?fM&RAjr-?(e4dZbs~{q{FaogM zZrJX?xieX(Y9;IH2z9@acNNTP*sEk{-uPs%8oF{!J<%Iy~cEtz>aVp`S})8`B4JYiu(BvN#u z=_>dXovD=E8sAqM*mF3m2P!CvREi#^IM?B%q1eDvLG+Ykh)_eJ->{^+8d5ERN|CI2o7cn!mq?$sqlA5IM9hl0ICRV?z$&VxFiD|hf`VKIPJb;zw zlG;&QTD1aPX*n*|DpV7uGvXRd7sLgGf=I#RF&D{}YjRvU_2e?FWX?*utazj>c|bM) zBgB7B#6DBVCk%Mjiu~mHNkN~8R8L*xoR_uoi)*+MN}jQj38AUo8eN{#<+|O)k$FdSzw)Em{YTEg0Cu5B|}!=y*7&I+m$$}{+S*C3%KjYqIjE#)(lhfikzB$pWkkg#KB#I0X0PZ{N!N{c)t=Czc? zc>*oDR2R~<0(DXgvRuy0^CPP1%TQ;}#`{U|&#d7y4EU#om!JHu9`fnD&)r3H6?83@ zYy~{6lNU;!uyyOq=OM6GY_+zpJI4z-pX>G<>DMUHJk#Qze7#y%$Fdwrc|@&z;q&&H z)K6`jXNy>s5`Jp#&-UKuh(d{@IXtR?!VdBu|@*ypm@m#*;t19DP;dwa_8>Hq)$ M07*qoM6N<$f<+x&WB>pF literal 0 HcmV?d00001 diff --git a/nerfstudio/engine/dreamfusion_trainer.py b/nerfstudio/engine/dreamfusion_trainer.py index de456152d..4aa9f1644 100644 --- a/nerfstudio/engine/dreamfusion_trainer.py +++ b/nerfstudio/engine/dreamfusion_trainer.py @@ -34,12 +34,16 @@ class DreamfusionTrainerConfig(TrainerConfig): _target: Type = field(default_factory=lambda: DreamfusionTrainer) """target class to instantiate""" pipeline: DreamfusionPipelineConfig = DreamfusionPipelineConfig() + """specifies the pipeline config""" + gradient_accumulation_steps: int = 4 + """number of gradient accumulation steps before optimizer step""" class DreamfusionTrainer(Trainer): """Dreamfusion trainer""" pipeline: DreamfusionPipeline + config: DreamfusionTrainerConfig def __init__(self, config: DreamfusionTrainerConfig, local_rank: int = 0, world_size: int = 1): assert isinstance(config, DreamfusionTrainerConfig) @@ -54,17 +58,26 @@ class DreamfusionTrainer(Trainer): Args: step: current iteration step to update sampler if using DDP (distributed) """ - self.optimizers.zero_grad_all() - model_outputs, loss_dict, metrics_dict = self.pipeline.get_train_loss_dict(step) - latents = model_outputs["latents"] - grad = model_outputs["grad"] - self.grad_scaler.scale(latents).backward(gradient=grad, retain_graph=True) - normals_loss = loss_dict["orientation_loss"] + loss_dict["pred_normal_loss"] - if self.pipeline.config.alphas_penalty: - normals_loss += loss_dict["alphas_loss"] - if self.pipeline.config.opacity_penalty: - normals_loss += loss_dict["opacity_loss"] - self.grad_scaler.scale(normals_loss).backward() # type: ignore + for i in range(self.config.gradient_accumulation_steps): + self.optimizers.zero_grad_all() + model_outputs, loss_dict, metrics_dict = self.pipeline.get_train_loss_dict(step) + latents = model_outputs["latents"] + grad = model_outputs["grad"] + self.grad_scaler.scale(latents).backward(gradient=grad, retain_graph=True) + normals_loss = loss_dict["orientation_loss"] + loss_dict["pred_normal_loss"] + if self.pipeline.config.alphas_penalty: + normals_loss += loss_dict["alphas_loss"] * self.config.pipeline.alphas_scale + if self.pipeline.config.opacity_penalty: + normals_loss += loss_dict["opacity_loss"] * self.config.pipeline.opacity_scale + self.grad_scaler.scale(normals_loss).backward() # type: ignore + # I noticed that without these del statements and the empty cache, I was getting OOM errors on the device + # running stable diffusion, since these variables aren't deleted (normally they would be freed when we + # exit the frame since they no longer are being referenced, but now we loop back instead of exiting the + # frame) + if i != self.config.gradient_accumulation_steps - 1: + del normals_loss, latents, grad, model_outputs, loss_dict, metrics_dict + torch.cuda.empty_cache() + self.optimizers.optimizer_scaler_step_all(self.grad_scaler) self.grad_scaler.update() diff --git a/nerfstudio/models/dreamfusion.py b/nerfstudio/models/dreamfusion.py index 38211bacd..e5f7cd341 100644 --- a/nerfstudio/models/dreamfusion.py +++ b/nerfstudio/models/dreamfusion.py @@ -76,6 +76,14 @@ class DreamFusionModelConfig(ModelConfig): """Initial strength of center density""" sphere_collider: bool = True """Use spherical collider instead of box""" + target_transmittance_start: float = 0.8 + """target transmittance for opacity penalty. This is the percent of the scene that is + background when rendered at the start of training""" + target_transmittance_end: float = 0.4 + """target transmittance for opacity penalty. This is the percent of the scene that is + background when rendered at the end of training""" + transmittance_end_schedule: int = 1000 + """number of iterations to reach target_transmittance_end""" class DreamFusionModel(Model): @@ -97,6 +105,7 @@ class DreamFusionModel(Model): self.train_normals = False self.train_shaded = False self.density_strength = 1.0 + self.target_transmittance = config.target_transmittance_start super().__init__(config=config, **kwargs) def get_training_callbacks( @@ -119,6 +128,13 @@ class DreamFusionModel(Model): ): self.train_shaded = True + def update_target_transmittance( + self, training_callback_attributes: TrainingCallbackAttributes, step: int # pylint: disable=unused-argument + ): + self.target_transmittance -= (1 / self.config.transmittance_end_schedule) * ( + self.config.target_transmittance_start - self.config.target_transmittance_end + ) + callbacks = [ TrainingCallback( where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], @@ -138,6 +154,12 @@ class DreamFusionModel(Model): func=start_shaded_training, args=[self, training_callback_attributes], ), + TrainingCallback( + where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION], + iters=(self.config.transmittance_end_schedule,), + func=update_target_transmittance, + args=[self, training_callback_attributes], + ), ] return callbacks @@ -260,7 +282,7 @@ class DreamFusionModel(Model): ) assert weights.shape[-1] == 1 - outputs["alphas_loss"] = torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01).mean() + outputs["alphas_loss"] = torch.sqrt(torch.sum(weights, dim=-2) ** 2 + 0.01) return outputs @@ -280,7 +302,15 @@ class DreamFusionModel(Model): outputs["rendered_pred_normal_loss"] ) - loss_dict["alphas_loss"] = outputs["alphas_loss"] + loss_dict["alphas_loss"] = outputs["alphas_loss"].mean() + + assert ( + self.target_transmittance >= self.config.target_transmittance_end + and self.target_transmittance <= self.config.target_transmittance_start + ) + loss_dict["opacity_loss"] = -torch.minimum( + (1 - outputs["accumulation"]).mean(), torch.tensor(self.target_transmittance) + ) return loss_dict diff --git a/nerfstudio/pipelines/dreamfusion_pipeline.py b/nerfstudio/pipelines/dreamfusion_pipeline.py index b8477ab3c..288dbc6fc 100644 --- a/nerfstudio/pipelines/dreamfusion_pipeline.py +++ b/nerfstudio/pipelines/dreamfusion_pipeline.py @@ -50,11 +50,12 @@ class DreamfusionPipelineConfig(VanillaPipelineConfig): """appended to prompt for back view""" alphas_penalty: bool = True """enables penalty to encourage sparse weights (penalizing for uniform density along ray)""" + alphas_scale: float = 1 + """scale for alphas penalty""" opacity_penalty: bool = True """enables penalty to encourage transparent scenes, as in "dreamfields" paper""" - target_transmittance: float = 0.8 - """target transmittance for opacity penalty""" - + opacity_scale: float = 1 + """scale for opacity penalty""" guidance_scale: float = 100 """guidance scale for sds loss""" @@ -145,9 +146,6 @@ class DreamfusionPipeline(VanillaPipeline): if self.config.opacity_penalty: accum_mean = np.mean(1.0 - accumulation) sds_loss *= np.min((0.5, accum_mean)) - loss_dict["opacity_loss"] = -torch.minimum( - (1 - model_outputs["accumulation"]).mean(), torch.tensor(self.config.target_transmittance) - ) model_outputs["latents"] = latents model_outputs["grad"] = grad