Update app.py

app.py

@@ -35,9 +35,9 @@ def load_sb_pipe(controlnet_version, sb_path="runwayml/stable-diffusion-v1-5"):
     
     pipe, params = FlaxStableDiffusionControlNetPipeline.from_pretrained(
         sb_path,
-        controlnet=controlnet,
-        dtype=jnp.float32,
-        from_pt=True 
+        controlnet=controlnet, 
+        revision="flax", 
+        dtype=jnp.bfloat16
     )
         
     pipe.scheduler = scheduler
@@ -56,9 +56,9 @@ high_threshold = 200
 
 pipe, params = load_sb_pipe(controlnet_version)
 
-pipe.enable_xformers_memory_efficient_attention()
-pipe.enable_model_cpu_offload()
-pipe.enable_attention_slicing()
+# pipe.enable_xformers_memory_efficient_attention()
+# pipe.enable_model_cpu_offload()
+# pipe.enable_attention_slicing()
 
 def pipe_inference(
     image,
@@ -78,18 +78,20 @@ def pipe_inference(
     resized_image = resize_image(image, resolution)
         
     if not is_canny:
-        resized_image = preprocess_canny(resized_image)
+        resized_image = preprocess_canny(resized_image, resolution)
 
     rng = create_key(seed)
-    # rng = jax.random.split(rng,)
-        
+    rng = jax.random.split(rng, jax.device_count())
+
     prompt_ids = pipe.prepare_text_inputs([prompt] * num_samples)
     negative_prompt_ids = pipe.prepare_text_inputs([negative_prompt] * num_samples)
     processed_image = pipe.prepare_image_inputs([resized_image] * num_samples)
+        
     p_params = replicate(params)
     prompt_ids = shard(prompt_ids)
     negative_prompt_ids = shard(negative_prompt_ids)
     processed_image = shard(processed_image)
+        
     output = pipe(
         prompt_ids=prompt_ids,
         image=processed_image,
@@ -122,15 +124,6 @@ def resize_image(image, resolution):
     
     
 def preprocess_canny(image, resolution=128):
-    h, w = image.shape
-    ratio = w/h
-    if ratio > 1 :
-        resized_image = cv2.resize(image, (int(resolution*ratio), resolution), interpolation=cv2.INTER_NEAREST)
-    elif ratio < 1 :
-        resized_image = cv2.resize(image, (resolution, int(resolution/ratio)), interpolation=cv2.INTER_NEAREST)
-    else:
-        resized_image = cv2.resize(image, (resolution, resolution), interpolation=cv2.INTER_NEAREST)
-
     processed_image = cv2.Canny(resized_image, low_threshold, high_threshold)
     processed_image = processed_image[:, :, None]
     processed_image = np.concatenate([processed_image, processed_image, processed_image], axis=2)
@@ -139,6 +132,7 @@ def preprocess_canny(image, resolution=128):
     processed_image = Image.fromarray(processed_image)
     return resized_image, processed_image
 
+
 def create_demo(process, max_images=12, default_num_images=4):
     with gr.Blocks() as demo:
         with gr.Row():
@@ -218,14 +212,12 @@ def create_demo(process, max_images=12, default_num_images=4):
                          inputs=inputs,
                          outputs=result,
                          api_name='canny')
-    return demo
 
 
 if __name__ == '__main__':
-    """
-    from model import Model
-    model = Model()
-    demo = create_demo(model.process_canny)
+
+    pipe_inference
+    demo = create_demo(pipe_inference)
     demo.queue().launch()
-    """
-    pass
+    # gr.Interface(create_demo).launch()
+