Update app.py

app.py

@@ -6,10 +6,23 @@ from flax.jax_utils import replicate
 from flax.training.common_utils import shard
 from share_btn import community_icon_html, loading_icon_html, share_js
 
+DTYPE = jnp.float16
+
 pipeline, pipeline_params = FlaxStableDiffusionPipeline.from_pretrained(
     "bguisard/stable-diffusion-nano-2-1",
-    dtype=jnp.float16
+    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, negative_prompt: str = "", inference_steps: int = 25, prng_seed: int = 0, guidance_scale: float = 9):
@@ -20,6 +33,7 @@ def generate_image(prompt: str, negative_prompt: str = "", inference_steps: int
     num_samples = 1
     prompt_ids = pipeline.prepare_inputs([prompt] * num_samples)
     prompt_ids = shard(prompt_ids)
+
     if negative_prompt == "":
         images = pipeline(
             prompt_ids=prompt_ids,