search

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
139 stars 12 forks source link

Real-CUGAN pro models are saved incorrectly #89

Closed RunDevelopment closed 9 months ago

RunDevelopment commented 9 months ago

As explained here, Real-CUGAN pro models have a weird trick to them. The authors added a pro key to the state dict, and that is the only difference to regular Real-CUGAN models. The issue with this pro key is that it is not part of the architecture itself. This means that saving a pro model will not include the pro key. The goal of this issue is to find a solution to this problem.

@joeyballentine suggested the following solution:

Maybe we could have something like the call API where we have descriptor methods for saving that are model specific?

My response:

[...] the issue with that is that we would have 2 (potentially) different state dicts per model. One state dict of the model itself (model_desc.model.state_dict()) and the modified state dict of the new API.

joeyballentine commented 9 months ago

The issue with this pro key is that it is not part of the architecture itself.

Could we just make it part of the architecture itself?

RunDevelopment commented 9 months ago

Good idea. The official repo sets pro to the int 1, but it doesn't actually check that value, it only checks whether the key exists. So we have a tensor (size [1]) and assign it to the pro field. PyTorch would then save this tensor just like any other tensor and add the pro key we need to the state dict.