Update app.py

app.py

@@ -1,15 +1,14 @@
 import os
 import sqlite3
 import base64
-from datetime import datetime, timedelta
+from datetime import datetime
 from typing import Optional
-from fastapi import FastAPI, Request, Response, Cookie, HTTPException
+from fastapi import FastAPI, Request, Cookie, HTTPException, WebSocket
 from fastapi.responses import HTMLResponse, RedirectResponse, JSONResponse
 from fastapi.templating import Jinja2Templates
-from fastapi.staticfiles import StaticFiles
 from huggingface_hub import whoami
-import secrets
 import httpx
+import websockets
 
 app = FastAPI()
 templates = Jinja2Templates(directory=".")
@@ -21,6 +20,9 @@ OAUTH_SCOPES = os.getenv("OAUTH_SCOPES", "openid profile")
 SPACE_HOST = os.getenv("SPACE_HOST", "multimodalart-krea-realtime-video.hf.space")
 OPENID_PROVIDER_URL = os.getenv("OPENID_PROVIDER_URL", "https://huggingface.co")
 
+# FAL API Key from environment
+FAL_API_KEY = os.getenv("FAL_API_KEY", "")
+
 # Database setup
 DB_PATH = "/data/usage.db"
 
@@ -70,7 +72,6 @@ def record_session(username: str, is_pro: bool):
         )
         conn.commit()
     except sqlite3.IntegrityError:
-        # Already recorded today
         pass
     finally:
         conn.close()
@@ -85,7 +86,6 @@ async def exchange_code_for_token(code: str, redirect_uri: str) -> dict:
     """Exchange OAuth code for access token"""
     token_url = f"{OPENID_PROVIDER_URL}/oauth/token"
     
-    # Prepare Basic Auth header
     credentials = f"{OAUTH_CLIENT_ID}:{OAUTH_CLIENT_SECRET}"
     b64_credentials = base64.b64encode(credentials.encode()).decode()
     
@@ -126,7 +126,6 @@ async def home(request: Request, access_token: Optional[str] = Cookie(None)):
     """Home page - check auth and show app or login"""
     
     if not access_token:
-        # Not logged in - show login button
         return templates.TemplateResponse("index.html", {
             "request": request,
             "authenticated": False,
@@ -135,11 +134,9 @@ async def home(request: Request, access_token: Optional[str] = Cookie(None)):
             "space_host": SPACE_HOST
         })
     
-    # Verify token and get user info
     try:
         user_info = await get_user_info(access_token)
     except:
-        # Invalid token, clear it
         response = templates.TemplateResponse("index.html", {
             "request": request,
             "authenticated": False,
@@ -151,7 +148,6 @@ async def home(request: Request, access_token: Optional[str] = Cookie(None)):
         response.delete_cookie("access_token")
         return response
     
-    # Check session limits
     can_start, used, limit = can_start_session(user_info["username"], user_info["is_pro"])
     
     return templates.TemplateResponse("index.html", {
@@ -172,14 +168,12 @@ async def oauth_callback(code: str, state: Optional[str] = None):
     redirect_uri = f"https://{SPACE_HOST}/oauth/callback"
     
     try:
-        # Exchange code for token
         token_data = await exchange_code_for_token(code, redirect_uri)
         access_token = token_data.get("access_token")
         
         if not access_token:
             raise HTTPException(status_code=400, detail="No access token received")
         
-        # Redirect to home with token as cookie
         response = RedirectResponse(url="/", status_code=302)
         response.set_cookie(
             key="access_token",
@@ -187,7 +181,7 @@ async def oauth_callback(code: str, state: Optional[str] = None):
             httponly=True,
             secure=True,
             samesite="lax",
-            max_age=30 * 24 * 60 * 60  # 30 days
+            max_age=30 * 24 * 60 * 60
         )
         
         return response
@@ -215,7 +209,6 @@ async def start_session(access_token: Optional[str] = Cookie(None)):
             detail=f"Daily limit reached. You've used {used}/{limit} sessions today."
         )
     
-    # Record the session
     record_session(user_info["username"], user_info["is_pro"])
     
     return {
@@ -254,4 +247,98 @@ async def logout():
 @app.get("/health")
 async def health():
     """Health check endpoint"""
-    return {"status": "ok", "oauth_enabled": bool(OAUTH_CLIENT_ID)}
+    return {
+        "status": "ok", 
+        "oauth_enabled": bool(OAUTH_CLIENT_ID),
+        "fal_api_key_configured": bool(FAL_API_KEY)
+    }
+
+@app.websocket("/ws/video-gen")
+async def websocket_video_gen(websocket: WebSocket):
+    """WebSocket proxy to FAL API - keeps API key secret"""
+    from fastapi import WebSocket
+    import websockets
+    import json
+    
+    await websocket.accept()
+    
+    # Get user from cookie
+    access_token = websocket.cookies.get("access_token")
+    if not access_token:
+        await websocket.close(code=1008, reason="Not authenticated")
+        return
+    
+    try:
+        user_info = await get_user_info(access_token)
+    except:
+        await websocket.close(code=1008, reason="Invalid session")
+        return
+    
+    # Check if user can start session
+    can_start, used, limit = can_start_session(user_info["username"], user_info["is_pro"])
+    if not can_start:
+        await websocket.close(code=1008, reason=f"Daily limit reached ({used}/{limit})")
+        return
+    
+    if not FAL_API_KEY:
+        await websocket.close(code=1011, reason="FAL API key not configured")
+        return
+    
+    # Fetch temporary FAL token
+    try:
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                "https://rest.alpha.fal.ai/tokens/",
+                headers={
+                    "Content-Type": "application/json",
+                    "Authorization": f"Key {FAL_API_KEY}"
+                },
+                json={
+                    "allowed_apps": ["krea-wan-14b"],
+                    "token_expiration": 5000
+                }
+            )
+            response.raise_for_status()
+            fal_token = response.json()
+    except Exception as e:
+        await websocket.close(code=1011, reason=f"Failed to get FAL token: {str(e)}")
+        return
+    
+    # Connect to FAL WebSocket
+    fal_ws_url = f"wss://fal.run/fal-ai/krea-wan-14b/ws?fal_jwt_token={fal_token}"
+    
+    try:
+        async with websockets.connect(fal_ws_url) as fal_ws:
+            # Relay messages between client and FAL
+            async def client_to_fal():
+                try:
+                    while True:
+                        # Receive from client
+                        data = await websocket.receive_bytes()
+                        # Forward to FAL
+                        await fal_ws.send(data)
+                except Exception as e:
+                    print(f"Client to FAL error: {e}")
+            
+            async def fal_to_client():
+                try:
+                    while True:
+                        # Receive from FAL
+                        message = await fal_ws.recv()
+                        # Forward to client
+                        if isinstance(message, str):
+                            await websocket.send_text(message)
+                        else:
+                            await websocket.send_bytes(message)
+                except Exception as e:
+                    print(f"FAL to client error: {e}")
+            
+            # Run both directions concurrently
+            import asyncio
+            await asyncio.gather(
+                client_to_fal(),
+                fal_to_client()
+            )
+    except Exception as e:
+        print(f"WebSocket proxy error: {e}")
+        await websocket.close(code=1011, reason=str(e))