Upload fluxai.py

fluxai.py

@@ -10,6 +10,10 @@ from pydantic import BaseModel
 from pymongo import MongoClient
 from models import *
 from huggingface_hub import InferenceClient
+from fastapi import UploadFile
+from fastapi.responses import JSONResponse
+import uuid
+from RyuzakiLib import GeminiLatest
 
 class FluxAI(BaseModel):
     user_id: int
@@ -24,6 +28,7 @@ router = APIRouter()
 load_dotenv()
 MONGO_URL = os.environ["MONGO_URL"]
 HUGGING_TOKEN = os.environ["HUGGING_TOKEN"]
+GOOGLE_API_KEY = os.environ["GOOGLE_API_KEY"]
 
 client_mongo = MongoClient(MONGO_URL)
 db = client_mongo["tiktokbot"]
@@ -85,17 +90,15 @@ 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:
-            # Generate the image from the flux AI model
             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)
@@ -107,7 +110,20 @@ async def fluxai_image(payload: FluxAI):
                     enhanced_image_bytes = io.BytesIO()
                     image.save(enhanced_image_bytes, format="JPEG", quality=95)
                     enhanced_image_bytes.seek(0)
-                return StreamingResponse(enhanced_image_bytes, media_type="image/jpeg")
+                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")