Spaces:
Running
Running
| from __future__ import annotations | |
| import os | |
| import pathlib | |
| import pickle | |
| import sys | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| from huggingface_hub import hf_hub_download | |
| app_dir = pathlib.Path(__file__).parent | |
| submodule_dir = app_dir / 'StyleGAN-Human' | |
| sys.path.insert(0, submodule_dir.as_posix()) | |
| HF_TOKEN = os.environ['HF_TOKEN'] | |
| class Model: | |
| def __init__(self, device: str | torch.device): | |
| self.device = torch.device(device) | |
| self.model = self.load_model('stylegan_human_v2_1024.pkl') | |
| def load_model(self, file_name: str) -> nn.Module: | |
| path = hf_hub_download('hysts/StyleGAN-Human', | |
| f'models/{file_name}', | |
| use_auth_token=HF_TOKEN) | |
| with open(path, 'rb') as f: | |
| model = pickle.load(f)['G_ema'] | |
| model.eval() | |
| model.to(self.device) | |
| with torch.inference_mode(): | |
| z = torch.zeros((1, model.z_dim)).to(self.device) | |
| label = torch.zeros([1, model.c_dim], device=self.device) | |
| model(z, label, force_fp32=True) | |
| return model | |
| def generate_z(self, z_dim: int, seed: int) -> torch.Tensor: | |
| return torch.from_numpy(np.random.RandomState(seed).randn( | |
| 1, z_dim)).to(self.device).float() | |
| def generate_single_image(self, seed: int, | |
| truncation_psi: float) -> np.ndarray: | |
| seed = int(np.clip(seed, 0, np.iinfo(np.uint32).max)) | |
| z = self.generate_z(self.model.z_dim, seed) | |
| label = torch.zeros([1, self.model.c_dim], device=self.device) | |
| out = self.model(z, | |
| label, | |
| truncation_psi=truncation_psi, | |
| force_fp32=True) | |
| out = (out.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to( | |
| torch.uint8) | |
| return out[0].cpu().numpy() | |
| def generate_interpolated_images( | |
| self, seed0: int, psi0: float, seed1: int, psi1: float, | |
| num_intermediate: int) -> list[np.ndarray]: | |
| seed0 = int(np.clip(seed0, 0, np.iinfo(np.uint32).max)) | |
| seed1 = int(np.clip(seed1, 0, np.iinfo(np.uint32).max)) | |
| z0 = self.generate_z(self.model.z_dim, seed0) | |
| z1 = self.generate_z(self.model.z_dim, seed1) | |
| vec = z1 - z0 | |
| dvec = vec / (num_intermediate + 1) | |
| zs = [z0 + dvec * i for i in range(num_intermediate + 2)] | |
| dpsi = (psi1 - psi0) / (num_intermediate + 1) | |
| psis = [psi0 + dpsi * i for i in range(num_intermediate + 2)] | |
| label = torch.zeros([1, self.model.c_dim], device=self.device) | |
| res = [] | |
| for z, psi in zip(zs, psis): | |
| out = self.model(z, label, truncation_psi=psi, force_fp32=True) | |
| out = (out.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to( | |
| torch.uint8) | |
| out = out[0].cpu().numpy() | |
| res.append(out) | |
| return res | |