feat: update model downloading from Azure Blob Storage and environment configuration

.env

@@ -0,0 +1,33 @@
+# OpenAI and API Keys
+OPENAI_API_KEY=your_openai_api_key
+TAVILY_API_KEY=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
+LANGCHAIN_API_KEY=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
+LANGCHAIN_TRACING_V2=true
+LANGCHAIN_PROJECT=P&ID agent
+LANGSMITH_API_KEY=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
+
+# Storage Configuration
+STORAGE_TYPE=local
+AZURE_STORAGE_CONTAINER_NAME='pnid-models'
+
+# Azure Storage Connection Strings
+AZURE_STORAGE_PRIMARY_CONNECTION='DefaultEndpointsProtocol=https;AccountName=pidmodelstorage019659c7a;AccountKey=PYE6QujpTMkciPKLMDsNQzDqCIn331on1cv4MvpRsBPC8ud84jeIVJQoargLuigoN6kdjAICi0u1+AStkTuRSg==;EndpointSuffix=core.windows.net'
+AZURE_STORAGE_SECONDARY_CONNECTION='DefaultEndpointsProtocol=https;AccountName=pidmodelstorage019659c7a;AccountKey=6DS5ApT8Z7pw7tSjavt+S83mX/DJquxf5NuEB2Q2svy9OyubBYh/Hog9jz1wWdKZID8cyul2ecrL+AStVNF9lQ==;EndpointSuffix=core.windows.net'
+
+# Runtime Settings
+USE_TORCH=1
+GRADIO_SERVER_PORT=7861
+
+# Model Paths
+MODEL_SDM_41_PATH=models/Intui_SDM_41.pt
+MODEL_SDM_30_PATH=models/Intui_SDM_30.pt
+MODEL_SDM_20_PATH=models/Intui_SDM_20.pt
+MODEL_DEEPLSD_PATH=models/deeplsd_md.tar
+MODEL_CRAFT_PATH=models/craft_mlt_25k.pth
+MODEL_ENGLISH_PATH=models/english_g2.pth
+MODEL_LDM_PATH=models/intui_LDM_01.pt
+
+# Legacy Model Paths (can be removed if not used)
+# YOLO_MODEL_PATH=models/yolo/yolov8n.pt
+# DEEPLSD_MODEL_PATH=models/deeplsd/deeplsd_md.tar
+# DOCTR_MODEL_PATH=models/doctr/ocr_predictor.pt

download_models.py

@@ -1,49 +1,58 @@
 import os
-import shutil
-import torch
-from doctr.models import ocr_predictor
-from ultralytics import YOLO
-from deeplsd.models.deeplsd_inference import DeepLSD
+from azure.storage.blob import BlobServiceClient
 from dotenv import load_dotenv
 
 # Load environment variables
 load_dotenv()
 
-def copy_local_models():
-    """Copy models from local directory to deployment"""
-    # Create model directories
-    os.makedirs('models/yolo', exist_ok=True)
-    os.makedirs('models/deeplsd', exist_ok=True)
-    os.makedirs('models/doctr', exist_ok=True)
-    
-    # Source paths (adjust these to your local paths)
-    local_models_dir = "../models"
-    
-    # Copy YOLO model
-    yolo_src = os.path.join(local_models_dir, "yolov8n.pt")
-    yolo_dst = "models/yolo/yolov8n.pt"
-    if os.path.exists(yolo_src):
-        shutil.copy2(yolo_src, yolo_dst)
-        print(f"Copied YOLO model to {yolo_dst}")
-    
-    # Copy DeepLSD model
-    deeplsd_src = os.path.join(local_models_dir, "deeplsd_md.tar")
-    deeplsd_dst = "models/deeplsd/deeplsd_md.tar"
-    if os.path.exists(deeplsd_src):
-        shutil.copy2(deeplsd_src, deeplsd_dst)
-        print(f"Copied DeepLSD model to {deeplsd_dst}")
-    
-    # Copy DocTR model if exists
-    doctr_src = os.path.join(local_models_dir, "ocr_predictor.pt")
-    doctr_dst = "models/doctr/ocr_predictor.pt"
-    if os.path.exists(doctr_src):
-        shutil.copy2(doctr_src, doctr_dst)
-        print(f"Copied DocTR model to {doctr_dst}")
-    else:
-        # Download DocTR model if not available locally
-        predictor = ocr_predictor(pretrained=True)
-        torch.save(predictor.state_dict(), doctr_dst)
-        print(f"Downloaded DocTR model to {doctr_dst}")
+def download_from_azure():
+    """Download models from Azure Blob Storage"""
+    try:
+        # Try primary connection first, fallback to secondary
+        connect_str = os.getenv('AZURE_STORAGE_PRIMARY_CONNECTION')
+        if not connect_str:
+            connect_str = os.getenv('AZURE_STORAGE_SECONDARY_CONNECTION')
+            
+        container_name = os.getenv('AZURE_STORAGE_CONTAINER_NAME', 'pnid-models')
+
+        # Create the BlobServiceClient
+        blob_service_client = BlobServiceClient.from_connection_string(connect_str)
+        container_client = blob_service_client.get_container_client(container_name)
+
+        # Create models directory
+        os.makedirs('models', exist_ok=True)
+
+        # Define model files to download
+        model_files = {
+            'Intui_SDM_41.pt': os.getenv('MODEL_SDM_41_PATH'),
+            'Intui_SDM_30.pt': os.getenv('MODEL_SDM_30_PATH'),
+            'Intui_SDM_20.pt': os.getenv('MODEL_SDM_20_PATH'),
+            'deeplsd_md.tar': os.getenv('MODEL_DEEPLSD_PATH'),
+            'craft_mlt_25k.pth': os.getenv('MODEL_CRAFT_PATH'),
+            'english_g2.pth': os.getenv('MODEL_ENGLISH_PATH'),
+            'intui_LDM_01.pt': os.getenv('MODEL_LDM_PATH')
+        }
+
+        # Download each model
+        for blob_name, local_path in model_files.items():
+            if not local_path:
+                continue
+                
+            print(f"Downloading {blob_name}...")
+            os.makedirs(os.path.dirname(local_path), exist_ok=True)
+            
+            if not os.path.exists(local_path) or os.path.getsize(local_path) == 0:
+                blob_client = container_client.get_blob_client(blob_name)
+                with open(local_path, "wb") as file:
+                    data = blob_client.download_blob()
+                    file.write(data.readall())
+                print(f"Downloaded {blob_name} to {local_path}")
+            else:
+                print(f"Skipping {blob_name}, already exists")
+
+    except Exception as e:
+        print(f"Error downloading models from Azure: {str(e)}")
+        raise
 
 if __name__ == "__main__":
-    copy_local_models()
+    download_from_azure()

gradioChatApp.py

@@ -25,7 +25,7 @@ import cv2
 import numpy as np
 import time
 from huggingface_hub import HfApi, login
-from download_models import download_models, copy_local_models
+from download_models import download_from_azure
 
 # Load environment variables from .env file
 load_dotenv()
@@ -789,9 +789,17 @@ def create_ui():
 
 
 def main():
-    # Download models if they don't exist
-    if not os.path.exists('models/yolo/yolov8n.pt'):
-        copy_local_models()
+    # Check for all required models
+    required_models = [
+        'models/yolo/yolov8n.pt',
+        'models/deeplsd/deeplsd_md.tar',
+        'models/doctr/craft_mlt_25k.pth',
+        'models/doctr/english_g2.pth',
+        'models/yolo/intui_LDM_01.pt'
+    ]
+    
+    if any(not os.path.exists(model) for model in required_models):
+        download_from_azure()
 
     demo = create_ui()
     # Remove HF Spaces conditional, just use local development settings