OpenWebUI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2025-06-26 18:26:48 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

feat: add comfyui api key

Browse Source
This commit is contained in:
Davide Del Popolo
2024-12-17 08:29:00 +01:00
parent 18fd3db3d5
commit 987b0f0dfd
5 changed files with 41 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
backend/open_webui/utils/images/comfyui.py
View File

@@ -16,14 +16,14 @@ log.setLevel(SRC_LOG_LEVELS["COMFYUI"])
default_headers = {"User-Agent": "Mozilla/5.0"}
def queue_prompt(prompt, client_id, base_url):
def queue_prompt(prompt, client_id, base_url,api_key):
log.info("queue_prompt")
p = {"prompt": prompt, "client_id": client_id}
data = json.dumps(p).encode("utf-8")
log.debug(f"queue_prompt data: {data}")
try:
req = urllib.request.Request(
f"{base_url}/prompt", data=data, headers=default_headers
f"{base_url}/prompt", data=data, headers={**default_headers, "Authorization": f"Bearer {api_key}"}
)
response = urllib.request.urlopen(req).read()
return json.loads(response)
@@ -32,12 +32,12 @@ def queue_prompt(prompt, client_id, base_url):
raise e
def get_image(filename, subfolder, folder_type, base_url):
def get_image(filename, subfolder, folder_type, base_url, api_key):
log.info("get_image")
data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
url_values = urllib.parse.urlencode(data)
req = urllib.request.Request(
f"{base_url}/view?{url_values}", headers=default_headers
f"{base_url}/view?{url_values}", headers={**default_headers, "Authorization": f"Bearer {api_key}"}
)
with urllib.request.urlopen(req) as response:
return response.read()
@@ -50,18 +50,18 @@ def get_image_url(filename, subfolder, folder_type, base_url):
return f"{base_url}/view?{url_values}"
def get_history(prompt_id, base_url):
def get_history(prompt_id, base_url, api_key):
log.info("get_history")
req = urllib.request.Request(
f"{base_url}/history/{prompt_id}", headers=default_headers
f"{base_url}/history/{prompt_id}", headers={**default_headers, "Authorization": f"Bearer {api_key}"}
)
with urllib.request.urlopen(req) as response:
return json.loads(response.read())
def get_images(ws, prompt, client_id, base_url):
prompt_id = queue_prompt(prompt, client_id, base_url)["prompt_id"]
def get_images(ws, prompt, client_id, base_url, api_key):
prompt_id = queue_prompt(prompt, client_id, base_url, api_key)["prompt_id"]
output_images = []
while True:
out = ws.recv()
@@ -74,7 +74,7 @@ def get_images(ws, prompt, client_id, base_url):
else:
continue # previews are binary data
history = get_history(prompt_id, base_url)[prompt_id]
history = get_history(prompt_id, base_url, api_key)[prompt_id]
for o in history["outputs"]:
for node_id in history["outputs"]:
node_output = history["outputs"][node_id]
@@ -113,7 +113,7 @@ class ComfyUIGenerateImageForm(BaseModel):
async def comfyui_generate_image(
model: str, payload: ComfyUIGenerateImageForm, client_id, base_url
model: str, payload: ComfyUIGenerateImageForm, client_id, base_url, api_key
):
ws_url = base_url.replace("http://", "ws://").replace("https://", "wss://")
workflow = json.loads(payload.workflow.workflow)
@@ -167,7 +167,8 @@ async def comfyui_generate_image(
try:
ws = websocket.WebSocket()
ws.connect(f"{ws_url}/ws?clientId={client_id}")
headers = {"Authorization": f"Bearer {api_key}"}
ws.connect(f"{ws_url}/ws?clientId={client_id}", header=headers)
log.info("WebSocket connection established.")
except Exception as e:
log.exception(f"Failed to connect to WebSocket server: {e}")
@@ -176,7 +177,7 @@ async def comfyui_generate_image(
try:
log.info("Sending workflow to WebSocket server.")
log.info(f"Workflow: {workflow}")
images = await asyncio.to_thread(get_images, ws, workflow, client_id, base_url)
images = await asyncio.to_thread(get_images, ws, workflow, client_id, base_url, api_key)
except Exception as e:
log.exception(f"Error while receiving images: {e}")
images = None