Blame for features/real_forged_classifier/model_loader.py

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)