Use CUDA to generate images faster

inference_code.py

@@ -5,13 +5,16 @@ 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)
+# pipeline, _params = FlaxStableDiffusionPipeline.from_pretrained(model_path, dtype=jax.numpy.bfloat16)
+pipeline = DiffusionPipeline.from_pretrained(
+    model_path,
+    from_flax=True, safety_checker=None).to("cuda")
 
 # 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_inference_steps = 20
 
     num_samples = jax.device_count()
     prompt = num_samples * [prompt]