ImportError: cannot import name 'enable_fused_na' from 'natten'

[justin4ai]

Hello, I'm trying to inference with your pretrained model. But I encounter ImportError: cannot import name 'enable_fused_na' from 'natten' error. I tried to change torch & natten versions but the same error occurs again and again.

python main.py type=inference

Traceback (most recent call last):
  File "/home/justin/workspace/StyleNAT/main.py", line 11, in <module>
    from models.generator import Generator
  File "/home/justin/workspace/StyleNAT/models/generator.py", line 22, in <module>
    from models.stylenat import HydraNeighborhoodAttention, MHSARPB, NeighborhoodAttentionSplitHead
  File "/home/justin/workspace/StyleNAT/models/stylenat.py", line 14, in <module>
    from natten import (
ImportError: cannot import name 'enable_fused_na' from 'natten' (/home/justin/anaconda3/envs/stylenat/lib/python3.10/site-packages/natten/__init__.py)

Also the following is my pip freeze output.

antlr4-python3-runtime==4.9.3
certifi==2024.2.2
charset-normalizer==3.3.2
click==8.1.7
cmake==3.29.3
contourpy==1.2.1
cycler==0.12.1
dill==0.3.8
docker-pycreds==0.4.0
einops==0.8.0
filelock==3.13.1
fonttools==4.52.4
fsspec==2024.2.0
ftfy==6.2.0
gitdb==4.0.11
GitPython==3.1.43
huggingface-hub==0.23.2
hydra-core==1.3.2
idna==3.7
imageio-ffmpeg==0.4.9
Jinja2==3.1.3
joblib==1.4.2
kiwisolver==1.4.5
lmdb==1.4.1
markdown-it-py==3.0.0
MarkupSafe==2.1.5
matplotlib==3.9.0
mdurl==0.1.2
mpmath==1.3.0
natten==0.17.1+torch220cu121
networkx==3.2.1
ninja==1.11.1.1
numpy==1.26.4
nvidia-cublas-cu12==12.1.3.1
nvidia-cuda-cupti-cu12==12.1.105
nvidia-cuda-nvrtc-cu12==12.1.105
nvidia-cuda-runtime-cu12==12.1.105
nvidia-cudnn-cu12==8.9.2.26
nvidia-cufft-cu12==11.0.2.54
nvidia-curand-cu12==10.3.2.106
nvidia-cusolver-cu12==11.4.5.107
nvidia-cusparse-cu12==12.1.0.106
nvidia-nccl-cu12==2.19.3
nvidia-nvjitlink-cu12==12.1.105
nvidia-nvtx-cu12==12.1.105
omegaconf==2.3.0
packaging==24.0
pillow==10.3.0
platformdirs==4.2.2
protobuf==4.25.3
psutil==5.9.8
Pygments==2.18.0
pyparsing==3.1.2
python-dateutil==2.9.0.post0
PyYAML==6.0.1
requests==2.31.0
rich==13.7.1
safetensors==0.4.3
scikit-learn==1.5.0
scipy==1.13.1
sentry-sdk==2.3.1
setproctitle==1.3.3
six==1.16.0
smmap==5.0.1
sympy==1.12
threadpoolctl==3.5.0
timm==1.0.3
torch==2.2.2+cu121
torchaudio==2.2.2+cu121
torchvision==0.17.2+cu121
tqdm==4.66.4
triton==2.2.0
typing_extensions==4.11.0
urllib3==2.2.1
wandb==0.17.0
wcwidth==0.2.13

The problematic code is as follows:

And I verified there are no such enable_fused_na() and so on.

Thanks for your great project, I hope I could hear from you.

Best, Junyeong Ahn

[alihassanijr]

Thank you for raising this issue; that's my fault, I changed the API name in 0.17.

Although now that I'm looking more closely, it looks like StyleNAT doesn't use FNA, even though the flags are being set. I'll push a temporary fix to unblock you right now and will wait for @stevenwalton to chime in.

StyleNAT should be able to use FNA for inference, but training gets complicated because FNA doesn't implement the backward pass for RPB.

[justin4ai]

I close this issue since at least the error is fixed although there are some others remaining. Please check conversation in your PR!

[stevenwalton]

@alihassanijr thanks for handling this so quickly.

I think FNA was still in alpha when we did this so it wasn't fully implemented or maybe there was miscommunication between us. We can fully switch to FNA. We can sync outside this chat.

@justin4ai, thanks for your interest in StyleNAT. We're always interested to see what people do with it, and don't hesitate to reach out if you're running into issues. Some of the param choices can be tricky. If you're doing face-swapping (as your profile suggests) the attention maps will be helpful for debugging and finding artifacts.

[justin4ai]

@stevenwalton Thanks for your suggestion! I really love your works. Actually I'm trying to replace mere attention module with your HydraNAT module in absorbing diffusion sampler for facial image generation! And this StyleNAT helps me a lot as both a reference for my trial and itself. I'd like to share my work if it finishes! Thank you so much anyways :)

[stevenwalton]

Great to hear! You'll need a bit more complex kerneling for things other than faces. Though this will play a bit differently in diffusion, but it should work fine.