Spaces:
Running
Running
| #!/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) interpolation | |
| Related Apps: | |
| - [TADNE](https://huggingface.co/spaces/hysts/TADNE) | |
| - [TADNE Image Viewer](https://huggingface.co/spaces/hysts/TADNE-image-viewer) | |
| - [TADNE Image Selector](https://huggingface.co/spaces/hysts/TADNE-image-selector) | |
| - [TADNE Image Search with DeepDanbooru](https://huggingface.co/spaces/hysts/TADNE-image-search-with-DeepDanbooru) | |
| ''' | |
| 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(model: nn.Module, z: torch.Tensor, truncation_psi: float, | |
| randomize_noise: bool) -> np.ndarray: | |
| 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() | |
| def generate_interpolated_images(seed0: int, seed1: int, num_intermediate: int, | |
| psi0: float, psi1: float, | |
| randomize_noise: bool, model: nn.Module, | |
| device: torch.device) -> list[np.ndarray]: | |
| seed0 = int(np.clip(seed0, 0, MAX_SEED)) | |
| seed1 = int(np.clip(seed1, 0, MAX_SEED)) | |
| z0 = generate_z(model.style_dim, seed0, device) | |
| z1 = generate_z(model.style_dim, seed1, device) | |
| 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)] | |
| res = [] | |
| for z, psi in zip(zs, psis): | |
| out = generate_image(model, z, psi, randomize_noise) | |
| res.append(out) | |
| return res | |
| device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') | |
| model = load_model(device) | |
| fn = functools.partial(generate_interpolated_images, | |
| model=model, | |
| device=device) | |
| examples = [ | |
| [29703, 55376, 3, 0.7, 0.7, False], | |
| [34141, 36864, 5, 0.7, 0.7, False], | |
| [74650, 88322, 7, 0.7, 0.7, False], | |
| [84314, 70317410, 9, 0.7, 0.7, False], | |
| [55376, 55376, 5, 0.3, 1.3, False], | |
| ] | |
| with gr.Blocks(css='style.css') as demo: | |
| gr.Markdown(DESCRIPTION) | |
| with gr.Row(): | |
| with gr.Column(): | |
| seed_1 = gr.Slider(label='Seed 1', | |
| minimum=0, | |
| maximum=MAX_SEED, | |
| step=1, | |
| value=29703) | |
| seed_2 = gr.Slider(label='Seed 2', | |
| minimum=0, | |
| maximum=MAX_SEED, | |
| step=1, | |
| value=55376) | |
| num_intermediate_frames = gr.Slider( | |
| label='Number of Intermediate Frames', | |
| minimum=1, | |
| maximum=21, | |
| step=1, | |
| value=3, | |
| ) | |
| psi_1 = gr.Slider(label='Truncation psi 1', | |
| minimum=0, | |
| maximum=2, | |
| step=0.05, | |
| value=0.7) | |
| psi_2 = gr.Slider(label='Truncation psi 2', | |
| 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.Gallery(label='Output') | |
| inputs = [ | |
| seed_1, | |
| seed_2, | |
| num_intermediate_frames, | |
| psi_1, | |
| psi_2, | |
| randomize_noise, | |
| ] | |
| gr.Examples( | |
| examples=examples, | |
| inputs=inputs, | |
| outputs=result, | |
| fn=fn, | |
| cache_examples=os.getenv('CACHE_EXAMPLES') == '1', | |
| ) | |
| run_button.click( | |
| fn=fn, | |
| inputs=inputs, | |
| outputs=result, | |
| api_name='run', | |
| ) | |
| demo.queue(max_size=10).launch() | |