Update main.py

main.py

@@ -1,27 +1,28 @@
 import os
 import time
 import uuid
-from typing import List, Optional, Literal, Any, Dict, Union
+from typing import List, Optional, Literal, Any, Dict
 
 import httpx
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from huggingface_hub import InferenceClient
-import asyncio
-
-
-# ---------------- Config (env) ----------------
-HF_TOKEN = os.getenv("HF_TOKEN")  # Hugging Face token (works for provider=fal-ai)
-WAN_MODEL = os.getenv("WAN_MODEL", "Wan-AI/Wan2.2-T2V-A14B")
 
+# ---------- Config via env ----------
+HF_TOKEN = os.getenv("HF_TOKEN")  # e.g. hf_jwt_...
+# Default matches your curl submit endpoint
+HF_SUBMIT_URL = os.getenv(
+    "HF_SUBMIT_URL",
+    "https://router.huggingface.co/fal-ai/fal-ai/wan/v2.2-a14b/text-to-video?_subdomain=queue",
+)
 UPLOAD_URL = os.getenv("UPLOAD_URL", "https://upload.snapzion.com/api/public-upload")
-UPLOAD_ACCESS_TOKEN = os.getenv("UPLOAD_ACCESS_TOKEN")  # your bearer token
+UPLOAD_ACCESS_TOKEN = os.getenv("UPLOAD_ACCESS_TOKEN")  # your Snapzion bearer
 
-# Optional tuning
-GEN_TIMEOUT_SEC = int(os.getenv("GEN_TIMEOUT_SEC", "900"))  # 15 min generation ceiling
+# Polling/backoff
+POLL_INTERVAL_SEC = float(os.getenv("POLL_INTERVAL_SEC", "3"))
+POLL_TIMEOUT_SEC = int(os.getenv("POLL_TIMEOUT_SEC", "900"))  # 15 min max
 
 
-# ---------------- OpenAI-compatible schemas ----------------
+# ---------- OpenAI-compatible schemas ----------
 class ChatMessage(BaseModel):
     role: Literal["system", "user", "assistant", "tool"]
     content: str
@@ -65,9 +66,8 @@ class ChatCompletionsResponse(BaseModel):
     }
 
 
-# ---------------- Helpers ----------------
+# ---------- Helpers ----------
 def extract_prompt(messages: List[ChatMessage]) -> str:
-    """Use the last user message as the prompt. Fallback to joining all user messages."""
     for m in reversed(messages):
         if m.role == "user" and m.content and m.content.strip():
             return m.content.strip()
@@ -77,42 +77,89 @@ def extract_prompt(messages: List[ChatMessage]) -> str:
     return "\n".join(user_texts).strip()
 
 
-async def generate_video_bytes(prompt: str) -> bytes:
-    """Calls huggingface_hub.InferenceClient with provider='fal-ai' (Wan T2V) and returns MP4 bytes."""
+async def hf_queue_submit(prompt: str) -> Dict[str, str]:
+    """
+    Submit a job to the HF router queue endpoint.
+    Returns a dict containing status_url and response_url (from HF).
+    """
+    if not HF_TOKEN:
+        raise HTTPException(status_code=500, detail="HF_TOKEN is not set.")
+
+    headers = {
+        "Authorization": f"Bearer {HF_TOKEN}",
+        "Content-Type": "application/json",
+        "Accept": "*/*",
+    }
+    payload = {"prompt": prompt}
+
+    async with httpx.AsyncClient(timeout=None) as client:
+        resp = await client.post(HF_SUBMIT_URL, headers=headers, json=payload)
+    if resp.status_code >= 400:
+        raise HTTPException(status_code=502, detail=f"HF submit error: {resp.text}")
+
+    data = resp.json()
+    status_url = data.get("status_url") or data.get("urls", {}).get("status_url")
+    response_url = data.get("response_url") or data.get("urls", {}).get("response_url")
+    if not status_url or not response_url:
+        raise HTTPException(status_code=502, detail=f"Unexpected HF submit response: {data}")
+
+    return {"status_url": status_url, "response_url": response_url}
+
+
+async def hf_queue_wait(status_url: str) -> None:
+    """
+    Polls the HF queue status_url until COMPLETED or error states.
+    """
     if not HF_TOKEN:
         raise HTTPException(status_code=500, detail="HF_TOKEN is not set.")
-    client = InferenceClient(provider="fal-ai", api_key=HF_TOKEN)
 
-    def _sync_generate() -> Union[bytes, Dict[str, Any]]:
-        # mirrors your Python example:
-        # video = client.text_to_video("prompt", model="Wan-AI/Wan2.2-T2V-A14B")
-        return client.text_to_video(prompt, model=WAN_MODEL)
+    headers = {"Authorization": f"Bearer {HF_TOKEN}", "Accept": "*/*"}
+    start = time.time()
+
+    async with httpx.AsyncClient(timeout=None) as client:
+        while True:
+            resp = await client.get(status_url, headers=headers)
+            if resp.status_code >= 400:
+                raise HTTPException(status_code=502, detail=f"HF status error: {resp.text}")
+
+            data = resp.json()
+            status = data.get("status")
+
+            if status in ("COMPLETED", "SUCCEEDED"):
+                return
+            if status in ("FAILED", "ERROR", "CANCELLED", "CANCELED"):
+                raise HTTPException(status_code=502, detail=f"HF job failed: {data}")
+
+            if time.time() - start > POLL_TIMEOUT_SEC:
+                raise HTTPException(status_code=504, detail="HF queue timed out.")
+
+            time.sleep(POLL_INTERVAL_SEC)
+
 
-    try:
-        result = await asyncio.wait_for(
-            asyncio.get_event_loop().run_in_executor(None, _sync_generate),
-            timeout=GEN_TIMEOUT_SEC,
-        )
-    except asyncio.TimeoutError:
-        raise HTTPException(status_code=504, detail="Video generation timed out.")
-    except Exception as e:
-        raise HTTPException(status_code=502, detail=f"Video generation failed: {e}")
+async def hf_queue_fetch_result(response_url: str) -> Dict[str, Any]:
+    """
+    Fetch the final response JSON, which includes {"video": {"url": ...}, ...}
+    """
+    if not HF_TOKEN:
+        raise HTTPException(status_code=500, detail="HF_TOKEN is not set.")
 
-    # fal-ai provider typically returns a dict with "video": bytes; sometimes raw bytes
-    if isinstance(result, (bytes, bytearray)):
-        return bytes(result)
+    headers = {"Authorization": f"Bearer {HF_TOKEN}", "Accept": "*/*"}
+    async with httpx.AsyncClient(timeout=None) as client:
+        resp = await client.get(response_url, headers=headers)
+    if resp.status_code >= 400:
+        raise HTTPException(status_code=502, detail=f"HF result error: {resp.text}")
+    return resp.json()
 
-    if isinstance(result, dict):
-        # common keys: "video" (bytes), "seed", etc.
-        vid = result.get("video")
-        if isinstance(vid, (bytes, bytearray)):
-            return bytes(vid)
 
-    raise HTTPException(status_code=502, detail=f"Unexpected generation result: {type(result)}")
+async def download_video(url: str) -> bytes:
+    async with httpx.AsyncClient(timeout=None) as client:
+        resp = await client.get(url)
+    if resp.status_code >= 400:
+        raise HTTPException(status_code=502, detail=f"Download failed: {resp.text}")
+    return resp.content
 
 
 async def upload_video_bytes(mp4_bytes: bytes) -> str:
-    """Uploads MP4 to Snapzion uploader and returns public URL."""
     if not UPLOAD_ACCESS_TOKEN:
         raise HTTPException(status_code=500, detail="UPLOAD_ACCESS_TOKEN is not set.")
     headers = {"Authorization": f"Bearer {UPLOAD_ACCESS_TOKEN}"}
@@ -120,12 +167,10 @@ async def upload_video_bytes(mp4_bytes: bytes) -> str:
 
     async with httpx.AsyncClient(timeout=None) as client:
         resp = await client.post(UPLOAD_URL, headers=headers, files=files)
-
     if resp.status_code >= 400:
         raise HTTPException(status_code=502, detail=f"Upload failed: {resp.text}")
 
     data = resp.json()
-    # Try common URL fields (adjust if your API returns a different shape)
     url = (
         data.get("url")
         or data.get("fileUrl")
@@ -137,31 +182,44 @@ async def upload_video_bytes(mp4_bytes: bytes) -> str:
     return url
 
 
-# ---------------- FastAPI app ----------------
-app = FastAPI(title="OpenAI-Compatible T2V Proxy (FAL via HF)")
-
+# ---------- FastAPI ----------
+app = FastAPI(title="OpenAI-Compatible T2V Proxy (HF Router Queue)")
 
 @app.get("/health")
 async def health():
-    return {"status": "ok", "model": WAN_MODEL}
-
+    return {"status": "ok", "submit_url": HF_SUBMIT_URL}
 
 @app.post("/v1/chat/completions", response_model=ChatCompletionsResponse)
 async def chat_completions(req: ChatCompletionsRequest):
     """
-    OpenAI-compatible endpoint:
-      - reads last user message as the T2V prompt
-      - generates a video with Wan-AI/Wan2.2-T2V-A14B via provider='fal-ai'
-      - uploads to your uploader
-      - returns the public URL inside the assistant message
+    1) submit to HF router queue (Bearer HF_TOKEN)
+    2) poll status_url until COMPLETED
+    3) fetch response_url -> video.url
+    4) download MP4, upload to Snapzion
+    5) return URL in OpenAI chat shape
     """
     prompt = extract_prompt(req.messages)
-    mp4 = await generate_video_bytes(prompt)
-    video_url = await upload_video_bytes(mp4)
 
+    # 1) Submit
+    urls = await hf_queue_submit(prompt)
+
+    # 2) Wait
+    await hf_queue_wait(urls["status_url"])
+
+    # 3) Fetch result JSON
+    result = await hf_queue_fetch_result(urls["response_url"])
+    video_url = (result.get("video") or {}).get("url")
+    if not video_url:
+        raise HTTPException(status_code=502, detail=f"HF result missing video.url: {result}")
+
+    # 4) Download + re-upload
+    mp4 = await download_video(video_url)
+    public_url = await upload_video_bytes(mp4)
+
+    # 5) Respond OpenAI-style
     now = int(time.time())
     completion_id = f"chatcmpl-{uuid.uuid4().hex}"
-    content = f"✅ Video generated & uploaded.\n**Prompt:** {prompt}\n**URL:** {video_url}"
+    content = f"✅ Video generated & uploaded.\n**Prompt:** {prompt}\n**URL:** {public_url}"
 
     return ChatCompletionsResponse(
         id=completion_id,