@@ -139,6 +139,27 @@ def load_upscalers():
139
139
key = lambda x : x .name .lower () if not isinstance (x .scaler , (UpscalerNone , UpscalerLanczos , UpscalerNearest )) else ""
140
140
)
141
141
142
+ # None: not loaded, False: failed to load, True: loaded
143
+ _spandrel_extra_init_state = None
144
+
145
+
146
+ def _init_spandrel_extra_archs () -> None :
147
+ """
148
+ Try to initialize `spandrel_extra_archs` (exactly once).
149
+ """
150
+ global _spandrel_extra_init_state
151
+ if _spandrel_extra_init_state is not None :
152
+ return
153
+
154
+ try :
155
+ import spandrel
156
+ import spandrel_extra_arches
157
+ spandrel .MAIN_REGISTRY .add (* spandrel_extra_arches .EXTRA_REGISTRY )
158
+ _spandrel_extra_init_state = True
159
+ except Exception :
160
+ logger .warning ("Failed to load spandrel_extra_arches" , exc_info = True )
161
+ _spandrel_extra_init_state = False
162
+
142
163
143
164
def load_spandrel_model (
144
165
path : str | os .PathLike ,
@@ -148,11 +169,16 @@ def load_spandrel_model(
148
169
dtype : str | torch .dtype | None = None ,
149
170
expected_architecture : str | None = None ,
150
171
) -> spandrel .ModelDescriptor :
172
+ global _spandrel_extra_init_state
173
+
151
174
import spandrel
175
+ _init_spandrel_extra_archs ()
176
+
152
177
model_descriptor = spandrel .ModelLoader (device = device ).load_from_file (str (path ))
153
- if expected_architecture and model_descriptor .architecture != expected_architecture :
178
+ arch = model_descriptor .architecture
179
+ if expected_architecture and arch .name != expected_architecture :
154
180
logger .warning (
155
- f"Model { path !r} is not a { expected_architecture !r} model (got { model_descriptor . architecture !r} )" ,
181
+ f"Model { path !r} is not a { expected_architecture !r} model (got { arch . name !r} )" ,
156
182
)
157
183
half = False
158
184
if prefer_half :
@@ -166,6 +192,6 @@ def load_spandrel_model(
166
192
model_descriptor .model .eval ()
167
193
logger .debug (
168
194
"Loaded %s from %s (device=%s, half=%s, dtype=%s)" ,
169
- model_descriptor , path , device , half , dtype ,
195
+ arch , path , device , half , dtype ,
170
196
)
171
197
return model_descriptor
0 commit comments