Using `hydra` to configure models fail to serialize

dasayan05 commented 8 months ago

Describe the bug

When using hydra to configure diffusers model instances, it fails to serialize (i.e. .save_config()). This is very similar to #3814 but the solution isn't really applicable to me.

Reason for this behavior: diffusers' s FrozenDict apparently doesn't take care of list/tuple arguments recursively.

PS: I am not sure if this issue belongs here, but still would like to see if someone has a solution.

Reproduction

I am trying to configure the UNet2DModel class using hydra's automatic instantiation mechanism, where I have the following config

model:
  _target_: diffusers.models.UNet2DModel
  sample_size: 32
  in_channels: 3
  down_block_types:
    - DownBlock2D
    - AttnDownBlock2D
    - AttnDownBlock2D
    - AttnDownBlock2D
  block_out_channels: ...
  ...

Now instantiate the model following the hydra way

> model = hydra.utils.instantiate(config.model)
<UNet2DModel>

.. and it works fine so far.

BUT, turns out that the FrozenDict that wraps over the internal config dictionary ..

> model.config
OrderedDict([('sample_size', 32), ('in_channels', 3), ('down_block_types', ['DownBlock2D', 'AttnDownBlock2D', 'AttnDownBlock2D', 'AttnDownBlock2D']), ...])

.. still contains some hydra-container inside, which is not serializable

> type(model.config['down_block_types'])
<class 'omegaconf.listconfig.ListConfig'>

.. and makes the .save_config() (and eventually .save_pretrained()) to fail

> model.save_config(...)
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/proj/gpu_mtk53576/miniconda3/envs/dl/lib/python3.9/site-packages/diffusers/configuration_utils.py", line 167, in save_config
    self.to_json_file(output_config_file)
  File "/proj/gpu_mtk53576/miniconda3/envs/dl/lib/python3.9/site-packages/diffusers/configuration_utils.py", line 601, in to_json_file
    writer.write(self.to_json_string())
  File "/proj/gpu_mtk53576/miniconda3/envs/dl/lib/python3.9/site-packages/diffusers/configuration_utils.py", line 590, in to_json_string
    return json.dumps(config_dict, indent=2, sort_keys=True) + "\n"
  File "/proj/gpu_mtk53576/miniconda3/envs/dl/lib/python3.9/json/__init__.py", line 234, in dumps
    return cls(
  File "/proj/gpu_mtk53576/miniconda3/envs/dl/lib/python3.9/json/encoder.py", line 201, in encode
    chunks = list(chunks)
  File "/proj/gpu_mtk53576/miniconda3/envs/dl/lib/python3.9/json/encoder.py", line 431, in _iterencode
    yield from _iterencode_dict(o, _current_indent_level)
  File "/proj/gpu_mtk53576/miniconda3/envs/dl/lib/python3.9/json/encoder.py", line 405, in _iterencode_dict
    yield from chunks
  File "/proj/gpu_mtk53576/miniconda3/envs/dl/lib/python3.9/json/encoder.py", line 438, in _iterencode
    o = _default(o)
  File "/proj/gpu_mtk53576/miniconda3/envs/dl/lib/python3.9/json/encoder.py", line 179, in default
    raise TypeError(f'Object of type {o.__class__.__name__} '
TypeError: Object of type ListConfig is not JSON serializable

Logs

No response

System Info

diffusers version: 0.25.1
Platform: Linux-5.4.0-26-generic-x86_64-with-glibc2.31
Python version: 3.9.12
PyTorch version (GPU?): 2.1.2+cu118 (True)
Huggingface_hub version: 0.20.2
Transformers version: 4.36.2
Accelerate version: 0.26.1
xFormers version: not installed
Using GPU in script?: True
Using distributed or parallel set-up in script?: True

Who can help?

@sayakpaul @patrickvonplaten

sayakpaul commented 8 months ago

Sorry, we are hydra users in the team and don't have enough expertise for you here.

yonetaniryo commented 4 months ago

Adding _convert_: all could resolve the same issue in my project.

for example:

model:
  _target_: diffusers.models.UNet2DModel
  _convert_: all
...

huggingface / diffusers