Ashrafb commited on
Commit
cb4762a
1 Parent(s): db9d314

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +58 -1
main.py CHANGED
@@ -1,4 +1,61 @@
 
 
 
 
 
 
 
 
 
1
  import os
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2
 
3
 
4
- exec(os.environ.get('API'))
 
1
+ from fastapi import FastAPI, File, UploadFile
2
+ from fastapi import FastAPI, File, UploadFile, Form, Request
3
+ from fastapi.responses import HTMLResponse, FileResponse
4
+ from fastapi.staticfiles import StaticFiles
5
+ from fastapi.templating import Jinja2Templates
6
+ from fastapi import FastAPI, File, UploadFile, HTTPException
7
+ from fastapi.responses import JSONResponse
8
+ from fastapi.responses import StreamingResponse
9
+ from gradio_client import Client
10
  import os
11
+ import io
12
+ app = FastAPI()
13
+
14
+ hf_token = os.environ.get('HF_TOKEN')
15
+ client = Client("https://ashrafb-ifur.hf.space/", hf_token=hf_token)
16
+
17
+ import tempfile
18
+
19
+
20
+
21
+ import base64
22
+ from fastapi.middleware.cors import CORSMiddleware
23
+
24
+ app.add_middleware(
25
+ CORSMiddleware,
26
+ allow_origins=["*"], # Adjust as needed, '*' allows requests from any origin
27
+ allow_credentials=True,
28
+ allow_methods=["*"],
29
+ allow_headers=["*"],
30
+ )
31
+
32
+
33
+ @app.post("/upload/")
34
+ async def upload_file(file: UploadFile = File(...), version: str = Form(...), scale: int = Form(...)):
35
+ with tempfile.NamedTemporaryFile(delete=False) as temp_file:
36
+ temp_file.write(await file.read())
37
+ temp_file_path = temp_file.name
38
+
39
+ try:
40
+ result = client.predict(temp_file_path, version, scale, api_name="/predict")
41
+
42
+
43
+ # Check if the result is valid
44
+ if result and len(result) == 2:
45
+ # Convert the image data to base64 string
46
+ with open(result[0], "rb") as image_file:
47
+ image_data = base64.b64encode(image_file.read()).decode("utf-8")
48
+
49
+ return {
50
+ "sketch_image_base64": f"data:image/png;base64,{image_data}",
51
+ "result_file": result[1]
52
+ }
53
+ else:
54
+ return {"error": "Invalid result from the prediction API."}
55
+ except Exception as e:
56
+ return {"error": str(e)}
57
+ finally:
58
+ if os.path.exists(temp_file_path):
59
+ os.unlink(temp_file_path)
60
 
61