Spaces:
Sleeping
Sleeping
| # Model loading utilities - Deployment Optimized | |
| import streamlit as st | |
| import whisper | |
| import easyocr | |
| import os | |
| from sentence_transformers import SentenceTransformer | |
| from transformers import CLIPProcessor, CLIPModel | |
| from gpt4all import GPT4All | |
| from config import MODELS_CONFIG, DEPLOYMENT_MODE | |
| from device_setup import setup_device | |
| def load_models(): | |
| try: | |
| # Writable caches first (prevent /.cache writes) | |
| import pathlib | |
| cache_root = os.environ.get("HF_HOME", "/tmp/hf") | |
| pathlib.Path(cache_root).mkdir(parents=True, exist_ok=True) | |
| # Make sure all relevant env vars point to the same place | |
| os.environ.setdefault("HF_HOME", cache_root) | |
| os.environ.setdefault("TRANSFORMERS_CACHE", cache_root) | |
| os.environ.setdefault("HUGGINGFACE_HUB_CACHE", cache_root) | |
| os.environ.setdefault("HF_DATASETS_CACHE", cache_root) | |
| os.environ.setdefault("SENTENCE_TRANSFORMERS_HOME", cache_root) | |
| os.environ.setdefault("XDG_CACHE_HOME", "/tmp") | |
| # Setup device | |
| device, has_gpu = setup_device() | |
| device_str = "cuda" if has_gpu else "cpu" | |
| progress = st.progress(0, text="Initializing models...") | |
| # Text embeddings (explicit cache folder) | |
| progress.progress(0.1, text="Loading text encoder...") | |
| st.write("📝 Loading text encoder...") | |
| text_encoder = SentenceTransformer( | |
| MODELS_CONFIG["text_encoder"], | |
| device=device, | |
| cache_folder=cache_root | |
| ) | |
| # CLIP (route Hugging Face cache explicitly) | |
| progress.progress(0.3, text="Loading CLIP model...") | |
| st.write("🖼️ Loading CLIP model...") | |
| clip_model_name = MODELS_CONFIG["clip_model"] | |
| clip_model = CLIPModel.from_pretrained(clip_model_name, cache_dir=cache_root).to(device) | |
| clip_processor = CLIPProcessor.from_pretrained(clip_model_name, cache_dir=cache_root) | |
| # Whisper (force CPU on free tier unless GPU present) | |
| progress.progress(0.5, text="Loading Whisper...") | |
| st.write("🎤 Loading Whisper model...") | |
| whisper_size = MODELS_CONFIG["whisper_base"] | |
| whisper_model = whisper.load_model(whisper_size, device=device_str) | |
| # EasyOCR (avoid CUDA probing on CPU Spaces) | |
| progress.progress(0.7, text="Loading OCR...") | |
| st.write("🔍 Loading OCR reader...") | |
| ocr_reader = easyocr.Reader(['en'], gpu=False if not has_gpu else True) | |
| # Local LLM | |
| progress.progress(0.9, text="Loading LLM...") | |
| st.write("🧠 Loading local LLM...") | |
| try: | |
| llm = GPT4All(MODELS_CONFIG["llm_model"], allow_download=True) | |
| st.write("✅ LLM loaded successfully!") | |
| except Exception as llm_error: | |
| st.warning(f"⚠️ LLM loading failed: {llm_error}") | |
| st.info("💡 The system will work without local LLM for document search") | |
| llm = None | |
| progress.progress(1.0, text="All models loaded!") | |
| st.success("🎉 All models loaded successfully!") | |
| return text_encoder, clip_model, clip_processor, whisper_model, ocr_reader, llm, device, has_gpu | |
| except Exception as e: | |
| st.error(f"❌ Error loading models: {e}") | |
| st.info("💡 Try refreshing the page or check your internet connection") | |
| return None, None, None, None, None, None, None, False | |
| def get_model_info(): | |
| '''Get information about loaded models''' | |
| return { | |
| "text_encoder": MODELS_CONFIG["text_encoder"], | |
| "clip_model": MODELS_CONFIG["clip_model"], | |
| "whisper_model": MODELS_CONFIG["whisper_base"], | |
| "llm_model": MODELS_CONFIG["llm_model"], | |
| "deployment_mode": DEPLOYMENT_MODE | |
| } | |