colab version not working

[joeaelkhoury]

Hello,

I was trying to test your code on google colab. but when i tried to run

from SD2.diff_pipe import StableDiffusionDiffImg2ImgPipeline

i got the below error.

---

AttributeError Traceback (most recent call last)

in () 12 import torch 13 from torchvision import transforms ---> 14 from SD2.diff_pipe import StableDiffusionDiffImg2ImgPipeline 15 from transformers import pipeline 16 6 frames /usr/local/lib/python3.10/dist-packages/jax/_src/deprecations.py in getattr(name) 52 warnings.warn(message, DeprecationWarning, stacklevel=2) 53 return fn ---> 54 raise AttributeError(f"module {module!r} has no attribute {name!r}") 55 56 return getattr AttributeError: module 'jax.random' has no attribute 'KeyArray'

[joeaelkhoury]

I just added the below

!pip install "jax[cuda12_pip]==0.4.23" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
!pip install --upgrade diffusers

and in the code diff_pipe.py I added the following from diffusers.utils.torch_utils import randn_tensor also i removed the randn_tensor from this line

from diffusers.utils import ( PIL_INTERPOLATION, deprecate, is_accelerate_available, is_accelerate_version, logging, replace_example_docstring, )