"""
Created By: ishwor subedi
Date: 2024-07-03
"""
import numpy as np
from PIL.Image import Image

from src.components.faceswap import FaceSwapper


class MainPipeline:
    def __init__(self, device='cpu'):
        self.device = device
        self.face_swapper = FaceSwapper(device=device)
        self.face_swapper.load_swapper_model(
            'https://drive.google.com/uc?id=1HvZ4MAtzlY74Dk4ASGIS9L6Rg5oZdqvu',
            'artifacts/inswapper/inswapper_128.onnx'
        )
        # self.face_swapper.load_enhancer_model('REAL-ESRGAN 2x', device)

    def face_swap(self, img1: np.array, img2: np.array) -> Image:
        result = self.face_swapper.face_swap(img1, img2)
        return result