how to get the minimun working example of FlaxStableDiffusionPipeline in google colab with tpu runtime

### Describe the bug

I try the code in google colab with tpu runtime
```
! python3 -m pip install -U diffusers[flax]
import diffusers, os
pipeline = diffusers.StableDiffusionPipeline.from_single_file('https://huggingface.co/chaowenguo/pal/blob/main/chilloutMix-Ni.safetensors')
pipeline.save_pretrained('chilloutMix', safe_serialization=False)
pipeline, params = diffusers.FlaxStableDiffusionPipeline.from_pretrained('./chilloutMix', from_pt=True, safety_checker=None)
```
I always get Your session crashed for an unknown reason. I want to get the mininum working example in google colab with tpu runtime

### Reproduction

N/A

### Logs

_No response_

### System Info

google colab with tpu runtime

### Who can help?

_No response_