본문 바로가기
Issue

HuggingFace) Error(s) in loading state_dict for '' Missing key(s) in state_dict

by 박채니 2024. 8. 30.
SMALL

안녕하세요, 코린이의 코딩 학습기 채니 입니다.
개인 포스팅용으로 내용에 오류 및 잘못된 정보가 있을 수 있습니다.

# 상황

HuggingFace에서 사전 학습된 모델을 가져와 프로젝트를 하나 만들려고 했다.

https://huggingface.co/TencentARC/PhotoMaker

 

TencentARC/PhotoMaker · Hugging Face

PhotoMaker Model Card Introduction Users can input one or a few face photos, along with a text prompt, to receive a customized photo or painting within seconds (no training required!). Additionally, this model can be adapted to any base model based on SDXL

huggingface.co

 

위 모델을 사용하려고 하고, 간단하게 말하자면 첨부한 이미지와 프롬프트를 참고하여 새로운 이미지를 생성해준다.

아래 코드 실행 시, state_dict 관련 오류가 발생했다.

pipe.load_photomaker_adapter(
    os.path.dirname(photomaker_ckpt),
    subfolder="",
    weight_name=os.path.basename(photomaker_ckpt),
    trigger_word="img"
)

🐛 오류 메세지

Error(s) in loading state_dict for PhotoMakerIDEncoder_CLIPInsightfaceExtendtoken:
    Missing key(s) in state_dict: "qformer_perceiver.token_proj.0.weight", "qformer_perceiver.token_proj.0.bias", "qformer_perceiver.token_proj.2.weight", "qformer_perceiver.token_proj.2.bias", "qformer_perceiver.token_norm.weight", "qformer_perceiver.token_norm.bias"
    ...
    ...

💡 해결방안

pipe.load_photomaker_adapter(
    os.path.dirname(photomaker_ckpt),
    subfolder="",
    weight_name=os.path.basename(photomaker_ckpt),
    trigger_word="img",
    pm_version="v1"
)

 

🚨 이 해결방안은 PhotoMaker를 사용하는 경우에만 해결할 수 있는 방안이기 때문에 엥..? 스러운 해결방안이다.

 

⭐️ 본인이 직접 checkpoint를 로드하는 경우⭐️ strict=False 옵션을 넣어주면 된다.

다만, 해당 옵션은 없는 키들을 무시한다는 거기 때문에 모델 품질의 저하가 올 수도 있다고 한다.

🔥 원인

일단 결론부터 말하자면.. 진짜 대박 삽질이였던 버그였다.

 

해당 오류의 원인은

사전학습된 모델과 추론하려는 모델이 일치하지 않아서, 혹은 버전이 일치하지 않아서 발생하는 오류이다.

 

쉽게 말해서 사전학습된 모델의 state_dict에는 존재하는 키들이 추론하려는 모델에 없다는 것!

 

📝 나의 삽질 기록

버그 원인은 알겠는데.. 도저히 이해가 안됐다.

base_model_path = 'SG161222/RealVisXL_V3.0'
device = torch.device("mps")
torch_dtype = torch.float16

photomaker_path = hf_hub_download(repo_id="TencentARC/PhotoMaker", filename="photomaker-v1.bin", repo_type="model")

# Load base model
pipe = PhotoMakerStableDiffusionXLPipeline.from_pretrained(
    base_model_path,  # can change to any base model based on SDXL
    torch_dtype=torch_dtype, 
    use_safetensors=True, 
    variant="fp16"
).to(device)

# Load PhotoMaker checkpoint
pipe.load_photomaker_adapter(
    os.path.dirname(photomaker_path),
    subfolder="",
    weight_name=os.path.basename(photomaker_path),
    trigger_word="img"  # define the trigger word
) 

 

코드를 보면,

 

1. hf_hub_download 함수를 통해 Hub에서 photomaker 모델 파일 다운

2. PhotoMakerStableDiffusionXLPipeline.from_pretrained 함수를 통해 사전학습된 모델 로드

3. Hub에서 가져온 photomaker checkpoint 결합

 

근데 왜 자꾸 key가 없다는 오류가 나는거지? 싶었다.

 

일단 버전이 일치하지 않아서 해당 오류가 발생된다는 아티클을 봤기 때문에 python 버전부터 하나하나 뜯어보았다.

$ python -V

3.12.X

 

나는 python 3.10 버전으로 가상환경을 만들었었는데, 3.12 버전을 사용하고 있었다.

 

$ which python

alias: 블라블라

 

현재 참조하고 있는 python 위치를 찾아보니 이전에 python 설치 시 설정했던 alias를 참조하고 있었다.

 

$ unalias python
$ python -V

3.10.X

 

alias를 해제하여 현재 가상환경의 python을 참조할 수 있도록 해주었다.

 

그리고 PhotoMaker github README에서 Mac GPU 사용법을 아래와 같이 정리해주었었는데,

$ xcode-select --install
$ brew install llvm libomp
$ pip install torch==2.1.2

 

여기서 계속 마음에 걸렸던 게 pip install torch==2.1.2 이 부분이다.
해당 스크립트를 실행하면 자꾸 Requirement torch==2.1.2 not found 버전이 없다는 오류가 발생했기 때문이다.

 

이 때까지만 해도 나는 "아! 요거구나! 이거 때문이다!" 했다.

(응.. 아니야..^^;)

 

torch 버전 관련 오류에 대해 구글링을 열심히 해봤는데, python 3.10 버전을 설치하면 된다고 한다.

...? 난 이미 3.10 버전인데요..


또 열심히 구글링을 해본 결과, 아래 명령어처럼 입력하라고 되어있었다.

$ python -m pip install torch==2.1.2 torchvision==0.16.2 --extra-index-url https://download.pytorch.org/whl/cu118

 

오! torch 버전 관련 오류가 안나고 설치가 잘 되었다!


이제 되겠거니~ 하고 실행해보니!

.... 여전히 똑같은 오류가 발생해서.. CUDA가 설치되어있는 Window 환경으로 줄행랑 쳤다.. 

(당연하다.. 왜냐면 이게 원인이 아니니까..ㅎ)

이 때까지만 해도 Mac에서 돌려서 그런거구나 했다..

(응.. 아니야..^^;)

 

참고 사이트
https://huggingface.co/TencentARC/PhotoMaker-V2/discussions/2
https://github.com/TencentARC/PhotoMaker/issues/175
https://github.com/Stability-AI/StableCascade/issues/23
https://mrnoobiest.tistory.com/entry/stable-diffusion-No-matching-distribution-found-for-torch212-%ED%95%B4%EA%B2%B0%EB%B2%95
https://github.com/Stability-AI/StableCascade/issues/23


 

Window 환경에서 세팅 다시 하고 실행했는데 이전에 포스팅 했던 Torch not compiled with CUDA enabled 오류가 또 발생했다.

 

Torch not compiled with CUDA enabled 오류 해결 포스팅은 아래 링크 참고
(https://chanychu.tistory.com/532)

 

해당 포스팅을 읽어보면 알겠지만, Mac 환경에서 실행 시 CUDA가 없어서 mps로 실행하여 오류를 처리해주었다.
하지만 지금은 CUDA가 있는 Window 환경이다. 따라서 torch와 CUDA 버전의 충돌일 것이라고 생각했다.

 

그래서 버전들을 쭉 확인해주었다.

# GPU 버전확인 
$ nvidia-smi
Wed Aug 28 17:54:09 2024
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 537.42                 Driver Version: 537.42       CUDA Version: 12.2     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                     TCC/WDDM  | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  NVIDIA GeForce RTX 3060 Ti   WDDM  | 00000000:2B:00.0  On |                  N/A |
|  0%   37C    P8              21W / 200W |    474MiB /  8192MiB |      2%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+

+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|    0   N/A  N/A       240    C+G   C:\Windows\explorer.exe                   N/A      |
|    0   N/A  N/A     29732    C+G   ...975_x64__8wekyb3d8bbwe\ms-teams.exe    N/A      |
|    0   N/A  N/A     45144    C+G   ...ekyb3d8bbwe\PhoneExperienceHost.exe    N/A      |
|    0   N/A  N/A     45176    C+G   ...GeForce Experience\NVIDIA Share.exe    N/A      |
|    0   N/A  N/A     47572    C+G   ...les\Microsoft OneDrive\OneDrive.exe    N/A      |
|    0   N/A  N/A     49004    C+G   ....Search_cw5n1h2txyewy\SearchApp.exe    N/A      |
|    0   N/A  N/A     51032    C+G   ...CBS_cw5n1h2txyewy\TextInputHost.exe    N/A      |
|    0   N/A  N/A     52416    C+G   ...e Stream\96.0.0.0\GoogleDriveFS.exe    N/A      |
|    0   N/A  N/A     55676    C+G   ...975_x64__8wekyb3d8bbwe\ms-teams.exe    N/A      |
|    0   N/A  N/A     56636    C+G   ...on\128.0.2739.42\msedgewebview2.exe    N/A      |
|    0   N/A  N/A     57756    C+G   ...Programs\Microsoft VS Code\Code.exe    N/A      |
+---------------------------------------------------------------------------------------+

# cuda 버전 확인
$ nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2020 NVIDIA Corporation
Built on Thu_Jun_11_22:26:48_Pacific_Daylight_Time_2020
Cuda compilation tools, release 11.0, V11.0.194
Build cuda_11.0_bu.relgpu_drvr445TC445_37.28540450_0

현재 설치되어 있는 CUDA는 v11.0인 것을 알 수 있다.


따라서 pytorch 공식 홈페이지(https://pytorch.org/get-started/previous-versions/)에서 CUDA 11.0과 호환되는 torch를 설치해주었다.

# Linux and Windows
# CUDA 11.0

$ pip install torch==1.7.1+cu110 torchvision==0.8.2+cu110 torchaudio==0.7.2 -f https://download.pytorch.org/whl/torch_stable.html

 

근데 또 ERROR: Could not find a version that satisfies the requirement torch==1.7.0+cpu 이런 오류가 발생했다..


찾아보니 torch 1.7.1 버전과 python 3.10 버전이랑 호환이 안된다고 한다.

(https://github.com/pytorch/pytorch/issues/47354)

 

.........ㅋ 🤦🏻‍♀️


그래서 python 버전을 3.9로 변경하려고 했는데 PhotoMaker github에 들어가보니 pytorch 버전이 2 이상이여야 한다길래 그냥 CUDA를 새로 설치해주었다. (ㅎㅎ)

 

이제 되겠다! 하고 스크립트를 실행했는데 하.. 또! state_dict 관련 오류가 발생했다.

결국엔.. PhotoMaker github에 들어가서 load_photomaker_adapter 함수를 뜯어보기로 했다...


해당 함수 코드

더보기
class PhotoMakerStableDiffusionXLPipeline(StableDiffusionXLPipeline):
    @validate_hf_hub_args
    def load_photomaker_adapter(
        self,
        pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
        weight_name: str,
        subfolder: str = '',
        trigger_word: str = 'img',
        pm_version: str = 'v2',
        **kwargs,
    ):
        """
        Parameters:
            pretrained_model_name_or_path_or_dict (`str` or `os.PathLike` or `dict`):
                Can be either:

                    - A string, the *model id* (for example `google/ddpm-celebahq-256`) of a pretrained model hosted on
                      the Hub.
                    - A path to a *directory* (for example `./my_model_directory`) containing the model weights saved
                      with [`ModelMixin.save_pretrained`].
                    - A [torch state
                      dict](https://pytorch.org/tutorials/beginner/saving_loading_models.html#what-is-a-state-dict).

            weight_name (`str`):
                The weight name NOT the path to the weight.

            subfolder (`str`, defaults to `""`):
                The subfolder location of a model file within a larger model repository on the Hub or locally.

            trigger_word (`str`, *optional*, defaults to `"img"`):
                The trigger word is used to identify the position of class word in the text prompt, 
                and it is recommended not to set it as a common word. 
                This trigger word must be placed after the class word when used, otherwise, it will affect the performance of the personalized generation.           
        """

        # Load the main state dict first.
        cache_dir = kwargs.pop("cache_dir", None)
        force_download = kwargs.pop("force_download", False)
        resume_download = kwargs.pop("resume_download", False)
        proxies = kwargs.pop("proxies", None)
        local_files_only = kwargs.pop("local_files_only", None)
        token = kwargs.pop("token", None)
        revision = kwargs.pop("revision", None)

        user_agent = {
            "file_type": "attn_procs_weights",
            "framework": "pytorch",
        }

        if not isinstance(pretrained_model_name_or_path_or_dict, dict):
            model_file = _get_model_file(
                pretrained_model_name_or_path_or_dict,
                weights_name=weight_name,
                cache_dir=cache_dir,
                force_download=force_download,
                resume_download=resume_download,
                proxies=proxies,
                local_files_only=local_files_only,
                token=token,
                revision=revision,
                subfolder=subfolder,
                user_agent=user_agent,
            )
            if weight_name.endswith(".safetensors"):
                state_dict = {"id_encoder": {}, "lora_weights": {}}
                with safe_open(model_file, framework="pt", device="cpu") as f:
                    for key in f.keys():
                        if key.startswith("id_encoder."):
                            state_dict["id_encoder"][key.replace("id_encoder.", "")] = f.get_tensor(key)
                        elif key.startswith("lora_weights."):
                            state_dict["lora_weights"][key.replace("lora_weights.", "")] = f.get_tensor(key)
            else:
                state_dict = torch.load(model_file, map_location="cpu")
        else:
            state_dict = pretrained_model_name_or_path_or_dict

        keys = list(state_dict.keys())
        if keys != ["id_encoder", "lora_weights"]:
            raise ValueError("Required keys are (`id_encoder` and `lora_weights`) missing from the state dict.")

        self.num_tokens =2
        self.pm_version = pm_version
        self.trigger_word = trigger_word
        # load finetuned CLIP image encoder and fuse module here if it has not been registered to the pipeline yet
        print(f"Loading PhotoMaker {pm_version} components [1] id_encoder from [{pretrained_model_name_or_path_or_dict}]...")
        self.id_image_processor = CLIPImageProcessor()
        if pm_version == "v1": # PhotoMaker v1 
            id_encoder = PhotoMakerIDEncoder()
        elif pm_version == "v2": # PhotoMaker v2
            id_encoder = PhotoMakerIDEncoder_CLIPInsightfaceExtendtoken()
        else:
            raise NotImplementedError(f"The PhotoMaker version [{pm_version}] does not support")

        id_encoder.load_state_dict(state_dict["id_encoder"], strict=True)
        id_encoder = id_encoder.to(self.device, dtype=self.unet.dtype)    
        self.id_encoder = id_encoder

        # load lora into models
        print(f"Loading PhotoMaker {pm_version} components [2] lora_weights from [{pretrained_model_name_or_path_or_dict}]")
        self.load_lora_weights(state_dict["lora_weights"], adapter_name="photomaker")

        # Add trigger word token
        if self.tokenizer is not None: 
            self.tokenizer.add_tokens([self.trigger_word], special_tokens=True)

        self.tokenizer_2.add_tokens([self.trigger_word], special_tokens=True)

 

큰 부분들만 봐보면 이런 흐름을 가지고 있다.

(아마 설명이 틀릴건데 대충 이런 흐름을 가지고 있다고 생각했다.)

 

1️⃣ 추론한 모델을 불러온 후

model_file = _get_model_file(
    pretrained_model_name_or_path_or_dict,
    weights_name=weight_name,
    cache_dir=cache_dir,
    force_download=force_download,
    resume_download=resume_download,
    proxies=proxies,
    local_files_only=local_files_only,
    token=token,
    revision=revision,
    subfolder=subfolder,
    user_agent=user_agent,
)
state_dict = torch.load(model_file, map_location="cpu")

 

2️⃣ 내부적으로 정의된 모델을 불러온다.

if pm_version == "v1": # PhotoMaker v1 
    id_encoder = PhotoMakerIDEncoder()
elif pm_version == "v2": # PhotoMaker v2
    id_encoder = PhotoMakerIDEncoder_CLIPInsightfaceExtendtoken()

 

3️⃣ 내부적으로 불러온 모델의 state_dict을 가져와서 우리가 로컬에서 추론하려는 모델의 id_encoder를 비교해준다.

id_encoder.load_state_dict(state_dict["id_encoder"], strict=True)

 

3번 코드에서 오류가 나는건데,

위에서 설명했듯이 state_dict 관련 Missing Keys 오류가 발생하면 대부분 strict=False 옵션을 통해 없는 키들을 무시해서 오류를 처리해준다고 한다.


근데 이건 내가 직접 로드하는 게 아니기 때문에 해당 방법을 사용할 수도 없을 뿐더러, 모델 품질이 저하되기 때문에 해당 방법을 사용하진 않았다.

 

이 때, 내 눈에 들어온 건 2번 코드인데 내가 참고한 github sample에 있는 코드에는 pm_version을 파라미터로 넘겨주지 않는다.

따라서 기본 값인 "v2"로 적용된다..

ㅋㅋㅋㅋㅋㅋ🤦🏻‍♀️..ㅋㅋ 장난...? 아니지.. 그냥 다 내 잘못이지...


... 내가 로드한 모델의 pm_version은 "v1"이였기 때문에..

당연히 다른 모델을 불러와 비교하니까 키가 없다는 오류가 발생했던 것.....

코드를 뜯어서 보니 당연히 오류가 날 수밖에 없더라..😂

어쨌든.. "v1"으로 파라미터를 넘겨주어서 해결했다.

 

진짜 역대급 대박 삽질이였다 ㅎ;; 이걸로 시간을 얼마나 날린 건지 모르겠다..ㅋ ㅠ....ㅋㅋ..ㅠㅠㅠ... 하하..핳하ㅏ하!!! 눙물..
정말 너무 허무했지만, 해결했으니 됐다. 그리고 얻은 것도 많다.. 그럼 됐지!
python이랑 안 친한데 덕분에 python 코드도 많이 읽었고, 구글링하면서 여러 지식도 얻게 되었으니 됐다!!!

 

 

 

LIST