model.py

from __future__ import annotations

import argparse
import os
import pathlib
import subprocess
import sys
from typing import Callable, Union

import dlib
import huggingface_hub
import numpy as np
import PIL.Image
import torch
import torch.nn as nn
import torchvision.transforms as T

if os.getenv("SYSTEM") == "spaces" and not torch.cuda.is_available():
    with open("patch.e4e") as f:
        subprocess.run("patch -p1".split(), cwd="encoder4editing", stdin=f)
    with open("patch.hairclip") as f:
        subprocess.run("patch -p1".split(), cwd="HairCLIP", stdin=f)

app_dir = pathlib.Path(__file__).parent

e4e_dir = app_dir / "encoder4editing"
sys.path.insert(0, e4e_dir.as_posix())

from models.psp import pSp
from utils.alignment import align_face

hairclip_dir = app_dir / "HairCLIP"
mapper_dir = hairclip_dir / "mapper"
sys.path.insert(0, hairclip_dir.as_posix())
sys.path.insert(0, mapper_dir.as_posix())

from mapper.datasets.latents_dataset_inference import LatentsDatasetInference
from mapper.hairclip_mapper import HairCLIPMapper


class Model:
    def __init__(self):
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.landmark_model = self._create_dlib_landmark_model()
        self.e4e = self._load_e4e()
        self.hairclip = self._load_hairclip()
        self.transform = self._create_transform()

    @staticmethod
    def _create_dlib_landmark_model():
        path = huggingface_hub.hf_hub_download(
            "public-data/dlib_face_landmark_model", "shape_predictor_68_face_landmarks.dat"
        )
        return dlib.shape_predictor(path)

    def _load_e4e(self) -> nn.Module:
        ckpt_path = huggingface_hub.hf_hub_download("public-data/e4e", "e4e_ffhq_encode.pt")
        ckpt = torch.load(ckpt_path, map_location="cpu")
        opts = ckpt["opts"]
        opts["device"] = self.device.type
        opts["checkpoint_path"] = ckpt_path
        opts = argparse.Namespace(**opts)
        model = pSp(opts)
        model.to(self.device)
        model.eval()
        return model

    def _load_hairclip(self) -> nn.Module:
        ckpt_path = huggingface_hub.hf_hub_download("public-data/HairCLIP", "hairclip.pt")
        ckpt = torch.load(ckpt_path, map_location="cpu")
        opts = ckpt["opts"]
        opts["device"] = self.device.type
        opts["checkpoint_path"] = ckpt_path
        opts["editing_type"] = "both"
        opts["input_type"] = "text"
        opts["hairstyle_description"] = "HairCLIP/mapper/hairstyle_list.txt"
        opts["color_description"] = "red"
        opts = argparse.Namespace(**opts)
        model = HairCLIPMapper(opts)
        model.to(self.device)
        model.eval()
        return model

    @staticmethod
    def _create_transform() -> Callable:
        transform = T.Compose(
            [
                T.Resize(256),
                T.CenterCrop(256),
                T.ToTensor(),
                T.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
            ]
        )
        return transform

    def detect_and_align_face(self, image: str) -> PIL.Image.Image:
        image = align_face(filepath=image, predictor=self.landmark_model)
        return image

    @staticmethod
    def denormalize(tensor: torch.Tensor) -> torch.Tensor:
        return torch.clamp((tensor + 1) / 2 * 255, 0, 255).to(torch.uint8)

    def postprocess(self, tensor: torch.Tensor) -> np.ndarray:
        tensor = self.denormalize(tensor)
        return tensor.cpu().numpy().transpose(1, 2, 0)

    @torch.inference_mode()
    def reconstruct_face(self, image: PIL.Image.Image) -> tuple[np.ndarray, torch.Tensor]:
        input_data = self.transform(image).unsqueeze(0).to(self.device)
        reconstructed_images, latents = self.e4e(input_data, randomize_noise=False, return_latents=True)
        reconstructed = torch.clamp(reconstructed_images[0].detach(), -1, 1)
        reconstructed = self.postprocess(reconstructed)
        return reconstructed, latents[0]

    @torch.inference_mode()
    def generate(
        self, editing_type: str, hairstyle_index: int, color_description: str, latent: torch.Tensor
    ) -> np.ndarray:
        opts = self.hairclip.opts
        opts.editing_type = editing_type
        opts.color_description = color_description

        if editing_type == "color":
            hairstyle_index = 0

        device = torch.device(opts.device)

        dataset = LatentsDatasetInference(latents=latent.unsqueeze(0).cpu(), opts=opts)
        w, hairstyle_text_inputs_list, color_text_inputs_list = dataset[0][:3]

        w = w.unsqueeze(0).to(device)
        hairstyle_text_inputs = hairstyle_text_inputs_list[hairstyle_index].unsqueeze(0).to(device)
        color_text_inputs = color_text_inputs_list[0].unsqueeze(0).to(device)

        hairstyle_tensor_hairmasked = torch.Tensor([0]).unsqueeze(0).to(device)
        color_tensor_hairmasked = torch.Tensor([0]).unsqueeze(0).to(device)

        w_hat = w + 0.1 * self.hairclip.mapper(
            w,
            hairstyle_text_inputs,
            color_text_inputs,
            hairstyle_tensor_hairmasked,
            color_tensor_hairmasked,
        )
        x_hat, _ = self.hairclip.decoder(
            [w_hat],
            input_is_latent=True,
            return_latents=True,
            randomize_noise=False,
            truncation=1,
        )
        res = torch.clamp(x_hat[0].detach(), -1, 1)
        res = self.postprocess(res)
        return res