Update app.py

app.py

@@ -13,12 +13,14 @@ from pathos.threading import ThreadPool as Pool
 #from diffusers.utils import export_to_gif
 #from huggingface_hub import hf_hub_download
 #from safetensors.torch import load_file
-from diffusers import FlaxStableDiffusionPipeline
+from diffusers import FlaxStableDiffusionPipeline, DiffusionPipeline
 #import jax
-import jax.numpy as jnp
+#import jax.numpy as jnp
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-pipe = FlaxStableDiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", dtype=jnp.bfloat16, token=os.getenv("hf_token")).to(device)
+
+#pipe = FlaxStableDiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", dtype=jnp.bfloat16, token=os.getenv("hf_token")).to(device)
+pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16, revision="refs/pr/1", token=os.getenv("hf_token")).to(device)
 
 def translate(text,lang):