Upload fluxai.py

fluxai.py

@@ -3,17 +3,15 @@ from fastapi.responses import StreamingResponse
 from PIL import Image, ImageEnhance
 from fastapi import HTTPException
 import io
-from io import BytesIO
 import requests
 import os
-import base64
 from dotenv import load_dotenv
 from pydantic import BaseModel
 from pymongo import MongoClient
 from models import *
 from huggingface_hub import InferenceClient
-from fastapi import UploadFile, File
-from fastapi.responses import JSONResponse, FileResponse
+from fastapi import UploadFile
+from fastapi.responses import JSONResponse
 import uuid
 from RyuzakiLib import GeminiLatest
 
@@ -78,6 +76,20 @@ def deduct_tokens_gpt(user_id, amount):
         return False
 
 
+@router.get("/akeno/gettoken")
+async def get_token_with_flux(user_id: int):
+    tokens = get_user_tokens_gpt(user_id)
+    if tokens:
+        return SuccessResponse(
+            status="True",
+            randydev={"tokens": f"Current tokens: {tokens}."}
+        )
+    else:
+        return SuccessResponse(
+            status="False",
+            randydev={"tokens": f"Not enough tokens. Current tokens: {tokens}."}
+        )
+
 @router.post("/akeno/mistralai", response_model=SuccessResponse, responses={422: {"model": SuccessResponse}})
 async def mistralai_(payload: MistralAI):
         try:
@@ -93,7 +105,7 @@ async def mistralai_(payload: MistralAI):
             )
 
 @router.post("/akeno/fluxai", response_model=SuccessResponse, responses={422: {"model": SuccessResponse}})
-async def fluxai_image(payload: FluxAI):
+async def fluxai_image(payload: FluxAI, file: UploadFile):
     if deduct_tokens_gpt(payload.user_id, amount=20):
         try:
             image_bytes = await schellwithflux(payload.args)
@@ -102,33 +114,34 @@ async def fluxai_image(payload: FluxAI):
                     status="False",
                     randydev={"error": "Failed to generate an image"}
                 )
-
             if payload.auto_enhancer:
-                with Image.open(BytesIO(image_bytes)) as image:
+                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 = "akeno.jpg"
+                    enhanced_image_bytes = io.BytesIO()
                     image.save(enhanced_image_bytes, format="JPEG", quality=95)
-                    with open(enhanced_image_bytes, "rb") as image_file:
-                        encoded_string = base64.b64encode(image_file.read())
-
-                example_test = "Explain how this picture looks like."
+                    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, enhanced_image_bytes)
-                
+                response = x.get_response_image(example_test, file_path)
+                url = f"https://randydev-ryuzaki-api.hf.space/{file_path}"
                 return SuccessResponse(
                     status="True",
-                    randydev={"image_data": encoded_string, "caption": response}
+                    randydev={"url": url, "caption": response}
                 )
-
             else:
-                return StreamingResponse(BytesIO(image_bytes), media_type="image/jpeg")
-
+                return StreamingResponse(io.BytesIO(image_bytes), media_type="image/jpeg")
+        
         except Exception as e:
             return SuccessResponse(
                 status="False",
@@ -138,5 +151,5 @@ async def fluxai_image(payload: FluxAI):
         tokens = get_user_tokens_gpt(payload.user_id)
         return SuccessResponse(
             status="False",
-            randydev={"error": f"Not enough tokens. Current tokens: {tokens}. Please support @xtdevs"}
+            randydev={"error": f"Not enough tokens. Current tokens: {tokens}."}
         )