Download models checks in automatic the existance of the models from app.py execution.

app.py

@@ -8,6 +8,9 @@ from pdf_processor import to_pdf, to_markdown
 from config import config
 from tts import text_to_speech  # Import TTS module
 
+
+
+
 # Set up logging with ANSI escape codes for colored output
 logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
 
@@ -117,4 +120,8 @@ with gr.Blocks() as demo:
     clear_button.add([file_input, md_render, pdf_display, md_text, output_file, is_ocr])
 
 if __name__ == "__main__":
+    import subprocess
+    print("Checking and downloading models if necessary...")
+    subprocess.run(["python", "download_models.py"])
+    print("Models are ready!")
     demo.launch(ssr_mode=True)

download_models_hf.py

@@ -6,61 +6,78 @@ from huggingface_hub import snapshot_download
 
 
 def download_json(url):
-    # 下载JSON文件
+    """Download JSON file from a URL."""
     response = requests.get(url)
-    response.raise_for_status()  # 检查请求是否成功
+    response.raise_for_status()
     return response.json()
 
 
 def download_and_modify_json(url, local_filename, modifications):
+    """Download and modify a JSON file if it doesn't exist or if it's outdated."""
     if os.path.exists(local_filename):
-        data = json.load(open(local_filename))
+        with open(local_filename, 'r', encoding='utf-8') as f:
+            data = json.load(f)
         config_version = data.get('config_version', '0.0.0')
+
+        # Only re-download if the version is outdated
         if config_version < '1.1.1':
             data = download_json(url)
     else:
         data = download_json(url)
 
-    # 修改内容
+    # Apply modifications
     for key, value in modifications.items():
         data[key] = value
 
-    # 保存修改后的内容
+    # Save modified JSON
     with open(local_filename, 'w', encoding='utf-8') as f:
         json.dump(data, f, ensure_ascii=False, indent=4)
 
 
-if __name__ == '__main__':
+def check_and_download_models():
+    """Download models only if they are not already present."""
+    model_base_path = os.path.expanduser("~/.cache/huggingface/hub")
+
+    mineru_models_path = os.path.join(model_base_path, "opendatalab__PDF-Extract-Kit-1.0")
+    layoutreader_models_path = os.path.join(model_base_path, "hantian__layoutreader")
+
+    # Check if models exist before downloading
+    if not os.path.exists(mineru_models_path):
+        mineru_patterns = [
+            "models/Layout/LayoutLMv3/*",
+            "models/Layout/YOLO/*",
+            "models/MFD/YOLO/*",
+            "models/MFR/unimernet_small_2501/*",
+            "models/TabRec/TableMaster/*",
+            "models/TabRec/StructEqTable/*",
+        ]
+        mineru_models_path = snapshot_download('opendatalab/PDF-Extract-Kit-1.0', allow_patterns=mineru_patterns)
+    
+    if not os.path.exists(layoutreader_models_path):
+        layoutreader_pattern = [
+            "*.json",
+            "*.safetensors",
+        ]
+        layoutreader_models_path = snapshot_download('hantian/layoutreader', allow_patterns=layoutreader_pattern)
 
-    mineru_patterns = [
-        "models/Layout/LayoutLMv3/*",
-        "models/Layout/YOLO/*",
-        "models/MFD/YOLO/*",
-        "models/MFR/unimernet_small_2501/*",
-        "models/TabRec/TableMaster/*",
-        "models/TabRec/StructEqTable/*",
-    ]
-    model_dir = snapshot_download('opendatalab/PDF-Extract-Kit-1.0', allow_patterns=mineru_patterns)
-
-    layoutreader_pattern = [
-        "*.json",
-        "*.safetensors",
-    ]
-    layoutreader_model_dir = snapshot_download('hantian/layoutreader', allow_patterns=layoutreader_pattern)
-
-    model_dir = model_dir + '/models'
-    print(f'model_dir is: {model_dir}')
-    print(f'layoutreader_model_dir is: {layoutreader_model_dir}')
+    # Print paths
+    print(f'model_dir is: {mineru_models_path}/models')
+    print(f'layoutreader_model_dir is: {layoutreader_models_path}')
 
+    # JSON configuration update
     json_url = 'https://github.com/opendatalab/MinerU/raw/master/magic-pdf.template.json'
     config_file_name = 'magic-pdf.json'
     home_dir = os.path.expanduser('~')
     config_file = os.path.join(home_dir, config_file_name)
 
     json_mods = {
-        'models-dir': model_dir,
-        'layoutreader-model-dir': layoutreader_model_dir,
+        'models-dir': f"{mineru_models_path}/models",
+        'layoutreader-model-dir': layoutreader_models_path,
     }
 
     download_and_modify_json(json_url, config_file, json_mods)
     print(f'The configuration file has been configured successfully, the path is: {config_file}')
+
+
+if __name__ == '__main__':
+    check_and_download_models()

tts.py

@@ -17,7 +17,7 @@ def text_to_speech(text: str, voice: str = "coral", model: str = "tts-1") -> str
         The file path to the generated audio file.
     """
     # Generate a unique filename using a hash of the text
-    output_file = Path(__file__).parent / f"speech_{abs(hash(text))}.pus"
+    output_file = Path(__file__).parent / 'output' / f"speech_{abs(hash(text))}.pus"
     try:
         response = openai.Audio.speech.create(
             model=model,
@@ -38,7 +38,7 @@ def text_to_speech_gtts(text: str) -> str:
     Returns:
         The file path to the generated audio file.
     """
-    output_file = Path(__file__).parent / f"speech_{abs(hash(text))}.mp3"
+    output_file = Path(__file__).parent / 'output' / f"speech_{abs(hash(text))}.mp3"
     try:
         tts = gTTS(text=text, lang='en')
         tts.save(str(output_file))