search

NVIDIA / modulus

Open-source deep-learning framework for building, training, and fine-tuning deep learning models using state-of-the-art Physics-ML methods
https://developer.nvidia.com/modulus
Apache License 2.0
795 stars 172 forks source link

🐛[BUG]: Diffusion network, SongUNetPosEmbed tries to apply positional embedding when its disabled. #504

Closed daviddpruitt closed 1 month ago

daviddpruitt commented 1 month ago

Version

0.6.0

On which installation method(s) does this occur?

Docker, Pip, Source

Describe the issue

When using SongUNetPosEmbed for diffusion there is an option to disable positional embedding by setting N_grid_channels to 0. If the positional embedding is disabled the forward pass will still apply the positional embedding resulting in a crash. See: https://github.com/NVIDIA/modulus/blob/f49ceae9b30df5fa065976d2fcca66e19f98ace1/modulus/models/diffusion/song_unet.py#L516

Minimum reproducible example

No response

Relevant log output

Traceback (most recent call last):
  File "/root/modulus_corr_patch/examples/generative/corrdiff/train.py", line 352, in <module>
    main()
  File "/usr/local/lib/python3.10/dist-packages/hydra/main.py", line 94, in decorated_main
    _run_hydra(
  File "/usr/local/lib/python3.10/dist-packages/hydra/_internal/utils.py", line 394, in _run_hydra
    _run_app(
  File "/usr/local/lib/python3.10/dist-packages/hydra/_internal/utils.py", line 457, in _run_app
    run_and_report(
  File "/usr/local/lib/python3.10/dist-packages/hydra/_internal/utils.py", line 223, in run_and_report
    raise ex
  File "/usr/local/lib/python3.10/dist-packages/hydra/_internal/utils.py", line 220, in run_and_report
    return func()
  File "/usr/local/lib/python3.10/dist-packages/hydra/_internal/utils.py", line 458, in <lambda>
    lambda: hydra.run(
  File "/usr/local/lib/python3.10/dist-packages/hydra/_internal/hydra.py", line 132, in run
    _ = ret.return_value
  File "/usr/local/lib/python3.10/dist-packages/hydra/core/utils.py", line 260, in return_value
    raise self._return_value
  File "/usr/local/lib/python3.10/dist-packages/hydra/core/utils.py", line 186, in run_job
    ret.return_value = task_function(task_cfg)
  File "/root/modulus_corr_patch/examples/generative/corrdiff/train.py", line 344, in main
    training_loop.training_loop(
  File "/root/modulus_corr_patch/examples/generative/corrdiff/training/training_loop.py", line 288, in training_loop                                                                                                    loss = loss_fn(
  File "/usr/local/lib/python3.10/dist-packages/modulus/metrics/diffusion/loss.py", line 411, in __call__                                                                                                               D_yn = net(input, y_lr, sigma, labels, augment_labels=augment_labels)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl                                                                                                           return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/modulus/models/diffusion/unet.py", line 152, in forward
    F_x = self.model(
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl                                                                                                           return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/nvtx/nvtx.py", line 116, in inner
    result = func(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/modulus/models/diffusion/song_unet.py", line 516, in forward                                                                                                            selected_pos_embd = self.positional_embedding_indexing(x, global_index)
  File "/usr/local/lib/python3.10/dist-packages/modulus/models/diffusion/song_unet.py", line 578, in positional_embedding_indexing                                                                                      self.pos_embd.to(x.dtype)
AttributeError: 'NoneType' object has no attribute 'to'

Environment details

No response