from pathlib import Path from typing import Any import shutil from .model import FFTCNN # Import the FFT CNN architecture (package-relative) from config import Config try: from huggingface_hub import hf_hub_download except Exception: hf_hub_download = None # NOTE: EfficientNet/nn imports are done lazily when torch is available. ELAForgeryNet = None # will be constructed dynamically when needed torch = None TORCH_AVAILABLE = False class ModelLoader: """A class to load and hold PyTorch models used by this feature. It loads: - an FFT-based CNN (downloaded from Hugging Face Hub) - an ELA-based document forgery detector (local .pth by default) """ def __init__(self, model_repo_id: str, model_filename: str, doc_model_path: str = None): # Try to import torch once and expose module-level variables global torch, TORCH_AVAILABLE try: import torch as _torch torch = _torch TORCH_AVAILABLE = True except Exception: torch = None TORCH_AVAILABLE = False print("[WARN] PyTorch not available; model loading will be skipped until torch is installed.") if TORCH_AVAILABLE: self.device = "cuda" if torch.cuda.is_available() else "cpu" else: self.device = "cpu" print(f"Using device: {self.device} (torch available: {TORCH_AVAILABLE})") # Load FFT CNN from HF Hub self.fft_model = None if TORCH_AVAILABLE: try: self.fft_model = self._load_fft_model(repo_id=model_repo_id, filename=model_filename) print("FFT CNN model loaded successfully from Hub.") except Exception: # Try local fallback path (if provided in config) self.fft_model = None local_path = Path(getattr(Config, 'REAL_FORGED_MODEL_LOCAL_PATH', '')) if local_path and local_path.exists(): try: print(f"Attempting to load FFT model from local path: {local_path}") model = FFTCNN() state = torch.load(str(local_path), map_location=torch.device(self.device)) state_dict = state.get('state_dict', state) if isinstance(state, dict) else state model.load_state_dict(state_dict, strict=False) model.to(self.device) model.eval() self.fft_model = model print("FFT CNN model loaded successfully from local path.") except Exception as e: print(f"Failed to load local FFT model: {e}") else: print("No local FFT model path configured or file missing; FFT model not loaded.") else: print("Skipping FFT model load because PyTorch is not installed.") # Load document forgery model (ELA CNN), downloading the checkpoint if needed. self.doc_model = None if doc_model_path is None: doc_model_path = Config.DOCUMENT_FORGERY_MODEL_PATH self.doc_model = None if TORCH_AVAILABLE: try: self.doc_model = self._load_document_forgery_model(Path(doc_model_path)) if self.doc_model is not None: print("Document forgery (ELA) model loaded successfully.") except Exception as e: print(f"Warning: failed to load document forgery model: {e}") else: print("Skipping document forgery model load because PyTorch is not installed.") def _load_fft_model(self, repo_id: str, filename: str): """Downloads and loads the FFT CNN model from a Hugging Face Hub repository.""" print(f"Attempting to download FFT CNN model from Hugging Face repo: {repo_id}") try: from huggingface_hub import hf_hub_download except Exception as e: raise RuntimeError(f"huggingface_hub not available: {e}") try: model_path = hf_hub_download(repo_id=repo_id, filename=filename, token=Config.HF_TOKEN) print(f"Model downloaded to: {model_path}") model = FFTCNN() model.load_state_dict(torch.load(model_path, map_location=torch.device(self.device))) model.to(self.device) model.eval() return model except Exception as e: print(f"Error downloading or loading FFT model from Hugging Face: {e}") raise def _load_document_forgery_model(self, path: Path): """Load the ELA-based document forgery model from a local .pth checkpoint. Returns the model instance or None if the file does not exist. """ # If the configured path doesn't exist, try sensible fallbacks in the repo. if not path.exists(): print(f"Document forgery model file not found at configured path: {path}") # 1) Try the configured document forgery checkpoint path relative to repo root repo_root = Path(__file__).resolve().parents[2] candidate = repo_root / 'features' / 'Model' / 'document_forgery' / path.name if candidate.exists(): path = candidate print(f"Found document forgery model at fallback path: {path}") else: # 2) Search the repo for any file with the configured checkpoint name print(f"Searching repository for '{path.name}'...") matches = list(repo_root.rglob(path.name)) if matches: path = matches[0] print(f"Found document forgery model at: {path}") else: try: path = self._download_document_forgery_model(path) except Exception as exc: print(f"Document forgery model not found in repository and download failed: {exc}") return None print(f"Loading document forgery model from: {path}") # Build the ELA model architecture lazily (requires torchvision & torch.nn) try: import torchvision.models as tv_models import torch.nn as nn except Exception as e: raise RuntimeError(f"Required packages for ELA model not available: {e}") backbone = tv_models.efficientnet_b0(weights='IMAGENET1K_V1') in_features = backbone.classifier[1].in_features backbone.classifier = nn.Sequential( nn.Dropout(p=0.4), nn.Linear(in_features, 256), nn.ReLU(inplace=True), nn.Dropout(p=0.2), nn.Linear(256, 2), ) model = backbone state = torch.load(str(path), map_location=torch.device(self.device)) # The checkpoint might be either a state_dict or a full checkpoint dict if isinstance(state, dict) and 'state_dict' in state: state_dict = state['state_dict'] else: state_dict = state # Attempt to load state dict; allow strict=False to be tolerant to minor key name differences model.load_state_dict(state_dict, strict=False) model.to(self.device) model.eval() return model def _download_document_forgery_model(self, target_path: Path) -> Path: """Download the document forgery checkpoint into the configured local path.""" if hf_hub_download is None: raise RuntimeError("huggingface_hub not available") repo_id = getattr(Config, "DOCUMENT_FORGERY_MODEL_REPO_ID", Config.REAL_FORGED_MODEL_REPO_ID) configured_name = getattr(Config, "DOCUMENT_FORGERY_MODEL_FILENAME", str(target_path)) candidate_filenames = [] for candidate in (configured_name, str(target_path), target_path.name): if candidate and candidate not in candidate_filenames: candidate_filenames.append(candidate) last_error = None for filename in candidate_filenames: try: print(f"Downloading document forgery model from Hugging Face repo: {repo_id} ({filename})") downloaded_path = hf_hub_download(repo_id=repo_id, filename=filename, token=Config.HF_TOKEN) target_path.parent.mkdir(parents=True, exist_ok=True) shutil.copy2(downloaded_path, target_path) print(f"Document forgery model downloaded to: {target_path}") return target_path except Exception as exc: last_error = exc raise RuntimeError(f"unable to download document forgery model: {last_error}") # --- Global Model Instance --- MODEL_REPO_ID = Config.REAL_FORGED_MODEL_REPO_ID MODEL_FILENAME = Config.REAL_FORGED_MODEL_FILENAME DOC_MODEL_PATH = Config.DOCUMENT_FORGERY_MODEL_PATH models = ModelLoader(model_repo_id=MODEL_REPO_ID, model_filename=MODEL_FILENAME, doc_model_path=DOC_MODEL_PATH)