File size: 2,715 Bytes
856c733 eb5387e 856c733 13df47b 856c733 13df47b fc5602d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | # main.py
import asyncio
import os
import uuid
from typing import Optional
from fastapi import FastAPI, HTTPException, BackgroundTasks
from contextlib import asynccontextmanager
from pydantic import BaseModel
from fastapi.responses import FileResponse
from playwright.async_api import async_playwright
from app.login_logic import GenApeLogin
from app.gen_logic import GenApeGen
playwright = None
browser = None
context = None
semaphore = asyncio.Semaphore(5)
BASE_URL = "https://app.genape.ai/ai-image-generator"
STORAGE_PATH = "genape_state.json"
tasks_status = {}
@asynccontextmanager
async def lifespan(app: FastAPI):
global playwright, browser, context
playwright = await async_playwright().start()
browser = await playwright.chromium.launch(
headless=True,
args=["--no-sandbox", "--disable-setuid-sandbox"]
)
storage = STORAGE_PATH if os.path.exists(STORAGE_PATH) else None
context = await browser.new_context(storage_state=storage)
yield
await browser.close()
await playwright.stop()
app = FastAPI(lifespan=lifespan)
class LoginRequest(BaseModel):
email: str
password: str
@app.post("/api/login")
async def login_api(req: LoginRequest):
login_handler = GenApeLogin(context, BASE_URL, STORAGE_PATH)
res = await login_handler.login(req.email, req.password)
if not res["ok"]:
raise HTTPException(500, res["error"])
return res
@app.post("/api/gen")
async def gen_image_async(prompt: str, image_link: str = None, background_tasks: BackgroundTasks = None):
task_id = str(uuid.uuid4())
tasks_status[task_id] = {"status": "processing"}
background_tasks.add_task(run_gen_task, task_id, prompt, image_link)
return {"task_id": task_id}
@app.get("/api/task/{task_id}")
async def get_task_status(task_id: str):
task = tasks_status.get(task_id)
if not task:
raise HTTPException(status_code=404, detail="Task ID không tồn tại")
return task
async def run_gen_task(task_id: str, prompt: str, image_link: Optional[str]):
gen_handler = GenApeGen(context, BASE_URL, semaphore)
try:
res = await gen_handler.generate_image(prompt, image_link)
if res.get("ok"):
tasks_status[task_id] = {
"status": "success",
"result": res["img_url"]
}
else:
tasks_status[task_id] = {
"status": "failed",
"error": res.get("error", "Unknown error")
}
except Exception as e:
tasks_status[task_id] = {
"status": "failed",
"error": str(e)
}
@app.get("/")
async def root():
return {"status": "running"} |