| from fastapi import FastAPI, HTTPException |
| from pydantic import BaseModel |
| from gradio_client import Client, handle_file |
| import uvicorn |
| import os |
| import time |
|
|
| app = FastAPI(title="Omni Editor API by Xalman") |
|
|
| |
| SOURCE_SPACE = "selfit-camera/omni-image-editor" |
|
|
| def get_client(): |
| """চেষ্টা করবে মূল AI স্পেসের সাথে কানেক্ট করার""" |
| max_retries = 3 |
| for i in range(max_retries): |
| try: |
| print(f"Connecting to AI Source (Attempt {i+1})...") |
| client = Client(SOURCE_SPACE) |
| return client |
| except Exception as e: |
| print(f"Connection failed: {e}") |
| if i < max_retries - 1: |
| time.sleep(5) |
| return None |
|
|
| class EditRequest(BaseModel): |
| imageUrl: str |
| prompt: str = "enhance quality" |
|
|
| @app.get("/") |
| def home(): |
| return { |
| "status": "Running", |
| "service": "Omni-Editor-API", |
| "author": "Xalman" |
| } |
|
|
| @app.post("/predict") |
| async def predict(data: EditRequest): |
| client = get_client() |
| |
| if client is None: |
| raise HTTPException( |
| status_code=503, |
| detail="AI Source Space is currently unavailable or sleeping. Please try again in a minute." |
| ) |
| |
| try: |
| |
| result = client.predict( |
| image=handle_file(data.imageUrl), |
| edit_command=data.prompt, |
| api_name="/predict" |
| ) |
| |
| return { |
| "status": "success", |
| "result_url": result |
| } |
| except Exception as e: |
| print(f"Prediction Error: {e}") |
| raise HTTPException(status_code=500, detail=str(e)) |
|
|
| if __name__ == "__main__": |
| port = int(os.environ.get("PORT", 7860)) |
| uvicorn.run(app, host="0.0.0.0", port=port) |
| |
