Could you please add a node for the "Face Color Enhancement and Restoration" feature, which allows for the repair and colorization of old photos, to the CodeFormer project?

[ericlee2023]

Could you please add a node for the "Face Color Enhancement and Restoration" feature, which allows for the repair and colorization of old photos, to the CodeFormer project? Currently, I only see a "Face Restoration" node available.

[zslefour]

In order to achieve what you want, you just need to replace the code in the _ init _.py like this：

class FaceRestoreModelLoader:
@classmethod
def INPUT_TYPES(s):
return {"required": { "model_name": (folder_paths.get_filename_list("facerestore_models"), ),
"model_type": (["codeformer", "colorization", "inpainting"],),
}}
RETURN_TYPES = ("FACERESTORE_MODEL",)
FUNCTION = "load_model"

CATEGORY = "facerestore_cf2"

# def load_model(self, model_name):
#     model_path = folder_paths.get_full_path("facerestore_models", model_name)
#     sd = comfy.utils.load_torch_file(model_path, safe_load=True)
#     out = model_loading.load_state_dict(sd).eval()
#     return (out, )

def load_model(self, model_name, model_type):
    
    
    if model_type == "codeformer":
        if "codeformer" in model_name.lower():
            print(f'\tLoading CodeFormer: {model_name}')
            model_path = folder_paths.get_full_path("facerestore_models", model_name)
            device = model_management.get_torch_device()
            codeformer_net = ARCH_REGISTRY.get("CodeFormer")(dim_embd=512, codebook_size=1024, n_head=8, n_layers=9,  
                                                            connect_list=["32", "64", "128", "256"],).to(device)
            checkpoint = torch.load(model_path)["params_ema"]
            codeformer_net.load_state_dict(checkpoint)
            out = codeformer_net.eval()  
            return (out, )
        else:
            model_path = folder_paths.get_full_path("facerestore_models", model_name)
            sd = comfy.utils.load_torch_file(model_path, safe_load=True)
            out = model_loading.load_state_dict(sd).eval()
            return (out, )

    elif model_type == "colorization":
        if "codeformer" in model_name.lower():
            print(f'\tLoading CodeFormer: {model_name}')
            model_path = folder_paths.get_full_path("facerestore_models", model_name)
            device = model_management.get_torch_device()
            codeformer_net = ARCH_REGISTRY.get("CodeFormer")(dim_embd=512, codebook_size=1024, n_head=8, n_layers=9,  
                                                            connect_list=["32", "64", "128"],).to(device)
            checkpoint = torch.load(model_path)["params_ema"]
            codeformer_net.load_state_dict(checkpoint)
            out = codeformer_net.eval()  
            return (out, )
        else:
            model_path = folder_paths.get_full_path("facerestore_models", model_name)
            sd = comfy.utils.load_torch_file(model_path, safe_load=True)
            out = model_loading.load_state_dict(sd).eval()
            return (out, )

    elif model_type == "inpainting":
        if "codeformer" in model_name.lower():
            print(f'\tLoading CodeFormer: {model_name}')
            model_path = folder_paths.get_full_path("facerestore_models", model_name)
            device = model_management.get_torch_device()
            codeformer_net = ARCH_REGISTRY.get("CodeFormer")(dim_embd=512, codebook_size=512, n_head=8, n_layers=9,  
                                                            connect_list=["32", "64", "128"],).to(device)
            checkpoint = torch.load(model_path)["params_ema"]
            codeformer_net.load_state_dict(checkpoint)
            out = codeformer_net.eval()  
            return (out, )
        else:
            model_path = folder_paths.get_full_path("facerestore_models", model_name)
            sd = comfy.utils.load_torch_file(model_path, safe_load=True)
            out = model_loading.load_state_dict(sd).eval()
            return (out, )

[zslefour]