Create main.py
Browse files
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")}
|