runner+worker: Add bearer auth

runner/app/routes/image_to_image.py

@@ -1,13 +1,14 @@
 from fastapi import Depends, APIRouter, UploadFile, File, Form
 from fastapi.responses import JSONResponse
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from app.pipelines import ImageToImagePipeline
 from app.dependencies import get_pipeline
 from app.routes.util import image_to_data_url, ImageResponse, HTTPError, http_error
 import PIL
 from typing import Annotated
 import logging
 import random
-from typing import List
+import os
 
 router = APIRouter()
 
@@ -35,7 +36,17 @@ async def image_to_image(
  seed: Annotated[int, Form()] = None,
  num_images_per_prompt: Annotated[int, Form()] = 1,
  pipeline: ImageToImagePipeline = Depends(get_pipeline),
+ token: HTTPAuthorizationCredentials = Depends(HTTPBearer(auto_error=False)),
 ):
+ auth_token = os.environ.get("AUTH_TOKEN")
+ if auth_token:
+ if not token or token.credentials != auth_token:
+ return JSONResponse(
+ status_code=401,
+ headers={"WWW-Authenticate": "Bearer"},
+ content=http_error("Invalid bearer token"),
+ )
+
  if model_id != "" and model_id != pipeline.model_id:
  return JSONResponse(
  status_code=400,

runner/app/routes/image_to_video.py

@@ -1,12 +1,14 @@
 from fastapi import Depends, APIRouter, UploadFile, File, Form
 from fastapi.responses import JSONResponse
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from app.pipelines import ImageToVideoPipeline
 from app.dependencies import get_pipeline
-from app.routes.util import image_to_data_url, VideoResponse, HTTPError
+from app.routes.util import image_to_data_url, VideoResponse, HTTPError, http_error
 import PIL
 from typing import Annotated
 import logging
 import random
+import os
 
 router = APIRouter()
 
@@ -34,15 +36,23 @@ async def image_to_video(
  noise_aug_strength: Annotated[float, Form()] = 0.02,
  seed: Annotated[int, Form()] = None,
  pipeline: ImageToVideoPipeline = Depends(get_pipeline),
+ token: HTTPAuthorizationCredentials = Depends(HTTPBearer(auto_error=False)),
 ):
+ auth_token = os.environ.get("AUTH_TOKEN")
+ if auth_token:
+ if not token or token.credentials != auth_token:
+ return JSONResponse(
+ status_code=401,
+ headers={"WWW-Authenticate": "Bearer"},
+ content=http_error("Invalid bearer token"),
+ )
+
  if model_id != "" and model_id != pipeline.model_id:
  return JSONResponse(
  status_code=400,
- content={
- "detail": {
- "msg": f"pipeline configured with {pipeline.model_id} but called with {model_id}"
- }
- },
+ content=http_error(
+ f"pipeline configured with {pipeline.model_id} but called with {model_id}"
+ ),
  )
 
  if seed is None:
@@ -62,7 +72,7 @@ async def image_to_video(
  logger.error(f"ImageToVideoPipeline error: {e}")
  logger.exception(e)
  return JSONResponse(
- status_code=500, content={"detail": {"msg": "ImageToVideoPipeline error"}}
+ status_code=500, content=http_error("ImageToVideoPipeline error")
  )
 
  output_frames = []

runner/app/routes/text_to_image.py

@@ -1,12 +1,13 @@
 from pydantic import BaseModel
 from fastapi import Depends, APIRouter
 from fastapi.responses import JSONResponse
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from app.pipelines import TextToImagePipeline
 from app.dependencies import get_pipeline
 from app.routes.util import image_to_data_url, ImageResponse, HTTPError, http_error
 import logging
-from typing import List
 import random
+import os
 
 router = APIRouter()
 
@@ -32,8 +33,19 @@ class TextToImageParams(BaseModel):
 @router.post("/text-to-image", response_model=ImageResponse, responses=responses)
 @router.post("/text-to-image/", response_model=ImageResponse, include_in_schema=False)
 async def text_to_image(
- params: TextToImageParams, pipeline: TextToImagePipeline = Depends(get_pipeline)
+ params: TextToImageParams,
+ pipeline: TextToImagePipeline = Depends(get_pipeline),
+ token: HTTPAuthorizationCredentials = Depends(HTTPBearer(auto_error=False)),
 ):
+ auth_token = os.environ.get("AUTH_TOKEN")
+ if auth_token:
+ if not token or token.credentials != auth_token:
+ return JSONResponse(
+ status_code=401,
+ headers={"WWW-Authenticate": "Bearer"},
+ content=http_error("Invalid bearer token"),
+ )
+
  if params.model_id != "" and params.model_id != pipeline.model_id:
  return JSONResponse(
  status_code=400,

runner/openapi.json

@@ -79,7 +79,12 @@
  }
  }
  }
- }
+ },
+ "security": [
+ {
+ "HTTPBearer": []
+ }
+ ]
  }
  },
  "/image-to-image": {
@@ -137,7 +142,12 @@
  }
  }
  }
- }
+ },
+ "security": [
+ {
+ "HTTPBearer": []
+ }
+ ]
  }
  },
  "/image-to-video": {
@@ -195,7 +205,12 @@
  }
  }
  }
- }
+ },
+ "security": [
+ {
+ "HTTPBearer": []
+ }
+ ]
  }
  }
  },
@@ -477,6 +492,12 @@
  ],
  "title": "VideoResponse"
  }
+ },
+ "securitySchemes": {
+ "HTTPBearer": {
+ "type": "http",
+ "scheme": "bearer"
+ }
  }
  }
 }