update

.cache/GeneratorV2_ffhq_Arcane_210624_e350.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:697e85aed0228a507ebcd834c393033598ae6ce3d855a4d7d9aa7cab5cf8319e
+size 8608740

.cache/GeneratorV2_gldv2_Hayao.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:24cda17c176fb9c052471770228d92aee16e1c36a9ca98e7b5d586dd2d99169d
+size 8607756

.cache/GeneratorV2_gldv2_Shinkai.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a3e74a1dc12a8050f52a4f11b526c0847b7bd6612d0adbdbed0b4713a75140c1
+size 8607920

.cache/GeneratorV2_train_photo_Hayao.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1e2485770ca36d073bc2b4992b8e58501be888559919e64d61c113ba6e6d5734
+size 8605752

app.py

@@ -12,18 +12,18 @@ def inference(
     image: np.ndarray,
     style,
     imgsz=None,
+    retain_color=False,
 ):
     if imgsz is not None:
         imgsz = int(imgsz)
 
-    retain_color = False
+    retain_color = retain_color
 
     weight = {
-        "AnimeGAN_Hayao": "hayao",
-        "AnimeGAN_Shinkai": "shinkai",
         "AnimeGANv2_Hayao": "hayao:v2",
         "AnimeGANv2_Shinkai": "shinkai:v2",
         "AnimeGANv2_Arcane": "arcane:v2",
+        "AnimeGANv2_Test": "GeneratorV2_train_photo_Hayao.pt",
     }[style]
     predictor = Predictor(
         weight,
@@ -39,29 +39,23 @@ def inference(
     return anime_image, save_path
 
 
-title = "AnimeGANv2: To produce your own animation."
-description = r"""Turn your photo into anime style 😊"""
-article = r"""
-[![GitHub Stars](https://img.shields.io/github/stars/ptran1203/pytorch-animeGAN?style=social)](https://github.com/ptran1203/pytorch-animeGAN)    
-### 🗻 Demo
-
-"""
+title = "图片动漫风格转换"
+description = r"""将图片转换成动漫风格"""
 
 gr.Interface(
     fn=inference,
     inputs=[
-        gr.components.Image(label="Input"),
+        gr.components.Image(label="输入图片"),
         gr.Dropdown(
             [
-                'AnimeGAN_Hayao',
-                'AnimeGAN_Shinkai',
                 'AnimeGANv2_Hayao',
                 'AnimeGANv2_Shinkai',
                 'AnimeGANv2_Arcane',
+                'AnimeGANv2_Test',
             ],
             type="value",
             value='AnimeGANv2_Hayao',
-            label='Style'
+            label='转换风格'
         ),
         gr.Dropdown(
             [
@@ -74,24 +68,22 @@ gr.Interface(
             ],
             type="value",
             value=None,
-            label='Image size'
-        )
+            label='图片大小'
+        ),
+        gr.Checkbox(value=False, label="保留原图颜色"),
     ],
     outputs=[
-        gr.components.Image(type="numpy", label="Output (The whole image)"),
-        gr.components.File(label="Download the output image")
+        gr.components.Image(type="numpy", label="转换后图片"),
+        gr.components.File(label="下载转换图片")
     ],
     title=title,
     description=description,
-    article=article,
     allow_flagging="never",
     examples=[
         ['example/face/girl4.jpg', 'AnimeGANv2_Arcane', None],
         ['example/face/leo.jpg', 'AnimeGANv2_Arcane', None],
         ['example/face/cap.jpg', 'AnimeGANv2_Arcane', None],
         ['example/face/anne.jpg', 'AnimeGANv2_Arcane', None],
-        # ['example/boy2.jpg', 'AnimeGANv3_Arcane', "No"],
-        # ['example/cap.jpg', 'AnimeGANv3_Arcane', "No"],
         ['example/landscape/pexels-camilacarneiro-6318793.jpg', 'AnimeGANv2_Hayao', None],
         ['example/landscape/pexels-nandhukumar-450441.jpg', 'AnimeGANv2_Hayao', None],
     ]

inference.py

@@ -41,7 +41,14 @@ def profile(func):
 
 def auto_load_weight(weight, version=None, map_location=None):
     """Auto load Generator version from weight."""
+    cache_dir = os.path.expanduser(".cache")
     weight_name = os.path.basename(weight).lower()
+    cached_weight = os.path.join(cache_dir, weight_name)
+
+    # Check if the cached weight file exists
+    if os.path.exists(cached_weight):
+        weight = cached_weight
+
     if version is not None:
         version = version.lower()
         assert version in {"v1", "v2", "v3"}, f"Version {version} does not exist"