do what coyo_1m space does

app.py

@@ -6,6 +6,8 @@ import torch
 torch.backends.cuda.matmul.allow_tf32 = True
 import torchvision
 import torchvision.transforms as T
+from flax.jax_utils import replicate
+from flax.training.common_utils import shard
 #from torchvision.transforms import v2 as T2
 import cv2
 import PIL
@@ -25,8 +27,8 @@ conditioning_image_transforms = T.Compose(
     ]
 )
 
-cnet = FlaxControlNetModel.from_pretrained("./models/catcon-controlnet-wd", dtype=jnp.bfloat16, from_flax=True)
-pipe = FlaxStableDiffusionControlNetPipeline.from_pretrained(
+cnet, cnet_params = FlaxControlNetModel.from_pretrained("./models/catcon-controlnet-wd", dtype=jnp.bfloat16, from_flax=True)
+pipe, params = FlaxStableDiffusionControlNetPipeline.from_pretrained(
         "./models/wd-1-5-b2",
         controlnet=cnet,
         dtype=jnp.bfloat16,
@@ -36,23 +38,38 @@ pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
 #pipe.enable_model_cpu_offload()
 #pipe.enable_xformers_memory_efficient_attention()
 
-generator = torch.manual_seed(0)
+def get_random(seed):
+    jax.random.PRNGKey(seed)
 
 # inference function takes prompt, negative prompt and image
 def infer(prompt, negative_prompt, image):
     # implement your inference function here
+    params["controlnet"] = cnet_params
+    num_samples = 1
+
     inp = Image.fromarray(image)
 
     cond_input = conditioning_image_transforms(inp)
     cond_input = T.ToPILImage()(cond_input)
 
-    cond_in = pipe.prepare_image_inputs([cond_input] * 4)
+    cond_img_in = pipe.prepare_image_inputs([cond_input] * num_samples)
+
+    prompt_in = pipe.prepare_text_inputs([prompt] * num_samples)
+    prompt_in = shard(prompt_in)
+
+    n_prompt_in = pipe.prepare_text_inputs([negative_prompt] * num_samples)
+    n_prompt_in = shard(n_prompt_in)
+
+    rng = get_random(0)
+    rng.random.split(rng, jax.device_count())
+
+    p_params = replicate(params)
     
     output = pipe(
-        prompt,
-        cond_in,
-        generator=generator,
-        num_images_per_prompt=4,
+        prompt_ids=prompts_in,
+        image=cond_img_in,
+        prng_seed=rng,
+        neg_prompt_ids=n_prompt_in,
         num_inference_steps=20,
         jit=True
             )