Added Segmentation to the Space

app.py

@@ -1,119 +1,161 @@
-import gradio as gr
-import jax
+from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
+from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
 from PIL import Image
-from flax.jax_utils import replicate
-from flax.training.common_utils import shard
-from diffusers import FlaxControlNetModel, FlaxStableDiffusionControlNetPipeline
-import jax.numpy as jnp
+import gradio as gr
 import numpy as np
+import requests
+import torch
 import gc
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Download and Create SAM Model
+
+print("[Downloading SAM Weights]")
+SAM_URL = "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"
+
+r = requests.get(SAM_URL, allow_redirects=True)
+
+print("[Writing SAM Weights]")
+
+with open("./sam_vit_h_4b8939.pth", "wb") as sam_weights:
+    sam_weights.write(r.content)
+    
+del r
+gc.collect()
+
+sam = sam_model_registry["vit_h"](checkpoint="./sam_vit_h_4b8939.pth").to(device)
+
+mask_generator = SamAutomaticMaskGenerator(sam)
+gc.collect()
+
+# Create ControlNet Pipeline
 
-controlnet, controlnet_params = FlaxControlNetModel.from_pretrained(
-    "mfidabel/controlnet-segment-anything", dtype=jnp.float32
-)
+print("Creating ControlNet Pipeline")
 
-pipe, params = FlaxStableDiffusionControlNetPipeline.from_pretrained(
-    "runwayml/stable-diffusion-v1-5", controlnet=controlnet, revision="flax", dtype=jnp.float32
-)
+controlnet = ControlNetModel.from_pretrained(
+    "mfidabel/controlnet-segment-anything", torch_dtype=torch.float16
+).to(device)
+
+pipe = StableDiffusionControlNetPipeline.from_pretrained(
+    "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16, safety_check=None
+).to(device)
 
-# Add ControlNet params and Replicate
-params["controlnet"] = controlnet_params
-p_params = replicate(params)
 
 # Description
 title = "# 🧨 ControlNet on Segment Anything 🤗"
 description = """This is a demo on 🧨 ControlNet based on Meta's [Segment Anything Model](https://segment-anything.com/).
 
-                Upload a Segment Anything Segmentation Map, write a prompt, and generate images 🤗 This demo is still a Work in Progress, so don't expect it to work well for now !! 
-
-                ⌛️ It takes about 30~ seconds to generate 4 samples, to get faster results, don't forget to reduce the Nº Samples to 1.
-
+                Upload an Image, Segment it with Segment Anything, write a prompt, and generate images 🤗
+                
+                ⌛️ It takes about 20~ 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.
-
+                A huge thanks goes out to @GoogleCloud, for providing us with powerful TPUs that enabled us to train this model; and to the @HuggingFace Team for organizing the sprint.
+                
                 Check out our [Model Card 🧨](https://huggingface.co/mfidabel/controlnet-segment-anything)
+                
               """
 
 about = """
-
-
         # 👨‍💻 About the model
-
+        
         This [model](https://huggingface.co/mfidabel/controlnet-segment-anything) is based on the [ControlNet Model](https://huggingface.co/blog/controlnet), which allow us to generate Images using some sort of condition image. For this model, we selected the segmentation maps produced by Meta's new segmentation model called [Segment Anything Model](https://github.com/facebookresearch/segment-anything) as the condition image. We then trained the model to generate images based on the structure of the segmentation maps and the text prompts given.
+        
 
-
+        
         # 💾 About the dataset
-
+        
         For the training, we generated a segmented dataset based on the [COYO-700M](https://huggingface.co/datasets/kakaobrain/coyo-700m) dataset. The dataset provided us with the images, and the text prompts. For the segmented images, we used [Segment Anything Model](https://github.com/facebookresearch/segment-anything). We then created 8k samples to train our model on, which isn't a lot, but as a team, we have been very busy with many other responsibilities and time constraints, which made it challenging to dedicate a lot of time to generating a larger dataset.  Despite the constraints we faced, we have still managed to achieve some nice results 🙌
-
+        
         You can check the generated datasets below ⬇️
         - [sam-coyo-2k](https://huggingface.co/datasets/mfidabel/sam-coyo-2k)
         - [sam-coyo-2.5k](https://huggingface.co/datasets/mfidabel/sam-coyo-2.5k)
         - [sam-coyo-3k](https://huggingface.co/datasets/mfidabel/sam-coyo-3k)
-
+        
 """
 
-examples = [["contemporary living room of a house", "low quality", "examples/condition_image_1.png"],
-            ["new york buildings,  Vincent Van Gogh starry night ", "low quality, monochrome", "examples/condition_image_2.png"],
-            ["contemporary living room,  high quality, 4k, realistic", "low quality, monochrome, low res", "examples/condition_image_3.png"],
-            ["internal stairs of a japanese house", "low quality,  low res,  people, kids", "examples/condition_image_4.png"],
-            ["a photo of a girl taking notes", "low quality,  low res,  painting", "examples/condition_image_5.png"],
-            ["painting of an hot air ballon flying over a valley, The Great Wave off Kanagawa style, blue and white colors", "low quality,  low res", "examples/condition_image_6.png"],
-            ["painting of families enjoying the sunset,  The Garden of Earthly Delights style,  joyful", "low quality,  low res", "examples/condition_image_7.png"]]
+gif_html = """ <img src="https://github.com/mfidabel/JAX_SPRINT_2023/blob/8632f0fde7388d7a4fc57225c96ef3b8411b3648/EX_1.gif?raw=true" alt= “” height="50%" class="about"> """
+
+examples = [["photo of a futuristic dining table, high quality, tricolor", "low quality, deformed, blurry, points", "examples/condition_image_1.jpeg"],
+            ["a monochrome photo of henry cavil using a shirt, high quality", "low quality,  low res, deformed", "examples/condition_image_2.jpeg"],
+            ["photo of a japanese living room, high quality, coherent", "low quality, colors, saturation, extreme brightness, blurry, low res", "examples/condition_image_3.jpeg"],
+            ["living room, detailed, high quality", "low quality, low resolution, render, oversaturated, low contrast", "examples/condition_image_4.jpeg"],
+            ["painting of the bodiam castle, Vicent Van Gogh style, Starry Night", "low quality, low resolution, render, oversaturated, low contrast", "examples/condition_image_5.jpeg"],
+            ["painting of food, olive oil can, purple wine, green cabbage, chili peppers, pablo picasso style, high quality", "low quality, low resolution, render, oversaturated, low contrast, realistic", "examples/condition_image_6.jpeg"],
+            ["Katsushika Hokusai painting of mountains, a sky and desert landscape,  The Great Wave off Kanagawa style, colorful", 
+             "low quality, low resolution, render, oversaturated, low contrast, realistic", "examples/condition_image_7.jpeg"]]
+
+default_example = examples[4]
+
+examples = examples[::-1]
 
 css = "h1 { text-align: center } .about { text-align: justify; padding-left: 10%; padding-right: 10%; }"
 
 # Inference Function
+def show_anns(anns):
+    if len(anns) == 0:
+        return
+    sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)
+    h, w =  anns[0]['segmentation'].shape
+    final_img = Image.fromarray(np.zeros((h, w, 3), dtype=np.uint8), mode="RGB")
+    for ann in sorted_anns:
+        m = ann['segmentation']
+        img = np.empty((m.shape[0], m.shape[1], 3), dtype=np.uint8)
+        for i in range(3):
+            img[:,:,i] = np.random.randint(255, dtype=np.uint8)
+        final_img.paste(Image.fromarray(img, mode="RGB"), (0, 0), Image.fromarray(np.uint8(m*255)))
+    
+    return final_img
+
+def segment_image(image, seed = 0):
+    # Generate Masks
+    np.random.seed(int(seed))
+    masks = mask_generator.generate(image)
+    torch.cuda.empty_cache()
+    # Create map
+    map = show_anns(masks)
+    del masks
+    gc.collect()
+    torch.cuda.empty_cache()
+    return map
+
 def infer(prompts, negative_prompts, image, num_inference_steps = 50, seed = 4, num_samples = 4):
     try:
-        rng = jax.random.PRNGKey(int(seed))
+        # Segment Image
+        print("Segmenting Everything")
+        segmented_map = segment_image(image, seed)
+        yield segmented_map, [Image.fromarray(np.zeros((512, 512, 3), dtype=np.uint8))] * num_samples
+        # Generate
+        rng = torch.Generator(device="cpu").manual_seed(seed)
         num_inference_steps = int(num_inference_steps)
-        image = Image.fromarray(image, mode="RGB")
-        num_samples = max(jax.device_count(), int(num_samples))
-        p_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)
-        processed_image = pipe.prepare_image_inputs([image] * num_samples)
-        
-        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,
-            params=p_params,
-            prng_seed=p_rng,
-            num_inference_steps=num_inference_steps,
-            neg_prompt_ids=negative_prompt_ids,
-            jit=True,
-        ).images
-
-        del negative_prompt_ids
-        del processed_image
-        del prompt_ids
-
-        output = output.reshape((num_samples,) + output.shape[-3:])
-        final_image = [np.array(x*255, dtype=np.uint8) for x in output]
-        print(output.shape)
+        print(f"Generating Prompt: {prompts} \nNegative Prompt: {negative_prompts} \nSamples:{num_samples}")
+        output = pipe([prompts] * num_samples, 
+                      [segmented_map] * num_samples,
+                      negative_prompt = [negative_prompts] * num_samples,
+                      generator = rng,
+                      num_inference_steps = num_inference_steps)
+    
+
+        final_image = output.images
         del output
         
     except Exception as e:
         print("Error: " + str(e))
-        final_image = [np.zeros((512, 512, 3), dtype=np.uint8)] * num_samples
+        final_image = segmented_map = [np.zeros((512, 512, 3), dtype=np.uint8)] * num_samples
     finally:
         gc.collect()
-        return final_image
+        torch.cuda.empty_cache()
+        yield segmented_map, final_image
     
 
-default_example = examples[5]
-
 cond_img = gr.Image(label="Input", shape=(512, 512), value=default_example[2])\
-                    .style(height=200)
+                    .style(height=400)
+
+segm_img = gr.Image(label="Segmented Image", shape=(512, 512), interactive=False)\
+                    .style(height=400)
 
 output = gr.Gallery(label="Generated images")\
                     .style(height=200, rows=[2], columns=[2], object_fit="contain")
@@ -132,17 +174,16 @@ with gr.Blocks(css=css) as demo:
 
         with gr.Column():
             # Examples
-            gr.Markdown("Try some of the examples below ⬇️")
-            gr.Examples(examples=examples,
-                    inputs=[prompt, negative_prompt, cond_img],
-                    outputs=output,
-                    fn=infer,
-                    examples_per_page=4)
+            gr.Markdown(gif_html)
 
     # Images
     with gr.Row(variant="panel"):
-        with gr.Column(scale=2):
+        with gr.Column(scale=1):
             cond_img.render()
+            
+        with gr.Column(scale=1):
+            segm_img.render()
+            
         with gr.Column(scale=1):
             output.render()
         
@@ -158,15 +199,29 @@ with gr.Blocks(css=css) as demo:
                 seed = gr.Slider(0, 1024, 4, step=1, label="Seed")
                 num_samples = gr.Slider(1, 4, 4, step=1, label="Nº Samples")
                 
-            submit = gr.Button("Generate")
+            segment_btn = gr.Button("Segment")
+            submit = gr.Button("Segment & Generate Images")
             # TODO: Download Button
 
     with gr.Row():
-        gr.Markdown(about, elem_classes="about")
+        with gr.Column():
+            gr.Markdown("Try some of the examples below ⬇️")
+            gr.Examples(examples=examples,
+                        inputs=[prompt, negative_prompt, cond_img],
+                        outputs=output,
+                        fn=infer,
+                        examples_per_page=4)
+            
+        with gr.Column():
+            gr.Markdown(about, elem_classes="about")
     
     submit.click(infer, 
                  inputs=[prompt, negative_prompt, cond_img, num_steps, seed, num_samples],
-                 outputs = output)
+                 outputs = [segm_img, output])
+    
+    segment_btn.click(segment_image,
+                     inputs=[cond_img, seed],
+                     outputs=segm_img)
     
 demo.queue()
 demo.launch()

requirements.txt

@@ -5,4 +5,7 @@ jax[cuda11_pip]
 jaxlib
 git+https://github.com/huggingface/diffusers@main
 opencv-python
-torch
+torch
+torchvision
+git+https://github.com/facebookresearch/segment-anything.git
+accelerate