Spaces:
Runtime error
Runtime error
File size: 3,722 Bytes
9c4c9e6 b4a6915 9c4c9e6 3744a88 b4a6915 3744a88 9c4c9e6 3744a88 9c4c9e6 3744a88 a258609 3744a88 9c4c9e6 3744a88 9c4c9e6 b4a6915 9c4c9e6 b4a6915 9c4c9e6 3744a88 9c4c9e6 3744a88 9c4c9e6 3744a88 9c4c9e6 59bef24 9c4c9e6 59bef24 9c4c9e6 59bef24 3744a88 9c4c9e6 3744a88 b4a6915 3744a88 59bef24 b4a6915 3744a88 b4a6915 59bef24 b4a6915 3744a88 9c4c9e6 59bef24 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 |
#!/usr/bin/env python
from __future__ import annotations
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
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = load_model(device)
def generate_z(z_dim: int, seed: int) -> torch.Tensor:
return torch.from_numpy(np.random.RandomState(seed).randn(1, z_dim)).float()
@torch.inference_mode()
def generate_image(seed: int, truncation_psi: float, randomize_noise: bool) -> np.ndarray:
seed = int(np.clip(seed, 0, np.iinfo(np.uint32).max))
z = generate_z(model.style_dim, seed)
z = z.to(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()
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()
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=generate_image,
inputs=[seed, psi, randomize_noise],
outputs=result,
api_name="run",
)
if __name__ == "__main__":
demo.queue(max_size=10).launch()
|