Create app.py

app.py

@@ -0,0 +1,64 @@
+from flax.jax_utils import replicate
+from jax import pmap
+from flax.training.common_utils import shard
+import jax
+import jax.numpy as jnp
+
+from pathlib import Path
+from PIL import Image
+import numpy as np
+
+from diffusers import FlaxStableDiffusionPipeline
+
+
+import os
+if 'TPU_NAME' in os.environ:
+    import requests
+    if 'TPU_DRIVER_MODE' not in globals():
+        url = 'http:' + os.environ['TPU_NAME'].split(':')[1] + ':8475/requestversion/tpu_driver_nightly'
+        resp = requests.post(url)
+        TPU_DRIVER_MODE = 1
+
+
+    from jax.config import config
+    config.FLAGS.jax_xla_backend = "tpu_driver"
+    config.FLAGS.jax_backend_target = os.environ['TPU_NAME']
+    print('Registered TPU:', config.FLAGS.jax_backend_target)
+else:
+    print('No TPU detected. Can be changed under "Runtime/Change runtime type".')
+
+import jax
+jax.local_devices()
+num_devices = jax.device_count()
+device_type = jax.devices()[0].device_kind
+
+print(f"Found {num_devices} JAX devices of type {device_type}.")
+
+def sd2_inference(pipeline, prompts, params, seed = 42, num_inference_steps = 50 ):
+    prng_seed = jax.random.PRNGKey(seed)
+    prompt_ids = pipeline.prepare_inputs(prompts)
+    params = replicate(params)
+    prng_seed = jax.random.split(prng_seed, jax.device_count())
+    prompt_ids = shard(prompt_ids)
+    images = pipeline(prompt_ids, params, prng_seed, num_inference_steps, jit=True).images
+    images = images.reshape((images.shape[0] * images.shape[1], ) + images.shape[-3:])
+    images = pipeline.numpy_to_pil(images)
+    return images
+def image_grid(imgs, rows, cols, down_sample = 1 ):
+    w,h = imgs[0].size
+    grid = Image.new('RGB', size=(cols*w, rows*h))
+    for i, img in enumerate(imgs): grid.paste(img, box=(i%cols*w, i//cols*h))
+    grid = grid.resize( (grid.size[0]//down_sample, grid.size[1]//down_sample) )
+    return grid
+
+
+HF_ACCESS_TOKEN = os.environ["HFAUTH"]
+
+# Load Model
+# - Reference: https://github.com/huggingface/diffusers/blob/main/README.md
+pipeline, params = FlaxStableDiffusionPipeline.from_pretrained(
+    "CompVis/stable-diffusion-v1-4",
+    use_auth_token = HF_ACCESS_TOKEN,
+    revision="bf16",
+    dtype=jnp.bfloat16,
+)