Upload app.py with huggingface_hub

app.py

@@ -1,42 +1,52 @@
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, Request, Response
 from fastapi.responses import StreamingResponse
 import requests
 import os
 
 app = FastAPI()
-
-# Optional: Add your HF_TOKEN to Space Secrets for private models
 HF_TOKEN = os.environ.get("HF_TOKEN")
+BASE_URL = "https://huggingface.co"
 
-@app.get("/")
-def home():
-    return {"status": "running", "message": "HF Storage Proxy"}
-
-@app.get("/resolve/{repo_id:path}/{filename:path}")
-async def proxy_resolve(repo_id: str, filename: str):
-    """
-    Proxies requests to huggingface.co/repo/resolve/main/file
-    """
-    hf_url = f"https://huggingface.co/{repo_id}/resolve/main/{filename}"
-    
-    headers = {}
-    if HF_TOKEN:
-        headers["Authorization"] = f"Bearer {HF_TOKEN}"
+def get_proxy_response(path: str, query: str = "", method: str = "GET", headers: dict = {}, body: bytes = None):
+    url = f"{BASE_URL}/{path}"
+    if query:
+        url += f"?{query}"
     
+    # Filter headers to avoid conflicts
+    filtered_headers = {k: v for k, v in headers.items() if k.lower() not in ["host", "content-length", "connection"]}
+    if HF_TOKEN and "authorization" not in [k.lower() for k in filtered_headers.keys()]:
+        filtered_headers["Authorization"] = f"Bearer {HF_TOKEN}"
+        
     try:
-        # We use stream=True to handle large model files without loading into memory
-        response = requests.get(hf_url, headers=headers, stream=True, allow_redirects=True)
-        response.raise_for_status()
+        resp = requests.request(
+            method=method,
+            url=url,
+            headers=filtered_headers,
+            data=body,
+            stream=True,
+            allow_redirects=True
+        )
+        
+        # Exclude headers that FastAPI/Uvicorn might handle or conflict with
+        excluded = ["content-encoding", "transfer-encoding", "connection", "keep-alive"]
+        resp_headers = {k: v for k, v in resp.headers.items() if k.lower() not in excluded}
         
         return StreamingResponse(
-            response.iter_content(chunk_size=1024*1024),
-            media_type=response.headers.get("content-type"),
-            headers={
-                "Content-Disposition": response.headers.get("Content-Disposition", f'attachment; filename="{filename}"'),
-                "Content-Length": response.headers.get("Content-Length")
-            }
+            resp.iter_content(chunk_size=128*1024),
+            status_code=resp.status_code,
+            headers=resp_headers,
+            media_type=resp.headers.get("content-type")
         )
-    except requests.exceptions.HTTPError as e:
-        raise HTTPException(status_code=e.response.status_code, detail=str(e))
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        return Response(content=str(e), status_code=500)
+
+@app.get("/resolve/{repo_id:path}/{filename:path}")
+async def old_resolve_compat(repo_id: str, filename: str):
+    """Keep compatibility with the previous /resolve structure."""
+    return get_proxy_response(f"{repo_id}/resolve/main/{filename}")
+
+@app.api_route("/{path:path}", methods=["GET", "HEAD", "OPTIONS", "POST"])
+async def catch_all(request: Request, path: str):
+    """Universal proxy to Hugging Face."""
+    body = await request.body() if request.method == "POST" else None
+    return get_proxy_response(path, request.url.query, request.method, dict(request.headers), body)