Celeskry commited on
Commit
856c733
·
verified ·
1 Parent(s): 61fc9df

Create main.py

Browse files
Files changed (1) hide show
  1. main.py +76 -0
main.py ADDED
@@ -0,0 +1,76 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # main.py
2
+ import asyncio
3
+ import os
4
+ import uuid
5
+ from typing import Optional
6
+ from fastapi import FastAPI, HTTPException, BackgroundTasks
7
+ from contextlib import asynccontextmanager
8
+ from pydantic import BaseModel
9
+ from fastapi.responses import FileResponse
10
+ from playwright.async_api import async_playwright
11
+
12
+ from app.login import GenApeLogin
13
+ from app.gen import GenApeGen
14
+
15
+ playwright = None
16
+ browser = None
17
+ context = None
18
+ semaphore = asyncio.Semaphore(5)
19
+
20
+ BASE_URL = "https://app.genape.ai/ai-image-generator"
21
+ STORAGE_PATH = "genape_state.json"
22
+
23
+ tasks_status = {}
24
+
25
+ @asynccontextmanager
26
+ async def lifespan(app: FastAPI):
27
+ global playwright, browser, context
28
+
29
+ playwright = await async_playwright().start()
30
+ browser = await playwright.chromium.launch(
31
+ headless=True,
32
+ args=["--no-sandbox", "--disable-setuid-sandbox"]
33
+ )
34
+
35
+ storage = STORAGE_PATH if os.path.exists(STORAGE_PATH) else None
36
+
37
+ context = await browser.new_context(storage_state=storage)
38
+
39
+ yield
40
+
41
+ await browser.close()
42
+ await playwright.stop()
43
+
44
+ app = FastAPI(lifespan=lifespan)
45
+
46
+ class LoginRequest(BaseModel):
47
+ email: str
48
+ password: str
49
+
50
+ @app.post("/api/login")
51
+ async def login_api(req: LoginRequest):
52
+ login_handler = GenApeLogin(context, BASE_URL, STORAGE_PATH)
53
+ res = await login_handler.login(req.email, req.password)
54
+
55
+ if not res["ok"]:
56
+ raise HTTPException(500, res["error"])
57
+
58
+ return res
59
+
60
+ @app.post("/api/gen")
61
+ async def gen_image_async(prompt: str, image_link: str = None, background_tasks: BackgroundTasks = None):
62
+ task_id = str(uuid.uuid4())
63
+ tasks_status[task_id] = {"status": "processing"}
64
+
65
+ background_tasks.add_task(run_gen_task, task_id, prompt, image_link)
66
+
67
+ return {"task_id": task_id}
68
+
69
+ async def run_gen_task(task_id: str, prompt: str, image_link: str):
70
+ gen_handler = GenApeGen(context, BASE_URL, semaphore)
71
+ res = await gen_handler.generate_image(prompt, image_link)
72
+
73
+ if res.get("ok"):
74
+ tasks_status[task_id] = {"status": "success", "result": res["img_url"]}
75
+ else:
76
+ tasks_status[task_id] = {"status": "failed", "error": res.get("error")}