Upload fluxai.py

fluxai.py

@@ -1,5 +1,6 @@
 from fastapi import APIRouter, Depends
 from fastapi.responses import StreamingResponse
+from PIL import Image, ImageEnhance
 import io
 import requests
 import os
@@ -10,7 +11,8 @@ from models import *
 
 class FluxAI(BaseModel):
     user_id: int
-    args: str
+    args: str,
+    auto_enhancer: bool
 
 router = APIRouter()
 
@@ -59,8 +61,19 @@ async def fluxai_image(payload: FluxAI):
                     status="False",
                     randydev={"error": "Failed to generate an image"}
                 )
-            
-            return StreamingResponse(io.BytesIO(image_bytes), media_type="image/jpeg")
+            if payload.enhancer:
+                image = Image.open(io.BytesIO(image_bytes))
+                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)
+                file_path = "response_image.jpg"
+                image.save(file_path, quality=95)
+                return StreamingResponse(file_path, media_type="image/jpeg")
+            else:
+                return StreamingResponse(io.BytesIO(image_bytes), media_type="image/jpeg")
         
         except Exception as e:
             return SuccessResponse(