use params

inference_code.py

@@ -8,7 +8,7 @@ from diffusers import FlaxStableDiffusionPipeline
 # prompt = "create a map with traffic signals, busway and residential buildings, in water color style"
 def generate_images(prompt):
     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)
 
     prng_seed = jax.random.PRNGKey(-1)
     num_inference_steps = 50
@@ -18,7 +18,7 @@ def generate_images(prompt):
     prompt_ids = pipeline.prepare_inputs(prompt)
 
     # shard inputs and rng
-    params = replicate(_params)
+    params = replicate(params)
     prng_seed = jax.random.split(prng_seed, jax.device_count())
     prompt_ids = shard(prompt_ids)