Update dtype and example

app.py

@@ -1,12 +1,27 @@
 import gradio as gr
 import jax
-from diffusers import FlaxStableDiffusionPipeline
+import jax.numpy as jnp
+from diffusers import FlaxPNDMScheduler, FlaxStableDiffusionPipeline
 from flax.jax_utils import replicate
 from flax.training.common_utils import shard
 
+DTYPE = jnp.bfloat16
+
 pipeline, pipeline_params = FlaxStableDiffusionPipeline.from_pretrained(
     "bguisard/stable-diffusion-nano-2-1",
+    dtype=DTYPE,
 )
+if DTYPE != jnp.float32:
+    # There is a known issue with schedulers when loading from a pre trained
+    # pipeline. We need the schedulers to always use float32.
+    # See: https://github.com/huggingface/diffusers/issues/2155
+    scheduler, scheduler_params = FlaxPNDMScheduler.from_pretrained(
+        pretrained_model_name_or_path="bguisard/stable-diffusion-nano-2-1",
+        subfolder="scheduler",
+        dtype=jnp.float32,
+    )
+    pipeline_params["scheduler"] = scheduler_params
+    pipeline.scheduler = scheduler
 
 
 def generate_image(prompt: str, inference_steps: int = 30, prng_seed: int = 0):
@@ -51,7 +66,17 @@ app = gr.Interface(
         "Stable Diffusion Nano allows for fast prototyping of diffusion models, "
         "enabling quick experimentation with easily available hardware."
     ),
-    examples=[["A watercolor painting of a bird", 30, 0]],
+    # Some examples were copied from hf.co/spaces/stabilityai/stable-diffusion
+    examples=[
+        # ["A watercolor painting of a bird", 30, 0],
+        [
+            "A small cabin on top of a snowy mountain in the style of Disney, artstation",
+            25,
+            3129302,
+        ],
+        # ["A mecha robot in a favela in expressionist style", 30, 827198341273],
+    ],
 )
 
 app.launch()
+