chaiNNer-org / spandrel

Spandrel gives your project support for various PyTorch architectures meant for AI Super-Resolution, restoration, and inpainting. Based on the model support implemented in chaiNNer.
MIT License
105 stars 7 forks source link

Move to architecture classes #153

Closed RunDevelopment closed 4 months ago

RunDevelopment commented 4 months ago

This a huge PR, because I fundamentally changed how we define architectures. Instead of having a single load function per architecture, each architecture now defines a <Name>Arch class that extends the new Architecture class. The basic idea is that an architecture (class) knows its ID, its name, a way to tell whether a state dict is of its own architecture (a detect method), and a way to load a state dict.

For users, this means that they can no longer just import a load function. E.g. to load a state dict manually, they now have to do this:

from spandrel.architectures.DAT import DATArch
model = DATArch().load(state_dict)

Other changes:

Note: the number of lines changed is overblown. I had to add one level of indentation to all load functions, so hide whitespace changes to see the actual diff.

I haven't updated the contributing docs yet, because I plan to add licensing information and an API for hyperparameters next. Docs will updated after that.