entrpn

import jax
import jax.numpy as jnp
import numpy as np

num_devices = jax.device_count()
device_type = jax.devices()[0].device_kind
assert "TPU" in device_type, "Available device is not a TPU, please select TPU from Edit > Notebook settings > Hardware accelerator"

from flax.jax_utils import replicate
from diffusers import FlaxStableDiffusionXLPipeline

entrpn

import gradio as gr

import jax
import jax.numpy as jnp
import numpy as np

from flax.jax_utils import replicate
from diffusers import FlaxStableDiffusionXLPipeline

from flax.training.common_utils import shard