Upload app.py

app.py

@@ -1,100 +1,75 @@
-#!/usr/bin/env python
-
-from __future__ import annotations
-
-import argparse
-import functools
-import os
-import pickle
-import sys
-
 import gradio as gr
+import os
 import numpy as np
 import torch
-import torch.nn as nn
-from huggingface_hub import hf_hub_download
+import pickle
+import types
+
+from huggingface_hub import hf_hub_url, cached_download
+
+TOKEN = os.environ['TOKEN']
+
+with open(cached_download(hf_hub_url('CorvaeOboro/gen_ability_icon', 'gen_ability_icon_stylegan2ada_20220801.pkl'), use_auth_token=TOKEN), 'rb') as f:
+    G = pickle.load(f)['G_ema']# torch.nn.Module
+
+device = torch.device("cpu")
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+    G = G.to(device)
+else:
+    _old_forward = G.forward
+
+    def _new_forward(self, *args, **kwargs):
+        kwargs["force_fp32"] = True
+        return _old_forward(*args, **kwargs)
+
+    G.forward = types.MethodType(_new_forward, G)
+
+    _old_synthesis_forward = G.synthesis.forward
+
+    def _new_synthesis_forward(self, *args, **kwargs):
+        kwargs["force_fp32"] = True
+        return _old_synthesis_forward(*args, **kwargs)
+
+    G.synthesis.forward = types.MethodType(_new_synthesis_forward, G.synthesis)
+
+
+def generate(num_images, interpolate):
+    if interpolate:
+        z1 = torch.randn([1, G.z_dim])# latent codes
+        z2 = torch.randn([1, G.z_dim])# latent codes
+        zs = torch.cat([z1 + (z2 - z1) * i / (num_images-1) for i in range(num_images)], 0)
+    else:
+        zs = torch.randn([num_images, G.z_dim])# latent codes
+    with torch.no_grad():
+        zs = zs.to(device)
+        img = G(zs, None, force_fp32=True, noise_mode='const') 
+        img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
+    return img.cpu().numpy()
+
+demo = gr.Blocks()
 
-sys.path.insert(0, 'gen_ability_icon')
+def infer(num_images, interpolate):
+    img = generate(round(num_images), interpolate)
+    imgs = list(img)
+    return imgs
 
-TITLE = 'gen_ability_icon'
-DESCRIPTION = '''creates circular magic ability icons from stylegan2ada model trained on synthetic dataset .
+with demo:
+    gr.Markdown(
+    """
+    # gen_ability_icon
+    creates circular magic ability icons from stylegan2ada model trained on synthetic dataset .
 more information here :  https://github.com/CorvaeOboro/gen_ability_icon.
-'''
-
-
-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 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()
-
-
-@torch.inference_mode()
-def generate_image(seed: int, truncation_psi: float, model: nn.Module,
-                   device: torch.device) -> np.ndarray:
-    seed = int(np.clip(seed, 0, np.iinfo(np.uint32).max))
-
-    z = generate_z(model.z_dim, seed, device)
-    label = torch.zeros([1, model.c_dim], device=device)
-
-    out = model(z, label, truncation_psi=truncation_psi, force_fp32=True)
-    out = (out.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
-    return out[0].cpu().numpy()
-
-
-def load_model(file_name: str, device: torch.device) -> nn.Module:
-    path = hf_hub_download(f'CorvaeOboro/gen_ability_icon' , f'{file_name}')
-    with open(path, 'rb') as f:
-        model = pickle.load(f)['G_ema']
-    model.eval()
-    model.to(device)
-    with torch.inference_mode():
-        z = torch.zeros((1, model.z_dim)).to(device)
-        label = torch.zeros([1, model.c_dim], device=device)
-        model(z, label, force_fp32=True)
-    return model
-
-
-def main():
-    args = parse_args()
-    device = torch.device(args.device)
-
-    model = load_model('gen_ability_icon_stylegan2ada_20220801.pkl', device)
-
-    func = functools.partial(generate_image, model=model, device=device)
-    func = functools.update_wrapper(func, generate_image)
-
-    gr.Interface(
-        func,
-        [
-            gr.inputs.Number(default=0, label='Seed'),
-            gr.inputs.Slider(
-                0, 2, step=0.05, default=0.7, label='Truncation psi'),
-        ],
-        gr.outputs.Image(type='numpy', label='Output'),
-        title=TITLE,
-        description=DESCRIPTION,
-        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()
+    """)
+    images_num = gr.inputs.Slider(default=1, label="Num Images", minimum=1, maximum=16, step=1)
+    interpolate = gr.inputs.Checkbox(default=False, label="Interpolate")
+    submit = gr.Button("Generate")
+    
+    
+    out = gr.Gallery()
+
+    submit.click(fn=infer, 
+               inputs=[images_num, interpolate], 
+               outputs=out)
+
+demo.launch()