update app.py && app_cli.py

app.py

@@ -15,7 +15,6 @@ if run_gradio:
     os.system("pip install gradio==3.50.2")
 
 import gradio as gr
-
 from loguru import logger
 
 os.environ["CUDA_VISIBLE_DEVICES"] = "0"
@@ -24,6 +23,8 @@ if os.environ.get('IS_MY_DEBUG') is None:
     result = subprocess.run(['pip', 'install', '-e', 'GroundingDINO'], check=True)
     print(f'pip install GroundingDINO = {result}')
 
+logger.info(f"Start app...")
+
 # result = subprocess.run(['pip', 'list'], check=True)
 # print(f'pip list = {result}')
 
@@ -188,7 +189,6 @@ def plot_boxes_to_image(image_pil, tgt):
 
         mask_draw.rectangle([x0, y0, x1, y1], fill=255, width=6)
 
-
     return image_pil, mask
 
 def load_image(image_path):
@@ -290,19 +290,20 @@ def mix_masks(imgs):
     re_img =  1 - re_img
     return  Image.fromarray(np.uint8(255*re_img))
 
-def set_device():
+def set_device(args):
+    global device
     if os.environ.get('IS_MY_DEBUG') is None:
-        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        device = args.cuda if torch.cuda.is_available() else 'cpu'
     else:
         device = 'cpu'
     print(f'device={device}')
-    return device
 
 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=device) #'cpu')
-    return groundingdino_model
+    logger.info(f"initialize groundingdino model...{type(groundingdino_model)}")
 
 def get_sam_vit_h_4b8939():
     if not os.path.exists('./sam_vit_h_4b8939.pth'):
@@ -1104,13 +1105,14 @@ if __name__ == "__main__":
     parser.add_argument("--debug", action="store_true", help="using debug mode")
     parser.add_argument("--share", action="store_true", help="share the app")
     parser.add_argument("--port", "-p", type=int, default=7860, help="port")
+    parser.add_argument("--cuda", "-c", type=str, default='cuda:0', help="cuda")
     args, _ = parser.parse_known_args()
     print(f'args = {args}')
 
     if os.environ.get('IS_MY_DEBUG') is None:
         os.system("pip list")
     
-    device = set_device()
+    set_device(args)
     if device == 'cpu':
         kosmos_enable = False
 
@@ -1118,7 +1120,7 @@ if __name__ == "__main__":
         kosmos_model, kosmos_processor = load_kosmos_model(device)
         
     if groundingdino_enable:
-        groundingdino_model = load_groundingdino_model('cpu')
+        load_groundingdino_model('cpu')
     
     if sam_enable:
         load_sam_model(device)

app_cli.py

@@ -53,7 +53,7 @@ from io import BytesIO
 from diffusers import StableDiffusionInpaintPipeline
 from huggingface_hub import hf_hub_download
 
-from  utils import computer_info
+from util_computer import computer_info
 # relate anything
 from ram_utils import iou, sort_and_deduplicate, relation_classes, MLP, show_anns, ram_show_mask
 from ram_train_eval import RamModel,RamPredictor
@@ -85,6 +85,7 @@ def get_args():
     argparser.add_argument("--input_image", "-i", type=str, default="", help="")
     argparser.add_argument("--text", "-t", type=str, default="", help="")
     argparser.add_argument("--output_image", "-o", type=str, default="", help="")
+    argparser.add_argument("--cuda", "-c", type=str, default='cpu', help="cuda")
     args = argparser.parse_args()
     return args
 
@@ -96,8 +97,8 @@ if __name__ == '__main__':
     logger.info(f'\nargs={args}\n')
 
     logger.info(f'loading models ... ')
-    # set_device()  # If you have enough GPUs, you can open this comment
-    groundingdino_model = load_groundingdino_model('cpu')
+    # set_device(args)  # If you have enough GPUs, you can open this comment
+    load_groundingdino_model('cpu')
     load_sam_model(device)
     # load_sd_model(device)
     load_lama_cleaner_model(device)
@@ -105,7 +106,7 @@ if __name__ == '__main__':
 
     input_image = Image.open(args.input_image)
 
-    output_images, _ = run_anything_task(input_image = input_image, 
+    run_rets = run_anything_task(input_image = input_image, 
                         text_prompt = args.text,  
                         task_type = 'remove', 
                         inpaint_prompt = '', 
@@ -120,6 +121,7 @@ if __name__ == '__main__':
                         kosmos_input = None,
                         cleaner_size_limit = -1,
                         )
+    output_images = run_rets[0]
     if len(output_images) > 0:
         logger.info(f'save result to {args.output_image} ... ')        
         output_images[-1].save(args.output_image)