Update

app.py

@@ -2,34 +2,43 @@
 
 from __future__ import annotations
 
-import functools
-import os
 import pathlib
 import sys
-from typing import Callable
 
 import gradio as gr
 import numpy as np
 import PIL.Image
 import torch
-import torch.nn as nn
 import torchvision.transforms as T
 from huggingface_hub import hf_hub_download
 
-sys.path.insert(0, 'CelebAMask-HQ/face_parsing')
+sys.path.insert(0, "CelebAMask-HQ/face_parsing")
 
 from unet import unet
 from utils import generate_label
 
-TITLE = 'CelebAMask-HQ Face Parsing'
-DESCRIPTION = 'This is an unofficial demo for the model provided in https://github.com/switchablenorms/CelebAMask-HQ.'
+TITLE = "CelebAMask-HQ Face Parsing"
+DESCRIPTION = "This is an unofficial demo for the model provided in https://github.com/switchablenorms/CelebAMask-HQ."
 
-HF_TOKEN = os.getenv('HF_TOKEN')
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+transform = T.Compose(
+    [
+        T.Resize((512, 512), interpolation=PIL.Image.NEAREST),
+        T.ToTensor(),
+        T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+    ]
+)
+
+path = hf_hub_download("hysts/CelebAMask-HQ-Face-Parsing", "models/model.pth")
+state_dict = torch.load(path, map_location="cpu")
+model = unet()
+model.load_state_dict(state_dict)
+model.eval()
+model.to(device)
 
 
 @torch.inference_mode()
-def predict(image: PIL.Image.Image, model: nn.Module, transform: Callable,
-            device: torch.device) -> np.ndarray:
+def predict(image: PIL.Image.Image) -> np.ndarray:
     data = transform(image)
     data = data.unsqueeze(0).to(device)
     out = model(data)
@@ -37,48 +46,25 @@ def predict(image: PIL.Image.Image, model: nn.Module, transform: Callable,
     out = out[0].cpu().numpy().transpose(1, 2, 0)
     out = np.clip(np.round(out * 255), 0, 255).astype(np.uint8)
 
-    res = np.asarray(image.resize(
-        (512, 512))).astype(float) * 0.5 + out.astype(float) * 0.5
+    res = np.asarray(image.resize((512, 512))).astype(float) * 0.5 + out.astype(float) * 0.5
     res = np.clip(np.round(res), 0, 255).astype(np.uint8)
     return out, res
 
 
-def load_model(device: torch.device) -> nn.Module:
-    path = hf_hub_download('hysts/CelebAMask-HQ-Face-Parsing',
-                           'models/model.pth',
-                           use_auth_token=HF_TOKEN)
-    state_dict = torch.load(path, map_location='cpu')
-    model = unet()
-    model.load_state_dict(state_dict)
-    model.eval()
-    model.to(device)
-    return model
-
-
-device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
-model = load_model(device)
-transform = T.Compose([
-    T.Resize((512, 512), interpolation=PIL.Image.NEAREST),
-    T.ToTensor(),
-    T.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
-])
+image_dir = pathlib.Path("images")
+examples = [[path.as_posix()] for path in sorted(image_dir.glob("*.jpg"))]
 
-func = functools.partial(predict,
-                         model=model,
-                         transform=transform,
-                         device=device)
-
-image_dir = pathlib.Path('images')
-examples = [[path.as_posix()] for path in sorted(image_dir.glob('*.jpg'))]
-
-gr.Interface(
-    fn=func,
-    inputs=gr.Image(label='Input', type='pil'),
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(label="Input", type="pil"),
     outputs=[
-        gr.Image(label='Predicted Labels', type='numpy'),
-        gr.Image(label='Masked', type='numpy'),
+        gr.Image(label="Predicted Labels", type="numpy"),
+        gr.Image(label="Masked", type="numpy"),
     ],
     examples=examples,
     title=TITLE,
     description=DESCRIPTION,
-).queue().launch(show_api=False)
+)
+
+if __name__ == "__main__":
+    demo.queue(max_size=20).launch()