Update app.py

app.py

@@ -16,7 +16,7 @@ import base64
 from subscriptions import fetch_subscription
 from typing import Optional
 from keywords import *
-
+from assets import save_base64_image, cleanup_image, is_base64_image, asset_router
 WAN_SPACE = "https://huggingface.co/spaces/Wan-AI/Wan2.1"
 WAN_API_BASE = f"{WAN_SPACE}/gradio_api"
 
@@ -28,7 +28,7 @@ app.add_middleware(
     allow_methods=["GET", "POST", "HEAD"],
     allow_headers=["*"],
 )
-
+app.include_router(asset_router)
 @app.get("/")
 async def reroute_to_status():
     return RedirectResponse(url="https://inference.js.org", status_code=status.HTTP_308_PERMANENT_REDIRECT)
@@ -1000,8 +1000,22 @@ async def genvideo_airforce(
         "seed": -1,
     }
 
+    temp_assets = []
+    
     if image_urls:
-        params["image"] = "|".join(image_urls[:2])
+        processed_urls = []
+    
+        for img in image_urls[:2]:
+            if is_base64_image(img):
+                image_id = save_base64_image(img)
+                temp_assets.append(image_id)
+    
+                served_url = f"{request.base_url}asset-cdn/assets/{image_id}"
+                processed_urls.append(served_url)
+            else:
+                processed_urls.append(img)
+    
+        params["image"] = "|".join(processed_urls)
 
     if inputMode == "fun":
         params["enhance"] = "true"
@@ -1011,10 +1025,15 @@ async def genvideo_airforce(
 
     print(f"[VIDEO GEN] Pollinations URL: {url}")
     url = url + f"&key={PKEY}"
-
-    async with httpx.AsyncClient(timeout=600) as client:
-        resp = await client.get(url)
-
+    resp = None
+    try: 
+        async with httpx.AsyncClient(timeout=600) as client:
+            resp = await client.get(url)
+    finally:
+        for aid in temp_assets:
+            cleanup_image(aid)
+    if resp is None:
+        raise HTTPException(502, "Video generation request failed")
     if resp.status_code != 200:
         body_text = ""
         try: