Hugging Face's logo

Spaces:
hysts
/
StyleSwin
Running

App Files Community
StyleSwin
File size: 4,679 Bytes 
20ed832
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8f39d59
 
 
 
 
 
 
20ed832
 
9250fe1
8f39d59
 
 
 
e5605f3
 
 
 
 
 
 
 
 
 
 
 
 
 
8f39d59
 
e5605f3
9250fe1
20ed832
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5d660ba
20ed832
 
 
 
9250fe1
 
 
20ed832
 
 
 
 
 
 
 
 
 
 
 
 
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
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
 
#!/usr/bin/env python

from __future__ import annotations

import argparse
import functools
import os
import sys

import gradio as gr
import huggingface_hub
import numpy as np
import PIL.Image
import torch
import torch.nn as nn

if os.environ.get('SYSTEM') == 'spaces':
    os.system("sed -i '14,21d' StyleSwin/op/fused_act.py")
    os.system("sed -i '12,19d' StyleSwin/op/upfirdn2d.py")

sys.path.insert(0, 'StyleSwin')

from models.generator import Generator

TITLE = 'microsoft/StyleSwin'
DESCRIPTION = '''This is an unofficial demo for https://github.com/microsoft/StyleSwin.

Expected execution time on Hugging Face Spaces: 3s (for 256x256 images), 7s (for 1024x1024 images)
'''
SAMPLE_IMAGE_DIR = 'https://huggingface.co/spaces/hysts/StyleSwin/resolve/main/samples'
ARTICLE = f'''## Generated images
### CelebA-HQ
- size: 1024x1024
- seed: 0-99
![CelebA-HQ samples]({SAMPLE_IMAGE_DIR}/celeba-hq.jpg)
### FFHQ
- size: 1024x1024
- seed: 0-99
![FFHQ samples]({SAMPLE_IMAGE_DIR}/ffhq.jpg)
### LSUN Church
- size: 256x256
- seed: 0-99
![LSUN Church samples]({SAMPLE_IMAGE_DIR}/lsun-church.jpg)

<center><img src="https://visitor-badge.glitch.me/badge?page_id=hysts.styleswin" alt="visitor badge"/></center>
'''

TOKEN = os.environ['TOKEN']

MODEL_REPO = 'hysts/StyleSwin'
MODEL_NAMES = [
    'CelebAHQ_256',
    'FFHQ_256',
    'LSUNChurch_256',
    'CelebAHQ_1024',
    'FFHQ_1024',
]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument('--device', type=str, default='cpu')
    parser.add_argument('--theme', type=str)
    parser.add_argument('--live', action='store_true')
    parser.add_argument('--share', action='store_true')
    parser.add_argument('--port', type=int)
    parser.add_argument('--disable-queue',
                        dest='enable_queue',
                        action='store_false')
    parser.add_argument('--allow-flagging', type=str, default='never')
    return parser.parse_args()


def load_model(model_name: str, device: torch.device) -> 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(MODEL_REPO,
                                                f'models/{model_name}.pt',
                                                use_auth_token=TOKEN)
    ckpt = torch.load(ckpt_path)
    model.load_state_dict(ckpt['g_ema'])
    model.to(device)
    model.eval()
    return model


def generate_z(seed: int, device: torch.device) -> torch.Tensor:
    return torch.from_numpy(np.random.RandomState(seed).randn(
        1, 512)).to(device).float()


def postprocess(tensors: torch.Tensor) -> torch.Tensor:
    assert tensors.dim() == 4
    tensors = tensors.cpu()
    std = torch.FloatTensor([0.229, 0.224, 0.225])[None, :, None, None]
    mean = torch.FloatTensor([0.485, 0.456, 0.406])[None, :, None, None]
    tensors = tensors * std + mean
    tensors = (tensors * 255).clamp(0, 255).to(torch.uint8)
    return tensors


@torch.inference_mode()
def generate_image(model_name: str, seed: int, model_dict: dict,
                   device: torch.device) -> PIL.Image.Image:
    model = model_dict[model_name]
    seed = int(np.clip(seed, 0, np.iinfo(np.uint32).max))
    z = generate_z(seed, device)
    out, _ = model(z)
    out = postprocess(out)
    out = out.numpy()[0].transpose(1, 2, 0)
    return PIL.Image.fromarray(out, 'RGB')


def main():
    gr.close_all()

    args = parse_args()
    device = torch.device(args.device)

    model_dict = {name: load_model(name, device) for name in MODEL_NAMES}

    func = functools.partial(generate_image,
                             model_dict=model_dict,
                             device=device)
    func = functools.update_wrapper(func, generate_image)

    gr.Interface(
        func,
        [
            gr.inputs.Radio(MODEL_NAMES,
                            type='value',
                            default='FFHQ_256',
                            label='Model',
                            optional=False),
            gr.inputs.Slider(0, 2147483647, step=1, default=0, label='Seed'),
        ],
        gr.outputs.Image(type='pil', label='Output'),
        title=TITLE,
        description=DESCRIPTION,
        article=ARTICLE,
        theme=args.theme,
        allow_flagging=args.allow_flagging,
        live=args.live,
    ).launch(
        enable_queue=args.enable_queue,
        server_port=args.port,
        share=args.share,
    )


if __name__ == '__main__':
    main()