update model

app.py

@@ -1,3 +1,4 @@
+import random
 import gradio as gr
 import imageio
 import numpy as np
@@ -93,14 +94,15 @@ class Model:
         detector_path = huggingface_hub.hf_hub_download("skytnt/fbanime-gan", "waifu_dect.onnx")
         anime_seg_path = huggingface_hub.hf_hub_download("skytnt/anime-seg", "isnetis.onnx")
 
-        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
+        providers = ['CPUExecutionProvider']
+        gpu_providers = ['CUDAExecutionProvider']
         g_mapping = onnx.load(g_mapping_path)
         w_avg = [x for x in g_mapping.graph.initializer if x.name == "w_avg"][0]
         w_avg = np.frombuffer(w_avg.raw_data, dtype=np.float32)[np.newaxis, :]
         w_avg = w_avg.repeat(16, axis=0)[np.newaxis, :]
         self.w_avg = w_avg
-        self.g_mapping = rt.InferenceSession(g_mapping_path, providers=providers)
-        self.g_synthesis = rt.InferenceSession(g_synthesis_path, providers=providers)
+        self.g_mapping = rt.InferenceSession(g_mapping_path, providers=gpu_providers + providers)
+        self.g_synthesis = rt.InferenceSession(g_synthesis_path, providers=gpu_providers + providers)
         self.encoder = rt.InferenceSession(encoder_path, providers=providers)
         self.detector = rt.InferenceSession(detector_path, providers=providers)
         detector_meta = self.detector.get_modelmeta().custom_metadata_map
@@ -130,7 +132,7 @@ class Model:
         mask = np.transpose(mask, (1, 2, 0))
         mask = mask[ph // 2:ph // 2 + h, pw // 2:pw // 2 + w]
         mask = transform.resize(mask, (h0, w0))
-        img0 = (img0*mask + 255*(1-mask)).astype(np.uint8)
+        img0 = (img0 * mask + 255 * (1 - mask)).astype(np.uint8)
         return img0
 
     def encode_img(self, img):
@@ -247,10 +249,12 @@ def get_thumbnail(img):
 
 
 def gen_fn(method, seed, psi1, psi2, noise):
-    z = RandomState(int(seed) + 2 ** 31).randn(1, 512) if method == 1 else np.random.randn(1, 512)
+    if method == 0:
+        seed = random.randint(0, 2 ** 32 - 1)
+    z = RandomState(int(seed)).randn(1, 1024)
     w = model.get_w(z.astype(dtype=np.float32), psi1, psi2)
     img_out = model.get_img(w, noise)
-    return img_out, w, get_thumbnail(img_out)
+    return img_out, seed, w, get_thumbnail(img_out)
 
 
 def encode_img_fn(img, noise):
@@ -259,7 +263,7 @@ def encode_img_fn(img, noise):
     img = model.remove_bg(img)
     imgs = model.detect(img, 0.2, 0.03)
     if len(imgs) == 0:
-        return "failed to detect waifu", None, None, None, None
+        return "failed to detect anime character", None, None, None, None
     w = model.encode_img(imgs[0])
     img_out = model.get_img(w, noise)
     return "success", imgs[0], img_out, w, get_thumbnail(img_out)
@@ -278,8 +282,7 @@ if __name__ == '__main__':
     app = gr.Blocks()
     with app:
         gr.Markdown("# full-body anime GAN\n\n"
-                    "![visitor badge](https://visitor-badge.glitch.me/badge?page_id=skytnt.full-body-anime-gan)\n\n"
-                    "the model is not well, just use for fun.")
+                    "![visitor badge](https://visitor-badge.glitch.me/badge?page_id=skytnt.full-body-anime-gan)\n\n")
         with gr.Tabs():
             with gr.TabItem("generate image"):
                 with gr.Row():
@@ -287,9 +290,9 @@ if __name__ == '__main__':
                         gr.Markdown("generate image randomly or by seed")
                         with gr.Row():
                             gen_input1 = gr.Radio(label="method", value="random",
-                                                  choices=["random", "use seed"], type="index")
-                            gen_input2 = gr.Number(value=1, label="seed ( int between -2^31 and 2^31 - 1 )")
-                        gen_input3 = gr.Slider(minimum=0, maximum=1, step=0.01, value=0.6, label="truncation psi 1")
+                                                  choices=["random", "seed"], type="index")
+                            gen_input2 = gr.Slider(minimum=0, maximum=2 ** 32 - 1, step=1, value=0, label="seed")
+                        gen_input3 = gr.Slider(minimum=0, maximum=1, step=0.01, value=1, label="truncation psi 1")
                         gen_input4 = gr.Slider(minimum=0, maximum=1, step=0.01, value=1, label="truncation psi 2")
                         gen_input5 = gr.Slider(minimum=0, maximum=1, step=0.01, value=1, label="noise strength")
                         with gr.Group():
@@ -304,7 +307,7 @@ if __name__ == '__main__':
                     with gr.Column():
                         gr.Markdown("you'd better upload a standing full-body image")
                         encode_img_input = gr.Image(label="input image")
-                        examples_data = [[f"examples/{x:02d}.png"] for x in range(1, 5)]
+                        examples_data = [[f"examples/{x:02d}.jpg"] for x in range(1, 5)]
                         encode_img_examples = gr.Dataset(components=[encode_img_input], samples=examples_data)
                         with gr.Group():
                             encode_img_submit = gr.Button("Run", variant="primary")
@@ -319,11 +322,10 @@ if __name__ == '__main__':
             with gr.TabItem("generate video"):
                 with gr.Row():
                     with gr.Column():
-                        gr.Markdown("## generate video between 2 images")
+                        gr.Markdown("generate video between 2 images")
                         with gr.Row():
                             with gr.Column():
-                                gr.Markdown("please select image 1")
-                                select_img1_dropdown = gr.Radio(label="source", value="current generated image",
+                                select_img1_dropdown = gr.Radio(label="Select image 1", value="current generated image",
                                                                 choices=["current generated image",
                                                                          "current encoded image"], type="index")
                                 with gr.Group():
@@ -331,8 +333,7 @@ if __name__ == '__main__':
                                 select_img1_output_img = gr.Image(label="selected image 1")
                                 select_img1_output_w = gr.Variable()
                             with gr.Column():
-                                gr.Markdown("please select image 2")
-                                select_img2_dropdown = gr.Radio(label="source", value="current generated image",
+                                select_img2_dropdown = gr.Radio(label="Select image 2", value="current generated image",
                                                                 choices=["current generated image",
                                                                          "current encoded image"], type="index")
                                 with gr.Group():
@@ -345,7 +346,7 @@ if __name__ == '__main__':
                     with gr.Column():
                         generate_video_output = gr.Video(label="output video")
         gen_submit.click(gen_fn, [gen_input1, gen_input2, gen_input3, gen_input4, gen_input5],
-                         [gen_output1, select_img_input_w1, select_img_input_img1])
+                         [gen_output1, gen_input2, select_img_input_w1, select_img_input_img1])
         encode_img_submit.click(encode_img_fn, [encode_img_input, gen_input5],
                                 [encode_img_output1, encode_img_output2, encode_img_output3, select_img_input_w2,
                                  select_img_input_img2])