Update app.py

app.py

@@ -2,45 +2,49 @@ from PIL import Image
 import torch
 import gradio as gr
 # Загрузка моделей с разными вариантами настроек
-model1_512 = torch.hub.load("susanli2016/animegan2-pytorch", "generator", pretrained="celeba_distill")
-model2_512 = torch.hub.load("susanli2016/animegan2-pytorch", "generator", pretrained="face_paint_512_v1")
-model3_512 = torch.hub.load("susanli2016/animegan2-pytorch", "generator", pretrained="face_paint_512_v2")
-model4_512 = torch.hub.load("susanli2016/animegan2-pytorch", "generator", pretrained="paprika")
-model1_1024 = torch.hub.load("susanli2016/animegan2-pytorch", "generator", pretrained="celeba_distill_1024")
-model2_1024 = torch.hub.load("susanli2016/animegan2-pytorch", "generator", pretrained="face_paint_1024_v1")
-model3_1024 = torch.hub.load("susanli2016/animegan2-pytorch", "generator", pretrained="face_paint_1024_v2")
-model4_1024 = torch.hub.load("susanli2016/animegan2-pytorch", "generator", pretrained="paprika_1024")
+model1 = torch.hub.load("bryandlee/animegan2-pytorch:main", "generator", pretrained="celeba_distill")
+model2 = torch.hub.load("bryandlee/animegan2-pytorch:main", "generator", pretrained="face_paint_512_v1")
+model3 = torch.hub.load("bryandlee/animegan2-pytorch:main", "generator", pretrained="face_paint_512_v2")
+model4 = torch.hub.load("bryandlee/animegan2-pytorch:main", "generator", pretrained="paprika")
+def load_face2paint_model(size):
+    global face2paint
+    if size == 512:
+        face2paint = torch.hub.load('bryandlee/animegan2-pytorch:main', 'face2paint',
+                                    size=size, device="cpu", side_by_side=False)
+    elif size == 1024:
+        face2paint = torch.hub.load("bryandlee/animegan2-pytorch:main", "face2paint",
+                                    size=size, device="cpu", side_by_side=False)
+load_face2paint_model(512)  # Загрузка модели с размером 512 по умолчанию
 def inference(img, ver, size):
-    if size == '512':
-        if ver == 'Celebrity Distill':
-            out = face2paint(model1_512, img)
-        elif ver == 'Face Paint v1':
-            out = face2paint(model2_512, img)
-        elif ver == 'Face Paint v2':
-            out = face2paint(model3_512, img)
-        elif ver == 'Paprika':
-            out = face2paint(model4_512, img)
-    elif size == '1024':
-        if ver == 'Celebrity Distill':
-            out = face2paint(model1_1024, img)
-        elif ver == 'Face Paint v1':
-            out = face2paint(model2_1024, img)
-        elif ver == 'Face Paint v2':
-            out = face2paint(model3_1024, img)
-        elif ver == 'Paprika':
-            out = face2paint(model4_1024, img)
+    if size != 512:
+        load_face2paint_model(size)  # Загрузка модели с выбранным размером
+    if ver == 'Celebrity Distill':
+        out = face2paint(model1, img)
+    elif ver == 'Face Paint v1':
+        out = face2paint(model2, img)
+    elif ver == 'Face Paint v2':
+        out = face2paint(model3, img)
+    elif ver == 'Paprika':
+        out = face2paint(model4, img)
     return out
-# Создание интерфейса Gradio с выбором размера и версии модели
-gr.Interface(inference, 
-              [gr.inputs.Image(type="pil"),
-               gr.inputs.Radio(['512', '1024'], type="value", default='512', label='Size'),
-               gr.inputs.Radio(['Celebrity Distill', 'Face Paint v1', 'Face Paint v2', 'Paprika'], 
-                               type="value", 
-                               default='Celebrity Distill', 
-                               label='Version')],
-              gr.outputs.Image(type="pil"),
-              title=title,
-              description=description,
-              article=article,
-              allow_flagging=False,
-              allow_screenshot=False).launch()
+title = "AnimeGANv2"
+description = "Gradio Demo for AnimeGanv2 Face Portrait. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below. Please use a cropped portrait picture for best results similar to the examples below."
+article = ""
+# Создание интерфейса Gradio с вариантами моделей и размеров
+interface = gr.Interface(inference,
+                         [gr.inputs.Image(type="pil"),
+                          gr.inputs.Radio(['Celebrity Distill', 'Face Paint v1', 'Face Paint v2', 'Paprika'],
+                                         type="value",
+                                         default='Celebrity Distill',
+                                         label='Version'),
+                          gr.inputs.Radio([512, 1024],
+                                         type="value",
+                                         default=512,
+                                         label='Model Size (pixels)')],
+                         gr.outputs.Image(type="pil"),
+                         title=title,
+                         description=description,
+                         article=article,
+                         allow_flagging=False,
+                         allow_screenshot=False)
+interface.launch()