Kosmos error: Memory limit exceeded

app.py

@@ -4,7 +4,7 @@ warnings.filterwarnings('ignore')
 
 import subprocess, io, os, sys, time
 # os.system("pip install gradio==3.36.1")
-os.system("pip install gradio==3.41.2")
+os.system("pip install gradio==3.40.1")
 import gradio as gr
 
 from loguru import logger
@@ -698,6 +698,8 @@ def run_anything_task(input_image, text_prompt, task_type, inpaint_prompt, box_t
     logger.info(f'run_anything_task_[{file_temp}]_9_9_')
     return output_images, gr.Gallery.update(label='result images'), None, None, None
 
+kosmos_enable = False
+
 def change_radio_display(task_type, mask_source_radio):
     text_prompt_visible = True
     inpaint_prompt_visible = False
@@ -710,21 +712,23 @@ def change_radio_display(task_type, mask_source_radio):
     kosmos_text_output_visible = False
 
     if task_type == "Kosmos-2":
-        text_prompt_visible = False
-        image_gallery_visible = False
-        kosmos_input_visible = True
-        kosmos_output_visible = True
-        kosmos_text_output_visible = True
-    else: 
-        if task_type == "inpainting":
-            inpaint_prompt_visible = True
-        if task_type == "inpainting" or task_type == "remove":
-            mask_source_radio_visible = True   
-            if mask_source_radio == mask_source_draw:
-                text_prompt_visible = False
-        if task_type == "relate anything":
+        if kosmos_enable:
+            text_prompt_visible = False
+            image_gallery_visible = False
+            kosmos_input_visible = True
+            kosmos_output_visible = True
+            kosmos_text_output_visible = True
+        
+
+    if task_type == "inpainting":
+        inpaint_prompt_visible = True
+    if task_type == "inpainting" or task_type == "remove":
+        mask_source_radio_visible = True   
+        if mask_source_radio == mask_source_draw:
             text_prompt_visible = False
-            num_relation_visible = True
+    if task_type == "relate anything":
+        text_prompt_visible = False
+        num_relation_visible = True
 
     return  (gr.Textbox.update(visible=text_prompt_visible), 
             gr.Textbox.update(visible=inpaint_prompt_visible), 
@@ -750,6 +754,10 @@ def get_model_device(module):
     except Exception as e:
         return 'Error'
 
+task_types = ["detection", "segment", "inpainting", "remove", "relate anything"]
+if kosmos_enable:
+    task_types.append("Kosmos-2")
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser("Grounded SAM demo", add_help=True)
     parser.add_argument("--debug", action="store_true", help="using debug mode")
@@ -768,7 +776,7 @@ if __name__ == "__main__":
     load_lama_cleaner_model()
     load_ram_model()
     
-    if os.environ.get('IS_MY_DEBUG') is None:
+    if os.environ.get('IS_MY_DEBUG') is None and kosmos_enable:
         kosmos_model, kosmos_processor = load_kosmos_model(device)
 
     if os.environ.get('IS_MY_DEBUG') is None:
@@ -779,13 +787,14 @@ if __name__ == "__main__":
     print(f'sd_model__{get_model_device(sd_model)}')
     print(f'lama_cleaner_model__{get_model_device(lama_cleaner_model)}')
     print(f'ram_model__{get_model_device(ram_model)}')
+    print(f'kosmos_model__{get_model_device(kosmos_model)}')
     
     block = gr.Blocks().queue()
     with block:
         with gr.Row():
             with gr.Column():
                 input_image = gr.Image(source='upload', elem_id="image_upload", tool='sketch', type='pil', label="Upload")    
-                task_type = gr.Radio(["detection", "segment", "inpainting", "remove", "relate anything", "Kosmos-2"],  value="detection", 
+                task_type = gr.Radio(task_types,  value="detection", 
                                                 label='Task type', visible=True) 
                 mask_source_radio = gr.Radio([mask_source_draw, mask_source_segment], 
                                     value=mask_source_segment, label="Mask from",
@@ -860,7 +869,8 @@ if __name__ == "__main__":
         DESCRIPTION = f'### This demo from [Grounded-Segment-Anything](https://github.com/IDEA-Research/Grounded-Segment-Anything). <br>'
         DESCRIPTION += f'RAM from [RelateAnything](https://github.com/Luodian/RelateAnything). <br>'
         DESCRIPTION += f'Remove(cleaner) from [lama-cleaner](https://github.com/Sanster/lama-cleaner). <br>'
-        DESCRIPTION += f'Kosmos-2 from [RelateAnything](https://huggingface.co/spaces/ydshieh/Kosmos-2). <br>'
+        if kosmos_enable:
+            DESCRIPTION += f'Kosmos-2 from [Kosmos-2](https://huggingface.co/spaces/ydshieh/Kosmos-2). <br>'
         DESCRIPTION += f'Thanks for their excellent work.'
         DESCRIPTION += f'<p>For faster inference without waiting in queue, you may duplicate the space and upgrade to GPU in settings. \
                         <a href="https://huggingface.co/spaces/yizhangliu/Grounded-Segment-Anything?duplicate=true"><img style="display: inline; margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space" /></a></p>'