Update main.py

main.py

@@ -594,10 +594,8 @@ async def get_completion(payload: Payload, request: Request, authenticated: bool
 @lru_cache(maxsize=256)
 def cached_url(url: str):
     return url
-
-router = APIRouter()
-
-@router.get("/images/{prompt:path}")
+    
+@app.get("/images/{prompt:path}")
 async def create_image(
     prompt: str,
     request: Request,
@@ -606,7 +604,6 @@ async def create_image(
     if not server_status:
         raise HTTPException(status_code=503, detail="Server is under maintenance.")
 
-    # forward all GET params
     query = request.url.query
 
     base = "https://image.pollinations.ai/prompt/"
@@ -614,12 +611,11 @@ async def create_image(
     if query:
         final_url += f"?{query}"
 
-    # caching
     final_url = cached_url(final_url)
 
     try:
-        client = httpx.AsyncClient(timeout=60)
-        resp = await client.get(final_url)
+        async with httpx.AsyncClient(timeout=60) as client:
+            resp = await client.get(final_url)
 
         if resp.status_code != 200:
             raise HTTPException(status_code=resp.status_code, detail="Image generation failed.")
@@ -634,7 +630,6 @@ async def create_image(
 
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Unexpected error: {e}")
-
 @app.get("/usage")
 async def get_usage_json(days: int = 7):
     return usage_tracker.get_usage_summary(days)