Upload fluxai.py

fluxai.py

@@ -17,6 +17,7 @@ from RyuzakiLib import GeminiLatest
 
 class FluxAI(BaseModel):
     user_id: int
+    api_key: str
     args: str
     auto_enhancer: bool = False
 
@@ -104,52 +105,67 @@ async def mistralai_(payload: MistralAI):
                 randydev={"error": f"An error occurred: {str(e)}"}
             )
 
+def get_all_api_keys():
+    user = collection.find({})
+    api_keys = []
+    for x in user:
+        api_key = x.get("ryuzaki_api_key")
+        if api_key:
+            api_keys.append(api_key)
+    return api_keys
+
 @router.post("/akeno/fluxai", response_model=SuccessResponse, responses={422: {"model": SuccessResponse}})
 async def fluxai_image(payload: FluxAI, file: UploadFile):
     if deduct_tokens_gpt(payload.user_id, amount=20):
-        try:
-            image_bytes = await schellwithflux(payload.args)
-            if image_bytes is None:
+        USERS_API_KEYS = get_all_api_keys()
+        if payload.api_key in USERS_API_KEYS:
+            try:
+                image_bytes = await schellwithflux(payload.args)
+                if image_bytes is None:
+                    return SuccessResponse(
+                        status="False",
+                        randydev={"error": "Failed to generate an image"}
+                    )
+                if payload.auto_enhancer:
+                    with Image.open(io.BytesIO(image_bytes)) as image:
+                        enhancer = ImageEnhance.Sharpness(image)
+                        image = enhancer.enhance(1.5)
+                        enhancer = ImageEnhance.Contrast(image)
+                        image = enhancer.enhance(1.2)
+                        enhancer = ImageEnhance.Color(image)
+                        image = enhancer.enhance(1.1)
+                        enhanced_image_bytes = io.BytesIO()
+                        image.save(enhanced_image_bytes, format="JPEG", quality=95)
+                        enhanced_image_bytes.seek(0)
+                    ext = file.filename.split(".")[-1]
+                    unique_filename = f"{uuid.uuid4().hex}.{ext}"
+                    file_path = os.path.join("uploads", unique_filename)
+                    os.makedirs(os.path.dirname(file_path), exist_ok=True)
+                    with open(file_path, "wb") as f:
+                        f.write(enhanced_image_bytes.getvalue())
+                    example_test = "Accurately identify the baked good in the image and provide an appropriate and recipe consistent with your analysis."
+                    x = GeminiLatest(api_keys=GOOGLE_API_KEY)
+                    response = x.get_response_image(example_test, file_path)
+                    url = f"https://randydev-ryuzaki-api.hf.space/{file_path}"
+                    return SuccessResponse(
+                        status="True",
+                        randydev={"url": url, "caption": response}
+                    )
+                else:
+                    return StreamingResponse(io.BytesIO(image_bytes), media_type="image/jpeg")
+            except Exception as e:
                 return SuccessResponse(
                     status="False",
-                    randydev={"error": "Failed to generate an image"}
-                )
-            if payload.auto_enhancer:
-                with Image.open(io.BytesIO(image_bytes)) as image:
-                    enhancer = ImageEnhance.Sharpness(image)
-                    image = enhancer.enhance(1.5)
-                    enhancer = ImageEnhance.Contrast(image)
-                    image = enhancer.enhance(1.2)
-                    enhancer = ImageEnhance.Color(image)
-                    image = enhancer.enhance(1.1)
-                    enhanced_image_bytes = io.BytesIO()
-                    image.save(enhanced_image_bytes, format="JPEG", quality=95)
-                    enhanced_image_bytes.seek(0)
-                ext = file.filename.split(".")[-1]
-                unique_filename = f"{uuid.uuid4().hex}.{ext}"
-                file_path = os.path.join("uploads", unique_filename)
-                os.makedirs(os.path.dirname(file_path), exist_ok=True)
-                with open(file_path, "wb") as f:
-                    f.write(enhanced_image_bytes.getvalue())
-                example_test = "Accurately identify the baked good in the image and provide an appropriate and recipe consistent with your analysis."
-                x = GeminiLatest(api_keys=GOOGLE_API_KEY)
-                response = x.get_response_image(example_test, file_path)
-                url = f"https://randydev-ryuzaki-api.hf.space/{file_path}"
-                return SuccessResponse(
-                    status="True",
-                    randydev={"url": url, "caption": response}
+                    randydev={"error": f"An error occurred: {str(e)}"}
                 )
-            else:
-                return StreamingResponse(io.BytesIO(image_bytes), media_type="image/jpeg")
-        
-        except Exception as e:
+        else:
             return SuccessResponse(
                 status="False",
-                randydev={"error": f"An error occurred: {str(e)}"}
+                randydev={"error": f"Error required api_key"}
             )
     else:
         tokens = get_user_tokens_gpt(payload.user_id)
         return SuccessResponse(
             status="False",
-            randydev={"error": f"Not enough tokens. Current tokens: {tokens}."}
+            randydev={"error": f"Not enough tokens. Current tokens: {tokens} and required api_key."}
         )