Update main.py

main.py

@@ -22,6 +22,38 @@ app = FastAPI(
 CAPTIONS_BASE_URL = "https://core.captions-web-api.xyz/proxy/v1/gen-ai/image"
 BEARER_TOKEN = os.getenv("CAPTIONS_BEARER_TOKEN", "eyJhbGciOiJSUzI1NiIsImtpZCI6IjU3YmZiMmExMWRkZmZjMGFkMmU2ODE0YzY4NzYzYjhjNjg3NTgxZDgiLCJ0eXAiOiJKV1QifQ.eyJnb29nbGUiOnRydWUsImlzcyI6Imh0dHBzOi8vc2VjdXJldG9rZW4uZ29vZ2xlLmNvbS9jYXB0aW9ucy1mNmRlOSIsImF1ZCI6ImNhcHRpb25zLWY2ZGU5IiwiYXV0aF90aW1lIjoxNzU1MzYyODEzLCJ1c2VyX2lkIjoic3hWek5XaUYyempXYmUxTjNjd3UiLCJzdWIiOiJzeFZ6TldpRjJ6aldiZTFOM2N3dSIsImlhdCI6MTc1NTM2MjgxMywiZXhwIjoxNzU1MzY2NDEzLCJmaXJlYmFzZSI6eyJpZGVudGl0aWVzIjp7fSwic2lnbl9pbl9wcm92aWRlciI6ImN1c3RvbSJ9fQ.jGuhWp-w8jlGy8xmMjqOyig_LVcr53udFgMjrQTJtKtE_J_iVkvMLncO2TnJ2BquoEp9pwVlZIG-imlFe6Uhtz95-t1oHENf5yzUWu3HocFsNVeAZh9avi_iObSYM_pFOT9lwRNzk1oMa6LbwViuVgTXvHDse9T4_nDfmCBbWngWksh1_JGtnrK2qPb5YD8Hr26itDRMx8mzUr2cQqtU9mU0R910CROqsNaQ9ovemeGe-2RT-hZku4VVYAMDOdvcFsgcf_BJTLRikmc3T7Ekx8T0KM6ZpTgr34wtnl7rpDBNOX0cOSYu3NEUDBnhNJKmPl5qL08gcYEur1ijP2mcTA")
 
+# Model mappings from OpenAI model names to Captions model IDs
+MODEL_MAPPINGS = {
+    "dall-e-3": "openai-dalle-3",
+    "dall-e-2": "openai-dalle-3",  # Fallback to dalle-3
+    "gpt-4o": "openai-gpt-4o-image",
+    "google-imagen-3": "google-imagen-3",
+    "imagen-3": "google-imagen-3",
+    "luma-photon": "luma-photon",
+    "photon": "luma-photon",
+    "flux-1-1-pro": "bfl-flux-1-1-pro",
+    "flux": "bfl-flux-1-1-pro",
+    "ideogram-v1": "ideogram-v1",
+    "ideogram": "ideogram-v1",
+    "recraft-v3": "recraft-v3",
+    "recraft": "recraft-v3",
+    "stable-diffusion-3-5": "stable-diffusion-3-5-large",
+    "sd-3-5": "stable-diffusion-3-5-large",
+    "stable-diffusion": "stable-diffusion-3-5-large"
+}
+
+# Available models information
+AVAILABLE_MODELS = {
+    "google-imagen-3": {"name": "Imagen 3", "provider": "Google"},
+    "openai-gpt-4o-image": {"name": "GPT-4o", "provider": "OpenAI"},
+    "luma-photon": {"name": "Photon", "provider": "Luma AI"},
+    "bfl-flux-1-1-pro": {"name": "Flux 1.1 Pro", "provider": "Black Forest Labs"},
+    "ideogram-v1": {"name": "Ideogram V1", "provider": "Ideogram"},
+    "openai-dalle-3": {"name": "DALL-E 3 HD", "provider": "OpenAI"},
+    "recraft-v3": {"name": "Recraft V3", "provider": "Recraft"},
+    "stable-diffusion-3-5-large": {"name": "SD 3.5", "provider": "Stability AI"}
+}
+
 # OpenAI-compatible request models
 class ImageGenerationRequest(BaseModel):
     prompt: str = Field(..., description="A text description of the desired image(s)")
@@ -57,6 +89,10 @@ class CaptionsStatusRequest(BaseModel):
 # In-memory storage for operation tracking (use Redis in production)
 operations_store = {}
 
+def get_captions_model_id(openai_model: str) -> str:
+    """Convert OpenAI model name to Captions model ID"""
+    return MODEL_MAPPINGS.get(openai_model, "openai-dalle-3")  # Default to DALL-E 3
+
 def get_aspect_ratio_from_size(size: str) -> int:
     """Convert OpenAI size format to Captions aspect ratio"""
     size_map = {
@@ -68,7 +104,7 @@ def get_aspect_ratio_from_size(size: str) -> int:
     }
     return size_map.get(size, 1)
 
-async def submit_image_generation(prompt: str, size: str = "1024x1024") -> str:
+async def submit_image_generation(prompt: str, model: str = "dall-e-3", size: str = "1024x1024") -> str:
     """Submit image generation request to Captions API"""
     headers = {
         "accept": "application/json, text/plain, */*",
@@ -82,7 +118,7 @@ async def submit_image_generation(prompt: str, size: str = "1024x1024") -> str:
     }
     
     payload = {
-        "modelId": "openai-gpt-4o-image",
+        "modelId": get_captions_model_id(model),
         "prompt": prompt,
         "aspectRatio": get_aspect_ratio_from_size(size),
         "magicPrompt": False,
@@ -155,21 +191,77 @@ async def check_generation_status(operation_id: str) -> dict:
 async def wait_for_completion(operation_id: str, max_wait_time: int = 300) -> dict:
     """Wait for image generation to complete with polling"""
     start_time = datetime.now()
+    retry_count = 0
+    max_retries = 3
     
     while True:
-        status_data = await check_generation_status(operation_id)
-        
-        # State 2 means completed
-        if status_data.get("state") == 2:
-            return status_data["complete"]
-        
-        # Check if we've exceeded max wait time
-        elapsed = (datetime.now() - start_time).total_seconds()
-        if elapsed > max_wait_time:
-            raise HTTPException(status_code=408, detail="Image generation timeout")
+        try:
+            status_data = await check_generation_status(operation_id)
+            retry_count = 0  # Reset retry count on successful request
+            
+            # State 2 means completed
+            if status_data.get("state") == 2:
+                if "complete" in status_data:
+                    return status_data["complete"]
+                else:
+                    raise HTTPException(status_code=500, detail="Generation completed but no result data")
+            
+            # State 3 means failed
+            if status_data.get("state") == 3:
+                raise HTTPException(status_code=500, detail="Image generation failed")
+            
+            # Check if we've exceeded max wait time
+            elapsed = (datetime.now() - start_time).total_seconds()
+            if elapsed > max_wait_time:
+                raise HTTPException(status_code=408, detail="Image generation timeout")
+            
+            # Log progress
+            if status_data.get("state") == 1:
+                logger.info(f"Operation {operation_id} still processing...")
+            
+            # Wait before next poll (progressive backoff)
+            wait_time = min(5, 2 + (elapsed / 60))  # Start at 2s, increase to max 5s
+            await asyncio.sleep(wait_time)
+            
+        except HTTPException:
+            raise
+        except Exception as e:
+            retry_count += 1
+            if retry_count >= max_retries:
+                logger.error(f"Max retries exceeded for operation {operation_id}: {e}")
+                raise HTTPException(status_code=500, detail="Failed to check generation status after multiple retries")
+            
+            logger.warning(f"Retry {retry_count}/{max_retries} for operation {operation_id}: {e}")
+            await asyncio.sleep(2 ** retry_count)  # Exponential backoff
+
+@app.get("/v1/models")
+async def list_models():
+    """List available models compatible with OpenAI format"""
+    models = []
+    for model_id, info in AVAILABLE_MODELS.items():
+        # Add both the Captions ID and common aliases
+        models.append({
+            "id": model_id,
+            "object": "model",
+            "created": 1234567890,  # Static timestamp
+            "owned_by": info["provider"].lower().replace(" ", "-"),
+            "name": info["name"],
+            "provider": info["provider"]
+        })
         
-        # Wait before next poll
-        await asyncio.sleep(2)
+        # Add OpenAI-style aliases
+        for alias, captions_id in MODEL_MAPPINGS.items():
+            if captions_id == model_id and alias not in [m["id"] for m in models]:
+                models.append({
+                    "id": alias,
+                    "object": "model", 
+                    "created": 1234567890,
+                    "owned_by": info["provider"].lower().replace(" ", "-"),
+                    "name": info["name"],
+                    "provider": info["provider"]
+                })
+    
+    return {"object": "list", "data": models}
 
 @app.post("/v1/images/generations", response_model=ImageGenerationResponse)
 async def create_image(request: ImageGenerationRequest):
@@ -178,14 +270,31 @@ async def create_image(request: ImageGenerationRequest):
     Compatible with OpenAI's image generation API.
     """
     try:
-        logger.info(f"Received image generation request: {request.prompt}")
+        logger.info(f"Received image generation request: prompt='{request.prompt[:100]}...', model='{request.model}', size='{request.size}'")
+        
+        # Validate model
+        captions_model_id = get_captions_model_id(request.model)
+        if captions_model_id not in AVAILABLE_MODELS:
+            raise HTTPException(status_code=400, detail=f"Model '{request.model}' is not supported")
+        
+        # Validate request parameters
+        if not request.prompt or len(request.prompt.strip()) == 0:
+            raise HTTPException(status_code=400, detail="Prompt cannot be empty")
+        
+        if len(request.prompt) > 1000:
+            raise HTTPException(status_code=400, detail="Prompt exceeds maximum length of 1000 characters")
         
         # Submit the image generation request
-        operation_id = await submit_image_generation(request.prompt, request.size)
+        operation_id = await submit_image_generation(request.prompt, request.model, request.size)
+        logger.info(f"Image generation submitted with operation ID: {operation_id}")
         
         # Wait for completion
         completion_data = await wait_for_completion(operation_id)
         
+        # Validate completion data
+        if not completion_data.get("assetResolvedUrl"):
+            raise HTTPException(status_code=500, detail="Generation completed but no image URL received")
+        
         # Format response in OpenAI format
         image_data = ImageData(
             url=completion_data.get("assetResolvedUrl"),
@@ -206,6 +315,111 @@ async def create_image(request: ImageGenerationRequest):
         logger.error(f"Unexpected error in image generation: {e}")
         raise HTTPException(status_code=500, detail="Internal server error")
 
+@app.post("/v1/images/generations/async")
+async def create_image_async(request: ImageGenerationRequest):
+    """
+    Starts an image generation request and returns operation ID for status checking.
+    Non-blocking version of the generation API.
+    """
+    try:
+        logger.info(f"Received async image generation request: prompt='{request.prompt[:100]}...', model='{request.model}', size='{request.size}'")
+        
+        # Validate model
+        captions_model_id = get_captions_model_id(request.model)
+        if captions_model_id not in AVAILABLE_MODELS:
+            raise HTTPException(status_code=400, detail=f"Model '{request.model}' is not supported")
+        
+        # Validate request parameters
+        if not request.prompt or len(request.prompt.strip()) == 0:
+            raise HTTPException(status_code=400, detail="Prompt cannot be empty")
+        
+        if len(request.prompt) > 1000:
+            raise HTTPException(status_code=400, detail="Prompt exceeds maximum length of 1000 characters")
+        
+        # Submit the image generation request
+        operation_id = await submit_image_generation(request.prompt, request.model, request.size)
+        
+        # Store request details for later retrieval
+        operations_store[operation_id] = {
+            "created": int(datetime.now().timestamp()),
+            "prompt": request.prompt,
+            "model": request.model,
+            "size": request.size,
+            "status": "processing"
+        }
+        
+        return {
+            "operation_id": operation_id,
+            "status": "submitted",
+            "created": int(datetime.now().timestamp()),
+            "status_url": f"/v1/images/generations/status/{operation_id}"
+        }
+        
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Unexpected error in async image generation: {e}")
+        raise HTTPException(status_code=500, detail="Internal server error")
+
+@app.get("/v1/images/generations/status/{operation_id}")
+async def get_generation_status(operation_id: str):
+    """
+    Check the status of an image generation operation.
+    """
+    try:
+        if operation_id not in operations_store:
+            raise HTTPException(status_code=404, detail="Operation ID not found")
+        
+        # Get current status from Captions API
+        status_data = await check_generation_status(operation_id)
+        operation_info = operations_store[operation_id]
+        
+        # State 1 = processing, State 2 = completed, State 3 = failed
+        if status_data.get("state") == 1:
+            return {
+                "operation_id": operation_id,
+                "status": "processing",
+                "created": operation_info["created"],
+                "estimated_completion": None
+            }
+        elif status_data.get("state") == 2:
+            # Update stored info
+            operations_store[operation_id]["status"] = "completed"
+            
+            # Format response in OpenAI format
+            image_data = ImageData(
+                url=status_data["complete"].get("assetResolvedUrl"),
+                revised_prompt=operation_info["prompt"]
+            )
+            
+            return {
+                "operation_id": operation_id,
+                "status": "completed",
+                "created": operation_info["created"],
+                "data": [image_data.dict()]
+            }
+        elif status_data.get("state") == 3:
+            operations_store[operation_id]["status"] = "failed"
+            return {
+                "operation_id": operation_id,
+                "status": "failed",
+                "created": operation_info["created"],
+                "error": "Image generation failed"
+            }
+        else:
+            return {
+                "operation_id": operation_id,
+                "status": "unknown",
+                "created": operation_info["created"],
+                "error": "Unknown status"
+            }
+            
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Error checking generation status: {e}")
+        raise HTTPException(status_code=500, detail="Failed to check generation status")
+
 @app.get("/health")
 async def health_check():
     """Health check endpoint"""
@@ -217,10 +431,19 @@ async def root():
     return {
         "message": "OpenAI Compatible Image Generation API",
         "version": "1.0.0",
+        "supported_models": list(AVAILABLE_MODELS.keys()),
+        "openai_aliases": list(MODEL_MAPPINGS.keys()),
         "endpoints": {
+            "models": "/v1/models",
             "image_generation": "/v1/images/generations",
+            "async_generation": "/v1/images/generations/async",
+            "status_check": "/v1/images/generations/status/{operation_id}",
             "health": "/health",
             "docs": "/docs"
+        },
+        "example_curl": {
+            "generate_image": "curl -X POST 'http://localhost:8000/v1/images/generations' -H 'Content-Type: application/json' -d '{\"prompt\": \"a cat\", \"model\": \"dall-e-3\", \"size\": \"1024x1024\"}'",
+            "list_models": "curl -X GET 'http://localhost:8000/v1/models'"
         }
     }