Loading model in google co-lab

[Vinaysukhesh98]

Loading pre-trained TextEncoder model
Some weights of CLIPTextModel were not initialized from the model checkpoint at openai/clip-vit-large-patch14 and are newly initialized: ['text_model.bias', 'text_model.masked_bias']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
Loading pre-trained UNET model

AttributeError Traceback (most recent call last)
in <cell line: 11>()
9
10 print("Loading pre-trained UNET model")
---> 11 unet = UNet2DConditionModel.from_pretrained('runwayml/stable-diffusion-v1-5',
12 subfolder="unet", revision='main', torch_dtype=dtype,
13 cache_dir=cache_dir).to(device)

17 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'