update app.py

app.py

@@ -58,6 +58,26 @@ from diffusers import StableDiffusionInpaintPipeline
 from huggingface_hub import hf_hub_download
 
 from  utils 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
+from mmengine.config import Config as mmengine_Config
+
+config_file = 'GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py'
+ckpt_repo_id = "ShilongLiu/GroundingDINO"
+ckpt_filenmae = "groundingdino_swint_ogc.pth"
+sam_checkpoint = './sam_vit_h_4b8939.pth' 
+output_dir = "outputs"
+device = 'cpu'
+
+os.makedirs(output_dir, exist_ok=True)
+groundingdino_model = None
+sam_model = None
+sam_predictor = None
+sam_mask_generator = None
+sd_pipe = None
+lama_cleaner_model= None
+ram_model = None
 
 def load_model_hf(model_config_path, repo_id, filename, device='cpu'):
     args = SLConfig.fromfile(model_config_path) 
@@ -224,55 +244,50 @@ def mix_masks(imgs):
     re_img =  1 - re_img
     return  Image.fromarray(np.uint8(255*re_img))
 
-config_file = 'GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py'
-ckpt_repo_id = "ShilongLiu/GroundingDINO"
-ckpt_filenmae = "groundingdino_swint_ogc.pth"
-sam_checkpoint = './sam_vit_h_4b8939.pth' 
-output_dir = "outputs"
-device = evice = 'cuda' if torch.cuda.is_available() else 'cpu'
-
-print(f'device={device}')
-
-# make dir
-os.makedirs(output_dir, exist_ok=True)
-
-# initialize groundingdino model
-logger.info(f"initialize groundingdino model...")
-groundingdino_model = load_model_hf(config_file, ckpt_repo_id, ckpt_filenmae)
-
-# initialize SAM
-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)
-
-os.system("pip list")
-# initialize stable-diffusion-inpainting
-logger.info(f"initialize stable-diffusion-inpainting...")
-sd_pipe = None
-if os.environ.get('IS_MY_DEBUG') is None:
-    sd_pipe = StableDiffusionInpaintPipeline.from_pretrained(
-            "runwayml/stable-diffusion-inpainting", 
-            # revision="fp16",
-            # "stabilityai/stable-diffusion-2-inpainting",
-            torch_dtype=torch.float16,
-    )
-    sd_pipe = sd_pipe.to(device)
-
-# initialize lama_cleaner
-logger.info(f"initialize lama_cleaner...")
-from lama_cleaner.helper import (
-    load_img,
-    numpy_to_bytes,
-    resize_max_size,
-)
-
-lama_cleaner_model = ModelManager(
-        name='lama',
-        device='cpu', # device,
+def set_device():
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    print(f'device={device}')
+
+def load_groundingdino_model():
+    # initialize groundingdino model
+    logger.info(f"initialize groundingdino model...")
+    groundingdino_model = load_model_hf(config_file, ckpt_repo_id, ckpt_filenmae)
+
+def load_sam_model():
+    # initialize SAM
+    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():
+    # initialize stable-diffusion-inpainting
+    logger.info(f"initialize stable-diffusion-inpainting...")
+    sd_pipe = None
+    if os.environ.get('IS_MY_DEBUG') is None:
+        sd_pipe = StableDiffusionInpaintPipeline.from_pretrained(
+                "runwayml/stable-diffusion-inpainting", 
+                # revision="fp16",
+                # "stabilityai/stable-diffusion-2-inpainting",
+                torch_dtype=torch.float16,
+        )
+        sd_pipe = sd_pipe.to(device)
+
+def load_lama_cleaner_model():
+    # initialize lama_cleaner
+    logger.info(f"initialize lama_cleaner...")
+    from lama_cleaner.helper import (
+        load_img,
+        numpy_to_bytes,
+        resize_max_size,
     )
 
+    lama_cleaner_model = ModelManager(
+            name='lama',
+            device='cpu', # device,
+        )
+
 def lama_cleaner_process(image, mask):
     ori_image = image
     if mask.shape[0] == image.shape[1] and mask.shape[1] == image.shape[0] and mask.shape[0] != mask.shape[1]:
@@ -330,32 +345,6 @@ def lama_cleaner_process(image, mask):
     image = Image.open(io.BytesIO(numpy_to_bytes(res_np_img, 'png')))
     return  image
 
-# 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
-from mmengine.config import Config as mmengine_Config
-input_size = 512
-hidden_size = 256
-num_classes = 56
-
-# load ram model
-model_path = "./checkpoints/ram_epoch12.pth"
-ram_config = dict(
-    model=dict(
-        pretrained_model_name_or_path='bert-base-uncased',
-        load_pretrained_weights=False,
-        num_transformer_layer=2,
-        input_feature_size=256,
-        output_feature_size=768,
-        cls_feature_size=512,
-        num_relation_classes=56,
-        pred_type='attention',
-        loss_type='multi_label_ce',
-    ),
-    load_from=model_path,
-)
-ram_config = mmengine_Config(ram_config)
-
 class Ram_Predictor(RamPredictor):
     def __init__(self, config, device='cpu'):
         self.config = config
@@ -368,7 +357,25 @@ class Ram_Predictor(RamPredictor):
             self.model.load_state_dict(torch.load(self.config.load_from, map_location=self.device))
         self.model.train()
 
-ram_model = Ram_Predictor(ram_config, device)
+def load_ram_model():
+    # load ram model
+    model_path = "./checkpoints/ram_epoch12.pth"
+    ram_config = dict(
+        model=dict(
+            pretrained_model_name_or_path='bert-base-uncased',
+            load_pretrained_weights=False,
+            num_transformer_layer=2,
+            input_feature_size=256,
+            output_feature_size=768,
+            cls_feature_size=512,
+            num_relation_classes=56,
+            pred_type='attention',
+            loss_type='multi_label_ce',
+        ),
+        load_from=model_path,
+    )
+    ram_config = mmengine_Config(ram_config)
+    ram_model = Ram_Predictor(ram_config, device)
 
 # visualization
 def draw_selected_mask(mask, draw):
@@ -531,7 +538,6 @@ def run_anything_task(input_image, text_prompt, task_type, inpaint_prompt, box_t
             except:
                 warnings.warn("Failed to load custom C++ ops. Running on CPU mode Only in groundingdino!")
 
-        groundingdino_device = 'cpu'
         boxes_filt, pred_phrases = get_grounding_output(
             groundingdino_model, image, text_prompt, box_threshold, text_threshold, device=groundingdino_device
         )
@@ -676,6 +682,15 @@ if __name__ == "__main__":
 
     print(f'args = {args}')
 
+    os.system("pip list")
+
+    set_device()
+    load_groundingdino_model()
+    load_sam_model()
+    load_sd_model()
+    load_lama_cleaner_model()
+    load_ram_model()
+
     block = gr.Blocks().queue()
     with block:
         with gr.Row():