o iS@s:UddlZddlZddlZddlZddlZddlZddlmZmZm Z ddl m Z ddl m Z mZddlZddlmZiZe eefed<dedefd d ZGd d d e ZGd ddZddZGdddeZGdddeZGdddZGdddeZGdddZGdddeZdZ dZ!Gdd d eZ"dS)!N)SequenceMappingDict) DynamicPrompt)ABCabstractmethod)is_linkNODE_CLASS_CONTAINS_UNIQUE_ID class_typereturncCs>|tvrt|Stj|}d|divt|<t|S)N UNIQUE_IDhidden)r nodesNODE_CLASS_MAPPINGS INPUT_TYPESgetvalues)r class_defr5/mnt/c/Users/fbmor/ComfyUI/comfy_execution/caching.pyinclude_unique_id_in_inputs  rc@sHeZdZddZeddZddZddZd d Zd d Z d dZ dS) CacheKeySetcCsi|_i|_dSN)keys subcache_keysself dynpromptnode_idsis_changed_cacherrr__init__s zCacheKeySet.__init__cstr)NotImplementedError)rrrrradd_keysszCacheKeySet.add_keyscCst|jSr)setrrrrr all_node_ids"zCacheKeySet.all_node_idscC |jSr)rrr$rrr get_used_keys% zCacheKeySet.get_used_keyscCr'r)rrr$rrrget_used_subcache_keys(r)z"CacheKeySet.get_used_subcache_keyscC|j|dSr)rrrnode_idrrr get_data_key+r&zCacheKeySet.get_data_keycCr+r)rrr,rrrget_subcache_key.r&zCacheKeySet.get_subcache_keyN) __name__ __module__ __qualname__r rr"r%r(r*r.r/rrrrrs  rc@seZdZddZdS) UnhashablecCstd|_dS)NNaN)floatvaluer$rrrr 2r&zUnhashable.__init__N)r0r1r2r rrrrr31s r3c Csnt|ttttttdfr|St|tr!tddt | DSt|t r4tt t dd|DStS)NcSs g|] \}}t|t|fqSr to_hashable).0kvrrr ; zto_hashable..cSsg|]}t|qSrr7)r9irrrr<=s) isinstanceintr5strboolbytestyper frozensetsorteditemsrzip itertoolscountr3)objrrrr85s  r8cs$eZdZfddZddZZS) CacheKeySetIDcst|||||_dSr)superr rr __class__rrr Cs zCacheKeySetID.__init__csZ|D]'}||jvr q|j|sq|j|}||df|j|<||df|j|<qdSNr )rrhas_nodeget_noderrrr-noderrrr"Gs   zCacheKeySetID.add_keys)r0r1r2r r" __classcell__rrrNrrLBs rLcsReZdZfddZdefddZddZdd Zd d Zd d Z ddZ Z S)CacheKeySetInputSignaturecs t|||||_||_dSr)rMr rrrrNrrr Rs z"CacheKeySetInputSignature.__init__r cCsdSNFrr$rrrinclude_node_id_in_inputWz2CacheKeySetInputSignature.include_node_id_in_inputcsb|D]+}||jvr q|j|sq|j|}||j|IdH|j|<||df|j|<qdSrP)rrrQrRget_node_signaturerrSrrrr"Zs   z"CacheKeySetInputSignature.add_keyscs\g}|||\}}|||||IdH|D]}|||||IdHqt|Sr)get_ordered_ancestryappendget_immediate_node_signaturer8)rrr- signature ancestors order_mapping ancestor_idrrrrZdsz,CacheKeySetInputSignature.get_node_signaturec s||s tdgS||}|d}tj|}||j|IdHg}|s4t|dr0|j s4t |r9| ||d}t | D]&} t|| r`|| \} } || } | | d| | ffqC| | || fqC|S)Nr4r NOT_IDEMPOTENTinputsANCESTOR)rQr5rRrrrrrXhasattrrbrr\rFrr) rrr-ancestor_order_mappingrTr rr^rckeyraancestor_socketancestor_indexrrrr]ls"        z6CacheKeySetInputSignature.get_immediate_node_signaturecCs g}i}|||||||fSr)get_ordered_ancestry_internal)rrr-r_r`rrrr[sz.CacheKeySetInputSignature.get_ordered_ancestryc Cs||sdS||d}t|}|D]'}t||r=||d}||vr=||t|d||<|||||qdS)Nrcr)rQrRrFrrr\lenrj) rrr-r_r`rc input_keysrgrarrrrjs     z7CacheKeySetInputSignature.get_ordered_ancestry_internal) r0r1r2r rBrXr"rZr]r[rjrUrrrNrrVQs  rVc@seZdZd+ddZddZddZdd Zd d Zd d ZddZ ddZ ddZ ddZ ddZ ddZeddZddZddZd d!Zd"d#Zd$d%Zd&d'Zd(d)Zd*S), BasicCacheFcCs2||_d|_||_||i|_i|_t|_dSrW) key_class initializedenable_providerscache subcachesr#_pending_store_tasksrrorqrrrr s zBasicCache.__init__cs:||_|||||_|j|IdH||_d|_dS)NT)rro cache_key_setr"rrprrrr set_prompts  zBasicCache.set_promptcCs6|jsJ|j}|jD] }||}q|Sr)rprvr%rsrunion)rrsubcacherrrr%s  zBasicCache.all_node_idscCFt|j}g}|jD] }||vr||q |D]}|j|=qdSr)r#rvr(rrr\)r preserve_keys to_removergrrr _clean_caches   zBasicCache._clean_cachecCrzr)r#rvr*rsr\)rpreserve_subcachesr|rgrrr_clean_subcachess   zBasicCache._clean_subcachescCs|jsJ||dSr)rpr}rr$rrr clean_unuseds  zBasicCache.clean_unusedcKdSrrrkwargsrrrpollrYzBasicCache.pollcCs.|jsdS|j|}||jvr|j|SdSrrprvr.rrrr- cache_keyrrr get_locals    zBasicCache.get_localcCs$|jsJ|j|}||j|<dSrrrr-r6rrrr set_locals  zBasicCache.set_localcs:|jsJ|j|}||j|<||||IdHdSr)rprvr.rr_notify_providers_storerrrr_set_immediates    zBasicCache._set_immediatecsX|jsdS|j|}||jvr|j|S|||IdH}|dur*||j|<|SdSr)rprvr.rr_check_providers_lookup)rr-rexternal_resultrrr_get_immediates    zBasicCache._get_immediatec sddlm}m}m}m}m}|jsdS|sdS||s dS||r&dS|||} | dur2dS||j |j d} |D]@} z| | | r]t || | | } |j| | |jjWq=ty}} z|d| jjd| WYd} ~ q=d} ~ wwdS)Nr_has_cache_providers_get_cache_providers CacheValue_contains_self_unequal_logger)outputsuiCache provider z error on store: )comfy_execution.cache_providerrrrrrrq_is_external_cacheable_value_build_contextrr should_cacheasyncio create_task_safe_provider_storertaddadd_done_callbackdiscard ExceptionwarningrOr0)rr-rr6rrrrrcontext cache_valueprovidertaskerrrrs4     &z"BasicCache._notify_providers_storec shddlm}z |||IdHWdSty3}z|d|jjd|WYd}~dSd}~ww)Nr)rrz async store error: )rron_storerrrOr0)rrrrrrrrr s (zBasicCache._safe_provider_storec sBddlm}m}m}m}m}|jsdS|sdS||rdS|||}|dur+dS|D]p} zO| |s8Wq.| |IdH} | dur~t | |sV| d| j j dWq.t | jttfsk| d| j j dWq.ddlm} | | jt| jdWSWq.ty} z| d| j j d | WYd} ~ q.d} ~ wwdS) Nrrz Provider z returned invalid typez returned invalid outputs) CacheEntry)rrrz error on lookup: )rrrrrrrqrr on_lookupr?rrOr0rlisttuple executionrrr) rr-rrrrrrrrresultrrrrrrs>      &z"BasicCache._check_providers_lookupcCst|do t|dS)Nrr)re)rr6rrrr5z'BasicCache._is_external_cacheable_valuecCs<|jr|jsdSz |j|ddWStyYdSw)Nr )rprrRrrr,rrr_get_class_type8s  zBasicCache._get_class_typec Cs~ddlm}m}m}z||}|durWdS|||||dWSty>}z|d|d|WYd}~dSd}~ww)Nr) CacheContext_serialize_cache_keyr)r-r cache_key_hashz'Failed to build cache context for node z: )rrrrrrr)rr-rrrrrrrrrr@szBasicCache._build_contextcsT|j|}|j|d}|durt|j}||j|<||j||jIdH|Sr) rvr/rsrrnrorwrr)rr- children_ids subcache_keyryrrr_ensure_subcacheOs   zBasicCache._ensure_subcachecCs.|jsJ|j|}||jvr|j|SdSr)rprvr/rs)rr-rrrr _get_subcacheXs    zBasicCache._get_subcachecCsPg}|jD] }|||j|dq|jD]}|||j|dq|S)N)rgr6)rry)rrr\rsrecursive_debug_dump)rrrgrrrr`s   zBasicCache.recursive_debug_dumpNF)r0r1r2r rwr%r}rrrrrrrr staticmethodrrrrrrrrrrrrrns,     $ rncsNeZdZdfdd ZddZddZdd Zd d Zd d ZddZ Z S)HierarchicalCacheFcstj||ddS)Nrq)rMr rurNrrr irzHierarchicalCache.__init__cCs~|jdusJ|j|}|dur|Sg}|dur(|||j|}|dus|}t|D]}||}|dur<dSq.|Sr)rget_parent_node_idr\reversedr)rr- parent_id hierarchyrrrrr_get_cache_forls      z HierarchicalCache._get_cache_forcs(||}|dur dS||IdHSr)rrrr-rrrrrr~s  zHierarchicalCache.getcCs"||}|dur dSt||Sr)rrnrrrrrrs  zHierarchicalCache.get_localcs.||}|dus J|||IdHdSr)rrrr-r6rrrrrr#s  zHierarchicalCache.setcCs(||}|dus Jt|||dSr)rrnrrrrrrs  zHierarchicalCache.set_localcs*||}|dus J|||IdHSr)rr)rr-rrrrrrensure_subcache_fors  z%HierarchicalCache.ensure_subcache_forr) r0r1r2r rrrr#rrrUrrrNrrhsrc@sTeZdZddZddZddZddZd d Zd d Zd dZ ddZ ddZ dS) NullCachecdSrrrrrrrwzNullCache.set_promptcCsgSrrr$rrrr%rYzNullCache.all_node_idscCrrrr$rrrrrYzNullCache.clean_unusedcKrrrrrrrrrYzNullCache.pollcrrrr,rrrrrz NullCache.getcCrrrr,rrrrrYzNullCache.get_localcrrrrr-r6rrrr#rz NullCache.setcCrrrrrrrrrYzNullCache.set_localcs|Srr)rr-rrrrrrzNullCache.ensure_subcache_forN) r0r1r2rwr%rrrrr#rrrrrrrs rcs^eZdZdfdd ZfddZddZd d Zd d Zd dZddZ fddZ Z S)LRUCachedFcs2tj||d||_d|_d|_i|_i|_dS)Nrr)rMr max_sizemin_generation generationused_generationchildren)rrorrqrNrrr s  zLRUCache.__init__cs>t|||IdH|jd7_|D]}||qdS)Nrk)rMrwr _mark_used)rrrrr-rNrrrws  zLRUCache.set_promptcstjjkrCjjkrCjd7_fddjD}|D]}j|=j|=|jvr4j|=q!tjjkrCjjksdS)Nrkcs g|] }j|jkr|qSr)rr)r9rgr$rrr<r=z)LRUCache.clean_unused..)rlrrrrrrrr)rr|rgrr$rrs  zLRUCache.clean_unusedcs||||IdHSr)rrr,rrrrs z LRUCache.getcCs(|j|}|dur|j|j|<dSdSr)rvr.rrrrrrrs zLRUCache._mark_usedcs|||||IdHSr)rrrrrrr#s z LRUCache.setcCs||t|||dSr)rrnrrrrrrs zLRUCache.set_localcsxt||IdH|j|IdH|||j|}g|j|<|D]}|||j||j|q&|Sr)rMrrvr"rr.rr\)rr-rrchild_idrNrrrs    zLRUCache.ensure_subcache_for)rF) r0r1r2r rwrrrr#rrrUrrrNrrs  rg?g?csReZdZdfdd ZddZfddZfdd Zfd d Zd d ZZ S)RAMPressureCacheFcstj|d|di|_dS)Nrr)rMr timestampsrurNrrr s zRAMPressureCache.__init__cCs |dSr)rr$rrrrs zRAMPressureCache.clean_unusedcs0t|j|j|<t||IdHdSr)timerrvr.rMr#rrNrrr#szRAMPressureCache.setcs*t|j|j|<t|IdHSr)rrrvr.rMrr,rNrrrszRAMPressureCache.getcs(t|j|j|<t||dSr)rrrvr.rMrrrNrrr szRAMPressureCache.set_localcstj|kr dSg}|jD],\}}t|j|j|}tfdd|j |9}t |||j ||fqtj|krq|ru| \}}}|j|=|j |d|j |d|j |dtj|krs|sFdSdSdSdS)Ncs^|durdS|D]$}t|ttfr|qt|tjr,|jjdkr,||7qdS)Ncpu) r?rrtorchTensordevicerDnumel element_size)routput ram_usagescan_list_for_ram_usagerrrs z=RAMPressureCache.ram_release..scan_list_for_ram_usage)psutilvirtual_memory availablerrrG%RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIERrrRAM_CACHE_DEFAULT_RAM_USAGErbisectinsortrpopr)rtarget clean_listrg cache_entry oom_score_rrr ram_releases" "zRAMPressureCache.ram_releaser) r0r1r2r rr#rrrrUrrrNrrs   r)#rrrIrrrtypingrrrcomfy_execution.graphrabcrrrcomfy_execution.graph_utilsrr rArB__annotations__rrr3r8rLrVrnrrrrrrrrrrs4    DT1<