[코랩 error] module 'jax.random' has no attribute 'KeyArray' 드림부스 dreambooth 에러

1. stable diffusion을 dreambooth로 파인튜닝 중 train에서 <AttributeError: module 'jax.random' has no attribute 'KeyArray'> 에러가 계속 발생했다.

2. 처음엔 GPU를 A100으로 돌려서, T4로 바꿔주니까 해결이 됐는데 그 다음 학습부터는 T4에서도 계속 같은 에러 발생.

3. 해결하려고 별 짓을 다했는데 안돼서 골머리를 앓던 중!

4. 다음 코드로 해결

 

!pip install "jax[cuda12_local]==0.4.23" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
!pip install diffusers==0.11.1
!pip install transformers scipy ftfy accelerate

 

train 전에 이걸 돌려주고,

train 시키면 학습이 잘된다!!

HuggingFace stable diffusion 코드가 돌아가는 버전을 맞춰줘야 하는 걸로 보임 

 

으하하..

이상입니다.