更新自动 florence 模型

util/download_weights.py

@@ -2,40 +2,27 @@ import os
 from pathlib import Path
 from modelscope import snapshot_download
 __WEIGHTS_DIR = Path("weights")
-MODEL_DIR = os.path.join(__WEIGHTS_DIR, "AI-ModelScope", "OmniParser-v2___0")
-def download():
+OMNI_PARSER_MODEL_DIR = os.path.join(__WEIGHTS_DIR, "AI-ModelScope", "OmniParser-v2___0")
+FLORENCE_MODEL_DIR = os.path.join(__WEIGHTS_DIR, "AI-ModelScope", "Florence-2-base-ft")
+
+def __download_omni_parser():
     # Create weights directory
    
     __WEIGHTS_DIR.mkdir(exist_ok=True)
     
-    # List of files to download
-    files = [
-        "icon_detect/train_args.yaml",
-        "icon_detect/model.pt",
-        "icon_detect/model.yaml",
-        "icon_caption/config.json",
-        "icon_caption/generation_config.json",
-        "icon_caption/model.safetensors"
-    ]
-    
-    # Check and download missing files
-    missing_files = []
-    for file in files:
-        file_path = os.path.join(MODEL_DIR, file)
-        if not os.path.exists(file_path):
-            missing_files.append(file)
-            break
-    
-    if not missing_files:
-        print("Model files already detected!")
-        return
-    
     snapshot_download(
         'AI-ModelScope/OmniParser-v2.0',
         cache_dir='weights'
         )
 
-    print("Download complete")
+def __download_florence_model():
+    snapshot_download('AI-ModelScope/Florence-2-base-ft',
+                      cache_dir='weights'
+                      )
+
+def download_models():
+    __download_omni_parser()
+    __download_florence_model()
 
 if __name__ == "__main__":
-    download()
+    download_models()