Spaces:
Sleeping
Sleeping
| 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) | |