Spaces:
Runtime error
Runtime error
| #!/usr/bin/env python | |
| from __future__ import annotations | |
| import functools | |
| import os | |
| import random | |
| import shlex | |
| import subprocess | |
| import sys | |
| import gradio as gr | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| from huggingface_hub import hf_hub_download | |
| if os.environ.get('SYSTEM') == 'spaces': | |
| with open('patch') as f: | |
| subprocess.run(shlex.split('patch -p1'), | |
| cwd='stylegan2-pytorch', | |
| stdin=f) | |
| if not torch.cuda.is_available(): | |
| with open('patch-cpu') as f: | |
| subprocess.run(shlex.split('patch -p1'), | |
| cwd='stylegan2-pytorch', | |
| stdin=f) | |
| sys.path.insert(0, 'stylegan2-pytorch') | |
| from model import Generator | |
| DESCRIPTION = '''# [TADNE](https://thisanimedoesnotexist.ai/) (This Anime Does Not Exist) | |
| Related Apps: | |
| - [TADNE Image Viewer](https://huggingface.co/spaces/hysts/TADNE-image-viewer) | |
| - [TADNE Image Selector](https://huggingface.co/spaces/hysts/TADNE-image-selector) | |
| - [TADNE Interpolation](https://huggingface.co/spaces/hysts/TADNE-interpolation) | |
| - [TADNE Image Search with DeepDanbooru](https://huggingface.co/spaces/hysts/TADNE-image-search-with-DeepDanbooru) | |
| ''' | |
| SAMPLE_IMAGE_DIR = 'https://huggingface.co/spaces/hysts/TADNE/resolve/main/samples' | |
| ARTICLE = f'''## Generated images | |
| - size: 512x512 | |
| - truncation: 0.7 | |
| - seed: 0-99 | |
|  | |
| ''' | |
| MAX_SEED = np.iinfo(np.int32).max | |
| def randomize_seed_fn(seed: int, randomize_seed: bool) -> int: | |
| if randomize_seed: | |
| seed = random.randint(0, MAX_SEED) | |
| return seed | |
| def load_model(device: torch.device) -> nn.Module: | |
| model = Generator(512, 1024, 4, channel_multiplier=2) | |
| path = hf_hub_download('public-data/TADNE', | |
| 'models/aydao-anime-danbooru2019s-512-5268480.pt') | |
| checkpoint = torch.load(path) | |
| model.load_state_dict(checkpoint['g_ema']) | |
| model.eval() | |
| model.to(device) | |
| model.latent_avg = checkpoint['latent_avg'].to(device) | |
| with torch.inference_mode(): | |
| z = torch.zeros((1, model.style_dim)).to(device) | |
| model([z], truncation=0.7, truncation_latent=model.latent_avg) | |
| return model | |
| def generate_z(z_dim: int, seed: int, device: torch.device) -> torch.Tensor: | |
| return torch.from_numpy(np.random.RandomState(seed).randn( | |
| 1, z_dim)).to(device).float() | |
| def generate_image(seed: int, truncation_psi: float, randomize_noise: bool, | |
| model: nn.Module, device: torch.device) -> np.ndarray: | |
| seed = int(np.clip(seed, 0, np.iinfo(np.uint32).max)) | |
| z = generate_z(model.style_dim, seed, device) | |
| out, _ = model([z], | |
| truncation=truncation_psi, | |
| truncation_latent=model.latent_avg, | |
| randomize_noise=randomize_noise) | |
| out = (out.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8) | |
| return out[0].cpu().numpy() | |
| device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') | |
| model = load_model(device) | |
| fn = functools.partial(generate_image, model=model, device=device) | |
| with gr.Blocks(css='style.css') as demo: | |
| gr.Markdown(DESCRIPTION) | |
| with gr.Row(): | |
| with gr.Column(): | |
| seed = gr.Slider(label='Seed', | |
| minimum=0, | |
| maximum=MAX_SEED, | |
| step=1, | |
| value=0) | |
| randomize_seed = gr.Checkbox(label='Randomize seed', value=True) | |
| psi = gr.Slider(label='Truncation psi', | |
| minimum=0, | |
| maximum=2, | |
| step=0.05, | |
| value=0.7) | |
| randomize_noise = gr.Checkbox(label='Randomize Noise', value=False) | |
| run_button = gr.Button('Run') | |
| with gr.Column(): | |
| result = gr.Image(label='Output') | |
| gr.Markdown(ARTICLE) | |
| run_button.click( | |
| fn=randomize_seed_fn, | |
| inputs=[seed, randomize_seed], | |
| outputs=seed, | |
| queue=False, | |
| api_name=False, | |
| ).then( | |
| fn=fn, | |
| inputs=[seed, psi, randomize_noise], | |
| outputs=result, | |
| api_name='run', | |
| ) | |
| demo.queue(max_size=10).launch() | |