from __future__ import annotations 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()) class Model: def __init__(self): self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") self.model = self.load_model("stylegan_human_v2_1024.pkl") def load_model(self, file_name: str) -> nn.Module: path = hf_hub_download("public-data/StyleGAN-Human", f"models/{file_name}") 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() @torch.inference_mode() 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() @torch.inference_mode() 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