184 lines
4.6 KiB
Python
184 lines
4.6 KiB
Python
"""
|
|
Meta AI FastAPI Service (v2.0)
|
|
|
|
Uses metaai-api library for Meta AI image generation.
|
|
See: https://github.com/mir-ashiq/metaai-api
|
|
"""
|
|
from contextlib import asynccontextmanager
|
|
from fastapi import FastAPI, BackgroundTasks, HTTPException
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
import asyncio
|
|
import uuid
|
|
|
|
|
|
|
|
from .models import (
|
|
GenerateRequest,
|
|
GenerateResponse,
|
|
ImageResult,
|
|
TaskStatusResponse,
|
|
HealthResponse,
|
|
GrokChatRequest,
|
|
GrokChatResponse
|
|
)
|
|
from .grok_client import GrokChatClient
|
|
from .meta_crawler import meta_crawler
|
|
|
|
# Initialize Grok client
|
|
grok_client = GrokChatClient()
|
|
|
|
|
|
# Task storage (in-memory for simplicity)
|
|
tasks: dict = {}
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
"""Startup and shutdown events"""
|
|
print("[MetaAI] Starting Meta AI service...")
|
|
yield
|
|
print("[MetaAI] Shutting down...")
|
|
|
|
|
|
app = FastAPI(
|
|
title="Meta AI Image Generation Service",
|
|
description="FastAPI wrapper for Meta AI image generation using metaai-api",
|
|
version="2.0.0",
|
|
lifespan=lifespan
|
|
)
|
|
|
|
# CORS middleware
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
|
|
@app.get("/health", response_model=HealthResponse)
|
|
async def health_check():
|
|
"""Health check endpoint"""
|
|
return HealthResponse(
|
|
status="healthy",
|
|
version="2.0.0",
|
|
browser_ready=True # metaai-api handles this internally
|
|
)
|
|
|
|
|
|
@app.get("/rate-limit")
|
|
async def get_rate_limit():
|
|
"""Get current rate limiting status"""
|
|
return meta_crawler.get_rate_limit_status()
|
|
|
|
|
|
@app.post("/generate/sync", response_model=GenerateResponse)
|
|
async def generate_sync(request: GenerateRequest):
|
|
"""
|
|
Synchronous image generation - returns when complete.
|
|
|
|
Requires:
|
|
- prompt: The image generation prompt
|
|
- cookies: Facebook/Meta cookies (JSON array or string format)
|
|
"""
|
|
try:
|
|
images = await meta_crawler.generate_images(
|
|
prompt=request.prompt,
|
|
cookies=request.cookies,
|
|
num_images=request.num_images
|
|
)
|
|
|
|
return GenerateResponse(
|
|
success=True,
|
|
images=images,
|
|
error=None
|
|
)
|
|
|
|
except Exception as e:
|
|
return GenerateResponse(
|
|
success=False,
|
|
images=[],
|
|
error=str(e)
|
|
)
|
|
|
|
|
|
@app.post("/generate", response_model=GenerateResponse)
|
|
async def generate_async(request: GenerateRequest, background_tasks: BackgroundTasks):
|
|
"""
|
|
Async image generation - returns immediately with task_id.
|
|
Poll /status/{task_id} for results.
|
|
"""
|
|
task_id = str(uuid.uuid4())
|
|
|
|
tasks[task_id] = {
|
|
"status": "pending",
|
|
"images": [],
|
|
"error": None
|
|
}
|
|
|
|
async def run_generation():
|
|
try:
|
|
images = await meta_crawler.generate_images(
|
|
prompt=request.prompt,
|
|
cookies=request.cookies,
|
|
num_images=request.num_images
|
|
)
|
|
tasks[task_id] = {
|
|
"status": "completed",
|
|
"images": images,
|
|
"error": None
|
|
}
|
|
except Exception as e:
|
|
tasks[task_id] = {
|
|
"status": "failed",
|
|
"images": [],
|
|
"error": str(e)
|
|
}
|
|
|
|
# Run in background
|
|
background_tasks.add_task(asyncio.create_task, run_generation())
|
|
|
|
return GenerateResponse(
|
|
success=True,
|
|
images=[],
|
|
error=None,
|
|
task_id=task_id
|
|
)
|
|
|
|
|
|
@app.get("/status/{task_id}", response_model=TaskStatusResponse)
|
|
async def get_task_status(task_id: str):
|
|
"""Get status of async generation task"""
|
|
if task_id not in tasks:
|
|
raise HTTPException(status_code=404, detail="Task not found")
|
|
|
|
task = tasks[task_id]
|
|
return TaskStatusResponse(
|
|
task_id=task_id,
|
|
status=task["status"],
|
|
images=task["images"],
|
|
error=task["error"]
|
|
)
|
|
|
|
|
|
@app.delete("/status/{task_id}")
|
|
async def delete_task(task_id: str):
|
|
"""Clean up completed task"""
|
|
if task_id in tasks:
|
|
del tasks[task_id]
|
|
return {"deleted": True}
|
|
raise HTTPException(status_code=404, detail="Task not found")
|
|
|
|
|
|
|
|
@app.post("/grok/chat", response_model=GrokChatResponse)
|
|
async def grok_chat(request: GrokChatRequest):
|
|
"""
|
|
Chat with Grok AI
|
|
"""
|
|
try:
|
|
response = await grok_client.chat(request.message, request.history, request.cookies, request.user_agent)
|
|
return GrokChatResponse(response=response)
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|