3j%SSKJr SSKJr SSKJrJrJrJr SSK r SSK J r SSK J r JrJrJr "SS \5r"S S \5r\"S S 9"SS55r"SS5rg)) annotations) dataclass)CallableLiteralMappingSequenceN)canonicalize_state_dict)ArchId ArchitectureModelDescriptor StateDictc\rSrSrSrSrg)UnsupportedModelError zl An error that will be thrown by `ArchRegistry` and `ModelLoader` if a model architecture is not supported. N__name__ __module__ __qualname____firstlineno____doc____static_attributes__rU/home/wildlama/miniconda3/lib/python3.13/site-packages/spandrel/__helpers/registry.pyrr rrc\rSrSrSrSrg)DuplicateArchitectureErrorzY An error that will be thrown by `ArchRegistry` if the same architecture is added twice. rNrrrrrrrrrT)frozencf\rSrSr%SrS\S'S\S'SrS\S '\S S S jj5rSr g ) ArchSupportz[ An entry in an `ArchRegistry` that describes how to detect and load a model architecture. Architecture[torch.nn.Module] architecturezCallable[[StateDict], bool]detectrtuple[ArchId, ...]beforec.[XRU5$)z` Creates an `ArchSupport` from an `Architecture` by using the architecture's ``detect`` method. )r"r&)archr(s rfrom_architectureArchSupport.from_architecture/s4f55rN)r)r*r$r(r'returnr") rrrrr__annotations__r( staticmethodr+rrrrr"r"sb0/ (' "$F # JL6+65G6 66rr"c\rSrSrSrSrSSjrSSjrSSjrSr SSjr SS jr SSS jjr S S .SS jjr \SSj5rSSjrSrg) ArchRegistry9zr A registry of architectures. Architectures are detected/loaded in insertion order unless `before` is specified. c./Ul/Ul0UlgN)_architectures_ordered_by_idselfs r__init__ArchRegistry.__init__@s57/1 46 rc[5nURUlURUlURUlU$)z! Returns a copy of the registry. )r1r5r6r7)r9news rcopyArchRegistry.copyFs4n!00}} [[  rcXR;$r4)r7r9ids r __contains__ArchRegistry.__contains__Ps[[  rc2UR[U5$r4)r7r rAs r __getitem__ArchRegistry.__getitem__Ss{{6":&&rc6[URS55$)z@ Returns an iterator over all architectures in insertion order. insertion)iter architecturesr8s r__iter__ArchRegistry.__iter__VsD&&{344rc,[UR5$r4)lenr5r8s r__len__ArchRegistry.__len__\s4&&''rcLURR[U5S5$r4)r7getr rAs rrSArchRegistry.get_s{{vbz400rcUS:Xa[UR5$US:Xa[UR5$[SU35e)z Returns a new list with all architectures in the registry. The order of architectures in the list is either insertion order or the order in which architectures are detected. rI detectionzInvalid order: )listr5r6 ValueError)r9orders rrKArchRegistry.architecturesbsF K ++, , k ! & &ug67 7rF)ignore_duplicatesc[UR5n[UR5n/nUHnURR U;a+U(aM&[ SURR 35eURU5 XdURR 'URU5 M [RU5nX0lXpl X@lU$)a Adds the given architectures to the registry. Throws an error if an architecture with the same ID already exists, unless `ignore_duplicates` is True, in which case the old architecture is retained. Throws an error if a circular dependency of `before` references is detected. If an error is thrown, the registry is left unchanged. Returns a list of architectures that were added. zDuplicate architecture: ) rWr5dictr7r%rBrappendr1 _get_orderedr6)r9r[rKnew_architectures new_by_idaddedr* new_ordereds raddArchRegistry.addrs$!!4!45% !D  ##y0$0.t/@/@/C/C.DE  $ $T *.2d''** + LL "#//0AB /#   rcB^^^^^^0m0mUHbnUTURR'URH6nUT;a/TU'TURURR5 M8 Md /m[ 5m/mSUUUUUU4SjjmUH nT"U5 M T$)Nc>URRT;a5[SSR/TQURRP535eUT;agTR U5 TR URR5 TR URR/5HnT"TU5 M TR U5 TR5 g)Nz/Circular dependency in architecture detection: z -> )r%rBrXjoinrdr^rSpop)r*r(by_id inv_beforeorderedseenstackvisits rro(ArchRegistry._get_ordered..visits  ##u, EfkkRpTYRp[_[l[l[o[oRpFqErst| HHTN LL**-- .$..):):)=)=rBeFm$C NN4 IIKr)r*r")r%rBr(r^set) rKr*r(rjrkrlrmrnros @@@@@@rr_ArchRegistry._get_ordereds13 +-!D*.E$##&& '+++)+Jv&6"))$*;*;*>*>?&"&(!$    "D $K"rc[U5nURH6nURU5(dMURR U5s $ [ e)z Detects the architecture of the given state dict and loads it. This will canonicalize the state dict if it isn't already. Throws an `UnsupportedModelError` if the model architecture is not supported. )r r6r&r%loadr)r9 state_dictr*s rrtArchRegistry.loadsJ-Z8 MMD{{:&&((--j99"$#r)r5r7r6N)r-r1)rBz ArchId | strr-bool)rB str | ArchIdr-r")r-int)rBrxr-zArchSupport | None)rI)rYz!Literal['insertion', 'detection']r-list[ArchSupport])rKr"r[rwr-rz)rKrzr-rz)rurr-r )rrrrrr:r>rCrFrLrPrSrKrdr/r_rtrrrrr1r19s 7 !'5 (1 4?808 8&#(&#& &  &P!!F$rr1) __future__r dataclassesrtypingrrrrtorch canonicalizer model_descriptorr r r r ExceptionrrXrr"r1rrrrs`"!77 1NNI   $666@T$T$r