From d989e2479fb019c46e4b5c5ce9988cf802e7b80a Mon Sep 17 00:00:00 2001
From: sudoskys <coldlando@hotmail.com>
Date: Sun, 17 Mar 2024 01:53:54 +0800
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20feat:=20add=20apply=5Flsb=20script,?=
 =?UTF-8?q?=20update=20enum=20with=20new=20model,=20and=20fix=20lsb=20extr?=
 =?UTF-8?q?actor?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Add apply_lsb script for image metadata
- Update enum with new model NAI_DIFFUSION_2
- Fix lsb_extractor in image metadata package
---
 .gitignore                                    |   1 +
 playground/image_metadata/apply_lsb.py        |  21 +++++
 playground/image_metadata/read_image.py       |   2 +
 playground/image_metadata/read_lsb.py         |   4 +-
 playground/image_metadata/sample-0316-out.png | Bin 0 -> 67086 bytes
 pyproject.toml                                |   2 +-
 .../sdk/ai/generate_image/_enum.py            |  17 ++++
 .../tool/image_metadata/__init__.py           |  75 ++++++++++++------
 8 files changed, 96 insertions(+), 26 deletions(-)
 create mode 100644 playground/image_metadata/apply_lsb.py
 create mode 100644 playground/image_metadata/sample-0316-out.png

diff --git a/.gitignore b/.gitignore
index d8c87a3..25918f9 100755
--- a/.gitignore
+++ b/.gitignore
@@ -175,3 +175,4 @@ cython_debug/
 /playground/newtag/
 /playground/oldtag/
 /playground/art_assert/
+/playground/unpack/
diff --git a/playground/image_metadata/apply_lsb.py b/playground/image_metadata/apply_lsb.py
new file mode 100644
index 0000000..fa4a716
--- /dev/null
+++ b/playground/image_metadata/apply_lsb.py
@@ -0,0 +1,21 @@
+from pathlib import Path
+
+from novelai_python.tool.image_metadata import ImageMetadata, ImageLsbDataExtractor
+from PIL import Image
+
+image = Path(__file__).parent.joinpath("sample-0316.png")
+write_out = Path(__file__).parent.joinpath("sample-0316-out.png")
+try:
+    meta = ImageMetadata.load_image(image)
+    with Image.open(image) as img:
+        new_io = meta.apply_to_image(img, inject_lsb=True)
+        with open(write_out, 'wb') as f:
+            f.write(new_io.getvalue())
+    new_meta = ImageMetadata.load_image(write_out)
+    data = ImageLsbDataExtractor().extract_data(write_out)
+except ValueError:
+    raise LookupError("Cant find a MetaData")
+print(data)
+print(new_meta)
+print(new_meta.used_model)
+print(ImageMetadata.verify_image_is_novelai(Image.open(write_out)))
diff --git a/playground/image_metadata/read_image.py b/playground/image_metadata/read_image.py
index ae9f74e..946e7a0 100644
--- a/playground/image_metadata/read_image.py
+++ b/playground/image_metadata/read_image.py
@@ -10,6 +10,8 @@
     title = img.info.get("Title", None)
     prompt = img.info.get("Description", None)
     comment = img.info.get("Comment", None)
+    print(img.info)
+
     assert isinstance(comment, str), ValueError("Comment Empty")
     try:
         comment = json.loads(comment)
diff --git a/playground/image_metadata/read_lsb.py b/playground/image_metadata/read_lsb.py
index 8d0bc24..9c2b833 100644
--- a/playground/image_metadata/read_lsb.py
+++ b/playground/image_metadata/read_lsb.py
@@ -1,11 +1,11 @@
 import os
 from pathlib import Path
 
-from novelai_python.tool.image_metadata.lsb_extractor import ImageDataExtractor
+from novelai_python.tool.image_metadata.lsb_extractor import ImageLsbDataExtractor
 
 image = Path(__file__).parent.joinpath("sample-0316.png")
 try:
-    data = ImageDataExtractor().extract_data(image)
+    data = ImageLsbDataExtractor().extract_data(image)
 except ValueError:
     raise LookupError("Cant find a MetaData")
 
diff --git a/playground/image_metadata/sample-0316-out.png b/playground/image_metadata/sample-0316-out.png
new file mode 100644
index 0000000000000000000000000000000000000000..8b366d106739981f660a23f33074e7e23cbca4af
GIT binary patch
literal 67086
zcmeI5&yQVKa^IVjU<5`SB)i1fbQok~6AG!3XQTiQEDz-wImUQKBiJ*DjZTYvBs*$$
zbGn<u(HMpY2E0fX+2t?r+S%j}@FJT4-e(tN7a%|unMHs&vDm3^RejE>d)|GoUpKq?
zL(-AntzX~zR-JS1z3;xqZt~xJ_dDPErC<EzU)<W-`lS!Qedi;SKhx`P{em6;(?8$+
zD-(bJ<bxlae7twEzu0>J!|ex)gT>MA$>QGj-lN?Ii+}g?kN(^+e^;3AAKpKCx_h+P
z`p)5x7yIvj_-B9k7yr#-|2Se#j_xkDzJId&$^K&d&ffj|PmcEv54L~sJKNv9{lV=!
zAN<<eckcYdU4wrt_?^Y^-J`w7CkENtKHl5^@#5&S?WYg-P8K(}KUo}~Z2#!V?*87%
zXE(QZAMO5h@8CgqICR%A4;LqUcMb5#@sp!_M~mZs`@jF!2U}Zj-W2~g4<9{R9Gq<Z
z(T&GPhmRhg-1ukPH!iErjhovyj!zbkkB#*0cg=XX*n99$C2qU<bno8D!)6`d-Q8cv
zy7ksA!#+8<dw6hfr+M9Z_x^(&n?NFOz4bM*9xoR6lz01`cfbDL+qZ6i-Fyc-$GeXn
z?=MvPwkq%4TkPCheDdVMj%~@&{_bZMxWBu9Y|D7#;BfDFVU->(?mh7qZXE2M?EQG5
z{`VITcJF?+^W(4WEPisbI8Y(dju!V9M~j2Ii=DlL`-ex5Ea&h*u%q3(c3K$n)?4p>
z!#Fzd_{0((#H4*f(dCh0Z+yD5u<A!DYBN7MI@~{4oa`)Rl-Q%gdy9QbJ$SOeZ&~*~
zJJ@}+cX#LHp)JwF!~J`9Hk5Yf!aI)_b}}BE?D@2Bz4hLE@2w!(d9){PmOeM>_>mPj
zIeMa1JbtueC5rv&-s2r@*e8qq!>2pPdk-G%S~PZd|L|yc=ZPlx)Xxd+Q%=J}clRxM
zDfHy7!EPKJ-+y{@+c<5SH@E-jlihoNd~<tobaZ%hbNl|{>B59>!p-eRd&kGN2sgK%
z9v;<8<zGKuJlNhnI@!B#J7pY3(Jz~)c1_!69({Uq`^mxmy@S1DJ1=f--`hL3^JzEu
z?mo1O#rFHIT~t=tcmAebHV^E?d2(c{;SE^MgFW%9)UnXn<&)#Zk?ob?Y&{QdZvRWW
zH2tjkL-Rb^-QU+qaJ2Z*6T5hKPCk3A%i_j&{rdb?-gQ2Ts9$1Mc5!d#;qI|6rW@aS
z^gH)Ie(=4IfA_<8e*dGpAMYRTec!I|_m95)Pk;FC@lSvI-tYbP?K{7J``ynLKl<T!
zKmF0uuN~h0>0*EXH{bcmH^28=KipsZ>aCL}zxCd~_`UDG`_ZT0-n;XiPrmnqqg!|0
z|Mg$L@yGx5?_Jxr`_ezseIxtOo$F-p(PHcCZ@u&0+wc9_t#{x4`mNiy>?-^x-~aiq
zY+c*h`tZ*CAOGaq)^GgmFaPA)U!R^n`-T5$6o1=Zr;6I~^wfmZIG9IT9xZt7ngwOH
z5jk?lXo{JUhsTh{6>Ob^DA|#dWGV0=!!2tZQ?@aBx%OQBi&mg-pCJZhIqjA8YM|tz
zX8P83<9owifAef>>)G|KYp4ImHuUuK)2;t|-L~bL!A?mrlx&Sgo-v&vrfMTg(Dcz(
z@|+L9q=pIzHx3HJ0cI6`%>M$|!_@1{Ssl%im02}o#lh$VP_8{mz3FGf|NGnwJ0ni*
z64tfp*Qibm%fJ&S#!N|fiMqtx%yOJWUS62=n>O9}A?D`5+iC;{(aPwwS9ln@Y#m);
z$q6-co;YU=*2M41)^zLJ0@vQO1lw=hd}$_4u<`6uT;xI<pk{l992&^5x=nffBX*kC
zPO=Jb`lNyMMniEd?X{e%Q`*e3t!9_buC98YtIxd@XDw7)`ebYDzAvc1?4bB{^VvD@
z=GN!7)6dL#%`SIsIxn3$&kXAC3Y-PbEOt?0NDqs*gt&*u8!>5pO!G#bz^M_-^>px(
zb9|log=(~Q2}PmO?9fGf=4>mt^x-^;C4u&k=QCaS(zVa_uK#8P1ol@x;N`|Ur=k6k
z&7Tbfa{3~L(wjyfT+yDlga;`u4fv{UK41YVI0cr14fQH3%y8_n&Zg-sb%Nzzl`lxk
z1CPTbs)t09M0OJBEkIX(Y@i<vt+S(_8M2#GX!4WUz<zQ#z4Ju{L$>#SVx!E~T7Guc
zD+Hgzh>61+9LYAaMOB$LN;)rSERfT5n%a8BfvG;~ny)PpK9f#GY{+xdIQ<Tev>d+^
z^qhDGaL%~b3*u&5CIiQ&#|_@&4#}(DE9oF7Cr5AuOau{&SvX<CN1jT?8;5GaC9e;;
zH2mHbYL-hwF?)f-8*N}Uxx(Gz7gbvMv6mJ&gL%YNTl&zrn|5gKst=_+-40N?)LQ^9
zzsi;5AW;@{t}ge9*vQx_r3f1*Nmv~HP7+~@CuiCX<1Mu*I8(&B4?Ausd6F>`3p(>)
zdrt8M$K@QKn|6+1&NsGu7RKTZVd>1cW)FMvaqajT`hv;n9|A$EnxOT{mQO=WnN{J5
z7t?fzCx$k@pDMz|E4DGBicd}u=T&~^!OHGrwlPI9i!r;wUwcLKt*@mCWzcuU^d{vm
z!~FVQ$c6`WV5>rJF@FY-Whw2=s9+qZGQ8Xog+tAwa`eikT<Zr;mXLv>qhOxf{+^FU
z7m`;ou9QEPhh<u7e{R*6K*iD4z!cv)w!ozQpU(h(b;}V~qD!6$T(S-0#)-{_$VDg9
z?=>nL*LXUGvJw97vjn-Ez>^aqU4y#OSq3At=1BG$$g3OjwX~V~x=h<-UPwiy&p=zE
z;^5;gw1xCJzdr-$djJ7q$+HVGjUsN*N|sY4DQkAaDU3lJPBL4C;AnU=ZOBw+zt+Qw
zY=_vPbK^Nn!ZOOU@?J*VXQyQxraAULgjcq@2B#Ci!p^LfPpEH-{Mkj`JHm2s7Ly*f
zU--h6u+$i#aXlj)?1_fTJpQ0*g0W!S1A-^+EuEbzbF!k)M>&@_=VMPK>k!}QtMLoz
zG&v;WMh;;&m-?Y`iz?RQ{zZ^l*T%7VjWMU>0*3T83jCH1&ZSCdQTlZpWy9-14wF1H
z+*S|{>ES%j#w$xhnZe}d#)sC%kTK_<^9^Y)%v?8DzrM3s;p~;q0BxBRg(9T{N!%fP
zIs}-$*`96(ce&9XJoG1x-0dec&ZR@vqnOj!!S4_j)7{ALM2Q{Xjb=uPPO?3VyG{NK
z{z5L1zJf#3^wx&{d=2^upw7)z?*hKy@@w%U+;>nmE7*d89sSrYvn*$VW%Lm-9s4kD
zci5o?q~Pz!vozA1L^&(*HH(ka;@x~Q4Yo`tK)*>_ss;Sr)~-5XwCZu&HMJ#~BGz_-
zIIsXBGrW;B^{B=0Pp@ia^oPwpk%G5)%gqjzb)J!xBYR>@S!U1CDdd&3>RNk)gTq%a
z7uSUSNQhAGsH?%o57}w@=k%7tI>-S>BlVy|UCT#1Xgxr5;&di?`YYlbyWfc~cA&K6
z__`S5OV3(04fpF{gPXFc`7uzd>E2jgfn1xs?`F1wlPfW??zJn}5*?Yh{CK|)UY_M0
z;Q>QBSaB(VGi7G<dsUo99b2=VZpiM)2#j`v7a~pl%mGniUBGc>@(RH7=C7fp-2#b+
zKn1j@0jRLl>+L>QCg!ZTG^D6g_tena+Xyki$a9?R%#+%Nd;m=PTaPlsM`ZSW4i0Ni
z=c2HlP7?7^X1rV_6NK{=kJ}^XXczpf$`TS>mmdoca#ngzvymldeh=WAp<rC24r@*2
z3Qlqr?42~Id#k0Sm`!aeGMS}W?l3!U7fPBL2J#f)BiGYF;+7L?#Zm{0B<;L%%pdLZ
z^0j5GMD#D8^eWF?U`P-jobhJ?Svf`2wqqzEhx;K{U{C|Bg+vyFT>V%*Ckv`jZhYK@
zf<12YU{a2kC8d!SE`5B}7SF50ecrr_;OSer)J0x|Z_Mdj<qp6iHgGunN5$ONv)b7q
zmtb~p^}WFs-gMso`GFIGIUxjz*-_XcUcw0t$tMxkAaS*j<CDhuHaL*80dxtnNp7^i
zlzP>UrLs$o(m{{uV_?oo!mjk`9L*9~hjR7BHYQ8#+T>HAr$nMxoo;1Bli1>3*m3g3
zm7m;AM9$b<CcRXVEyi@Fi}zCcij}g`Y~A`U>}LYo6$O$h#a6i2pPoc&F~8cuC>nZu
zTTq<MH(hsuSQp_Kb<^=hyssgNc*sGL_i`oO<s_Nx3vY@(>grT|MMKwH-*W)p5@(hc
z8-xga`le0`MLT}wn<_cEvOPI4jW-C>RFWr}R1IN~0AB;i@?5ml3*>K{5<q<s#Nw(;
zd1dR!4YyK83{01v;#p{wA+q2ZrJ}r`Fbpo;eY`K2l`r3GFrK%~RrxI8^DjYs4-h;o
z?7*%oancBJbLP`~KD}i-@W=!!9TI%9rB@pt*inK-iuX%6a#>idaRpel(T}aw^WP89
z$G|CAA$jcO0it!m@re!`@hU?H&4;@rR`QCW4CTVL0`_Hgm!)BSXP2YkYP{r>JHYT+
z-+ximHlc%8usJg*Bv{#1rw#eQBJm;Tu<}(rn6D%&!0NPRConJdL!oSt5Yl<cZqM`f
zOa?Vlal^I~r}BWQK6oXGBgG{h{G23@m~y_LgDE2GG}BjCZYEq^>vSFhg%NY5jY{z-
zh@VZuR*G=sMvggb2*V3oQUe?^ym2htlYYP`K9@)Q3oh&O%uKa~r8Hv(7%n{nXfZ*u
zuw-K!#2@&uCG`SFzPWtN;S`v|S8QYl4nMw^BgdJI_NuPH)fuCYIbf@Q&Zqb)f6YCI
z9?t-s6phUlRq7rhj(;qrc+FM3`s${RT;M9E>f#!0`CiPRyYp5XJFc^aYs9SEWyj8u
zB$?uC!jwenw}b5-pw9$8XoYX{m^KfC#0AbW{#73(DPM6_k1G90o;)>|x*$g!{;QXB
z-JhqM^I+?2=ISc$E5lONgp29+7NGZVZDOlfTh5a0K~FJWip_Nb5jgnpD~~by;6n<h
zeAyKXr{b^9xU{b7R5iWPKBLeGA6KZv?NfkReCOAJKvGT$%7@<q9s@4}Fg|2w@0jYF
z)8Qfp7ZU%gDb>9?eeQ0{GQfqr@XUi#9|N;Z|CmQ5WTj7BbuB#@)?mfMCy72zftdUi
zGvZW(xbzn{=Bm;)F(h4$HvTHl#LO*0CAg@5s$UMZ>cOvLzG_ERS0SA!d?zqDX%rkg
zC0TxSaHB1M9W_VT`l`4Wr;oE+Rk_Yg80;$NW<z|5cV%D+xv?X!UEKjRYCER;OsOh^
zsVi1({7#aG(Fc#hb1{6#kxTi+BzaZb)$4R{Rh8;0E!DV4@y&AVQ-Cih<t?kI-6$;)
zcR4W%NUAy@PDu9XufDPCbul<!A34r^W#=p#awTv%z=|y`{W4UZ@ehN{r7r!czflP|
z3}SV|MxE~QJF)SQ3uTR0U&Vn@9A2<RR#<VLOFe@_;R>nxrghfk`0M~6HzlsUZI+wA
z7~%z#!Q$3z>DWjmfR!#kHW<n37fzTO$49PgjX8Vt1#)uS05ZYP85?piJ5h{!_0<@D
zBr75NBI9}Ut){*S)Qm-iG<9jzRjImww1FgEIpM~)E|22yRV=X;veK`BT&0_ik+s&G
zZ=-lkl+IG-GPf|XN_w_#KNPBmI)9pPRGBt16334%IGZlL`ia#Xbg<;k<&7BG8mvNC
zuk&$k*3CCA4ffoG$Jl7k<;3p>f<(?CTT8aT6<2lT72M>c5Kb}4xg7H5;uQ}LX>@f|
zX!XBF`ihxd)eQI;DE?ZyZbbFZ^&u=6b^r;l;-gP?^;N9&Rb%8@g?o+t<FcG%D(6fs
zCjfisn}UJV9~wh#KD}*tsK=*0;zueM9pq@kj}%Ur`ckaMg?Wv~xCAe7D!GdMJb+c6
ztSg4eNt5ztY~{Cr1>+NvEfOSz*bd5rd<=3%ANk;3V_B!RZW<$Qv`cPVou-svswF+{
z0O2U^cGX!%x<iOnFSFV=g2mU)mg>u@E~p`yJT<(e3-gi&Q$VVX%Wh{=z7dCw9C~Aw
zH&{8H00r(3cw)t8dh>}c{SvN%ldO68<<pCSUa>VsT)xTyruwVz3UrmPud67{6whjZ
z6>f#UuHuAxIiM53yb5@*B69YVA{qCqzfzM|C5cF4$?NIzVI%R);cLw=fX)=pI~A6+
z<`rL2D`Kk0a+f=RXU(OuA<EA<e;U{*gDuR82S6NtW-c3DSaOm0z>N8de>HpsSRylA
zWvSLLPdo!OOiO|dDMZoz4ggrP#5hSBDN4muANh)-6GzIwDhB2?l$z2d3}4wzxRx<-
z0>swT4Pvn+t6sYmo|pX6!NP^DxW?4qe=o7t3FH)uuX0r173#En6(uf_d+uF%i9Q|0
z&u<15mv5qGTQ{%j8B@ijw}QG#;$&kpr-dUn{gLz*%shI62Pqt}nsY_Erf^l`lDM5_
zhD)x=N+qU1bpq6dtur;hBY`6m<7UmBlqy7E<|MG1JJbQ<5}YvT6Tp|!U}YzhU$czD
z=j!pNfl7tOq^@jz{PU~4(q@{wsDMv49OMuirQ&n>HIF#HD&ecXFRwng@>-3vR>|V|
zF>vy?gg|Eu?G`U8=q%p+DkHpJ3~;h*3|+Efz)5P1G1=8uYk9&~2dzL|qFj-sl}b$U
z^9;bIlu%PG%fohPzXzxa?#xoyEvfj=qpHKoF*B>VnVzS_7ZA%TAvMLSw<+h+D&}hZ
zn?Ws{X)T<#gy81)w*}L&oklUI>-biFY(qBuHHP29ZPYmYqio<WwX(D-wsiQ`E(|{w
z9-E>hrRl}ur|oR%cbDSsELc;0kPD_tFp|mFYStJz)d#-%zS6Y<zT#|0SrT{6_zT3I
z0Wv#qGTh-Ztg!n}56gxXJ&-}H7_Z3feraQ_eBczToD~a3_^-y=6|iN~^SiIu8O!z?
zKl{r+v9rLR>I99C=D{u+Y2gn&cf{TU`Zo{qCB_pI<31-aFPMr$uNV<vONUE(9q?7(
zSNgaDTBTJh=-625{3u`^y|+h#j&{+vY#{^fM}Ha&Pro%d<@DhBjbs{Yt&MTndQnpn
z-$3z#%$Ru;tNIHnf37mMKkZ+&mxwHGTdBV;RAc;|@OUK92J_pa{i{1NpsZ!8J{Mo*
z6KC*(2z%r+j7~s5J57=znDi0e$u6u>LIhR6FyR#@&h<C9B#yIy{cpb<U`O9vknJr&
zj{dgLe=f|tt&$XsIXyV~oGH>hkzf5aN1VipUQ3Bewnp{ii^v!HxOS=pk*tK-wvYNt
z0*FY~#SUad>-Ouu08Uaba?vB#qxqKM7}W~G$f>m(PPy<(+VC9H9Z#(Y%^s;C$!iw)
z#c?ua%_vuG<>cDiyS@zQ4|Imt0=q*LLFt=qekZj}3i4V`Iwge*49H0rlVVC!UcRiW
zW$w%-k23u%n$TahLAO_3=OX-VX*Ni1yVt)AFqdC}Mx$uEGR2k~Y$A%tmOi`19q*?Z
zPY{NWG*iR^M{K(8FR%r#;v`i4JkvCaOB}Uw)i1F`pYjDrS*_XBAz1n^zad(@Y_*>h
z?ba-SF)fj4%d~+65@#MiO`O+a$O%sED$j#5Yy?ka*I4)+!X3elmv^XLUrHA%2v@kD
zt7NXGy#=^tC7zk0E1#>sZM?Q8;jN@v;_5c3FfTBkVh8)*C*ZPP=K|KLDFIuOSHTG)
zzf0-o7#GPY^hzHoM!Ue5RS4E45knRhR;IY-)tTH8C`j_<F9+apZ!M)w?i<}Vg@lm$
zM}MXDA=^CeB)^#n!>@Y5L%u9eb&k;D<6k0bizIC#JMkvyS<Hng)HAP<xsv8|fiTs<
zO*ka^1vFQRuwt3Q9|?+wQqmQkZYqyx#E?#Ew}5yf$W+UBHm}50A0^|)68upybX5&b
zt8A`Q9vZma*Y>i%$T{QaJtStraYo*Y4AYmEq6}YxjK!K*ur3!^!w5E{6K#?u<tA?d
z{WT%J+L;UmbXnagGFBeqT1kU>w51!X>7MNeyV_fx9J=}gj1{u=YD(Tk9OopS*GO^I
zlu>W}jDc6w9_FNrrMXa$bqY4m@{aA1*{Lg+K;-Aoh1sd$XP5`65Rz@csyoYf2vc6V
zk&JqZcODf9w!GA_9Mg~nTjFiV!Vwc$eH3()ku7Zkbqp=g=X6S~cjn;hTgdUtRdpaq
zhv{+-2f$<5iwIl2LewrwOG-Ci1)JHTeZ~Cdtt&TOx@sdau5eoKZvaW_34P0HQbe4o
zih90MoIh>+JqbWcB~WJq+T3`~bT4EaqUrO=X(V?pSxij{!lzS$xI#-KIhOq9s2pBj
zz683&D9NgYyXmIb*5)&SZN47lysZ>5J44wFNp?;*8cm5!5`)VVVzR+uM-^_<@26Df
zCW$wRw4))6=aXvx5(YGVHDa!1QKE_gz|NIjiV0tq6V-ELAzjrqJ6DSu1?u(P&SIW=
zs%L;?8J9bl106A?=at`>w3T*(w7bVxI;)wx^RxZ52gdBljEIX<7;us`rufSrbmyH+
zjEnWA8R{fprpQWNZq6%<IFCiSb4-BkPXpT<hQg&(X6qi%HjuDO+x|8C5qY&_5;j+z
z&W!GNr38DJ;%kaF?h7ilMaqe)wkI20#mIl2F>)$5`3POcOxrR-nmt-6dkbi%k*;M8
zWTJBB5sRz09orQPb9N7hadW+krr$}7yPR~L+dwj$&J>r{kE`N>gWfJlinxx^Rt%r_
zlwkAWuIZcV^$Pc;BvH$-f_QX~HWbL#u7KjrKuWy@gtKoIDP*KfHw5BYZc-p2Ky@yE
zPm3M3c;fu_fdy6@PQF2XSa<FMV@6UeNm8y#mez{;{MC|7BG&2xlw(7-Mp`#kt!_(R
zyL=BIkeBYqrK_8Ua^*Q?Ythy@yJNxJ<OcI@?7pCoC|izf@=C2?2JMcbRD%ac;IT&*
zkG2NjK10ruoRM?RaBXRoK1-?IG-^_KVz;Zm&O`pR{b#}S5Xgv_9kq(~s#l>-4n633
zF61sFD>|`grpH&qS;@PUZUdK~>FUmagpID8223XnEYF_^val@CPGYN!vGgk3CVrX{
zr*yc{k#k%%H{<StLq0G_c+74t;q4{+nH}ufn?3`yoO;!n&Sp7LyJBUD5vlC7C0f~v
zQB`n)JEIgiZB(xC#L@gftNSW-+P<tBOi_G;?z67GvF1!nE>^28jXB-l93xS&1@t?B
zDSlh^jrR2Duid)icJ1YKXUMXlMGK)TH@OlctYR?&n?|~#-o!UC$dxnrXQn;LxHO}d
z-h^~!dYhyP!^}3tgCoyn8s3KV?7y}pUDvh`H~bD@4d{Mg!|%4X@h#t!xJac4VH+xt
z-ysa=GIOcuFb7*o$mmA8PKfFgIQ&__G4jS(7X$N(NT#)-9j($CkQJHXj#cifwFJ%r
ztwt)^E~zcwjHde^23<ehdL!B51R%j+%5<j*-|WKVEU-bDL|#I%>JL5cAi{~9xO5GZ
zs4W(eIJRJ<%7miAq!TMC6Ue$efWlOsYJ;hARjtG<m7L?N^(uGa@U%)sPWl>c`gIZe
zR@lD~c1^z!q?dF(1Nes6Nx-e$++Ir8Y8$2|ciW)2wS(bJAS)`@E1LA`k?&|fH0G=;
zTY=Vk+VX2uzwCiTu2*m{qfJ_KtWA>}tft5;+u9c*e}|q6WKlikv!uGg3wljgx^4U$
z7QC*LApXj5fU1~Cb47Kf(MG;C$Ys71rjuhnpym=<m`dp<Rl6%VgLBy-2Lwnk3a`@1
z-B4C)>6LavicMkXp$1!*$9)9oo8svwjN7H0_TnVb$>6v$-#pQ-o&L?~v#rk)klX5;
z=R`Va2hHOQt`n&pC1&Mnf?b89GjP=KwlTiglBLlsS52E`5~k>+dRp-drXiQz<U=v3
ztKYWWZUy!hAist-KOXQx53~hz8foJb7{z2zHnMeWpf&?=5#?<qsFKjJohv>$um;oo
zC}}=jCN(`Pi$#MhBo0681yFj@C^k8-6_!sDLet&WK|}sgPtF!2*AqP}4%kR}rB@pt
z9I|t?yKLG@qZik)=`m&g3!z@6aKs%zCxd1cQha$k_Z6_5#%!>v&(1cu99UOS?v`A;
z^FTj1zYc~GOgZku#^=Gr$hTs6?uuA-CtYii7=GE=mWR+s{#AJRYF@<>e-7C;w(^>;
zl5g1Ny4wsr13cpnAh-8qmqh!SfYyXXcneUvnIg%jfkZj2d)?8lPW&12L_5%>iu&;C
z7$#dFp`IB2cB#&BcZ$wHzY!~0G3TCA%eog1s#hrMAPKI;M+bH(?y4OKwcJqwK(5(E
zofvSA6wMxTwM}*o@Kvy)qEN*w!A7~e`-F(j3s#PVC>*22oy&=O%uf_&vs{yt%?&@U
zpymw({y{F)?_Ic7OUN<W0be>sC+Vh!$>wj$fLUVd@iAo;G_STER$Y?mMbWBv^h7(=
z`@VARb=~lJx7P_EZ!m<`ZTx`|pKvuU-;CYv0*W&X*G@9|WD_&#`pA_yhw+HI+yhyX
zD@Uy*&}gqrJR@*fJgdr#{G(m+Scl@8s<}Wq5mMC&0Ao$haNAsA<V*jM5+)z~0_%Hi
zc$mOMJ$VBs85T%l2Jrz78yt3@LDXnncEASOSY?BhR|>+aAX(nFCZ=#lRZJ&=3El#<
zXh4~lY|-gudc9>gNW`YMrSy^KGFV7kmi<JDSk`R7+M8mRKdy$=`eYf^GJ2KrkHv{A
zrh$zN@7klktIi>;Rw<{i2*MP1=q2nvJGAEkp#n*sOLyZ}5*#MoU1N%4BjwYmc5=uy
zd6`jd&pYEFyaqC?ew_1DqoXdaRW52F#aiRF%80R;uS)e-e9?uAyKr~v4!~W&^|%AT
zrh5Rol5zt)Uw-VOi^WehFXzZkkzAy;K*2W7NSz)tUl_V}KqH*O*IA?>N9>$`&Nt?5
zV$boggiKL96^D*2b!tj+Pt<m^R=y>S<<hMj-2wXO^?;1!pi+*QGed6aZj%EWk~(mV
z{u-BF_*xxus6oyZ%F8YVAClIqU3KCqeFa#CAvT)pU>WPr0KTDq8jjVmkt$PlY|2S0
zj;**xwQC%%8aKG#zvR;^G3`!Jd=|aZF|_sEwKm}!F~#<ksLMCgP<tlJSV<qz(>p+0
z=n+;oZ?Z7}DN#2QT^PkjTmEq-@F->jbDYtHE1<9Pk!jQ^U8vu>gd#Aza<Af4+L)Ty
z`aMAX91w6Fdx5USCjwF6gI;#kh0(x=4i4#{DGus()q^2WIf=bml1|mtoPO>Z7>V;h
zmQq4o`O$8&+{DzwqzQNUGk~X>DalJ|$?ib;HB44_IPHiN9Cqe3Tvdb-UJ^VyJY!t>
z#H7F2L6@6WRi~HTEu2X8LJgM~+2q8W7!KeWV2sG-vpLgFj0^c{4tm9cL8_o|U}V=(
zImkg)+^YPFd47FWmC;YvwIW@Q*ou)}<8wUg@HN(WjAX)N2E8==1mIr;we-Y*kR(=g
zcY#q#iA{9H$d}Hjx%f7baMgO7<XsGN(JobuTqK>3oe&@WtL*dQmv?`OSr37tg&HZK
z+Qw!I6^n1gk%JVM26%MYqfapn<-%Jr_#O1xAL`65I+^ou_>NdY*%{sfRK6~*bnF6W
z!~By##@j>M46bpgP54x695q&XgsZXhOn0ntKKranbTxKMm9z5K;ns1FE4ivZ{HtVs
zWuw0&N|d}rH&QYcB<;yYR~h2?s}CI9lJzBQ4u=?rPBL_2OTD$gr7LtdN9=7`0J#dY
z6kl@8Wv=$I!mAO>k#!|4mOmBjd2uR40;Mt?lPNcwf^rCplpk9<ex!VL;5R8%xh^Xv
zdztXo_|@oCy(s%6Kj*YGwHfkBUdF4h>SMZMEMtphQS$N*pq~jG+mQ`S>BOrK+KNT5
zcv?Y{ha^XK%_A;ZbHUX3>!;7H<fYA`ENT5_-ru$Jr$D^E;p~4s14KKqw7emVe=rp7
z1$P|#hDGLdqnIKwa>Ro^Cy7hWao~qn_Uk*&nguJs3@7+pzPicakO>LXHox&%K)Q+e
zDp=dj4_q~nZxmDk2h%sdqx;5IA9~~mTh-;kmeflF8CJJPKk?U7GUH8W1zxkQdYgT2
zthR29^GH?wrt@b3Y;A=59}64PWHS`kfX3qOQ-HYz*U)#8()~Mp!#3Nab}0GHij$BE
zNkBrPC{WxlN*zm&_H*{ZiXhLTn=Nn2&7TT<x>LAb{HDgA1pwmr0)j<eWo3ra*-!_x
ztbkeE?<By%xGiiBw{^gWZqnbV>_ouTYARRG=~(u3)+ut?qQuD2x=Ka8H0%g&0pk~b
z&Q@tirLP(TcT0h{<oPnm9&^$rm$2rbUkt$ym0kXP4^XNYuq9fGPKtI8GzzFdDkt>r
z;a5B;H|8YL-O`Cmn%#-Qh9)H$qa0u>&01|1WW-u%+HjezReWEY@PcoNa?44<suPo3
zH5y?7m~W)1^+?e#gZS?Q*spw8d3Es2wo^Y8?6-Ww>0_WqrnnKF`0L|@lx=Y?1+!vn
zl@nM~Qk(kl1_d$8DO?THrleXhk|W!OOpT;S6_?8kA5;^5Q@P5#m5~(e7&N=Nni9#4
zdhqvzvB-?wNJ<@G;<dO$dzo#&1{PWIpr_a~`|k}Y_J7;~l%T=JdWKd&e}%lC;->3P
zkf^xY$qN1`D*(21!vJGNiQ^Z{haAd9!h<wFy$C<%*VHj4yhg=@r<bcXx+Gks4G-h5
zSoE3LjEj{ZYqg~B)re_C#q^OcZ{Yb`!TJYco_W!ni~WA!_!iJN*;mV-5yd8d$Y~`y
z8&{|L=ccaa%<nKe=s|J<n;2d{>SoRq+(NF5sS#a|AqQD=(AUMlND9Yt;IFwQR^32L
zxN`wS0<MzSRQ5~8Z~E4qWWR#r1W-}sg^aBZ?#>o>5S>znO?QaeGEAK<cq%57MWo@v
z4)xQ7{JnJ<78h{j74Q`sIju=obG`|n8@W2+D0xNg;M68L=4)6fb|<i<eWu?9wCh=0
zXa7n_fXs+z0Dqe>@b%o;CBov6n)#;94a$>*Eo||5tr~+%KB~x<c1Z>KN=o{qxw$D4
zyDXVCbNd>r+`f8t1NTn>y0+sE;P(SR0RruIJn`3!$>t*}WNt?COiu-7$_jvu3LUGs
zy?{&aMuN$0AV?y&Sg)Kq*YY`=TsQk!DiSj2J75QUg!8}ICmshZ=z+G0vbh5o*n$za
zJAvPhjn(h>lq8zfxMj;N-NTA$DAxojUGZof0F;~-b<!jt-pCqiNr7@HFgP~F4QXr5
z3vlu-IPx-6q%YsdFFQxQ@>`C@q}%J7eeSF8>#;ZS9iV;KvqwRlA$_-CRytKf#<PQi
z3tslBoWM@2gNHvb)B1{O)cUf>*^DtbuJ%YrS6gMCy{@_ttQu0!{HLF9o$7A{DT}9n
zEzmy(>a;lh?>~R`U;e@0x%Ml6?KOhZ{{b;$X+$2YHg5{k4Vv2)5J^#(<|9*_td8p^
z;Mm!6*7JD8-m8Xw7GLd1xplx-eOIBYj-H)$N4OqWwKp$YcbmBAd*$nP8vJLw1AODJ
zY^&cff4o<mx?}<R-JyO&tS19VNsaL|ARf^;f>nRQ2PQTcJGeStDDpJb*LM8HH;(?J
zHUIBFyz~CYmokaD750<C*ojHj6JTZq{WSU0r;mR3gFpX+|Dnb{{NOux{=;wnt3Un!
DRryH6

literal 0
HcmV?d00001

diff --git a/pyproject.toml b/pyproject.toml
index ecf858f..1021699 100755
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "novelai-python"
-version = "0.4.0"
+version = "0.4.1"
 description = "NovelAI Python Binding With Pydantic"
 authors = [
     { name = "sudoskys", email = "coldlando@hotmail.com" },
diff --git a/src/novelai_python/sdk/ai/generate_image/_enum.py b/src/novelai_python/sdk/ai/generate_image/_enum.py
index 71c8ace..3d0b084 100755
--- a/src/novelai_python/sdk/ai/generate_image/_enum.py
+++ b/src/novelai_python/sdk/ai/generate_image/_enum.py
@@ -70,6 +70,7 @@ class Model(Enum):
     NAI_DIFFUSION_3_INPAINTING = "nai-diffusion-3-inpainting"
 
     NAI_DIFFUSION = "nai-diffusion"
+    NAI_DIFFUSION_2 = "nai-diffusion-2"
     NAI_DIFFUSION_INPAINTING = "nai-diffusion-inpainting"
 
     SAFE_DIFFUSION = "safe-diffusion"
@@ -84,3 +85,19 @@ class Model(Enum):
                          Model.SAFE_DIFFUSION_INPAINTING,
                          Model.FURRY_DIFFUSION_INPAINTING
                          ]
+
+PROMOTION = {
+    "Stable Diffusion 1D44365E": Model.SAFE_DIFFUSION,
+    "Stable Diffusion F4D50568": Model.SAFE_DIFFUSION,
+    "Stable Diffusion 81274D13": Model.NAI_DIFFUSION,
+    "Stable Diffusion 3B3287AF": Model.NAI_DIFFUSION,
+    "Stable Diffusion 4CC42576": Model.NAI_DIFFUSION_FURRY,
+    "Stable Diffusion 1D09C008": Model.NAI_DIFFUSION_FURRY,
+    "Stable Diffusion 1D09D794": Model.NAI_DIFFUSION_FURRY,
+    "Stable Diffusion F64BA557": Model.NAI_DIFFUSION_FURRY,
+    "Stable Diffusion 49BFAF6A": Model.NAI_DIFFUSION_2,
+    "Stable Diffusion F1022D28": Model.NAI_DIFFUSION_2,
+    "Stable Diffusion XL B0BDF6C1": Model.NAI_DIFFUSION_3,
+    "Stable Diffusion XL C1E1DE52": Model.NAI_DIFFUSION_3,
+    "Stable Diffusion XL 8BA2AF87": Model.NAI_DIFFUSION_3,
+}
diff --git a/src/novelai_python/tool/image_metadata/__init__.py b/src/novelai_python/tool/image_metadata/__init__.py
index 9dd8ed9..68619e4 100644
--- a/src/novelai_python/tool/image_metadata/__init__.py
+++ b/src/novelai_python/tool/image_metadata/__init__.py
@@ -14,6 +14,8 @@
 from pydantic import BaseModel, ConfigDict
 
 from .lsb_extractor import ImageLsbDataExtractor
+from .lsb_injector import inject_data
+from ...sdk.ai.generate_image._enum import PROMOTION, Model
 
 
 class ImageMetadata(BaseModel):
@@ -69,11 +71,20 @@ def negative_prompt(self):
     model_config = ConfigDict(extra="allow")
 
     @property
-    def jsonfy(self):
-        return {"Title": self.title, "Description": self.description, "Comment": self.comment.model_dump()}
+    def used_model(self) -> Union[Model, None]:
+        """
+        Get the model used to generate the image from the Model code
+        :return:  Model or None
+        """
+        return PROMOTION.get(self.Source, None)
 
     @staticmethod
     def reset_alpha(input_img: BytesIO) -> BytesIO:
+        """
+        Remove LSB from the image, set the alpha channel to 254
+        :param input_img:
+        :return:
+        """
         image = Image.open(input_img).convert('RGBA')
         data = np.array(image)
         data[..., 3] = 254
@@ -82,23 +93,32 @@ def reset_alpha(input_img: BytesIO) -> BytesIO:
         new_image.save(_new_img_io, format="PNG")
         return _new_img_io
 
-    def write_out(self,
-                  img_io: BytesIO,
-                  *,
-                  remove_lsb: bool = False
-                  ):
-        if remove_lsb:
-            img_io = self.reset_alpha(img_io)
-        with Image.open(img_io) as img:
-            metadata = PngInfo()
-            for k, v in self.jsonfy.items():
-                if isinstance(v, dict):
-                    v = json.dumps(v)
-                metadata.add_text(k, v)
-            _new_img = BytesIO()
-            # Save original image with metadata
-            img.save(_new_img, format="PNG", pnginfo=metadata, optimize=False, compress_level=0)
-        return _new_img
+    def apply_to_image(self,
+                       origin_image: Image.Image,
+                       *,
+                       inject_lsb: bool = True
+                       ) -> BytesIO:
+        """
+        Write metadata to origin_image
+        If you set inject_lsb to True, the image will be injected with metadata using LSB.
+
+        **But if you set inject_lsb to False, the image will be reset to the 254 alpha channel**
+        :param origin_image:  BytesIO
+        :param inject_lsb:  Inject metadata using LSB
+        :return:  BytesIO
+        """
+        metadata = PngInfo()
+        for k, v in self.model_dump(mode="json").items():
+            if isinstance(v, dict):
+                v = json.dumps(v)
+            metadata.add_text(k, v)
+        if inject_lsb:
+            # Inject metadata using LSB
+            origin_image = inject_data(origin_image, metadata)
+        # Save original image with metadata (and LSB if inject_lsb is True)
+        new_img = BytesIO()
+        origin_image.save(new_img, format="PNG", pnginfo=metadata, optimize=False, compress_level=0)
+        return new_img
 
     @classmethod
     def load_image(cls,
@@ -106,8 +126,9 @@ def load_image(cls,
                    ):
         """
         Load image and extract metadata using LSB/Metadata
-        :param image_io:
-        :return:
+        :param image_io: str, bytes, Path, BytesIO
+        :return: ImageMetadata
+        :raises ValidationError:  Data extraction failed
         """
         try:
             image_data = ImageLsbDataExtractor().extract_data(image_io)
@@ -115,7 +136,6 @@ def load_image(cls,
         except Exception as e:
             logger.debug(f"Error trying extracting data in LSB: {e}")
         else:
-            print(model)
             return model
         with Image.open(image_io) as img:
             title = img.info.get("Title", None)
@@ -127,13 +147,22 @@ def load_image(cls,
             except Exception as e:
                 logger.debug(f"Error loading comment: {e}")
                 comment = {}
-            return cls(Title=title, Description=prompt, Comment=cls.CommentModel(**comment))
+        return cls(Title=title, Description=prompt, Comment=cls.CommentModel(**comment))
 
     @staticmethod
     def verify_image_is_novelai(
             image: Union[Image.Image, np.ndarray],
             verify_key_hex: str = "Y2JcQAOhLwzwSDUJPNgL04nS0Tbqm7cSRc4xk0vRMic="
     ) -> bool:
+        """
+        Verify if the image is a NovelAI generated image
+        :param image:  Image.Image or np.ndarray
+        :param verify_key_hex:
+        :return:  bool
+        :raises RuntimeError:  No metadata found in image
+        :raises RuntimeError:  Comment not in metadata
+        :raises RuntimeError:  signed_hash not in comment
+        """
         # MIT:https://github.com/NovelAI/novelai-image-metadata/blob/main/nai_sig.py
         if isinstance(image, Image.Image):
             image = np.array(image)