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

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

please help

[jsrgb]

Same error, TPU runtime:

!pip install -q jax jaxlib flax transformers ftfy
!pip install -q diffusers
from diffusers import FlaxStableDiffusionPipeline

Pkg versions:

jax: 0.4.33
jaxlib: 0.4.33
flax: 0.8.5
transformers: 4.44.2
ftfy: 6.2.3
diffusers: 0.30.3

Tail of app.log:

Timestamp	Level	Message
Oct 10, 2024, 2:41:11 PM	WARNING	WARNING:root:kernel f2407886-2e9a-44b3-ba3e-64a7df75644d restarted
Oct 10, 2024, 2:41:11 PM	INFO	KernelRestarter: restarting kernel (1/5), keep random ports
Oct 10, 2024, 2:41:08 PM	WARNING	E1010 21:41:08.711022 6384 process_state.cc:808] RAW: Raising signal 6 with default behavior
Oct 10, 2024, 2:41:06 PM	WARNING	E1010 21:41:05.999826 6384 coredump_hook.cc:603] RAW: Dumping core locally.
Oct 10, 2024, 2:41:06 PM	WARNING	E1010 21:41:05.999823 6384 coredump_hook.cc:551] RAW: Cannot send fingerprint to Coroner: [NOT_FOUND] stat failed on crash reporting socket /var/google/services/logmanagerd/remote_coredump.socket (Is the listener running?): No such file or directory
Oct 10, 2024, 2:41:06 PM	WARNING	E1010 21:41:05.999805 6384 coredump_hook.cc:542] RAW: Sending fingerprint to remote end.
Oct 10, 2024, 2:41:06 PM	WARNING	E1010 21:41:05.999802 6384 client.cc:272] RAW: Coroner client retries enabled (b/136286901), will retry for up to 30 sec.
Oct 10, 2024, 2:41:06 PM	WARNING	E1010 21:41:05.999788 6384 coredump_hook.cc:447] RAW: Remote crash data gathering hook invoked.
Oct 10, 2024, 2:41:05 PM	WARNING	https://symbolize.stripped_domain/r/?trace=796e3eaec9fc,796d190214f8,796e3ea9851f&map=5edeb7d86db111100e979a74159a3982:796d09400000-796d19240ba0
Oct 10, 2024, 2:41:05 PM	WARNING	@ 0x796e3ea98520 (unknown) (unknown)
Oct 10, 2024, 2:41:05 PM	WARNING	@ 0x796d190214f9 928 (unknown)
Oct 10, 2024, 2:41:05 PM	WARNING	PC: @ 0x796e3eaec9fc (unknown) pthread_kill
Oct 10, 2024, 2:41:05 PM	WARNING	*** SIGABRT received by PID 6384 (TID 6384) on cpu 35 from PID 6384; stack trace: ***
Oct 10, 2024, 2:41:05 PM	WARNING	https://symbolize.stripped_domain/r/?trace=796e3eaec9fc,796e3ea9851f&map=
Oct 10, 2024, 2:41:05 PM	WARNING	what(): File is already registered: xla/service/cpu/backend_config.proto
Oct 10, 2024, 2:41:05 PM	WARNING	terminate called after throwing an instance of 'google::protobuf::FatalException'
Oct 10, 2024, 2:41:05 PM	WARNING	[libprotobuf FATAL external/com_google_protobuf/src/google/protobuf/message.cc:258] File is already registered: xla/service/cpu/backend_config.proto
Oct 10, 2024, 2:40:29 PM	INFO	Kernel started: f2407886-2e9a-44b3-ba3e-64a7df75644d, name: python3

The above works fine on CPU / T4 GPU runtime.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.