Update fluxai.py

fluxai.py

@@ -1,6 +1,7 @@
 from fastapi import APIRouter, Depends
 from fastapi.responses import StreamingResponse
 from PIL import Image, ImageEnhance
+from fastapi import HTTPException
 import io
 import requests
 import os
@@ -55,12 +56,14 @@ def deduct_tokens_gpt(user_id, amount):
 async def fluxai_image(payload: FluxAI):
     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.enhancer:
                 image = Image.open(io.BytesIO(image_bytes))
                 enhancer = ImageEnhance.Sharpness(image)
@@ -69,9 +72,10 @@ async def fluxai_image(payload: FluxAI):
                 image = enhancer.enhance(1.2)
                 enhancer = ImageEnhance.Color(image)
                 image = enhancer.enhance(1.1)
-                file_path = "response_image.jpg"
-                image.save(file_path, quality=95)
-                return StreamingResponse(file_path, media_type="image/jpeg")
+                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")
             else:
                 return StreamingResponse(io.BytesIO(image_bytes), media_type="image/jpeg")