Update main.py

main.py

@@ -2,6 +2,7 @@ from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from huggingface_hub import InferenceClient
+import streamlit as st
 
 app = FastAPI()
 
@@ -53,4 +54,16 @@ def generate(item: Item):
 
 @app.post("/generate/")
 async def generate_text(item: Item):
-    return {"response": generate(item)}
+    # reject if key not the same 
+    apiKey = st.secrets["API_KEY"]
+    # if apiKey != key:
+    #     return jsonify({
+    #         "image_url": image_url,
+    #         "model": model,
+    #         "result": "Invalid API Key",
+    #     }), 400
+    
+    return {
+        "response": generate(item),
+        "key": apiKey
+    }