content: # https://github.com/eladrich/pixel2style2pixel
  import torch
  from torch import nn
  from ldm.thirdp.psp.model_irse import Backbone
  
  
  class IDFeatures(nn.Module):
      def __init__(self, model_path):
          super(IDFeatures, self).__init__()
          print('Loading ResNet ArcFace')
          self.facenet = Backbone(input_size=112, num_layers=50, drop_ratio=0.6, mode='ir_se')
          self.facenet.load_state_dict(torch.load(model_path, map_location="cpu"))
          self.face_pool = torch.nn.AdaptiveAvgPool2d((112, 112))
          self.facenet.eval()
  
      def forward(self, x, crop=False):
          # Not sure of the image range here
          if crop:
              x = torch.nn.functional.interpolate(x, (256, 256), mode="area")
              x = x[:, :, 35:223, 32:220]
          x = self.face_pool(x)
          x_feats = self.facenet(x)
          return x_feats
  