Colab TPUv2 session crash on diffusers import

[dan-crdll]

To reproduce the problem:

!pip install transformers
!pip install diffusers
!pip install tensorboard
!pip install --upgrade torch torchvision torchaudio
!pip install git+https://github.com/openai/CLIP.git@a1d071733d7111c9c014f024669f959182114e33

import jax.tools.colab_tpu

jax.tools.colab_tpu.setup_tpu()

import torch
import torchvision
import torch_xla
import torch_xla.core.xla_model as xm

from diffusers import UNet2DConditionModel, AutoencoderKL, PNDMScheduler
import clip

device = xm.xla_device()
print(f'TPU Device: {device}')

Session crash with unknown problem when importing diffusers

[tolgacangoz]

It seems that classes of diffusers and torch_xla cannot be imported in the same Colab TPUv2 session currently 🤔.
Cc: @pcuenca

[pcuenca]

Interesting. I could replicate a similar crash by importing diffusers first (it works) and then importing torch_xla (it crashed). This sounds like an issue for Colab, wouldn't you think?

!pip install diffusers

import jax.tools.colab_tpu
jax.tools.colab_tpu.setup_tpu()

from diffusers import UNet2DConditionModel, AutoencoderKL, PNDMScheduler
import torch_xla

(The same crash happens even if you don't call jax.tools.colab_tpu.setup_tpu())

[suzukimain]

I had the same problem, but doing this solved it.
huggingface/transformers#30516 (comment)

[dan-crdll]

Thanks, this solved it!