Update app.py

app.py

@@ -1,7 +1,5 @@
 import argparse
 import os
-from typing import Optional
-import io
 
 import gradio as gr
 import huggingface_hub
@@ -9,9 +7,6 @@ import numpy as np
 import onnxruntime as rt
 import pandas as pd
 from PIL import Image
-from fastapi import FastAPI, File, UploadFile, Form
-from fastapi.responses import JSONResponse
-import uvicorn
 
 TITLE = "WaifuDiffusion Tagger"
 DESCRIPTION = """
@@ -62,9 +57,6 @@ def load_labels(dataframe) -> list[str]:
     return tag_names, rating_indexes, general_indexes, character_indexes
 
 def mcut_threshold(probs):
-    """
-    Maximum Cut Thresholding (MCut)
-    """
     sorted_probs = probs[probs.argsort()[::-1]]
     difs = sorted_probs[:-1] - sorted_probs[1:]
     t = difs.argmax()
@@ -127,11 +119,11 @@ class Predictor:
     def predict(
         self,
         image,
-        model_repo=SWINV2_MODEL_DSV3_REPO,
-        general_thresh=0.35,
-        general_mcut_enabled=False,
-        character_thresh=0.85,
-        character_mcut_enabled=False,
+        model_repo,
+        general_thresh,
+        general_mcut_enabled,
+        character_thresh,
+        character_mcut_enabled,
     ):
         self.load_model(model_repo)
         
@@ -168,25 +160,10 @@ class Predictor:
 
         return sorted_general_strings, rating, character_res, general_res
 
-predictor = Predictor()
-
-@app.post("/tagging")
-async def tagging_endpoint(
-    image: UploadFile = File(...),
-    threshold: Optional[float] = Form(0.05)
-):
-    image_data = await image.read()
-    pil_image = Image.open(io.BytesIO(image_data)).convert("RGBA")
-    sorted_general_strings, _, _, _ = predictor.predict(
-        pil_image, 
-        general_thresh=threshold
-    )
-    tags = sorted_general_strings.split(", ")
-    return JSONResponse(content={"tags": tags})
-
-def create_demo():
+def main():
     args = parse_args()
-    
+    predictor = Predictor()
+
     dropdown_list = [
         SWINV2_MODEL_DSV3_REPO,
         CONV_MODEL_DSV3_REPO,
@@ -287,10 +264,8 @@ def create_demo():
         )
         
         demo.queue(max_size=10)
-        return demo
-
-app = FastAPI()
-app = gr.mount_gradio_app(app, create_demo(), path="/")
+    
+    demo.launch()
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    main()