Update app.py

app.py

@@ -8,9 +8,6 @@ import gradio as gr
 
 from loguru import logger
 
-# os.system("pip install diffuser==0.6.0")
-# os.system("pip install transformers==4.29.1")
-
 os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 
 if os.environ.get('IS_MY_DEBUG') is None:
@@ -69,7 +66,10 @@ ckpt_repo_id = "ShilongLiu/GroundingDINO"
 ckpt_filenmae = "groundingdino_swint_ogc.pth"
 sam_checkpoint = './sam_vit_h_4b8939.pth' 
 output_dir = "outputs"
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
+if os.environ.get('IS_MY_DEBUG') is None:
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+else:
+    device = 'cpu'
 
 os.makedirs(output_dir, exist_ok=True)
 groundingdino_model = None
@@ -77,8 +77,9 @@ sam_device = None
 sam_model = None
 sam_predictor = None
 sam_mask_generator = None
-sd_pipe = None
+sd_model = None
 lama_cleaner_model= None
+lama_cleaner_model_device = device
 ram_model = None
 
 def get_sam_vit_h_4b8939():
@@ -165,16 +166,6 @@ def load_image(image_path):
     image, _ = transform(image_pil, None)  # 3, h, w
     return image_pil, image
 
-def load_model(model_config_path, model_checkpoint_path, device):
-    args = SLConfig.fromfile(model_config_path)
-    args.device = device
-    model = build_model(args)
-    checkpoint = torch.load(model_checkpoint_path, map_location=device) #"cpu")
-    load_res = model.load_state_dict(clean_state_dict(checkpoint["model"]), strict=False)
-    print(load_res)
-    _ = model.eval()
-    return model
-
 def get_grounding_output(model, image, caption, box_threshold, text_threshold, with_logits=True, device="cpu"):
     caption = caption.lower()
     caption = caption.strip()
@@ -258,18 +249,21 @@ def mix_masks(imgs):
     return  Image.fromarray(np.uint8(255*re_img))
 
 def set_device():
-    device = 'cuda' if torch.cuda.is_available() else 'cpu'
-    print(f'device={device}')
+    global device
+    if os.environ.get('IS_MY_DEBUG') is None:
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    else:
+        device = 'cpu'
 
 def load_groundingdino_model():
     # initialize groundingdino model
     global groundingdino_model
     logger.info(f"initialize groundingdino model...")
-    groundingdino_model = load_model_hf(config_file, ckpt_repo_id, ckpt_filenmae)
+    groundingdino_model = load_model_hf(config_file, ckpt_repo_id, ckpt_filenmae, device='cpu')
 
 def load_sam_model():
     # initialize SAM
-    global sam_model, sam_predictor, sam_mask_generator, sam_device
+    global sam_model, sam_predictor, sam_mask_generator, sam_device, device
     logger.info(f"initialize SAM model...")
     sam_device = device
     sam_model = build_sam(checkpoint=sam_checkpoint).to(sam_device)
@@ -278,26 +272,26 @@ def load_sam_model():
 
 def load_sd_model():
     # initialize stable-diffusion-inpainting
-    global sd_pipe
+    global sd_model, device
     logger.info(f"initialize stable-diffusion-inpainting...")
-    sd_pipe = None
+    sd_model = None
     if os.environ.get('IS_MY_DEBUG') is None:
-        sd_pipe = StableDiffusionInpaintPipeline.from_pretrained(
+        sd_model = StableDiffusionInpaintPipeline.from_pretrained(
                 "runwayml/stable-diffusion-inpainting", 
                 revision="fp16",
                 # "stabilityai/stable-diffusion-2-inpainting",
                 torch_dtype=torch.float16,
         )
-        sd_pipe = sd_pipe.to(device)
+        sd_model = sd_model.to(device)
 
 def load_lama_cleaner_model():
     # initialize lama_cleaner
-    global lama_cleaner_model
+    global lama_cleaner_model, device
     logger.info(f"initialize lama_cleaner...")
 
     lama_cleaner_model = ModelManager(
             name='lama',
-            device='cpu', # device,
+            device=lama_cleaner_model_device,
         )
 
 def lama_cleaner_process(image, mask, cleaner_size_limit=1080):
@@ -517,6 +511,7 @@ mask_source_segment = "type what to detect below"
 
 def run_anything_task(input_image, text_prompt, task_type, inpaint_prompt, box_threshold, text_threshold, 
             iou_threshold, inpaint_mode, mask_source_radio, remove_mode, remove_mask_extend, num_relation, cleaner_size_limit=1080):
+    
     if (task_type == 'relate anything'):
         output_images = relate_anything(input_image['image'], num_relation)
         return output_images, gr.Gallery.update(label='relate images')
@@ -566,7 +561,7 @@ def run_anything_task(input_image, text_prompt, task_type, inpaint_prompt, box_t
             groundingdino_model, image, text_prompt, box_threshold, text_threshold, device=groundingdino_device
         )
         if boxes_filt.size(0) == 0:
-            logger.info(f'run_anything_task_[{file_temp}]_{task_type}_[{text_prompt}]_1_[No objects detected, please try others.]_')
+            logger.info(f'run_anything_task_[{file_temp}]_{task_type}_[{text_prompt}]_1___{groundingdino_device}/[No objects detected, please try others.]_')
             return [], gr.Gallery.update(label='No objects detected, please try others.😂😂😂😂')
         boxes_filt_ori = copy.deepcopy(boxes_filt)
 
@@ -640,7 +635,7 @@ def run_anything_task(input_image, text_prompt, task_type, inpaint_prompt, box_t
             # inpainting pipeline
             image_source_for_inpaint = image_pil.resize((512, 512))
             image_mask_for_inpaint = mask_pil.resize((512, 512))
-            image_inpainting = sd_pipe(prompt=inpaint_prompt, image=image_source_for_inpaint, mask_image=image_mask_for_inpaint).images[0]
+            image_inpainting = sd_model(prompt=inpaint_prompt, image=image_source_for_inpaint, mask_image=image_mask_for_inpaint).images[0]
         else:
             # remove from mask
             logger.info(f'run_anything_task_[{file_temp}]_{task_type}_5_')
@@ -707,6 +702,8 @@ def change_radio_display(task_type, mask_source_radio):
 
 def get_model_device(module):
     try:
+        if module is None:
+            return 'None'
         if isinstance(module, torch.nn.DataParallel):
             module = module.module
         for submodule in module.children():
@@ -714,8 +711,9 @@ def get_model_device(module):
                 parameters = submodule._parameters
                 if "weight" in parameters:
                     return parameters["weight"].device
+        return 'UnKnown'
     except Exception as e:
-        return 'ohoh'
+        return 'Error'
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser("Grounded SAM demo", add_help=True)
@@ -732,10 +730,12 @@ if __name__ == "__main__":
     load_lama_cleaner_model()
     load_ram_model()
 
-    os.system("pip list")
+    if os.environ.get('IS_MY_DEBUG') is None:
+        os.system("pip list")
+
     print(f'groundingdino_model__{get_model_device(groundingdino_model)}')
     print(f'sam_model__{get_model_device(sam_model)}')
-    print(f'sd_model__{get_model_device(sd_pipe)}')
+    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)}')
     
@@ -790,3 +790,4 @@ if __name__ == "__main__":
 
     computer_info()
     block.launch(server_name='0.0.0.0', debug=args.debug, share=args.share)
+