added basic auth middleware

.gitignore

@@ -8,6 +8,7 @@ __pycache__
 /venv
 /tmp
 /model.ckpt
+.env
 # /models/**/*
 /GFPGANv1.3.pth
 /gfpgan/weights/*.pth

basic_auth_middleware.py

@@ -0,0 +1,39 @@
+import base64
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.responses import Response
+from fastapi import FastAPI
+
+class BasicAuthMiddleware(BaseHTTPMiddleware):
+    def __init__(self, app, username: str, password: str):
+        super().__init__(app)
+        self.username = username
+        self.password = password
+
+    async def dispatch(self, request, call_next):
+        # Extract the Authorization header
+        auth_header = request.headers.get("Authorization")
+        if not auth_header or not auth_header.startswith("Basic "):
+            return self._unauthorized_response()
+
+        try:
+            # Decode and split the credentials
+            encoded_credentials = auth_header.split(" ")[1]
+            # We should add a step to santize the input here
+            decoded_credentials = base64.b64decode(encoded_credentials).decode("utf-8")
+            provided_username, provided_password = decoded_credentials.split(":")
+
+            # Check credentials
+            if provided_username == self.username and provided_password == self.password:
+                response = await call_next(request)
+                return response
+            else:
+                return self._unauthorized_response()
+        except Exception:
+            return self._unauthorized_response()
+
+    def _unauthorized_response(self):
+        return Response(
+            content="Unauthorized",
+            status_code=401,
+            headers={"WWW-Authenticate": "Basic realm='api'"}
+        )

extensions/sd-webui-controlnet

@@ -1 +1 @@
-Subproject commit 8bbbd0e55ef6e5d71b09c2de2727b36e7bc825b0
+Subproject commit 56cec5b2958edf3b1807b7e7b2b1b5186dbd2f81

webui.py

@@ -30,6 +30,13 @@ def api_only():
     initialize.initialize()
 
     app = FastAPI()
+
+    # Initialize the Basic Authentication middleware
+    from basic_auth_middleware import BasicAuthMiddleware
+    USERNAME = os.getenv('SDAPI_USERNAME')
+    PASSWORD = os.getenv('SDAPI_PASSWORD')
+    app.add_middleware(BasicAuthMiddleware, username=USERNAME, password=PASSWORD)
+
     initialize_util.setup_middleware(app)
     api = create_api(app)