From 55f0b45b236478ffd7d369dfaecf54a88912b1cc Mon Sep 17 00:00:00 2001
From: John Zheng <bluevisor@gmail.com>
Date: Sat, 6 Jul 2024 21:32:57 -0700
Subject: [PATCH] Added MPS support for macOS

---
 __pycache__/ps_blend_node.cpython-310.pyc | Bin 5157 -> 5214 bytes
 ps_blend_node.py                          |  24 ++++++++++++++--------
 2 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/__pycache__/ps_blend_node.cpython-310.pyc b/__pycache__/ps_blend_node.cpython-310.pyc
index f180ceb51d114f5a7a28c54784545badacbe46c6..8bcc0778f1661c3ddc6240915a4a99c0c9b6c817 100644
GIT binary patch
delta 2315
zcma)6O>7%Q6yBL#uh;919Xn3#ByM7wQsO2J>3@>8B5l)_pEd;bM<Sr6u4j`v>#U=-
zQ)sltQWNC<s8Z2Lh*PnF9zbxYgv0>^XG(=QZ~z~Asss{JIUpo12;LjpaY0nVE5Cj7
zz3=_b%-$ZjKM=^ZwdoT4lHJ$l#{1q1#Jjs;#Uw_u(u{OWVwA~i5|dZy47qdsM@kgF
zMb4IPc_gVIEszSWN=~pKGqOquAQgqtDk*yGF0gtQRAKJ}yS#2!gxv?2XB|@!b4kVP
zdO%ZM+R(hNx}dq<1<jJ3{RJQMu9C~dRar@c3Vm<VIoc(`Qe!@**H@Gl)grh6)?(o5
zQHHnZV}53=NmmJLV*z;6vQ#E>lm%J)ns!w#OLK&USQvB#aKbuR1a#n(X<05Sr8$Ll
zvJGomQFmn)wWxE=Rg3<c)Q-VA2J2fe3*CjjT*Z=#2CV$9;i|5;Dhr+hye{B<1eQ8Z
z0Vnn!fOiA_DRLu{z#anD1K1BDipt_viMrm#Vxy(7IOuIG?s{V|jUI_@WbvvT!|ez_
zLQ<7QQexkQK)n#?0S1psEFP;T_!96w!2iUUS75-aq!MI_DghlvUe_&%-6)@w?Aze-
zghg)>M+c{NVXAlV4p=q`OAm`fQEChlVQCSmk**Ake!&>Sgs@q03yZ}mg~ZIDNA&2C
zs<eqUTF8L!X6PRhPu9B-BRUK4lOp6<4Efj|lvWS+B=~F*J}&w|8BbXf8*tT#grkSz
z2kPac-fW;=#+<zcoRh-&vxf6li?BiAI0V~(0raSNgFRf|hGLuGs*nf<iwQ7UgT0;t
zECOe7{+-%c+9tAm-}3xFCEDLYeZ#`@_jONfNn~lne>w^i#_W$YDllx){sbb7ShRGL
zI&!n5=9#hI1><%IPn|?7JZMIfDK@UXxx$lHr6TNQ@Z2Fh#~S+=5D@}!2t1&P0|4&;
z9t7|x`EA_uox<fP7&_1m!V#$PRj}+5mNWHWt=ZsNvyX#)x3Hgc;aeiGLxAlOSQarR
zItK0@!1oILH5YC!9DF#>_8?yHQ2}4&S#9HqF`>VGMRwKt?GQ8rSmVO-9$38f^M@Ar
z3`ShV^{{<6Xr&W30H3N-48m;U{GrxWin~F)98swl0`>;C-3`rCmSFpDfWPaXBPXTg
zfvJ`F^vT4btZB2vRF0X6Og^!k&)D;c!>1<?4%d|QGD|J@!MFP{v|r!<Dsf`^<Zw-1
zT*|KuOyv^CrcWnQ>2$xDHJwx;ld}`4i>XXDbv|njHDgbouF2`8Wj^iunvC#ny+Zcz
z&-Ed)lmDz=BBT7ge`{<nBz3l;7y?oA=FP%vc6o7>f8f948wKdphtwJ8XN}mYeaPxZ
zu^Gh#h@|QqM12s%j3+(2n6aG$O+EE`9b!qi?#TLAntZBofZsL#+9c9AN0CBM>_YJ@
z2nauc`ab?kTdyiI9pb@2H1f}?Mf6EN9;g^KC9^ocUC15$L7+b=n9pF4BPgCj(W=9D
zVcDF8k=g779|^wFDKd*1{75uVjPgowknHC7f@VjO)Kq3(%%n|c3~b3Ll=en@wAZub
zWbA@-3b2}XK9#m0V%{0SB`xIm$L)Rb7ciEX-t!3SE2fjn&t@~0=}hpS+Nb>J>`b9}
zkxzso+g?I-8oA11DsMT5L7nmDmltQD#hE~EDV5F?F3lE9JD+p-#Zd1-G35Vh7`N6L
z1`&aO{vHUOYQ%sCrABq4%Q`Vs{!M6t=-dcfMz<%TMC7nXr7JOh--z*Q_{fBF1_Qke
zA}Kp*)X#yaDRyentWUKy+OQZ&qrBPth18Py1@G#(J~@wmg4YtFKneyK{`oKnO`=36
cimniaYE&v2NeAso>fkZtzjX{~ilBw*-;<v=?f?J)

delta 2320
zcmZuxOKcNY6n*c_c*Yrz?Zh~7;?R(WPh*l&5(**wMED97=$A^E3d`h;!Hhkzo}CJ!
zu`H1whgLzLQKhPz+A6J7p-M$+mvmF9f=yR#rOu+ON?lY)RX3=Mir)KdlTeY!ckVm)
zp8I>B-$p+ljm(9^27$ky%@1;YL$@RAM@C@938w{;ZhdK*9VeV|Wr=WQk)@S;M}B2A
z;0-B-S|dP6g*0e|TxJcWNkkW?4&|&#s{xlbR5t*c=CZb?yQC3tHLf%nB~sV9Ds}C$
z=IWO<TM^?God*`_H41(j6rx{eIW|D#O6S4Wim|Fz4Z1b`85dlwb<wL`yo(!4<Oby-
zZo-$<NuA~x5A(=U@P<++Im)BF8*~+L%6oVWbnvOOx>8qjDqqLrOTnt)D!kWbV!_p_
zp<B#Jz`D<7HE|QV54u{Dx~fg8Cak(#(}iA|igZ-KF#*I+z+y<80ZyIAfcFD_2eCMj
z*dxFO0DB;#Xgp~%S8wZsw#t*Bhj`M}`(YY`gb(ssvZ?gr2AGhLHWitaIu9YxdI<Cr
z22T*4>~AHw5BLj!Kf<8PFyKYn2=k#P1zqpxHUw{%Pn|o!*Z_=wfDu_!{>Ok1L;ts^
zFnAJb()u9{Ru(7iWOZ<P(PrQ=is@V(`oU5S_(kQ!phIK~k|yh5?FRC|9g*A}trp~b
z&I7tpqNgw-6kU6M2kM#*jsR(sB-u!UGB(?UZ*jFa!NEiETUzC#-f5s+#unhdB)MO-
zxuZ7aDaqRko1oW#6imoK;LzRLglbqiYmf-WiU}|oy01+E7K!8LxZzgf1i{44pOAJ;
zru}|3(sN6${jX-{N1%>N>hDjf*rV*Mo$%!Zj2z2Z)49ZC@~<GmHW}gXPKlip<a`RY
zSHQ+Zv=zOE!vxjnxXPW3bXM3~lgY5xfVy2$N839Zks%^*DgvO%69HeiJ9v;`3U~%4
zBx5JwJ;;J+6fEz7Wl~ydM_a*Ghl7XDISCFsq{Hj3jz<%)Lx8<1v2%!V*)ia|0N*L`
zoGUNrB3w{s4A5N?y|Sz<tFFGx@U|?n)=iN<1Fn+eQ_|`Bvf^rMOeWBGOUrGr=wOlL
zYDm6$AwG?79=w?`7hXudXN5J^;U?jKQ>{j&SCsFSFC&~F@2}v%iKH9F+m#^Io_3Ud
z?I>}#+wHl?_z>T}0>!(b9DSRl4xIMY*?FGzmD$p~ujY#vC#0P8?sh#iQt--`8Hi2Z
z8==kKXNI6tUSDXOj(dkgSLw8OKeV-fH`v8i6yqR#y=+x7h56!y_at=h{4_xCe96RK
zX|gS=#Ph{JRvvxZc>KfH2U4oohs2F2HldgSk<!E=)W<;hf!Rzk?}!5(J$ofBA)cb*
zs22(M(Our=@RJdlNE}BVMX?jb2@ntsgNl9LvB=Pb415^YK6QM0CEYcfnJwf?B@5#{
z+le>|Zu`BDA`L@KV1!8!-b{33O8Ol{=opHVC{_p1b!pzpTqsx$Kjy7O&-KbmWKAJ-
zG*L`@|3=54@>I8VYN1zH<@rKo_=1(26V~utDLb35Tuo75<JSB6S<8>Kp-jbc%5!4D
z%;lYIVYo#Tdmvz{w>8K!Ub8!qV3j#PRuXxq(q?4Jm$D_x`=xus`qP+R&ifQ-zOi75
zxpJnEw=FT__4J&cIfD*Q`%P<jBr;j~oy=s#TWEe0scJD>w#5-p(|UQnm;p!eF08z}
zJwsDCyVg7pqP99GJR@Qoh+g=WAAsm*hGI~YnbcG~BX*Dmz1L!P#AHduq;VEkOcqnT
z?_${<en2)Q(!fe7g2P8-LHMeZEn2NDc)AbbELgtrGV4C1XT9z5#FW5;tirDx0}&*Q
m8dNn@s<I#>HFH?n+b7ZwIP=BQRi8N}Z(DrS&?F;BW9&c9Bs%B-

diff --git a/ps_blend_node.py b/ps_blend_node.py
index 7954fd2..b3faf7d 100644
--- a/ps_blend_node.py
+++ b/ps_blend_node.py
@@ -1,5 +1,4 @@
 import torch
-import numpy as np
 
 class PSBlendNode:
     @classmethod
@@ -19,15 +18,23 @@ class PSBlendNode:
     CATEGORY = "image/blending"
 
     def blend_images(self, image1, image2, blend_mode, opacity):
-        # Determine the device (CPU or GPU)
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        # GPU availability check
+        if torch.cuda.is_available():
+            print("PS Blend Node is using CUDA.")
+            device = torch.device("cuda")
+        elif torch.backends.mps.is_available():
+            print("PS Blend Node is using  MPS.")
+            device = torch.device("mps")
+        else:
+            print("No GPU acceleration available, using CPU")
+            device = torch.device("cpu")
 
         # Move tensors to the appropriate device
         img1 = image1.to(device)
         img2 = image2.to(device)
 
-        self.print_tensor_shape(img1, "img1")
-        self.print_tensor_shape(img2, "img2")
+        # self.print_tensor_shape(img1, "img1")
+        # self.print_tensor_shape(img2, "img2")
 
         # Ensure images have the same shape
         if img1.shape != img2.shape:
@@ -158,18 +165,19 @@ class PSBlendNode:
         elif blend_mode == "Luminosity":
             result = set_lum(img2, lum(img1))
         
-        self.print_tensor_shape(result, "result before opacity")
+        # self.print_tensor_shape(result, "result before opacity")
 
         # Apply opacity
         opacity_tensor = torch.tensor(opacity, device=device).view(1, 1, 1, 1)
-        self.print_tensor_shape(opacity_tensor, "opacity_tensor")
+
+        # self.print_tensor_shape(opacity_tensor, "opacity_tensor")
         
         result_rgb = img2 * (1 - opacity_tensor) + result * opacity_tensor
         
         # Blend alpha channels
         result_a = img1_a * opacity_tensor + img2_a * (1 - opacity_tensor)
 
-        self.print_tensor_shape(result, "final result")
+        # self.print_tensor_shape(result, "final result")
         
         # Combine RGB and alpha channels
         result = torch.cat([result_rgb, result_a], dim=-1)