Update app.py

app.py

@@ -841,50 +841,91 @@ AIRFORCE_KEY = os.getenv("AIRFORCE")
 AIRFORCE_VIDEO_MODEL = "grok-imagine-video"
 AIRFORCE_API_URL = "https://api.airforce/v1/images/generations"
 
+valid_ratios = {"3:2", "2:3", "1:1", "", None}
+ratios = {"3:2", "2:3", "1:1"}
+
+valid_modes = {"normal", "fun", "", None}
+modes = {"normal", "fun"}
+
 MAX_VIDEO_RETRIES = 6
 
+
 @app.get("/gen/video/airforce/{prompt}")
 @app.post("/gen/video/airforce")
 async def genvideo_airforce(request: Request, prompt: str = None):
+    aspectRatio = "3:2"
+    inputMode = "normal"
+
+    user_body = {}
     if prompt is None:
-        body = await request.json()
-        prompt = body.get("prompt")
-        duration = body.get("duration", 5)
+        user_body = await request.json()
+        prompt = user_body.get("prompt")
+        ratio = user_body.get("ratio")
+        mode = user_body.get("mode")
+        image_urls = user_body.get("image_urls")
+
+        if not prompt:
+            raise HTTPException(400, "Prompt is required")
+
+        if ratio not in valid_ratios:
+            raise HTTPException(
+                status_code=400,
+                detail=f"Invalid aspect ratio {ratio}. Must be one of 3:2, 2:3, or 1:1. Default is 3:2"
+            )
+        if ratio in ratios:
+            aspectRatio = ratio
+
+        if mode not in valid_modes:
+            raise HTTPException(
+                status_code=400,
+                detail=f"Invalid mode {mode}. Must be 'normal' or 'fun'. Default is normal"
+            )
+        if mode in modes:
+            inputMode = mode
+
+        if image_urls:
+            if not isinstance(image_urls, list):
+                raise HTTPException(400, "image_urls must be a list")
 
-    body = {
-        "model": "grok-imagine-video",
+            if len(image_urls) > 2:
+                raise HTTPException(400, "You may provide at most two image URLs")
+
+    payload = {
+        "model": AIRFORCE_VIDEO_MODEL,
         "prompt": prompt,
         "n": 1,
         "size": "1024x1024",
         "response_format": "b64_json",
         "sse": False,
-        "mode": "normal",
-        "aspectRatio": "3:2"
+        "mode": inputMode,
+        "aspectRatio": aspectRatio
     }
 
+    if prompt is None:
+        if image_urls:
+            payload["image_urls"] = image_urls
+
     async with httpx.AsyncClient(timeout=600) as client:
         resp = await client.post(
-            "https://api.airforce/v1/images/generations",
-            headers={"Authorization": f"Bearer {AIRFORCE_KEY}", "Content-Type": "application/json"},
-            json=body
+            AIRFORCE_API_URL,
+            headers={
+                "Authorization": f"Bearer {AIRFORCE_KEY}",
+                "Content-Type": "application/json"
+            },
+            json=payload
         )
 
     if resp.status_code != 200:
         return JSONResponse(status_code=resp.status_code, content=resp.json())
+
     if not resp.content:
-        raise HTTPException(
-            status_code=502,
-            detail="api.airforce returned empty response"
-        )
+        raise HTTPException(502, "api.airforce returned empty response")
 
     try:
         result = resp.json()
         b64_video = result["data"][0]["b64_json"]
     except Exception:
-        raise HTTPException(
-            502,
-            f"Invalid api.airforce response: {resp.text[:500]}"
-        )
+        raise HTTPException(502, f"Invalid api.airforce response: {resp.text[:500]}")
 
     if not b64_video:
         raise HTTPException(502, "Airforce returned empty b64_json")