Update app.py

app.py

@@ -8,6 +8,13 @@ import numpy as np
 import onnxruntime as ort
 from PIL import Image
 from huggingface_hub import hf_hub_download
+import io
+
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import JSONResponse
+import uvicorn
+
+app = FastAPI()
 
 def _yield_tags_from_txt_file(txt_file: str):
     with open(txt_file, 'r') as f:
@@ -39,17 +46,26 @@ def image_preprocess(image: Image.Image) -> np.ndarray:
     height_pad_right = 512 - f_height - height_pad_left
     width_pad_left = (512 - f_width) // 2
     width_pad_right = 512 - f_width - width_pad_left
-    data = np.pad(data, ((height_pad_left, height_pad_right), (width_pad_left, width_pad_right), (0, 0)),
-                  mode='constant', constant_values=0.0)
+    data = np.pad(
+        data,
+        ((height_pad_left, height_pad_right), (width_pad_left, width_pad_right), (0, 0)),
+        mode='constant',
+        constant_values=0.0
+    )
 
     assert data.shape == (512, 512, 3), f'Shape (512, 512, 3) expected, but {data.shape!r} found.'
     return data.reshape((1, 512, 512, 3))  # B x H x W x C
 
 RE_SPECIAL = re.compile(r'([\\()])')
 
-def image_to_deepdanbooru_tags(image: Image.Image, threshold: float,
-                               use_spaces: bool, use_escape: bool, include_ranks: bool, score_descend: bool) \
-        -> Tuple[str, Mapping[str, float]]:
+def image_to_deepdanbooru_tags(
+    image: Image.Image,
+    threshold: float,
+    use_spaces: bool,
+    use_escape: bool,
+    include_ranks: bool,
+    score_descend: bool
+) -> Tuple[str, Mapping[str, float]]:
     tags = get_deepdanbooru_tags()
     session = get_deepdanbooru_onnx()
     input_name = session.get_inputs()[0].name
@@ -78,64 +94,62 @@ def image_to_deepdanbooru_tags(image: Image.Image, threshold: float,
 
     return output_text, filtered_tags
 
-if __name__ == '__main__':
-    import io
-    from fastapi import FastAPI, File, UploadFile
-    from fastapi.responses import JSONResponse
-    from fastapi.middleware.cors import CORSMiddleware
-
-    with gr.Blocks() as demo:
-        with gr.Row():
-            with gr.Column():
-                gr_input_image = gr.Image(type='pil', label='Original Image')
-                gr_threshold = gr.Slider(0.0, 1.0, 0.5, label='Tagging Confidence Threshold')
-                with gr.Row():
-                    gr_space = gr.Checkbox(value=False, label='Use Space Instead Of _')
-                    gr_escape = gr.Checkbox(value=True, label='Use Text Escape')
-                    gr_confidence = gr.Checkbox(value=False, label='Keep Confidences')
-                    gr_order = gr.Checkbox(value=True, label='Descend By Confidence')
-
-                gr_btn_submit = gr.Button(value='Tagging', variant='primary')
-
-            with gr.Column():
-                with gr.Tabs():
-                    with gr.Tab("Tags"):
-                        gr_tags = gr.Label(label='Tags')
-                    with gr.Tab("Exported Text"):
-                        gr_output_text = gr.TextArea(label='Exported Text')
-
-        gr_btn_submit.click(
-            image_to_deepdanbooru_tags,
-            inputs=[gr_input_image, gr_threshold, gr_space, gr_escape, gr_confidence, gr_order],
-            outputs=[gr_output_text, gr_tags],
-        )
-
-    # Get the FastAPI app from Gradio Blocks
-    app = demo.app
-
-    # Allow cross-origin requests (optional, useful for testing)
-    origins = ["*"]
-    app.add_middleware(
-        CORSMiddleware,
-        allow_origins=origins,
-        allow_methods=["*"],
-        allow_headers=["*"],
+@app.post("/tagging")
+async def tagging_endpoint(image: UploadFile = File(...)):
+    image_data = await image.read()
+    pil_image = Image.open(io.BytesIO(image_data)).convert("RGB")
+    # Use default parameters or customize as needed
+    output_text, filtered_tags = image_to_deepdanbooru_tags(
+        pil_image,
+        threshold=0.5,
+        use_spaces=False,
+        use_escape=False,
+        include_ranks=False,
+        score_descend=True
+    )
+    tags = list(filtered_tags.keys())
+    return JSONResponse(content={"tags": tags})
+
+def gradio_interface(
+    image: Image.Image,
+    threshold: float,
+    use_spaces: bool,
+    use_escape: bool,
+    include_ranks: bool,
+    score_descend: bool
+):
+    output_text, filtered_tags = image_to_deepdanbooru_tags(
+        image, threshold, use_spaces, use_escape, include_ranks, score_descend
+    )
+    return output_text, filtered_tags
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        with gr.Column():
+            gr_input_image = gr.Image(type='pil', label='Original Image')
+            gr_threshold = gr.Slider(0.0, 1.0, 0.5, label='Tagging Confidence Threshold')
+            with gr.Row():
+                gr_space = gr.Checkbox(value=False, label='Use Space Instead Of _')
+                gr_escape = gr.Checkbox(value=True, label='Use Text Escape')
+                gr_confidence = gr.Checkbox(value=False, label='Keep Confidences')
+                gr_order = gr.Checkbox(value=True, label='Descend By Confidence')
+
+            gr_btn_submit = gr.Button(value='Tagging', variant='primary')
+
+        with gr.Column():
+            with gr.Tabs():
+                with gr.Tab("Tags"):
+                    gr_tags = gr.Label(label='Tags')
+                with gr.Tab("Exported Text"):
+                    gr_output_text = gr.TextArea(label='Exported Text')
+
+    gr_btn_submit.click(
+        gradio_interface,
+        inputs=[gr_input_image, gr_threshold, gr_space, gr_escape, gr_confidence, gr_order],
+        outputs=[gr_output_text, gr_tags],
     )
 
-    @app.post("/api/analyze_image")
-    async def analyze_image(file: UploadFile = File(...)):
-        contents = await file.read()
-        image = Image.open(io.BytesIO(contents))
-        output_text, filtered_tags = image_to_deepdanbooru_tags(
-            image,
-            threshold=0.5,
-            use_spaces=False,
-            use_escape=True,
-            include_ranks=False,
-            score_descend=True
-        )
-        return JSONResponse(content=filtered_tags)
-
-    # Launch the Gradio app
-    demo.queue(concurrency_count=os.cpu_count()).launch(server_name="0.0.0.0")
+app = gr.mount_gradio_app(app, demo, path="/")
 
+if __name__ == '__main__':
+    uvicorn.run(app, host='0.0.0.0', port=7860)