Added Segmentation Map Notebook

app.py

@@ -4,7 +4,7 @@ from PIL import Image
 from flax.jax_utils import replicate
 from flax.training.common_utils import shard
 from diffusers import FlaxControlNetModel, FlaxStableDiffusionControlNetPipeline
-from diffusers.utils import load_image
+import jax.profiler
 import jax.numpy as jnp
 import numpy as np
 import gc
@@ -22,6 +22,8 @@ pipe, params = FlaxStableDiffusionControlNetPipeline.from_pretrained(
 params["controlnet"] = controlnet_params
 p_params = replicate(params)
 
+jax.profiler.save_device_memory_profile("memory.prof")
+
 # Description
 title = "# 🧨 ControlNet on Segment Anything 🤗"
 description = """This is a demo on 🧨 ControlNet based on Meta's [Segment Anything Model](https://segment-anything.com/).
@@ -30,6 +32,9 @@ description = """This is a demo on 🧨 ControlNet based on Meta's [Segment Anyt
 
                 ⌛️ It takes about 30~ seconds to generate 4 samples, to get faster results, don't forget to reduce the Nº Samples to 1.
 
+                You can obtain the Segmentation Map of any Image through this Colab: [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/mfidabel/JAX_SPRINT_2023/blob/main/Segment_Anything_JAX_SPRINT.ipynb)
+
+                
                 A huge thanks goes out to @Google Cloud, for providing us with powerful TPUs that enabled us to train this model; and to the @HuggingFace Team for organizing the sprint.
               """