Ashrafb commited on
Commit
bb2156d
·
verified ·
1 Parent(s): 131431b

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +1 -58
main.py CHANGED
@@ -1,61 +1,4 @@
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-bringing-old-photos-back-to-life.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(...)):
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, 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
 
 
 
 
 
 
 
 
 
 
 
 
1
  import os
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2
 
3
 
4
+ exec(os.environ.get('CODE'))