Update app.py

app.py

@@ -23,8 +23,8 @@ def infer(prompts, negative_prompts):
     rng = create_key(0)
     rng = jax.random.split(rng, jax.device_count())
     
-    prompt_ids = pipe.prepare_text_inputs([prompts] * num_samples)
-    negative_prompt_ids = pipe.prepare_text_inputs([negative_prompts] * num_samples)
+    prompt_ids = pipe.prepare_inputs([prompts] * num_samples)
+    negative_prompt_ids = pipe.prepare_inputs([negative_prompts] * num_samples)
     
     p_params = replicate(params)
     prompt_ids = shard(prompt_ids)