Update app.py

app.py

@@ -1,15 +1,8 @@
 import gradio as gr
 import numpy as np
 import torch
-import jax
-import jax.numpy as jnp
-from flax.jax_utils import replicate
-from flax.training.common_utils import shard
 from PIL import Image
 from segment_anything import SamPredictor, sam_model_registry, SamAutomaticMaskGenerator
-from diffusers import (
-    FlaxStableDiffusionControlNetPipeline,
-    FlaxControlNetModel,
 )
 from transformers import pipeline
 
@@ -49,11 +42,6 @@ with gr.Blocks() as demo:
         """
         ## Work in Progress
         ### About
-        We have trained a JAX ControlNet model for semantic segmentation on Wildlife Animal Images.
-        
-        For the training data creation we used the [Wildlife Animals Images](https://www.kaggle.com/datasets/anshulmehtakaggl/wildlife-animals-images) dataset.
-        We created segmentation masks with the help of [Grounded SAM](https://github.com/IDEA-Research/Grounded-Segment-Anything) where we used the animals names 
-        as input prompts for detection and more accurate segmentation.
         
         ### How To Use
         
@@ -64,10 +52,6 @@ with gr.Blocks() as demo:
         mask_img = gr.Image(label="Mask", interactive=False)
         output_img = gr.Image(label="Output", interactive=False)
 
-    with gr.Row():
-        prompt_text = gr.Textbox(lines=1, label="Prompt")
-        negative_prompt_text = gr.Textbox(lines=1, label="Negative Prompt")
-
     with gr.Row():
         submit = gr.Button("Submit")
         clear = gr.Button("Clear")
@@ -84,7 +68,8 @@ with gr.Blocks() as demo:
             pil_img = Image.fromarray(np_img, 'RGB')
             mask_images.append(pil_img)
         
-        return np.stack(mask_images)
+        #return np.stack(mask_images)
+        return image
 
     # def infer(
     #     image, prompts, negative_prompts, num_inference_steps=50, seed=4, num_samples=4