Makhinur commited on
Commit
ee3e372
·
verified ·
1 Parent(s): 6b5f754

Create main.py

Browse files
Files changed (1) hide show
  1. main.py +52 -0
main.py ADDED
@@ -0,0 +1,52 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import FastAPI, File, UploadFile, Form
2
+ from fastapi.responses import JSONResponse
3
+ from fastapi.middleware.cors import CORSMiddleware
4
+ from gradio_client import Client
5
+ import os
6
+ import tempfile
7
+ import base64
8
+
9
+ app = FastAPI()
10
+
11
+ # Setup CORS
12
+ app.add_middleware(
13
+ CORSMiddleware,
14
+ allow_origins=["*"], # Adjust as needed
15
+ allow_credentials=True,
16
+ allow_methods=["*"],
17
+ allow_headers=["*"],
18
+ )
19
+
20
+ # Initialize Gradio client
21
+ hf_token = os.environ.get('HF_TOKEN')
22
+ client = Client("https://Makhinur/Image_Face_Upscale_Restoration-GFPGAN.hf.space/", hf_token=hf_token)
23
+
24
+ @app.post("/upload/")
25
+ async def upload_file(file: UploadFile = File(...), version: str = Form(...), scale: int = Form(...)):
26
+ # Save the uploaded file temporarily
27
+ with tempfile.NamedTemporaryFile(delete=False) as temp_file:
28
+ temp_file.write(await file.read())
29
+ temp_file_path = temp_file.name
30
+
31
+ try:
32
+ # Call the Gradio API
33
+ result = client.predict(temp_file_path, version, scale, api_name="/predict")
34
+
35
+ # Check if the result is valid
36
+ if result and len(result) == 2:
37
+ # Convert the image data to a base64 string
38
+ with open(result[0], "rb") as image_file:
39
+ image_data = base64.b64encode(image_file.read()).decode("utf-8")
40
+
41
+ return JSONResponse({
42
+ "sketch_image_base64": f"data:image/png;base64,{image_data}",
43
+ "result_file": result[1]
44
+ })
45
+ else:
46
+ return JSONResponse({"error": "Invalid result from the prediction API."}, status_code=500)
47
+ except Exception as e:
48
+ return JSONResponse({"error": str(e)}, status_code=500)
49
+ finally:
50
+ # Clean up the temporary file
51
+ if os.path.exists(temp_file_path):
52
+ os.unlink(temp_file_path)