AttributeError: module 'jax.tree_util' has no attribute 'register_pytree_with_keys_class'

[emmanueldufourq]

Hello!

I've tried to run the code on Google cola but encountering the following issue:

AttributeError: module 'jax.tree_util' has no attribute 'register_pytree_with_keys_class'

This code produced the error:

! git clone https://github.com/Zeqiang-Lai/Anything2Image.git
%cd Anything2Image
! pip install -r requirements.txt
import anything2image.imagebind as ib
import torch
from diffusers import StableUnCLIPImg2ImgPipeline

# construct models
device = "cuda:0" if torch.cuda.is_available() else "cpu"
pipe = StableUnCLIPImg2ImgPipeline.from_pretrained(
    "stabilityai/stable-diffusion-2-1-unclip", torch_dtype=torch.float16
).to(device)
model = ib.imagebind_huge(pretrained=True).eval().to(device)

This was however on the free Colab (T4 GPU). Not sure if switching to Pro would get rid of that error?

I tried adding this line of code:
!pip install "jax<=0.3.16" "jaxlib<=0.3.16"
but it didn't resolve the issue.

Any suggestions?
Thanks!!

[Zeqiang-Lai]

I don't know the reasons as well. However, even if it is solved, the free T4 GPU could not support the demo since it does not have enough memory. We need at least about 22 Gb GPU memory to run it.

BTW
This notebook works fine locally on my computer. Maybe you could try it locally in your computer.
https://colab.research.google.com/github/Zeqiang-Lai/Anything2Image/blob/main/colab.ipynb