Upload 2 files

fluxai.py

@@ -0,0 +1,76 @@
+from fastapi import APIRouter, Depends
+from fastapi.responses import StreamingResponse
+import io
+import requests
+import os
+from dotenv import load_dotenv
+from pydantic import BaseModel
+
+class FluxAI(BaseModel):
+    user_id: int
+    args: str
+
+router = APIRouter()
+
+load_dotenv()
+MONGO_URL = os.environ["MONGO_URL"]
+HUGGING_TOKEN = os.environ["HUGGING_TOKEN"]
+
+client_mongo = MongoClient(MONGO_URL)
+db = client_mongo["tiktokbot"]
+collection = db["users"]
+
+async def schellwithflux(args):
+    API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell"
+    headers = {"Authorization": f"Bearer {HUGGING_TOKEN}"}
+    payload = {"inputs": args}
+    response = requests.post(API_URL, headers=headers, json=payload)
+    if response.status_code != 200:
+        print(f"Error status {response.status_code}")
+        return None
+    return response.content
+
+def get_user_tokens_gpt(user_id):
+    user = collection.find_one({"user_id": user_id})
+    if not user:
+        return 0
+    return user.get("tokens", 0)
+
+def deduct_tokens_gpt(user_id, amount):
+    tokens = get_user_tokens_gpt(user_id)
+    if tokens >= amount:
+        collection.update_one(
+            {"user_id": user_id},
+            {"$inc": {"tokens": -amount}}
+        )
+        return True
+    else:
+        return False
+
+@router.post("/akeno/fluxai", response_model=SuccessResponse, responses={422: {"model": SuccessResponse}})
+async def fluxai_image(payload: FluxAI):
+    if payload.user_id == 1191668125:
+        return SuccessResponse(status="False", randydev={"message": "Only Developer."})
+    
+    if deduct_tokens_gpt(payload.user_id, amount=20):
+        try:
+            image_bytes = await schellwithflux(payload.args)
+            if image_bytes is None:
+                return SuccessResponse(
+                    status="False",
+                    randydev={"error": "Failed to generate an image"}
+                )
+            
+            return StreamingResponse(io.BytesIO(image_bytes), media_type="image/jpeg")
+        
+        except Exception as e:
+            return SuccessResponse(
+                status="False",
+                randydev={"error": f"An error occurred: {str(e)}"}
+            )
+    else:
+        tokens = get_user_tokens_gpt(payload.user_id)
+        return SuccessResponse(
+            status="False",
+            randydev={"error": f"Not enough tokens. Current tokens: {tokens}. Please support @xtdevs"}
+        )

main.py

@@ -82,6 +82,7 @@ from gpytranslate import SyncTranslator
 
 import logging
 import functions as code
+from fluxai import router as fluxai_router
 
 logging.basicConfig(level=logging.ERROR)
 logging.basicConfig(level=logging.INFO)
@@ -125,6 +126,7 @@ collection = db["users"]
 trans = SyncTranslator()
 
 app = FastAPI(docs_url=None, redoc_url="/")
+app.include_router(fluxai_router, prefix="/api/v1")
 
 timeout = 100
 
@@ -195,7 +197,6 @@ RAMDOM_STATUS = [
     "spammer",
 ]
 
-
 def remove_sibyl_system_banned(user_id):
     update_doc = {
         "sibyl_ban": None,
@@ -532,7 +533,7 @@ def getfedbans_(payload: GetsaFedBans, api_key: str = Depends(validate_api_key_f
                 "error": str(e)
             }
         ) 
-                
+
 @app.post("/user/fedban", response_model=SuccessResponse, responses={422: {"model": SuccessResponse}})
 def fedbans_(payload: FedBans, api_key: str = Depends(validate_api_key_fedbans)):
     if payload.user_id == 1191668125:
@@ -541,7 +542,6 @@ def fedbans_(payload: FedBans, api_key: str = Depends(validate_api_key_fedbans))
     date_joined = str(dt.now())
     if not payload.hashtag.startswith("#"):
         return SuccessResponse(status="False", randydev={"message": "Invalid hashtag."})
-
     try:
         new_user_spammers(
             user_id=payload.user_id,