Additions to the last commit

2024-11-18 22:26:13 +08:00
parent 2166920cf0
commit e76a8e634c
1 changed files with 52 additions and 0 deletions
@@ -0,0 +1,52 @@
+import json
+import os
+import yaml
+import requests
+import pathlib
+from aiohttp import web
+
+root_path = pathlib.Path(__file__).parent.parent.parent
+config_path = os.path.join(root_path,'config.yaml')
+class FluxAIAPI:
+    def __init__(self):
+        self.api_url = "https://fluxaiimagegenerator.com/api"
+        self.origin = "https://fluxaiimagegenerator.com"
+        self.user_agent = None
+        self.cookie = None
+
+    def promptGenerate(self, text, cookies=None):
+        cookie = self.cookie if cookies is None else cookies
+        if cookie is None:
+            if os.path.isfile(config_path):
+                with open(config_path, 'r') as f:
+                    data = yaml.load(f, Loader=yaml.FullLoader)
+                    if 'FLUXAI_COOKIE' not in data:
+                        raise Exception("Please add FLUXAI_COOKIE to config.yaml")
+                    if "FLUXAI_USER_AGENT" in data:
+                        self.user_agent = data["FLUXAI_USER_AGENT"]
+                    self.cookie = cookie = data['FLUXAI_COOKIE']
+
+        headers = {
+            "Cookie": cookie,
+            "Referer": "https://fluxaiimagegenerator.com/flux-prompt-generator",
+            "Origin": self.origin,
+            "Content-Type": "application/json",
+        }
+        if self.user_agent is not None:
+            headers['User-Agent'] = self.user_agent
+
+        url = self.api_url + '/prompt'
+        json = {
+            "prompt": text
+        }
+
+        response = requests.post(url, json=json, headers=headers)
+        res = response.json()
+        data = res['data']
+        if "error" in data:
+            return data['error']
+        elif "prompt" in data:
+            return data['prompt']
+
+fluxaiAPI = FluxAIAPI()
+