|
from __future__ import annotations |
|
|
|
import os |
|
import pathlib |
|
import sys |
|
|
|
import huggingface_hub |
|
import numpy as np |
|
import torch |
|
import torch.nn as nn |
|
|
|
if os.getenv('SYSTEM') == 'spaces': |
|
os.system("sed -i '14,21d' StyleSwin/op/fused_act.py") |
|
os.system("sed -i '12,19d' StyleSwin/op/upfirdn2d.py") |
|
|
|
current_dir = pathlib.Path(__file__).parent |
|
submodule_dir = current_dir / 'StyleSwin' |
|
sys.path.insert(0, submodule_dir.as_posix()) |
|
|
|
from models.generator import Generator |
|
|
|
|
|
class Model: |
|
MODEL_NAMES = [ |
|
'CelebAHQ_256', |
|
'FFHQ_256', |
|
'LSUNChurch_256', |
|
'CelebAHQ_1024', |
|
'FFHQ_1024', |
|
] |
|
|
|
def __init__(self): |
|
self.device = torch.device( |
|
'cuda:0' if torch.cuda.is_available() else 'cpu') |
|
self._download_all_models() |
|
self.model_name = self.MODEL_NAMES[3] |
|
self.model = self._load_model(self.model_name) |
|
|
|
self.std = torch.FloatTensor([0.229, 0.224, |
|
0.225])[None, :, None, |
|
None].to(self.device) |
|
self.mean = torch.FloatTensor([0.485, 0.456, |
|
0.406])[None, :, None, |
|
None].to(self.device) |
|
|
|
def _load_model(self, model_name: str) -> nn.Module: |
|
size = int(model_name.split('_')[1]) |
|
channel_multiplier = 1 if size == 1024 else 2 |
|
model = Generator(size, |
|
style_dim=512, |
|
n_mlp=8, |
|
channel_multiplier=channel_multiplier) |
|
ckpt_path = huggingface_hub.hf_hub_download('public-data/StyleSwin', |
|
f'models/{model_name}.pt') |
|
ckpt = torch.load(ckpt_path) |
|
model.load_state_dict(ckpt['g_ema']) |
|
model.to(self.device) |
|
model.eval() |
|
return model |
|
|
|
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) |
|
|
|
def _download_all_models(self): |
|
for name in self.MODEL_NAMES: |
|
self._load_model(name) |
|
|
|
def generate_z(self, seed: int) -> torch.Tensor: |
|
seed = int(np.clip(seed, 0, np.iinfo(np.uint32).max)) |
|
z = np.random.RandomState(seed).randn(1, 512) |
|
return torch.from_numpy(z).float().to(self.device) |
|
|
|
def postprocess(self, tensors: torch.Tensor) -> np.ndarray: |
|
assert tensors.dim() == 4 |
|
tensors = tensors * self.std + self.mean |
|
tensors = (tensors * 255).clamp(0, 255).to(torch.uint8) |
|
return tensors.permute(0, 2, 3, 1).cpu().numpy() |
|
|
|
@torch.inference_mode() |
|
def generate_image(self, seed: int) -> np.ndarray: |
|
z = self.generate_z(seed) |
|
out, _ = self.model(z) |
|
out = self.postprocess(out) |
|
return out[0] |
|
|
|
def set_model_and_generate_image(self, model_name: str, |
|
seed: int) -> np.ndarray: |
|
self.set_model(model_name) |
|
return self.generate_image(seed) |
|
|
