1
import torch
2
import numpy as np
3

4

5
def faceshifter_batch(source_emb: torch.tensor, 
6
                      target: torch.tensor,
7
                      G: torch.nn.Module) -> np.ndarray:
8
    """
9
    Apply faceshifter model for batch of target images
10
    """
11
    
12
    bs = target.shape[0]
13
    assert target.ndim == 4, "target should have 4 dimentions -- B x C x H x W"
14
    
15
    if bs > 1:
16
        source_emb = torch.cat([source_emb]*bs)
17
    
18
    with torch.no_grad():
19
        Y_st, _ = G(target, source_emb)
20
        Y_st = (Y_st.permute(0, 2, 3, 1)*0.5 + 0.5)*255
21
        Y_st = Y_st[:, :, :, [2,1,0]].type(torch.uint8)
22
        Y_st = Y_st.cpu().detach().numpy()    
23
    return Y_st
24