update app.py

app.py

@@ -38,6 +38,7 @@ import cv2
 import numpy as np
 import matplotlib.pyplot as plt
 
+groundingdino_enable = True
 sam_enable = True
 inpainting_enable = True
 ram_enable = True
@@ -103,16 +104,10 @@ sam_predictor = None
 sam_mask_generator = None
 sd_model = None
 lama_cleaner_model= None
-lama_cleaner_model_device = device
 ram_model = None
 kosmos_model = None
 kosmos_processor = None
 
-def get_sam_vit_h_4b8939():
-    if not os.path.exists('./sam_vit_h_4b8939.pth'):
-        logger.info(f"get sam_vit_h_4b8939.pth...")
-        result = subprocess.run(['wget', 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth'], check=True)
-        print(f'wget sam_vit_h_4b8939.pth result = {result}') 
 
 def load_model_hf(model_config_path, repo_id, filename, device='cpu'):
     args = SLConfig.fromfile(model_config_path) 
@@ -282,24 +277,31 @@ def set_device():
         device = 'cpu'
     print(f'device={device}')
 
-def load_groundingdino_model():
+def load_groundingdino_model(device):
     # initialize groundingdino model
-    global groundingdino_model
     logger.info(f"initialize groundingdino model...")
-    groundingdino_model = load_model_hf(config_file, ckpt_repo_id, ckpt_filenmae, device='cpu')
+    groundingdino_model = load_model_hf(config_file, ckpt_repo_id, ckpt_filenmae, device=device) #'cpu')
+    return groundingdino_model
+
+def get_sam_vit_h_4b8939():
+    if not os.path.exists('./sam_vit_h_4b8939.pth'):
+        logger.info(f"get sam_vit_h_4b8939.pth...")
+        result = subprocess.run(['wget', 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth'], check=True)
+        print(f'wget sam_vit_h_4b8939.pth result = {result}')
 
-def load_sam_model():
+def load_sam_model(device):
     # initialize SAM
-    global sam_model, sam_predictor, sam_mask_generator, sam_device, device
+    global sam_model, sam_predictor, sam_mask_generator, sam_device
+    get_sam_vit_h_4b8939()
     logger.info(f"initialize SAM model...")
     sam_device = device
     sam_model = build_sam(checkpoint=sam_checkpoint).to(sam_device)
     sam_predictor = SamPredictor(sam_model)
     sam_mask_generator = SamAutomaticMaskGenerator(sam_model)
 
-def load_sd_model():
+def load_sd_model(device):
     # initialize stable-diffusion-inpainting
-    global sd_model, device
+    global sd_model
     logger.info(f"initialize stable-diffusion-inpainting...")
     sd_model = None
     if os.environ.get('IS_MY_DEBUG') is None:
@@ -311,14 +313,14 @@ def load_sd_model():
         )
         sd_model = sd_model.to(device)
 
-def load_lama_cleaner_model():
+def load_lama_cleaner_model(device):
     # initialize lama_cleaner
-    global lama_cleaner_model, device
+    global lama_cleaner_model
     logger.info(f"initialize lama_cleaner...")
 
     lama_cleaner_model = ModelManager(
             name='lama',
-            device=lama_cleaner_model_device,
+            device=device,
         )
 
 def lama_cleaner_process(image, mask, cleaner_size_limit=1080):
@@ -390,7 +392,7 @@ class Ram_Predictor(RamPredictor):
             self.model.load_state_dict(torch.load(self.config.load_from, map_location=self.device))
         self.model.train()
 
-def load_ram_model():
+def load_ram_model(device):
     # load ram model
     global ram_model
     if os.environ.get('IS_MY_DEBUG') is not None:
@@ -830,20 +832,20 @@ if __name__ == "__main__":
     if kosmos_enable:
         kosmos_model, kosmos_processor = load_kosmos_model(device)
         
-    load_groundingdino_model()
+    if groundingdino_enable:
+        groundingdino_model = load_groundingdino_model('cpu')
     
     if sam_enable:
-        get_sam_vit_h_4b8939()
-        load_sam_model()
+        load_sam_model(device)
 
     if inpainting_enable:
-        load_sd_model()
+        load_sd_model(device)
 
     if lama_cleaner_enable:
-        load_lama_cleaner_model()
+        load_lama_cleaner_model(device)
 
     if ram_enable:
-        load_ram_model()
+        load_ram_model(device)
     
     if os.environ.get('IS_MY_DEBUG') is None:
         os.system("pip list")
@@ -865,7 +867,7 @@ if __name__ == "__main__":
                 mask_source_radio = gr.Radio([mask_source_draw, mask_source_segment], 
                                     value=mask_source_segment, label="Mask from",
                                     visible=False) 
-                text_prompt = gr.Textbox(label="Detection Prompt[To detect multiple objects, seperating each name with '.', like this: cat . dog . chair ]", placeholder="Cannot be empty")                                                
+                text_prompt = gr.Textbox(label="Detection Prompt[To detect multiple objects, seperating each with '.', like this: cat . dog . chair ]", placeholder="Cannot be empty")                                                
                 inpaint_prompt = gr.Textbox(label="Inpaint Prompt (if this is empty, then remove)", visible=False)
                 num_relation = gr.Slider(label="How many relations do you want to see", minimum=1, maximum=20, value=5, step=1, visible=False)
                 
@@ -946,6 +948,7 @@ if __name__ == "__main__":
                         <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>'
         gr.Markdown(DESCRIPTION)
 
+    print(f'device={device}')
     computer_info()
     block.launch(server_name='0.0.0.0', debug=args.debug, share=args.share)