Update main.py

main.py

@@ -13,6 +13,9 @@ from google.generativeai.types import HarmCategory, HarmBlockThreshold
 import google.ai.generativelanguage as glm
 from PIL import Image
 import io
+import random
+
+api_keys = os.getenv('GEMINI_API_KEYS').split(',')
 
 from inference_sdk import InferenceHTTPClient
 
@@ -104,22 +107,7 @@ async def predict_wheat(image: UploadFile = File(...), model_id: str = "grian/1"
 
 
 
-genai.configure(api_key="AIzaSyDReyKulYGTIj-8s-Nw38hGJDasjNpZhHA")
-
-generation_config = {
-    "temperature": 1,
-    "top_p": 0.95,
-    "top_k": 64,
-    "max_output_tokens": 8192,
-    "response_mime_type": "text/plain",
-}
 
-safety_settings = {
-    HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
-    HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
-    HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
-    HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
-}
 
 def process_image(file: UploadFile):
     image = Image.open(file.file)
@@ -142,6 +130,24 @@ def process_image(file: UploadFile):
 
 @app.post("/analyze-image")
 async def analyze_image(file: UploadFile = File(...)):
+    selected_api_key = random.choice(api_keys)
+    print(f"Selected API Key: {selected_api_key}")
+    genai.configure(api_key=selected_api_key)
+
+    generation_config = {
+        "temperature": 1,
+        "top_p": 0.95,
+        "top_k": 64,
+        "max_output_tokens": 8192,
+        "response_mime_type": "text/plain",
+    }
+    
+    safety_settings = {
+        HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
+        HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
+        HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
+        HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
+    }
     # Process the image
     blob = process_image(file)