replicate params once

app.py

@@ -24,10 +24,10 @@ controlnet, controlnet_params = FlaxControlNetModel.from_pretrained(
 pipe, params = FlaxStableDiffusionControlNetPipeline.from_pretrained(
     "runwayml/stable-diffusion-v1-5", controlnet=controlnet, revision="flax", dtype=jnp.bfloat16
 )
+params["controlnet"] = controlnet_params
+p_params = replicate(params)
 
 def infer(prompts, negative_prompts, image):
-    params["controlnet"] = controlnet_params
-    
     num_samples = 1 #jax.device_count()
     rng = create_key(0)
     rng = jax.random.split(rng, jax.device_count())
@@ -38,7 +38,6 @@ def infer(prompts, negative_prompts, image):
     negative_prompt_ids = pipe.prepare_text_inputs([negative_prompts] * num_samples)
     processed_image = pipe.prepare_image_inputs([canny_image] * num_samples)
     
-    p_params = replicate(params)
     prompt_ids = shard(prompt_ids)
     negative_prompt_ids = shard(negative_prompt_ids)
     processed_image = shard(processed_image)