shard conditioning image input

app.py

@@ -54,6 +54,7 @@ def infer(prompt, negative_prompt, image):
     cond_input = T.ToPILImage()(cond_input)
 
     cond_img_in = pipe.prepare_image_inputs([cond_input] * num_samples)
+    cond_img_in = shard(cond_img_in)
 
     prompt_in = pipe.prepare_text_inputs([prompt] * num_samples)
     prompt_in = shard(prompt_in)
@@ -69,11 +70,11 @@ def infer(prompt, negative_prompt, image):
     output = pipe(
         prompt_ids=prompt_in,
         image=cond_img_in,
+        params=p_params,
         prng_seed=rng,
+        num_inference_steps=50,
         neg_prompt_ids=n_prompt_in,
-        params=p_params,
-        num_inference_steps=20,
-        jit=True
+        jit=True,
             ).images
 
     output_images = pipe.numpy_to_pil(np.asarray(output.reshape((num_samples,) + output.shape[-3:])))