Update modeling_omniparser.py

modeling_omniparser.py

@@ -48,14 +48,12 @@ class OmniparserModel(PreTrainedModel):
             repo_id="ThomasDh-C/OmniParser-v2.0",
             allow_patterns="icon_detect/**",
         )
-        print(icon_detect_path)
-        self.icon_detect_path = os.path.join(icon_detect_path, 'model.pt')
+        self.icon_detect_path = os.path.join(icon_detect_path, 'icon_detect', 'model.pt')
         icon_caption_path = snapshot_download(
             repo_id="ThomasDh-C/OmniParser-v2.0",
             allow_patterns="icon_caption/**",
         )
-        print(icon_caption_path)
-        caption_model_processor = self.get_caption_model_processor(model_name_or_path=icon_caption_path)
+        caption_model_processor = self.get_caption_model_processor(model_name_or_path=os.path.join(icon_caption_path, 'icon_caption'))
         self.caption_model, self.caption_model_processor = caption_model_processor['model'], caption_model_processor['processor']
         self.reader = easyocr.Reader(['en'])
         self.to_pil = ToPILImage()