[WIP] Load text_model.embeddings.position_ids outsude state_dict

[StAlKeR7779]

What type of PR is this? (check all applicable)

• Refactor
• Feature
• Bug Fix
• Optimization
• Documentation Update

Description

In transformers 4.31.0 text_model.embeddings.position_ids no longer part of state_dict.
Fix untested as can't run right now but should be correct. Also need to check how transformers 4.30.2 works with this fix.

Related Tickets & Documents

huggingface/transformers@8e5d161#diff-7f53db5caa73a4cbeb0dca3b396e3d52f30f025b8c48d4daf51eb7abb6e2b949R191
https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.register_buffer

QA Instructions, Screenshots, Recordings

  File "C:\Users\artis\Documents\invokeai\.venv\lib\site-packages\invokeai\backend\model_management\convert_ckpt_to_diffusers.py", line 844, in convert_ldm_clip_checkpoint
    text_model.load_state_dict(text_model_dict)
  File "C:\Users\artis\Documents\invokeai\.venv\lib\site-packages\torch\nn\modules\module.py", line 2041, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for CLIPTextModel:
        Unexpected key(s) in state_dict: "text_model.embeddings.position_ids".

[lstein]

The patch seems to work with 4.31.0. However there is a reproducible error in the logs during the conversion process:

Task exception was never retrieved
future: <Task finished name='Task-1053' coro=<AsyncServer._emit_internal() done, defined at /home/lstein/invokeai-main/.venv/lib/python3.10/site-packages/socketio/asyncio_server.py:427> exception=TypeError('Object of type PosixPath is not JSON serializable')>
Traceback (most recent call last):
  File "/home/lstein/invokeai-main/.venv/lib/python3.10/site-packages/socketio/asyncio_server.py", line 437, in _emit_internal
    await self._send_packet(sid, self.packet_class(
  File "/home/lstein/invokeai-main/.venv/lib/python3.10/site-packages/socketio/asyncio_server.py", line 442, in _send_packet
    encoded_packet = pkt.encode()
  File "/home/lstein/invokeai-main/.venv/lib/python3.10/site-packages/socketio/packet.py", line 64, in encode
    encoded_packet += self.json.dumps(data, separators=(',', ':'))
  File "/usr/lib/python3.10/json/__init__.py", line 238, in dumps
    **kw).encode(obj)
  File "/usr/lib/python3.10/json/encoder.py", line 199, in encode
    chunks = self.iterencode(o, _one_shot=True)
  File "/usr/lib/python3.10/json/encoder.py", line 257, in iterencode
    return _iterencode(o, 0)
  File "/usr/lib/python3.10/json/encoder.py", line 179, in default
    raise TypeError(f'Object of type {o.__class__.__name__} '
TypeError: Object of type PosixPath is not JSON serializable

My guess is that this is related to a bad model manager event. I will investigate.

[lstein]

As expected, the patch does not work when used with transformers 4.30.2:

  File "/home/lstein/Projects/InvokeAI-main/invokeai/app/services/processor.py", line 70, in __process
    outputs = invocation.invoke(
  File "/home/lstein/invokeai-main/.venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/lstein/Projects/InvokeAI-main/invokeai/app/invocations/compel.py", line 88, in invoke
    tokenizer_info = context.services.model_manager.get_model(
  File "/home/lstein/Projects/InvokeAI-main/invokeai/app/services/model_manager_service.py", line 359, in get_model
    model_info = self.mgr.get_model(
  File "/home/lstein/Projects/InvokeAI-main/invokeai/backend/model_management/model_manager.py", line 472, in get_model
    model_path = model_class.convert_if_required(
  File "/home/lstein/Projects/InvokeAI-main/invokeai/backend/model_management/models/stable_diffusion.py", line 120, in convert_if_required
    return _convert_ckpt_and_cache(
  File "/home/lstein/Projects/InvokeAI-main/invokeai/backend/model_management/models/stable_diffusion.py", line 293, in _convert_ckpt_and_cache
    convert_ckpt_to_diffusers(
  File "/home/lstein/Projects/InvokeAI-main/invokeai/backend/model_management/convert_ckpt_to_diffusers.py", line 1149, in convert_ckpt_to_diffusers
    pipe = load_pipeline_from_original_stable_diffusion_ckpt(checkpoint_path, **kwargs)
  File "/home/lstein/Projects/InvokeAI-main/invokeai/backend/model_management/convert_ckpt_to_diffusers.py", line 1109, in load_pipeline_from_original_stable_diffusion_ckpt
    text_model = convert_ldm_clip_checkpoint(checkpoint)
  File "/home/lstein/Projects/InvokeAI-main/invokeai/backend/model_management/convert_ckpt_to_diffusers.py", line 845, in convert_ldm_clip_checkpoint
    text_model.load_state_dict(text_model_dict)
  File "/home/lstein/invokeai-main/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 2041, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for CLIPTextModel:
	Missing key(s) in state_dict: "text_model.embeddings.position_ids". 

I've put in a try: block which tries fetching the state dict first without the key and then with the key added back in. I think this is marginally better than checking the module version number, since the transformers team might put the key back into the next version.

[lstein]

Tested with both 4.31.0 and 4.30.2 and seems to be working as expected.