Стабильное распространение: AttributeError: модуль «jax.random» не имеет атрибута «KeyArray»
Когда я запускаю стабильную диффузию в Colab https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/stable_diffusion.ipynb
без каких-либо изменений он терпит неудачу на линии
from diffusers import StableDiffusionPipeline
Журнал ошибок есть
AttributeError: module 'jax.random' has no attribute 'KeyArray'
Как я могу исправить это или любую подсказку?
Импорт должен работать, ipynb должен работать без ошибок.
Ответы 3
В конце концов, нам нужно понизить версию jax, Попробуйте каждый вариант, от последнего к более раннему, и, к счастью, он работает для
jax==0.4.23 jaxlib==0.4.23
Обновление @hunsnowboarder, оно по-прежнему актуально? У меня только что была такая же ошибка, и при установке 0.4.23, по крайней мере, мой проект запускается в Colab, он работает медленно на T4, но показывает «Использование графического процессора» начиная с 10 ГБ и продолжает расти, как и должно быть, я здесь новичок, так что Я не понимаю, работает ли он так, как должен.
@SKB_BGPL - странно... Я получил ошибку, когда попытался запустить ее на машине с графическим процессором... и я решил свою проблему, используя позднюю модель диффузии, чтобы больше не было конфликтов с библиотеками jax
jax.random.KeyArrayустарел в JAX v0.4.16 и удален в JAX v0.4.24. Учитывая это, похоже, что стабильный код распространения HuggingFace работает только с JAX v0.4.23 или более ранней версии.
Вы можете установить JAX v0.4.23 с поддержкой графического процессора следующим образом:
pip install "jax[cuda12_pip]==0.4.23" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
или, если вы предпочитаете ориентироваться на локальную установку CUDA, например:
pip install "jax[cuda12_local]==0.4.23" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
Дополнительную информацию об установке графического процессора см. в разделе Установка JAX: NVIDIA GPU.
Из руководства Colab обновите второй сегмент так:
!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
# Change this
# !pip install diffusers==0.11.1
# To just
!pip install diffusers
Если вы уже запускали pip install в среде выполнения Colab, вам нужно либо отключиться и открыть новую среду выполнения (моя рекомендация), либо использовать --upgrade.
Диффузору v0.11.1 уже более 18 месяцев, и ноутбук работает с текущей версией v0.29.0 без каких-либо других изменений. Вместо использования старой версии диффузоров, требующей старой версии jax, мы можем использовать последние версии.
Эта версия jaxlib не поддерживает CUDA. Как установить эти версии jax и jaxlib, если мы хотим, чтобы в jaxlib была поддержка cuda 12 12.2? Потому что нет смысла использовать Colab без графического процессора для стабильных диффузионных проектов....