add app

app.py

@@ -0,0 +1,17 @@
+import gradio as gr
+from inference_code import generate_images
+
+def generate_image_predictions(prompt):
+    images = generate_images(prompt)
+    return images
+
+iface = gr.Interface(
+    fn=generate_image_predictions,
+    inputs=gr.inputs.Textbox("Enter a text prompt here"),
+    outputs=[gr.outputs.Image() for i in range(4)],
+    title="Map Diffuser",
+    description="Generates four images from a given text prompt.",
+    examples=[["Satellite image of amsterdam with industrial area and highways"], ["Satellite image with forests and residential, no water"], ["A person playing guitar on a stage"]]
+)
+
+iface.launch()

inference_code.py

@@ -0,0 +1,27 @@
+import jax
+import numpy as np
+from flax.jax_utils import replicate
+from flax.training.common_utils import shard
+from diffusers import FlaxStableDiffusionPipeline
+
+model_path = "sabman/map-diffuser-v3"
+pipeline, params = FlaxStableDiffusionPipeline.from_pretrained(model_path, dtype=jax.numpy.bfloat16)
+
+#prompt = "create a map with traffic signals, busway and residential buildings, in water color style"
+def generate_images(prompt):
+    prng_seed = jax.random.PRNGKey(-1)
+    num_inference_steps = 50
+    
+    num_samples = jax.device_count()
+    prompt = num_samples * [prompt]
+    prompt_ids = pipeline.prepare_inputs(prompt)
+    
+    # shard inputs and rng
+    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 = pipeline.numpy_to_pil(np.asarray(images.reshape((num_samples,) + images.shape[-3:])))
+    return images
+

requirements.txt

@@ -0,0 +1,9 @@
+transformers>=4.25.1
+datasets
+flax
+optax
+torch
+torchvision
+ftfy
+tensorboard
+Jinja2