How to finetune stage 2?

[tintin-py]

Excellent work and thanks for sharing the repository.

Problem

I am trying to finetune the CodeFormer model. I am focused on finetuning the encoder and transformer and hence, I am planning to train only Stage 2 using my own dataset of HQ images.
As per the config file provided in options/CodeFormer_stage2.yml the model that is being used is CodeFormerIdxModel. However, I keep getting AttributeError when I try to load the weights to the model.

Code

Here is the code snippet to reproduce the problem.

root_path = '.'
opt_path = './options/dummy.yml'
opt = parse(opt_path, root_path)

ckpt_path = './weights/CodeFormer/codeformer_stage2.pth'
ckpt = torch.load(ckpt_path)['params_ema']
model = CodeFormerIdxModel(opt)
model.load_state_dict(ckpt)

Error

I get the following error when I execute the above code.
AttributeError: 'CodeFormerIdxModel' object has no attribute 'load_state_dict'.

Queries

I have the following queries

1. The CodeFormerModel class has the load_state_dict method. However, despite having the same inheritance pattern i.e. BaseModel -> SRModel -> CodeFormerIdxModel why am I not getting the AttributeError when I do the same for CodeFormerIdxModel?
2. There are several types of CodeFormer models in this repository, namely CodeFormerModel, CodeFormerIdxModel and CodeFormerJointModel. What is the puspose of each model?
3. Can I finetune CodeFormerModel in stage 2 using the pre-trained weights codeformer_stage2.pth?