Update app.py

app.py

@@ -432,6 +432,13 @@ async def generate_image(request: Request, prompt: str = None):
     else:
         chosen_model = "zimage"
 
+    if prompt is None:
+        if (await request.json()).get("mode"):
+            if mode == "fantasy":
+                chosen_model = "flux"
+            elif mode == "realistic":
+                chosen_model = "zimage"
+
     print(f"[IMAGE GEN] Routing to model: {chosen_model}")
 
     url = f"https://gen.pollinations.ai/image/{prompt}?model={chosen_model}&key={PKEY2}"
@@ -760,91 +767,6 @@ async def genvideo(request: Request, prompt: str = None):
         }
     )
 
-@app.get("/gen/video/exp/{prompt}")
-@app.post("/gen/video/exp")
-async def genvideo_wan(request: Request, prompt: str = None):
-    client_ip = request.client.host
-    check_rate_limit(client_ip)
-
-    if prompt is None:
-        body = await request.json()
-        prompt = body.get("prompt")
-
-    if not prompt:
-        raise HTTPException(400, "Prompt is required")
-
-    session_hash = uuid.uuid4().hex
-
-    join_payload = {
-        "fn_index": 0,
-        "session_hash": session_hash,
-        "data": [prompt],
-    }
-
-    async with httpx.AsyncClient(timeout=60.0) as client:
-        # 1️⃣ Join queue
-        join = await client.post(
-            f"{WAN_API_BASE}/queue/join",
-            json=join_payload,
-        )
-
-        if join.status_code != 200:
-            raise HTTPException(
-                join.status_code,
-                f"Failed to join WAN queue: {join.text}"
-            )
-
-        for _ in range(120):
-            await asyncio.sleep(1)
-
-            poll = await client.get(
-                f"{WAN_API_BASE}/queue/data",
-                params={"session_hash": session_hash},
-            )
-
-            if poll.status_code != 200:
-                continue
-
-            data = poll.json()
-
-            if data.get("status") == "complete":
-                try:
-                    video_url = data["data"][0]
-                except Exception:
-                    raise HTTPException(
-                        500,
-                        "WAN completed but returned unexpected format"
-                    )
-
-                # 3️⃣ Fetch video
-                video = await client.get(video_url)
-
-                if video.status_code != 200:
-                    raise HTTPException(
-                        502,
-                        "Failed to fetch generated video"
-                    )
-
-                return Response(
-                    content=video.content,
-                    media_type="video/mp4",
-                    headers={
-                        "Content-Length": str(len(video.content)),
-                        "Accept-Ranges": "bytes",
-                    }
-                )
-
-            if data.get("status") == "error":
-                raise HTTPException(
-                    500,
-                    f"WAN generation error: {data}"
-                )
-
-    raise HTTPException(
-        504,
-        "WAN video generation timed out"
-    )
-
 AIRFORCE_KEY = os.getenv("AIRFORCE")
 AIRFORCE_VIDEO_MODEL = "grok-imagine-video"
 AIRFORCE_API_URL = "https://api.airforce/v1/images/generations"