+jZUdZddlZddlZddlmZmZmZddlm Z ddl m Z ddl Z ddlZ ddlZ ddlZddlmZmZddlmZddlmZejd Zejd Zd Zd ZGd dZdejdejfdZ de!de"e#e#e#ffdZ$de%de%fdZ&de%de%de'de%fdZ(de)e%de)e%de)e"e#e#ffdZ*de)e)e%de!ddfdZ+de,e"e#e#fde)e)e#fdZ-Gd d!ej.Z/d"Z0Gd#d$ej.Z1ged%d&RZ2e"e!d'fe3d(<d)Z4e"e"e!e5fd'fe3d*<Gd+d,ej.Z6 d=d-ejd.e)d/e"e#e#e#fd0e#d1e#d2e)e)e#dzdejfd3Z7d4Z8e"e!d'fe3d5<d6Z9e"e"e!e5fd'fe3d7<Gd8d9ej.Z:Gd:d;eZ;de;fd<Zu,ComfyUI nodes for the pure-PyTorch MediaPipe Face Landmarker port. Custom IO types: FACE_LANDMARKER — FaceLandmarkerModel wrapper (ModelPatcher inside) FACE_LANDMARKS — {"frames": List[List[face_dict]], "image_size": (H, W), "connection_sets": dict[str, frozenset[(int, int)]]} face_dict: bbox_xyxy, blendshapes, landmarks_xy, landmarks_3d, presence, score, transformation_matrix MediaPipeFaceLandmarker also emits the core BOUNDING_BOX type — pair with DrawBBoxes. N)Image ImageColor ImageDraw)tqdm)override)ComfyExtensionio)FaceLandmarker)$transformation_matrix_from_detectionFACE_DETECTION_MODELFACE_LANDMARKS)canonical_verticesprocrustes_indicesprocrustes_weights) face_ovalleft_eye right_eye left_eyebrow right_eyebrowlipsc2eZdZdZdefdZdededefdZ dS) FaceLandmarkerModelaLoaded FaceLandmarker variants + ModelPatcher per variant. Safetensors layout: `detector_short.*` / `detector_full.*` plus shared `mesh.*`, `blendshapes.*`, `canonical_*`, and `topology.*`. PReLU forces plain-nn / fp32 (manual_cast strands buffers across devices). state_dictc ~ tj|_tj}t j|_idDD]a}ttt | |tdd<btjfdtDd<ddzdzd<|_fdt"D|_d D}i|_i|_d D]}d |d  t-|}| fd Dt1||jd|}||d||j|<tj||j|tj||j|<dS)Nc<g|]}|d|S) topology. startswith).0ks z0FaceLandmarkerModel.__init__..4s)EEE1<< +D+DE!EEErc3(K|] }|V dSN)rpbases r! z/FaceLandmarkerModel.__init__..6s'.O.O1tAw.O.O.O.O.O.Or#contoursirisesnoseallc`i|]*}||+Sr&)popnumpy)rr rs r! z0FaceLandmarkerModel.__init__..:s45l5l5lWXa9J9J9P9P9R9R5l5l5lr#cDi|]\}}|d||S))zmesh.z blendshapes.r)rr vs r!r1z0FaceLandmarkerModel.__init__..<s0aaa41aF_9`9`a!Qaaar#)shortfull detector_.cvi|]5\}}|d|td|6S)z detector.N)rlen)rr r3prefixs r!r1z0FaceLandmarkerModel.__init__..CsKpppTQ[\[g[ghn[o[op5Ac&kkllO55qpppr#)devicedtype operationsdetector_variantF)strict) load_deviceoffload_devicesize)comfymodel_managementtext_encoder_devicer@text_encoder_offload_devicetorchfloat32r< frozensetmaptupler/tolistr9union_CONTOUR_PARTSconnection_sets_CANONICAL_KEYScanonical_dataitemsmodelspatchersdictupdater evalload_state_dict model_patcherCoreModelPatcher module_size) selfrrAr sharedvariantsubflr(r:s ` @@r!__init__zFaceLandmarkerModel.__init__-s@ 1EEGG/KKMM] &(EEZEEE [ [A)23ujnnQ>O>O>V>V>X>X3Y3Y)Z)ZD3{##$$% & &,9;;,.O.O.O.O.O.O.OPZ:&h7$v,FU 595l5l5l5l\k5l5l5laa:#3#3#5#5aaa13 EG (  G++++Fv,,C JJppppZEUEUEWEWppp q q q~TZTXkrsssxxzzB   s5  1 1 1#%DK %*%8%I%I 0+77;;&J&&DM' " "  r# num_faces score_threshr^ctj|j||j||||S)N)rbrc)rCrDload_model_gpurTrS detect_batch)r\imagesrbrcr^s r!rfz FaceLandmarkerModel.detect_batchMsC --dmG.DEEE{7#009[g0hhhr#N) __name__ __module__ __qualname____doc__rUraintfloatstrrfr&r#r!rr%sm4@iciiQTiiiiiir#rimagereturnc|dddfddddtjS)N.o@?r)muladd_clamp_torGuint8cpur0)ros r!_image_to_uint8r|Rsh bqb>  e $ $ ) )# . . 5 5a = = @ @ M M Q Q S S Y Y [ [[r#colorc^ tj|ddS#t$rYdSwxYw)Nrr)rrur)rgetrgb ValueError)r}s r! _parse_colorrVsC ''++ {{s  ,,facec|dt|d|d|d|d|ddS)zKShallow copy of a face_dict with array-fields cloned so callers can mutate. bbox_xyxy blendshapes landmarks_xy landmarks_3dpresencescorerrrrrr)copyrU)rs r! _copy_facer]sn[)..00T-011^,1133^,1133Z(W   r#abtc 2dz dzdzzfddDdz dzdzzdz dzdzzdz dzdzzdz dzdzzd S) Nrc^i|])}|dz d|zd|zz*S)rrr&)rr rrrs r!r1z_lerp_face..lsDnnnXYQUa &6q&99A-@PQR@S.yoVVVNO!K."3a nQ6G"GH!K."3a nQ6G"GHJVVVr#cg|]B}d|dd|ddzzd|dd|ddzzfCSrr&rs r!r"z _match_faces..{rr#N)axisc3pK|]0}ttD]}||f||fV1dSr%)ranger9)riaibrdistss r!r)z_match_faces..sR``BRWX[\]X^X^R_R_``Br2vB/```````r#) nparraylinalgnormsetsortedrr9appendadd) rr centers_a centers_bpairsused_aused_b candidates_rrrs ` @r! _match_facesrts[ A VVSTVVVWWIVVSTVVVWWI INN9QQQW- $?bN I IE#%EuuFuuF`````%A--`````J 2r <<2<<  b"X 2 2 Lr#framesmodec^ |dkrdSdtD}|sdS|dkr/g}tD]\}|r|} |rd|D<dSttD]r tfd|Ddt fd|Ddd D<]d D<uz z z t } fd |D<dS) zIn-place fill empty frame slots from neighbouring detections. Multi-face aware: pairs faces across bracketing frames by greedy bbox-centre NN. When counts differ, unmatched faces are dropped from the synthesised frame.emptyNcg|] \}}|| Sr&r&)rifrs r!r"z(_fill_missing_frames..s! 4 4 451b 4Q 4 4 4r#previousc,g|]}t|Sr&rrs r!r"z(_fill_missing_frames..s999qZ]]999r#c3(K|] }|k|V dSr%r&rr3rs r!r)z'_fill_missing_frames..'00A!a%%a%%%%00r#defaultc3(K|] }|k|V dSr%r&rs r!r)z'_fill_missing_frames..rr#c,g|]}t|Sr&rrs r!r"z(_fill_missing_frames..???1A???r#c,g|]}t|Sr&rrs r!r"z(_fill_missing_frames..rr#chg|].\}}t||/Sr&)r)rrrrnext_iprev_irs r!r"z(_fill_missing_frames..s>___QUQRTUF6N1$5vf~a7H!LL___r#) enumeraterr9maxminr) rrvalidlastrrrrrrs ` @@@@r!_fill_missing_framesrs w 4 4If-- 4 4 4E  zv&& : :EAr : :99D999q  3v;;   ` ` !9  0000000$???0000000$??? >??v???F1II ^??v???F1IIV0A @@E_______Y^___F1II ` `r#edgesc@ i}|D]o\}}||t|||t|pt}g}|D]}||vr|g}||d|c } t fd||Dd} | | |krn0|| || || c }[|||S)zWalk an unordered edge set into one or more closed-loop vertex rings (handles multi-loop sets like FACEMESH_LIPS: outer + inner).rTc3(K|] }|k|V dSr%r&)rr3prevs r!r)z!_ordered_rings..s'99aqDyyyyyy99r#N) setdefaultrrnextr) radjrrvisitedringsstartringcurnxtrs @r!_ordered_ringsrsH "C((1 q#%%  $$Q''' q#%%  $$Q''''GE   G   w E c !99993s89994@@C{cUll KK    KK   SID#  !  T Lr#cNeZdZdZedZedejfdZdS)LoadMediaPipeFaceLandmarkerzLoad MediaPipe Face Landmarker v2 weights. Contains both detector variants (short / full), shared mesh, blendshapes, and canonical geometry.c tjdgdddtjdt jddgt g S) Nrrfacial mediapipe face landmark face mesh blazefaceface detectionz%Load Face Detection Model (MediaPipe)z model/loaders model_name detectionz,Face detection model from models/detection/.)optionstooltip)node_idsearch_aliases display_namecategoryinputsoutputs)r SchemaComboInput folder_pathsget_filename_listFaceDetectionTypeOutputclss r! define_schemaz)LoadMediaPipeFaceLandmarker.define_schemassy1www@$|\5ST_5`5`'UWW'--//0    r#rpctjtjd|d}t |}t j|S)NrT) safe_load)rCutilsload_torch_filerget_full_path_or_raiserr NodeOutput)rrsdwrappers r!executez#LoadMediaPipeFaceLandmarker.executesH [ ( ()L[Zd)e)equ ( v v%b))}W%%%r#N rhrirjrk classmethodrr rrr&r#r!rrsdII   [  &BM&&&[&&&r#r)rr interpolatecNeZdZdZedZedejfdZdS)MediaPipeFaceLandmarkeruBlazeFace → FaceMesh v2 → ARKit-52 blendshapes, batched across the input. Also emits a BOUNDING_BOX list (landmark-extent bbox per face) — pair with DrawBBoxes for detector-only viz or MediaPipeFaceMeshVisualize for the mesh overlay.c^tjdgddddtdtjdtjdgd d d tjd dddddtjdddddddtjdttdddgt d tj d!g"S)#Nrrz!Detect Face Landmarks (MediaPipe)image/detectionz/Detects facial landmarks using MediaPipe model.face_detection_modelror>)r4r5bothr4uFace detector range. 'short' is tuned for close-up faces (within ~2 m of the camera); 'full' covers farther / smaller faces (up to ~5 m) but is slower. 'both' runs both detectors and keeps whichever found more faces per frame (~2× detection cost).)rrrrbrrzDMaximum faces to return per frame. 0 = no cap (return all detected).rrrsteprmin_confidencertgg?g{Gz?Tz?BlazeFace score threshold. Lower to catch small/occluded faces.)rrrradvancedrmissing_frame_fallbackra$Per-frame behaviour when detection fails in a batch. 'empty' leaves the frame faceless. 'previous' copies the most recent successful detection. 'interpolate' lerps landmarks/bbox/blendshapes between bracketing successful frames. Multi-face: pairs faces across frames by greedy bbox-centre NN.)rrrrface_landmarks)rbboxesrrrr descriptionrr) r rrrrrIntFloatlist_FALLBACK_MODESFaceLandmarksTyper BoundingBoxrs r!rz%MediaPipeFaceLandmarker.define_schemas>y-www<&I!''(>??w''1;T;T;T^e(kll  [!%kmm/#3UYdh'hjj7oAVAV`grv(|}}&"((6F(GG%%h//1    r#rpc j}t|jdd\}} d|dk} | rdzntjdt dtttffd } | r3| d| d fd tD} n | |} t| |g} | D]}g}|D]c}t|| |||d <d |d D\}}}}| ||||z ||z dt|ddd| |tj| || fjd| S)Nrrr r rr^rpc g}td|d5}tdD]}t|z}|fdt||Dt t | t j||z z | ||z  dddn #1swxYwY|S)NzMediaPipe Face Landmarker ())totaldescrc g|] }| Sr&r&)rbiimg_nps r!r"zAMediaPipeFaceLandmarker.execute.._run..s<<<<<._runs]$&CA$L'$L$L$LMMM 'QSq!U++ ' 'Aa%i++CJJ3@@<<<%:%: ' A   ((T\S1W-Ez)R)RSSSIIcAg&&&& ' ' ' ' ' ' ' ' ' ' ' ' ' ' ' 'JsCC44C8;C8r4r5cg|]>}t|t|kr|n|?Sr&)r9)rr#full_res short_ress r!r"z3MediaPipeFaceLandmarker.execute..'sX((("%Yr]!3!3s8B<7H7H!H!H " hWYl(((r#transformation_matrixc34K|]}t|VdSr%rm)rr3s r!r)z2MediaPipeFaceLandmarker.execute..3s(!C!Cq%((!C!C!C!C!C!Cr#rrr)xywidthheightlabelr)r image_sizerO)rQr|shaperCr ProgressBarrnrrUrrr rrmr rrO)rr ror>rbrr canonicalHWis_bothr/rr per_frameper_bbrx1y1x2y2r+r,r1r$r-r2r.s ` `` @@@@@@@r!rzMediaPipeFaceLandmarker.execute sO)7  '',rr"1a"f,%,QUU1 {&&z22 # $tDz"2               ,W ItF||H((((((((((FF T*++FV%;<<< " "IF D D-QRSUVXY[d-e-e)*!C!CAkN!C!C!CBB BR"r'RRTW_epuvwxwAqBqBCCDDDD MM& ! ! ! !}q!f1E1UWWX^`` `r#Nrr&r#r!rrsj   [ >-`+-=-`-`-`[-`-`-`r#rr+r,._ALL_CONNECTION_PARTS) rT)rT)rT)rT)rT)rT)r+T)r,T) tesselationF_CUSTOM_FEATUREScPeZdZdZedZeddejfdZdS)MediaPipeFaceMeshVisualizezuDraw a FACEMESH_* subset over an image. Topology travels with the FACE_LANDMARKS payload (set at detection time).ctjdgddddtdtjddd tjd d tjd gtjdgtjddtDgtjddtj ddddddtj ddddddgtj gS) NrM)rrrrrrr visualizez$Visualize Face Landmarks (MediaPipe)r z-Draws face landmarks mesh on the input image.rroTz.If not connected, a black canvas will be used.)optionalr connectionsz'all' = oval+eyes+brows+lips+irises+nose. 'fill' = solid face_oval polygon (silhouette mask). 'custom' = toggle each feature individually (including 'tesselation', the full 2547-edge wireframe).r-fillcustomc `g|]+\}}tj||d|d,S)z Draw the 'z' connection set.rrr Booleanr)rfeatrs r!r"z.^sY:::!.gJ,,T75Y$5Y5Y5Y-[[:::r#rrr}z#00ff00r thicknessrrz7Edge line thickness in pixels. 0 disables edge drawing.r point_sizerr z8Landmark dot radius in pixels. 0 disables point drawing.r) r rrrr DynamicComboOptionrKColorrrrs r!rz(MediaPipeFaceMeshVisualize.define_schemaMsiy0EEE?&G!''(899w?opp%%!a..ub99..vr::..x::2B:::&  w :: [!%^`` \1!!%_aa',X__&&'9    r#Nrpc |dd}d}|dkr$td}t} nc|dkr6fdtD} tjfd| D} n'tjfdtD} t |t |t |} } } |d }|=|d \}}tjt|||d ftj }nt|}|j d }t|}tj|}tj|}t#|D]H}||kr||ng}t%|||| | | | |||<||dzIt)jt-j|tjtjdS)NrOrQrRrrScFg|]\}}|d|SFget)rrXrrQs r!r"z6MediaPipeFaceMeshVisualize.execute..ws1ZZZgdA[__TSX=Y=YZTZZZr#c3(K|] }|V dSr%r&rr'setss r!r)z5MediaPipeFaceMeshVisualize.execute..xs''?'?AQ'?'?'?'?'?'?r#c3(K|] }|V dSr%r&rfs r!r)z5MediaPipeFaceMeshVisualize.execute..zs''O'OAQ'O'O'O'O'O'Or#rr;rrr<rrr;r<rs)rrIrKrMrHrrlrzerosr9rzr|r<rCrr= empty_liker _draw_meshr&r rrG from_numpyryrDintermediate_deviceintermediate_dtypediv_)rrrQr}rZr\rosel fill_ringsrpartsrgbthickpsizerr?r@r$r+n_framesr-outr#facesrgs ` @r!rz"MediaPipeFaceMeshVisualize.executens3/0-(-1 &=='[(9::JKKEE H__ZZZZ)9ZZZE%IKK%'?'?'?'?'?'?'?@EE%IKK%'O'O'O'O9N'O'O'OPE(//YZEU) =!,/DAqXs6{{Aq!4BHEEEFF$U++F LOv;;{&&q))mF##(( ) )B"$x--F2JJRE UE3ujYYCG  a ( ( ( (}U-c2255)==??(;;==6   $u++ r#r%rr&r#r!rMrMIsj77  [ @_a_l[r#rM image_rgbrzrurZr\rsc  |dko|}|r ||s|dkr|Stj|}tj|} |dz} |I|D]2} | d|D]%} | fd| D|&3t j|S|D]K} | djd} |r|D]\}}|| krx|| krr| t|dft|dfft|dft|dffg|||dkr<|  ||dkrcD]`\}}| t|| z t|| z t|| zt|| zf|aMt j|S)Nrrtrclg|]0}t|dft|dff1Srrr5rrlmkss r!r"z_draw_mesh..s>SSSuT!Q$Z00%QT 2C2CDSSSr#rRr)rRr8)rr fromarrayrDrawpolygonrasarrayr<linermpointflattenrLellipse)r{rzrrurZr\rs draw_edgespildrawrrrnrrr6r7rs @r!rmrmsXQ(5J  Z' 'zQ~~ /) $ $C >#  DSA _ _A^$D" _ _ SSSSdSSSZ] ^^^^ _z#  a a  JqM  c c c1q55QUUIId1a4j 1 15ad3D3DE %d1a4j 1 15ad3D3DEGMPXaccc ?? JJt||~~,,..SJ 9 9 9 9 !^^ a a1 eAhhlE!HHqL%((Q,aST U\_ ```` :c??r#)rrrrr+ _MASK_REGIONS)rI)rF)rF)rF)r+F_MASK_CUSTOM_FEATUREScNeZdZdZedZedejfdZdS)MediaPipeFaceMaskzBinary mask from face landmarks, filled polygon per face. One mask per frame in the batch; faces in the same frame composite (union).ctjdgddddtdtjddtjd gtjd d t Dg gtjg S)Nr)rrrz face maskrrrOzDraw Face Mask (MediaPipe)r z!Draws a mask from face landmarks.rregionsz'all' = union of face_oval+lips+eyes+irises (which collapses to face_oval since it encloses the rest). 'custom' = toggle each region individually for combos like lips+eyes.r-rSc `g|]+\}}tj||d|d,S)z Include the 'z' region in the mask.rUrV)rregrs r!r"z3MediaPipeFaceMask.define_schema..sY:::!-WJ,,S'5_S5_5_5_-aa:::r#rYr) r rrrr]r^rMaskrrs r!rzMediaPipeFaceMask.define_schemasy'sss5&;!''(899%%K..ub99..x::1F:::&  W^^%%&+    r#rpc|dd}|dkrfdtD}ntt}fd|D}|d}|d\}}tjt |||ftj} tj t |} t|D]\} } | rxtj d ||fd } tj| }| D]2}|d |D]%}|fd |Dd &3tj| | | <| | dzt%jt)j| tjtjdS)NrOrrScFg|]\}}|d|Srbrc)rrrrs r!r"z-MediaPipeFaceMask.execute..s1YYYfc1SRWAXAXYcYYYr#cDg|]}t|D]}|Sr&)r)rrrrgs r!r"z-MediaPipeFaceMask.execute..s2HHHsnT#Y.G.GHHHHHHr#rr;riLrrclg|]0}t|dft|dff1Sr~r5rs r!r"z-MediaPipeFaceMask.execute..s?%[%[%[QRuT!Q$Z'8'8%QT :K:K&L%[%[%[r#rurrrjrs)rrrrrkr9rzrCrr=rrnewrrrrr&r rrGrnryrDrorprq)rrrrrpickedrrr?r@masksr-r#rBrrrrrrgs ` @@r!rzMediaPipeFaceMask.executes/0i  (??YYYY(=YYYFF-((FHHHHfHHH)l+1#f++q!,BH==={&&s6{{33&v.. ) )MB  ,iaVQ// ~c**"ggA^,D %gg %[%[%[%[VZ%[%[%[be ffffgJsOOb  a ( ( ( (}U-e4477)==??(;;==8   $u++ r#Nrr&r#r!rrsdFF  [ 2[r#rcLeZdZedeeejfdZdS)MediaPipeFaceExtensionrpc:KttttgSr%)rrrMr)r\s r! get_node_listz$MediaPipeFaceExtension.get_node_lists+-DF`bsttr#N) rhrirjrrtyper ComfyNoderr&r#r!rrsL uT$r|*<%=uuuXuuur#rc"KtSr%)rr&r#r!comfy_entrypointrs ! # ##r#r%)=rkr0rrGPILrrr tqdm.autortyping_extensionsrcomfy.model_managementrCcomfy.model_patcher comfy.utilsrcomfy_api.latestrr &comfy_extras.mediapipe.face_landmarkerr $comfy_extras.mediapipe.face_geometryr CustomrrrPrNrTensorndarrayr|rnrKrlrrUrrmrrrrrIrrrrrrH__annotations__rKboolrMrmrrrrrr&r#r!rs    ,,,,,,,,,,&&&&&&////////AAAAAAUUUUUUBI455BI.//T`*i*i*i*i*i*i*i*iZ\5<\BJ\\\\c3m 4 T d    $4EdDJ4:$uS#X2G.`d4j!1``````B)E#s(O4d3i4&&&&&",&&&47T`T`T`T`T`blT`T`T`p*M>)L8)LV)L)LuS#XLLL 2%c4i(#-.   DDDDDDDDRFJ"*T#sC-(58,0cOd,BNPj@"Z uS#XYYY7uU39-s2388888 888vuuuuu^uuu $ 6$$$$$$r#