Update app.py

app.py

@@ -9,26 +9,22 @@ import onnxruntime as ort
 from PIL import Image
 from huggingface_hub import hf_hub_download
 
-
 def _yield_tags_from_txt_file(txt_file: str):
     with open(txt_file, 'r') as f:
         for line in f:
             if line:
                 yield line.strip()
 
-
 @lru_cache()
 def get_deepdanbooru_tags() -> List[str]:
     tags_file = hf_hub_download('chinoll/deepdanbooru', 'tags.txt')
     return list(_yield_tags_from_txt_file(tags_file))
 
-
 @lru_cache()
 def get_deepdanbooru_onnx() -> ort.InferenceSession:
     onnx_file = hf_hub_download('chinoll/deepdanbooru', 'deepdanbooru.onnx')
     return ort.InferenceSession(onnx_file)
 
-
 def image_preprocess(image: Image.Image) -> np.ndarray:
     if image.mode != 'RGB':
         image = image.convert('RGB')
@@ -49,10 +45,8 @@ def image_preprocess(image: Image.Image) -> np.ndarray:
     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]]:
@@ -84,8 +78,12 @@ 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():
@@ -111,4 +109,33 @@ if __name__ == '__main__':
             inputs=[gr_input_image, gr_threshold, gr_space, gr_escape, gr_confidence, gr_order],
             outputs=[gr_output_text, gr_tags],
         )
-    demo.queue(os.cpu_count()).launch()
+
+    # 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("/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")
+