from __future__ import annotations import os import pathlib import pickle import sys import lpips import numpy as np import torch import torch.nn as nn from huggingface_hub import hf_hub_download current_dir = pathlib.Path(__file__).parent submodule_dir = current_dir / 'stylegan3' sys.path.insert(0, submodule_dir.as_posix()) HF_TOKEN = os.environ['HF_TOKEN'] class LPIPS(lpips.LPIPS): @staticmethod def preprocess(image: np.ndarray) -> torch.Tensor: data = torch.from_numpy(image).float() / 255 data = data * 2 - 1 return data.permute(2, 0, 1).unsqueeze(0) @torch.inference_mode() def compute_features(self, data: torch.Tensor) -> list[torch.Tensor]: data = self.scaling_layer(data) data = self.net(data) return [lpips.normalize_tensor(x) for x in data] @torch.inference_mode() def compute_distance(self, features0: list[torch.Tensor], features1: list[torch.Tensor]) -> float: res = 0 for lin, x0, x1 in zip(self.lins, features0, features1): d = (x0 - x1)**2 y = lin(d) y = lpips.lpips.spatial_average(y) res += y.item() return res class Model: MODEL_NAMES = [ 'dogs_1024', 'elephants_512', 'horses_256', 'bicycles_256', 'lions_512', 'giraffes_512', 'parrots_512', ] TRUNCATION_TYPES = [ 'Multimodal (LPIPS)', 'Multimodal (L2)', 'Global', ] def __init__(self, device: str | torch.device): self.device = torch.device(device) self._download_all_models() self._download_all_cluster_centers() self._download_all_cluster_center_images() self.model_name = self.MODEL_NAMES[0] self.model = self._load_model(self.model_name) self.cluster_centers = self._load_cluster_centers(self.model_name) self.cluster_center_images = self._load_cluster_center_images( self.model_name) self.lpips = LPIPS() self.cluster_center_lpips_feature_dict = self._compute_cluster_center_lpips_features( ) def _load_model(self, model_name: str) -> nn.Module: path = hf_hub_download('hysts/Self-Distilled-StyleGAN', f'models/{model_name}_pytorch.pkl', use_auth_token=HF_TOKEN) with open(path, 'rb') as f: model = pickle.load(f)['G_ema'] model.eval() model.to(self.device) return model def _load_cluster_centers(self, model_name: str) -> torch.Tensor: path = hf_hub_download('hysts/Self-Distilled-StyleGAN', f'cluster_centers/{model_name}.npy', use_auth_token=HF_TOKEN) centers = np.load(path) centers = torch.from_numpy(centers).float().to(self.device) return centers def _load_cluster_center_images(self, model_name: str) -> np.ndarray: path = hf_hub_download('hysts/Self-Distilled-StyleGAN', f'cluster_center_images/{model_name}.npy', use_auth_token=HF_TOKEN) return np.load(path) def set_model(self, model_name: str) -> None: if model_name == self.model_name: return self.model_name = model_name self.model = self._load_model(model_name) self.cluster_centers = self._load_cluster_centers(model_name) self.cluster_center_images = self._load_cluster_center_images( model_name) def _download_all_models(self): for name in self.MODEL_NAMES: self._load_model(name) def _download_all_cluster_centers(self): for name in self.MODEL_NAMES: self._load_cluster_centers(name) def _download_all_cluster_center_images(self): for name in self.MODEL_NAMES: self._load_cluster_center_images(name) def generate_z(self, seed: int) -> torch.Tensor: seed = int(np.clip(seed, 0, np.iinfo(np.uint32).max)) return torch.from_numpy( np.random.RandomState(seed).randn(1, self.model.z_dim)).float().to( self.device) def compute_w(self, z: torch.Tensor) -> torch.Tensor: label = torch.zeros((1, self.model.c_dim), device=self.device) w = self.model.mapping(z, label) return w @staticmethod def truncate_w(w_center: torch.Tensor, w: torch.Tensor, psi: float) -> torch.Tensor: if psi == 1: return w return w_center.lerp(w, psi) @torch.inference_mode() def synthesize(self, w: torch.Tensor) -> torch.Tensor: return self.model.synthesis(w) def postprocess(self, tensor: torch.Tensor) -> np.ndarray: tensor = (tensor.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to( torch.uint8) return tensor.cpu().numpy() def compute_lpips_features(self, image: np.ndarray) -> list[torch.Tensor]: data = self.lpips.preprocess(image) return self.lpips.compute_features(data) def _compute_cluster_center_lpips_features( self) -> dict[str, list[list[torch.Tensor]]]: res = dict() for name in self.MODEL_NAMES: images = self._load_cluster_center_images(name) res[name] = [ self.compute_lpips_features(image) for image in images ] return res def compute_distance_to_cluster_centers( self, ws: torch.Tensor, distance_type: str) -> list[torch.Tensor]: if distance_type == 'l2': return self._compute_l2_distance_to_cluster_centers(ws) elif distance_type == 'lpips': return self._compute_lpips_distance_to_cluster_centers(ws) else: raise ValueError def _compute_l2_distance_to_cluster_centers( self, ws: torch.Tensor) -> np.ndarray: dist2 = ((self.cluster_centers - ws[0, 0])**2).sum(dim=1) return dist2.cpu().numpy() def _compute_lpips_distance_to_cluster_centers( self, ws: torch.Tensor) -> np.ndarray: x = self.synthesize(ws) x = self.postprocess(x)[0] feat0 = self.compute_lpips_features(x) cluster_center_features = self.cluster_center_lpips_feature_dict[ self.model_name] distances = [ self.lpips.compute_distance(feat0, feat1) for feat1 in cluster_center_features ] return np.asarray(distances) def find_nearest_cluster_center(self, ws: torch.Tensor, distance_type: str) -> int: distances = self.compute_distance_to_cluster_centers(ws, distance_type) return int(np.argmin(distances)) def generate_image(self, seed: int, truncation_psi: float, truncation_type: str) -> np.ndarray: z = self.generate_z(seed) ws = self.compute_w(z) if truncation_type == self.TRUNCATION_TYPES[2]: w0 = self.model.mapping.w_avg else: if truncation_type == self.TRUNCATION_TYPES[0]: distance_type = 'lpips' elif truncation_type == self.TRUNCATION_TYPES[1]: distance_type = 'l2' else: raise ValueError cluster_index = self.find_nearest_cluster_center(ws, distance_type) w0 = self.cluster_centers[cluster_index] new_ws = self.truncate_w(w0, ws, truncation_psi) out = self.synthesize(new_ws) out = self.postprocess(out) return out[0] def set_model_and_generate_image(self, model_name: str, seed: int, truncation_psi: float, truncation_type: str) -> np.ndarray: self.set_model(model_name) return self.generate_image(seed, truncation_psi, truncation_type)