Cleanup

app.py

@@ -1,11 +1,15 @@
 import gradio as gr
 import jax
+import jax.numpy as jnp
 from diffusers import FlaxStableDiffusionPipeline
 from flax.jax_utils import replicate
 from flax.training.common_utils import shard
 
 pipeline, pipeline_params = FlaxStableDiffusionPipeline.from_pretrained(
     "bguisard/stable-diffusion-nano",
+    dtype=jnp.float16,
+    resume_download=True,
+    use_memory_efficient_attention=True
 )
 
 

pre-requirements.txt

@@ -0,0 +1,5 @@
+pip
+setuptools
+wheel
+ninja
+cmake

requirements.txt

@@ -1,4 +1,6 @@
 transformers
-diffusers
-jax[cuda11_pip] #jax[cuda11_cudnn82] #jax[cuda11_cudnn86] #jax[cuda11_cudnn805]
-flax
+flax
+jax[cuda11_pip]
+-f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+jaxlib
+git+https://github.com/huggingface/diffusers@main