Upload 7 files

app.py

@@ -1,94 +1,7 @@
+import fn
 import gradio as gr
-import os
-import torch
-import numpy as np
-from PIL import Image
-import cv2
-import pytorch_lightning as pl
-from model import ISNetDIS, ISNetGTEncoder, U2NET, U2NET_full2, U2NET_lite2, MODNet
 
-def get_mask(model, input_img):
-    h, w = input_img.shape[0], input_img.shape[1]
-    ph, pw = 0, 0
-    tmpImg = np.zeros([h, w, 3], dtype=np.float16)
-    tmpImg[ph // 2:ph // 2 + h, pw // 2:pw // 2 + w] = cv2.resize(input_img, (w, h)) / 255
-    tmpImg = tmpImg.transpose((2, 0, 1))
-    tmpImg = torch.from_numpy(tmpImg).unsqueeze(0).type(torch.FloatTensor).to(model.device)
-    with torch.no_grad():
-        pred = model(tmpImg)
-        pred = pred[0, :, ph // 2:ph // 2 + h, pw // 2:pw // 2 + w]
-        pred = cv2.resize(pred.cpu().numpy().transpose((1, 2, 0)), (w, h))[:, :, np.newaxis]
-        return pred
-
-def get_net(net_name):
-    if net_name == "isnet":
-        return ISNetDIS()
-    elif net_name == "isnet_is":
-        return ISNetDIS()
-    elif net_name == "isnet_gt":
-        return ISNetGTEncoder()
-    elif net_name == "u2net":
-        return U2NET_full2()
-    elif net_name == "u2netl":
-        return U2NET_lite2()
-    elif net_name == "modnet":
-        return MODNet()
-    raise NotImplemented
-
-# from anime-segmentation.train
-class AnimeSegmentation(pl.LightningModule):
-    def __init__(self, net_name):
-        super().__init__()
-        assert net_name in ["isnet_is", "isnet", "isnet_gt", "u2net", "u2netl", "modnet"]
-        self.net = get_net(net_name)
-        if net_name == "isnet_is":
-            self.gt_encoder = get_net("isnet_gt")
-            for param in self.gt_encoder.parameters():
-                param.requires_grad = False
-        else:
-            self.gt_encoder = None
-
-    @classmethod
-    def try_load(cls, net_name, ckpt_path, map_location=None):
-        state_dict = torch.load(ckpt_path, map_location=map_location)
-        if "epoch" in state_dict:
-            return cls.load_from_checkpoint(ckpt_path, net_name=net_name, map_location=map_location)
-        else:
-            model = cls(net_name)
-            if any([k.startswith("net.") for k, v in state_dict.items()]):
-                model.load_state_dict(state_dict)
-            else:
-                model.net.load_state_dict(state_dict)
-            return model
-
-    def forward(self, x):
-        if isinstance(self.net, ISNetDIS):
-            return self.net(x)[0][0].sigmoid()
-        if isinstance(self.net, ISNetGTEncoder):
-            return self.net(x)[0][0].sigmoid()
-        elif isinstance(self.net, U2NET):
-            return self.net(x)[0].sigmoid()
-        elif isinstance(self.net, MODNet):
-            return self.net(x, True)[2]
-        raise NotImplemented
-    
-def animeseg(image):
-    if not image:
-        return None
-
-    if torch.cuda.is_available():
-        device = 'cuda'
-    else:
-        device = 'cpu'
-    
-    model = AnimeSegmentation.try_load('isnet_is', 'anime-seg/isnetis.ckpt', device)
-    model.eval()
-    model.to(device)
-
-    img = np.array(image, dtype=np.uint8)
-    mask = get_mask(model, img)
-    img = np.concatenate((mask * img + 1 - mask, mask * 255), axis=2).astype(np.uint8)
-    return img
+fn.load_model()
 
 with gr.Blocks() as demo:
     title = gr.Markdown('# katanuki')
@@ -97,9 +10,10 @@ with gr.Blocks() as demo:
         dst_image = gr.Image(label="Result", interactive=False, type="numpy")
 
     src_image.change(
-        fn=animeseg,
+        fn=fn.animeseg,
         inputs=[src_image],
         outputs=[dst_image],
     )
 
-demo.launch()
+if __name__ == '__main__':
+    demo.launch()

fn.py

@@ -0,0 +1,101 @@
+import gradio as gr
+import os
+import torch
+import numpy as np
+from PIL import Image
+import cv2
+import pytorch_lightning as pl
+from model import ISNetDIS, ISNetGTEncoder, U2NET, U2NET_full2, U2NET_lite2, MODNet
+
+model = None
+
+def get_mask(model, input_img):
+    h, w = input_img.shape[0], input_img.shape[1]
+    ph, pw = 0, 0
+    tmpImg = np.zeros([h, w, 3], dtype=np.float16)
+    tmpImg[ph // 2:ph // 2 + h, pw // 2:pw // 2 + w] = cv2.resize(input_img, (w, h)) / 255
+    tmpImg = tmpImg.transpose((2, 0, 1))
+    tmpImg = torch.from_numpy(tmpImg).unsqueeze(0).type(torch.FloatTensor).to(model.device)
+    with torch.no_grad():
+        pred = model(tmpImg)
+        pred = pred[0, :, ph // 2:ph // 2 + h, pw // 2:pw // 2 + w]
+        pred = cv2.resize(pred.cpu().numpy().transpose((1, 2, 0)), (w, h))[:, :, np.newaxis]
+        return pred
+
+def get_net(net_name):
+    if net_name == "isnet":
+        return ISNetDIS()
+    elif net_name == "isnet_is":
+        return ISNetDIS()
+    elif net_name == "isnet_gt":
+        return ISNetGTEncoder()
+    elif net_name == "u2net":
+        return U2NET_full2()
+    elif net_name == "u2netl":
+        return U2NET_lite2()
+    elif net_name == "modnet":
+        return MODNet()
+    raise NotImplemented
+
+# from anime-segmentation.train
+class AnimeSegmentation(pl.LightningModule):
+    def __init__(self, net_name):
+        super().__init__()
+        assert net_name in ["isnet_is", "isnet", "isnet_gt", "u2net", "u2netl", "modnet"]
+        self.net = get_net(net_name)
+        if net_name == "isnet_is":
+            self.gt_encoder = get_net("isnet_gt")
+            for param in self.gt_encoder.parameters():
+                param.requires_grad = False
+        else:
+            self.gt_encoder = None
+
+    @classmethod
+    def try_load(cls, net_name, ckpt_path, map_location=None):
+        state_dict = torch.load(ckpt_path, map_location=map_location)
+        if "epoch" in state_dict:
+            return cls.load_from_checkpoint(ckpt_path, net_name=net_name, map_location=map_location)
+        else:
+            model = cls(net_name)
+            if any([k.startswith("net.") for k, v in state_dict.items()]):
+                model.load_state_dict(state_dict)
+            else:
+                model.net.load_state_dict(state_dict)
+            return model
+
+    def forward(self, x):
+        if isinstance(self.net, ISNetDIS):
+            return self.net(x)[0][0].sigmoid()
+        if isinstance(self.net, ISNetGTEncoder):
+            return self.net(x)[0][0].sigmoid()
+        elif isinstance(self.net, U2NET):
+            return self.net(x)[0].sigmoid()
+        elif isinstance(self.net, MODNet):
+            return self.net(x, True)[2]
+        raise NotImplemented
+
+def load_model():
+    global model
+
+    if torch.cuda.is_available():
+        device = 'cuda'
+    else:
+        device = 'cpu'
+    
+    model = AnimeSegmentation.try_load('isnet_is', 'anime-seg/isnetis.ckpt', device)
+    model.eval()
+    model.to(device)
+    
+def animeseg(image):
+    global model
+
+    if not image:
+        return None
+
+    if not model:
+        model = load_model()
+
+    img = np.array(image, dtype=np.uint8)
+    mask = get_mask(model, img)
+    img = np.concatenate((mask * img + 1 - mask, mask * 255), axis=2).astype(np.uint8)
+    return img

install.bat

@@ -0,0 +1,56 @@
+@echo off
+
+rem -------------------------------------------
+rem NOT guaranteed to work on Windows
+
+set REPOS=https://huggingface.co/spaces/aka7774/katanuki
+set APPDIR=katanuki
+set VENV=venv
+
+rem -------------------------------------------
+
+set INSTALL_DIR=%~dp0
+cd /d %INSTALL_DIR%
+
+:git_clone
+set DL_URL=%REPOS%
+set DL_DST=%APPDIR%
+git clone %DL_URL% %APPDIR%
+if exist %DL_DST% goto install_python
+
+set DL_URL=https://github.com/git-for-windows/git/releases/download/v2.41.0.windows.3/PortableGit-2.41.0.3-64-bit.7z.exe
+set DL_DST=PortableGit-2.41.0.3-64-bit.7z.exe
+curl -L -o %DL_DST% %DL_URL%
+if not exist %DL_DST% bitsadmin /transfer dl %DL_URL% %DL_DST%
+%DL_DST% -y
+del %DL_DST%
+
+set GIT=%INSTALL_DIR%PortableGit\bin\git
+%GIT% clone %REPOS%
+
+:install_python
+set DL_URL=https://github.com/indygreg/python-build-standalone/releases/download/20240107/cpython-3.10.13+20240107-i686-pc-windows-msvc-shared-install_only.tar.gz
+set DL_DST="%INSTALL_DIR%python.tar.gz"
+curl -L -o %DL_DST% %DL_URL%
+if not exist %DL_DST% bitsadmin /transfer dl %DL_URL% %DL_DST%
+tar -xzf %DL_DST%
+
+set PYTHON=%INSTALL_DIR%python\python.exe
+set PATH=%PATH%;%INSTALL_DIR%python310\Scripts
+
+:install_venv
+cd %APPDIR%
+%PYTHON% -m venv %VENV%
+set PYTHON=%VENV%\Scripts\python.exe
+
+:install_pip
+set DL_URL=https://bootstrap.pypa.io/get-pip.py
+set DL_DST=%INSTALL_DIR%get-pip.py
+curl -o %DL_DST% %DL_URL%
+if not exist %DL_DST% bitsadmin /transfer dl %DL_URL% %DL_DST%
+%PYTHON% %DL_DST%
+
+%PYTHON% -m pip install gradio
+%PYTHON% -m pip install -r requirements.txt
+
+pause

main.py

@@ -0,0 +1,43 @@
+import os
+import sys
+import time
+import signal
+import psutil
+import io
+
+from fastapi import FastAPI, Request, status, Form, UploadFile
+from fastapi.staticfiles import StaticFiles
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel, Field
+from fastapi.exceptions import RequestValidationError
+from fastapi.responses import JSONResponse
+
+import fn
+import gradio as gr
+from app import demo
+
+app = FastAPI()
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=['*'],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+gr.mount_gradio_app(app, demo, path="/gradio")
+
+fn.load_model()
+
+@app.post("/katanuki")
+async def katanuki_image(file: UploadFile = Form(...)):
+    try:
+        file_content = await file.read()
+        file_stream = io.BytesIO(file_content)
+
+		dst_image = fn.animeseg(src_image)
+
+        return Response(content=dst_image, media_type="image/webp")
+    except Exception as e:
+        return {"error": str(e)}

requirements.txt

@@ -1,6 +1,9 @@
+fastapi
+uvicorn
 opencv-python
 pytorch_lightning
 torch
 torchvision
 numpy
 scipy
+python-multipart

venv.sh

@@ -0,0 +1,7 @@
+#!/usr/bin/bash
+
+python3 -m venv venv
+curl -kL https://bootstrap.pypa.io/get-pip.py | venv/bin/python
+
+venv/bin/python -m pip install gradio
+venv/bin/python -m pip install -r requirements.txt