|
from r_chainner.archs.face.gfpganv1_clean_arch import GFPGANv1Clean |
|
from r_chainner.types import PyTorchModel |
|
|
|
|
|
class UnsupportedModel(Exception): |
|
pass |
|
|
|
|
|
def load_state_dict(state_dict) -> PyTorchModel: |
|
|
|
state_dict_keys = list(state_dict.keys()) |
|
|
|
if "params_ema" in state_dict_keys: |
|
state_dict = state_dict["params_ema"] |
|
elif "params-ema" in state_dict_keys: |
|
state_dict = state_dict["params-ema"] |
|
elif "params" in state_dict_keys: |
|
state_dict = state_dict["params"] |
|
|
|
state_dict_keys = list(state_dict.keys()) |
|
|
|
|
|
if ( |
|
"toRGB.0.weight" in state_dict_keys |
|
and "stylegan_decoder.style_mlp.1.weight" in state_dict_keys |
|
): |
|
model = GFPGANv1Clean(state_dict) |
|
return model |
|
|
