tools/models_handler.py

import torch
from ultralytics import YOLO
from faster_whisper import WhisperModel
from pyannote.audio import Pipeline
from google.cloud import translate_v2 as translate
from speechbrain.inference.speaker import SpeakerRecognition
from tools.tools import UniverseToolsClass

ModelsUBL_process = 'cpu' # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device
ModelsBAT_process = 'cpu'
ModelFace_process = 'cpu'
ModelBATHotspot_process = 'cpu'
ModelBevarageIndrasty_process = 'cpu'
ModelMaterial_process = 'cpu'
ModelMemoOcr_process = 'cpu'
ModelIntImgAna_process = None
ModelSpeechAnalyzer_process = None


class ModelsUBL:
    def __init__(self,device=ModelsUBL_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device
        self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu')
        self.daModel = YOLO("AI_Models/UBL_Models/ublDA_v8.1.pt").to(self.device)
        self.qpdsModel = YOLO("AI_Models/UBL_Models/ublQPDS_2_v1.2.pt").to(self.device)
        self.sosModel = YOLO("AI_Models/UBL_Models/ublSOS_v4.5.pt").to(self.device)
        self.mtSOSModel = YOLO("AI_Models/UBL_Models/ublMTSOS_v1.5.pt").to(self.device)
        self.sachetModel = YOLO("AI_Models/UBL_Models/sachetModel_v4.2.pt").to(self.device)
        self.megaHangerModel = YOLO("AI_Models/UBL_Models/megaHanger_v1.2.pt").to(self.device)
        self.sovmModel = YOLO("AI_Models/UBL_Models/sovmPOSM_v1.3.pt").to(self.device)
        self.mPOSM = YOLO("AI_Models/UBL_Models/POSM1_POSM2_FAT_v15.1.pt").to(self.device)
        self.st_orientation = YOLO("AI_Models/UBL_Models/st_orientation_v5.1.pt").to(self.device)
        self.exclusivity_model = YOLO("AI_Models/UBL_Models/exclusivity_v6.3.pt").to(self.device)
        self.sosModel_square = YOLO("AI_Models/UBL_Models/sosSquare_v1.2.pt").to(self.device)

        if 'cuda' in self.device:
            torch.cuda.empty_cache()
class ModelsBAT:
    def __init__(self,device=ModelsBAT_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device
        self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu')
        self.blanksModel = YOLO('AI_Models/BAT_Models/blanks_ghw_v8.1.pt').to(self.device)
        self.posmModel = YOLO('AI_Models/BAT_Models/batPOSMall_v3.pt').to(self.device)

        if 'cuda' in self.device:
            torch.cuda.empty_cache()

class ModelFace:
    def __init__(self,device=ModelFace_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device
        self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu')
        self.faceModel = YOLO('AI_Models/Face_Models/faceModel.pt').to(self.device)

        if 'cuda' in self.device:
            torch.cuda.empty_cache()

class ModelBATHotspot:
    def __init__(self,device=ModelBATHotspot_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device
        self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu')
        self.BATHotspotModel = YOLO('AI_Models/BAT_Models/batPOSMall_v3.pt').to(self.device)

        if 'cuda' in self.device:
            torch.cuda.empty_cache()
        
        
class ModelBevarageIndrasty:
    def __init__(self,device=ModelBevarageIndrasty_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device
        self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu')
        self.FridgeModel_fdz = YOLO('AI_Models/Fridge_Models/fdzModel.pt').to(self.device)
        self.FridgeModel_count = YOLO('AI_Models/Fridge_Models/countModel.pt').to(self.device)

        if 'cuda' in self.device:
            torch.cuda.empty_cache()
        
        

class ModelMaterial:
    def __init__(self,device=ModelMaterial_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device
        self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu')
        self.MaterialGpModel_gp = YOLO('AI_Models/Material_Models/gp_v1.pt').to(self.device)
        self.MaterialGpModel_marchent = YOLO('AI_Models/Material_Models/marchent_v13.1.pt').to(self.device)
        self.MaterialGpModel_uddokta = YOLO('AI_Models/Material_Models/uddokta_v13.1.pt').to(self.device)

        if 'cuda' in self.device:
            torch.cuda.empty_cache()


class ModelMemoOcr:
    def __init__(self,device=ModelMemoOcr_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device
        self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu')
        self.MemoOcrModel_crop = YOLO('AI_Models/MemoOCR_Models/cropModel.pt').to(self.device)
        self.MemoOcrModel_n_item = YOLO('AI_Models/MemoOCR_Models/n_item.pt').to(self.device)
        self.MemoOcrModel_n_price = YOLO('AI_Models/MemoOCR_Models/n_price.pt').to(self.device)
        self.MemoOcrModel_n_qty = YOLO('AI_Models/MemoOCR_Models/n_qty.pt').to(self.device)

        if 'cuda' in self.device:
            torch.cuda.empty_cache()

class ModelIntImgAna:
    def __init__(self,device=ModelIntImgAna_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device
        self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu')
        self.ClasfiModel = YOLO('AI_Models/IntImgAna_Models/classification.pt').to(self.device)
        self.IntImgAnaModel = YOLO('AI_Models/IntImgAna_Models/prego_v4.1.pt').to(self.device)
        self.IntImgAnaFrameModel = YOLO('AI_Models/IntImgAna_Models/prego_frame_v3.1.pt').to(self.device)
        self.HeinekenModel = YOLO('AI_Models/IntImgAna_Models/heineken_sku.pt').to(self.device)
        self.HeinekenFrameModel = YOLO('AI_Models/IntImgAna_Models/heineken_frame.pt').to(self.device)
        self.HeinekenPlanogramModel = YOLO('AI_Models/IntImgAna_Models/heineken_planogram.pt').to(self.device)

        if 'cuda' in self.device:
            torch.cuda.empty_cache()

class ModelMthouseOcr:
    def __init__(self,device=ModelIntImgAna_process):
        self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu')
        self.PageCropModel = YOLO('AI_Models/MtHouseOCR_Model/page.pt').to(self.device)

class ModelSpeechAnalyzer:
    def __init__(self,device=ModelSpeechAnalyzer_process):
        # Determine target device
        self.device = (
            "cpu" if device == "cpu"
            else device if device
            else ("cuda" if torch.cuda.is_available() else "cpu")
        )
        self.torch_device = torch.device(self.device)

        # Hugging Face token: environment override if set
        self.tools_universe = UniverseToolsClass()
        self.hf_token = self.tools_universe.decrypt_token()
        if not self.hf_token:
            print("⚠️ No Hugging Face token provided — gated models may fail to load.")

        # -----------------------
        # 1️⃣ Whisper model (Speech-to-Text)
        # -----------------------
        compute_type = "int8" if self.device == "cpu" else "int8_float32"
        self.whisper_model = WhisperModel(
            "large-v3",
            device=self.device,
            compute_type=compute_type
        )

        # -----------------------
        # 2️⃣ Speaker Verification (Speaker Embeddings)
        # -----------------------
        self.speaker_verification_model = SpeakerRecognition.from_hparams(
            source="speechbrain/spkrec-ecapa-voxceleb",
            savedir="AI_Models/SpeechAnalyzer_Models/pretrained_models/spkrec",
            run_opts={"device": self.device}
        )

        # -----------------------
        # 3️⃣ Speaker Diarization (Who Spoke When)
        # -----------------------
        try:
            self.diarization_pipeline = Pipeline.from_pretrained(
                "pyannote/speaker-diarization-community-1",
                token=self.hf_token
            ).to(self.torch_device)
            print(f"✅ Diarization pipeline loaded on {self.device}")
        except Exception as e:
            print(f"❌ Diarization pipeline failed: {e}")
            self.diarization_pipeline = None
        
        # -----------------------
        # 4️⃣ translate
        # -----------------------
        self.translate_client = translate.Client.from_service_account_json("tools/service_account.json")
        
        # -----------------------
        # 5️⃣ Optional GPU memory cleanup
        # -----------------------
        if self.device == "cuda":
            torch.cuda.empty_cache()
            print("🧠 GPU memory cache cleared.")
            print("🧠 GPU memory cache cleared.")






# class ModelsUBL:
#     def __init__(self,device=ModelsUBL_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device
#         self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu')
#         self.daModel = None
#         self.qpdsModel = None
#         self.sosModel = None
#         self.mtSOSModel = None
#         self.sachetModel = None
#         self.megaHangerModel = None
#         self.sovmModel = None
#         self.mPOSM = None
#         self.st_orientation = None
#         self.exclusivity_model = None
#         self.sosModel_square = None

#         if 'cuda' in self.device:
#             torch.cuda.empty_cache()
# class ModelsBAT:
#     def __init__(self,device=ModelsBAT_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device
#         self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu')
#         self.blanksModel = None
#         self.posmModel = None

#         if 'cuda' in self.device:
#             torch.cuda.empty_cache()

# class ModelFace:
#     def __init__(self,device=ModelFace_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device
#         self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu')
#         self.faceModel = None

#         if 'cuda' in self.device:
#             torch.cuda.empty_cache()

# class ModelBATHotspot:
#     def __init__(self,device=ModelBATHotspot_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device
#         self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu')
#         self.BATHotspotModel = None

#         if 'cuda' in self.device:
#             torch.cuda.empty_cache()
        
        
# class ModelBevarageIndrasty:
#     def __init__(self,device=ModelBevarageIndrasty_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device
#         self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu')
#         self.FridgeModel_fdz = None
#         self.FridgeModel_count = None

#         if 'cuda' in self.device:
#             torch.cuda.empty_cache()
        
        

# class ModelMaterial:
#     def __init__(self,device=ModelMaterial_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device
#         self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu')
#         self.MaterialGpModel_gp = None
#         self.MaterialGpModel_marchent = None
#         self.MaterialGpModel_uddokta = None

#         if 'cuda' in self.device:
#             torch.cuda.empty_cache()


# class ModelMemoOcr:
#     def __init__(self,device=ModelMemoOcr_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device
#         self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu')
#         self.MemoOcrModel_crop = None
#         self.MemoOcrModel_n_item = None
#         self.MemoOcrModel_n_price = None
#         self.MemoOcrModel_n_qty = None

#         if 'cuda' in self.device:
#             torch.cuda.empty_cache()

# class ModelIntImgAna:
#     def __init__(self,device=ModelIntImgAna_process): # If None [it will automatically select the device based on system configuration] else 'cpu' or 'cuda' to force select the device
#         self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu')
#         self.ClasfiModel = YOLO('AI_Models/IntImgAna_Models/classification.pt').to(self.device)
#         self.IntImgAnaModel = YOLO('AI_Models/IntImgAna_Models/prego_v4.1.pt').to(self.device)
#         self.IntImgAnaFrameModel = YOLO('AI_Models/IntImgAna_Models/prego_frame_v3.0.pt').to(self.device)
#         self.HeinekenModel = YOLO('AI_Models/IntImgAna_Models/heineken_sku.pt').to(self.device)
#         self.HeinekenFrameModel = YOLO('AI_Models/IntImgAna_Models/heineken_frame.pt').to(self.device)
#         self.HeinekenPlanogramModel = YOLO('AI_Models/IntImgAna_Models/heineken_planogram.pt').to(self.device)

#         if 'cuda' in self.device:
#             torch.cuda.empty_cache()

# class ModelMthouseOcr:
#     def __init__(self,device=ModelIntImgAna_process):
#         self.device = 'cpu' if device=='cpu' else device if device else ('cuda' if torch.cuda.is_available() else 'cpu')
#         self.PageCropModel = None

# class ModelSpeechAnalyzer:
#     def __init__(self,device=ModelSpeechAnalyzer_process):
#         self.whisper_model = None
#         self.speaker_verification_model = None
#         self.diarization_pipeline = None
#         self.translate_client = None