mirror of
https://github.com/open-webui/open-webui
synced 2024-11-07 17:19:53 +00:00
ad6e8edcd3
this commit adds three environment variables: - COMFYUI_FLUX: determines whether Flux is used, the workflow is completely different so this is necessary. - COMFYUI_FLUX_WEIGHT_DTYPE: sets the weight precision for Flux. you will probably want to set this to "fp8_e4m3fn" as the fp16 weights take up about 24GB of VRAM. optional, defaults to "default". - COMFYUI_FLUX_FP8_CLIP: Flux requires two CLIP models downloaded, one of which is available in fp8 and fp16. set to true if you are using the fp8 CLIP weights.
406 lines
9.7 KiB
Python
406 lines
9.7 KiB
Python
import websocket # NOTE: websocket-client (https://github.com/websocket-client/websocket-client)
|
|
import uuid
|
|
import json
|
|
import urllib.request
|
|
import urllib.parse
|
|
import random
|
|
import logging
|
|
|
|
from config import SRC_LOG_LEVELS
|
|
|
|
log = logging.getLogger(__name__)
|
|
log.setLevel(SRC_LOG_LEVELS["COMFYUI"])
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from typing import Optional
|
|
|
|
COMFYUI_DEFAULT_PROMPT = """
|
|
{
|
|
"3": {
|
|
"inputs": {
|
|
"seed": 0,
|
|
"steps": 20,
|
|
"cfg": 8,
|
|
"sampler_name": "euler",
|
|
"scheduler": "normal",
|
|
"denoise": 1,
|
|
"model": [
|
|
"4",
|
|
0
|
|
],
|
|
"positive": [
|
|
"6",
|
|
0
|
|
],
|
|
"negative": [
|
|
"7",
|
|
0
|
|
],
|
|
"latent_image": [
|
|
"5",
|
|
0
|
|
]
|
|
},
|
|
"class_type": "KSampler",
|
|
"_meta": {
|
|
"title": "KSampler"
|
|
}
|
|
},
|
|
"4": {
|
|
"inputs": {
|
|
"ckpt_name": "model.safetensors"
|
|
},
|
|
"class_type": "CheckpointLoaderSimple",
|
|
"_meta": {
|
|
"title": "Load Checkpoint"
|
|
}
|
|
},
|
|
"5": {
|
|
"inputs": {
|
|
"width": 512,
|
|
"height": 512,
|
|
"batch_size": 1
|
|
},
|
|
"class_type": "EmptyLatentImage",
|
|
"_meta": {
|
|
"title": "Empty Latent Image"
|
|
}
|
|
},
|
|
"6": {
|
|
"inputs": {
|
|
"text": "Prompt",
|
|
"clip": [
|
|
"4",
|
|
1
|
|
]
|
|
},
|
|
"class_type": "CLIPTextEncode",
|
|
"_meta": {
|
|
"title": "CLIP Text Encode (Prompt)"
|
|
}
|
|
},
|
|
"7": {
|
|
"inputs": {
|
|
"text": "Negative Prompt",
|
|
"clip": [
|
|
"4",
|
|
1
|
|
]
|
|
},
|
|
"class_type": "CLIPTextEncode",
|
|
"_meta": {
|
|
"title": "CLIP Text Encode (Prompt)"
|
|
}
|
|
},
|
|
"8": {
|
|
"inputs": {
|
|
"samples": [
|
|
"3",
|
|
0
|
|
],
|
|
"vae": [
|
|
"4",
|
|
2
|
|
]
|
|
},
|
|
"class_type": "VAEDecode",
|
|
"_meta": {
|
|
"title": "VAE Decode"
|
|
}
|
|
},
|
|
"9": {
|
|
"inputs": {
|
|
"filename_prefix": "ComfyUI",
|
|
"images": [
|
|
"8",
|
|
0
|
|
]
|
|
},
|
|
"class_type": "SaveImage",
|
|
"_meta": {
|
|
"title": "Save Image"
|
|
}
|
|
}
|
|
}
|
|
"""
|
|
|
|
FLUX_DEFAULT_PROMPT = """
|
|
{
|
|
"5": {
|
|
"inputs": {
|
|
"width": 1024,
|
|
"height": 1024,
|
|
"batch_size": 1
|
|
},
|
|
"class_type": "EmptyLatentImage"
|
|
},
|
|
"6": {
|
|
"inputs": {
|
|
"text": "Input Text Here",
|
|
"clip": [
|
|
"11",
|
|
0
|
|
]
|
|
},
|
|
"class_type": "CLIPTextEncode"
|
|
},
|
|
"8": {
|
|
"inputs": {
|
|
"samples": [
|
|
"13",
|
|
0
|
|
],
|
|
"vae": [
|
|
"10",
|
|
0
|
|
]
|
|
},
|
|
"class_type": "VAEDecode"
|
|
},
|
|
"9": {
|
|
"inputs": {
|
|
"filename_prefix": "ComfyUI",
|
|
"images": [
|
|
"8",
|
|
0
|
|
]
|
|
},
|
|
"class_type": "SaveImage"
|
|
},
|
|
"10": {
|
|
"inputs": {
|
|
"vae_name": "ae.sft"
|
|
},
|
|
"class_type": "VAELoader"
|
|
},
|
|
"11": {
|
|
"inputs": {
|
|
"clip_name1": "clip_l.safetensors",
|
|
"clip_name2": "t5xxl_fp16.safetensors",
|
|
"type": "flux"
|
|
},
|
|
"class_type": "DualCLIPLoader"
|
|
},
|
|
"12": {
|
|
"inputs": {
|
|
"unet_name": "flux1-dev.sft",
|
|
"weight_dtype": "default"
|
|
},
|
|
"class_type": "UNETLoader"
|
|
},
|
|
"13": {
|
|
"inputs": {
|
|
"noise": [
|
|
"25",
|
|
0
|
|
],
|
|
"guider": [
|
|
"22",
|
|
0
|
|
],
|
|
"sampler": [
|
|
"16",
|
|
0
|
|
],
|
|
"sigmas": [
|
|
"17",
|
|
0
|
|
],
|
|
"latent_image": [
|
|
"5",
|
|
0
|
|
]
|
|
},
|
|
"class_type": "SamplerCustomAdvanced"
|
|
},
|
|
"16": {
|
|
"inputs": {
|
|
"sampler_name": "euler"
|
|
},
|
|
"class_type": "KSamplerSelect"
|
|
},
|
|
"17": {
|
|
"inputs": {
|
|
"scheduler": "simple",
|
|
"steps": 20,
|
|
"denoise": 1,
|
|
"model": [
|
|
"12",
|
|
0
|
|
]
|
|
},
|
|
"class_type": "BasicScheduler"
|
|
},
|
|
"22": {
|
|
"inputs": {
|
|
"model": [
|
|
"12",
|
|
0
|
|
],
|
|
"conditioning": [
|
|
"6",
|
|
0
|
|
]
|
|
},
|
|
"class_type": "BasicGuider"
|
|
},
|
|
"25": {
|
|
"inputs": {
|
|
"noise_seed": 778937779713005
|
|
},
|
|
"class_type": "RandomNoise"
|
|
}
|
|
}
|
|
"""
|
|
|
|
|
|
def queue_prompt(prompt, client_id, base_url):
|
|
log.info("queue_prompt")
|
|
p = {"prompt": prompt, "client_id": client_id}
|
|
data = json.dumps(p).encode("utf-8")
|
|
req = urllib.request.Request(f"{base_url}/prompt", data=data)
|
|
return json.loads(urllib.request.urlopen(req).read())
|
|
|
|
|
|
def get_image(filename, subfolder, folder_type, base_url):
|
|
log.info("get_image")
|
|
data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
|
|
url_values = urllib.parse.urlencode(data)
|
|
with urllib.request.urlopen(f"{base_url}/view?{url_values}") as response:
|
|
return response.read()
|
|
|
|
|
|
def get_image_url(filename, subfolder, folder_type, base_url):
|
|
log.info("get_image")
|
|
data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
|
|
url_values = urllib.parse.urlencode(data)
|
|
return f"{base_url}/view?{url_values}"
|
|
|
|
|
|
def get_history(prompt_id, base_url):
|
|
log.info("get_history")
|
|
with urllib.request.urlopen(f"{base_url}/history/{prompt_id}") 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"]
|
|
output_images = []
|
|
while True:
|
|
out = ws.recv()
|
|
if isinstance(out, str):
|
|
message = json.loads(out)
|
|
if message["type"] == "executing":
|
|
data = message["data"]
|
|
if data["node"] is None and data["prompt_id"] == prompt_id:
|
|
break # Execution is done
|
|
else:
|
|
continue # previews are binary data
|
|
|
|
history = get_history(prompt_id, base_url)[prompt_id]
|
|
for o in history["outputs"]:
|
|
for node_id in history["outputs"]:
|
|
node_output = history["outputs"][node_id]
|
|
if "images" in node_output:
|
|
for image in node_output["images"]:
|
|
url = get_image_url(
|
|
image["filename"], image["subfolder"], image["type"], base_url
|
|
)
|
|
output_images.append({"url": url})
|
|
return {"data": output_images}
|
|
|
|
|
|
class ImageGenerationPayload(BaseModel):
|
|
prompt: str
|
|
negative_prompt: Optional[str] = ""
|
|
steps: Optional[int] = None
|
|
seed: Optional[int] = None
|
|
width: int
|
|
height: int
|
|
n: int = 1
|
|
cfg_scale: Optional[float] = None
|
|
sampler: Optional[str] = None
|
|
scheduler: Optional[str] = None
|
|
sd3: Optional[bool] = None
|
|
flux: Optional[bool] = None
|
|
flux_weight_dtype: Optional[str] = None
|
|
flux_fp8_clip: Optional[bool] = None
|
|
|
|
|
|
def comfyui_generate_image(
|
|
model: str, payload: ImageGenerationPayload, client_id, base_url
|
|
):
|
|
ws_url = base_url.replace("http://", "ws://").replace("https://", "wss://")
|
|
|
|
comfyui_prompt = json.loads(COMFYUI_DEFAULT_PROMPT)
|
|
|
|
if payload.cfg_scale:
|
|
comfyui_prompt["3"]["inputs"]["cfg"] = payload.cfg_scale
|
|
|
|
if payload.sampler:
|
|
comfyui_prompt["3"]["inputs"]["sampler"] = payload.sampler
|
|
|
|
if payload.scheduler:
|
|
comfyui_prompt["3"]["inputs"]["scheduler"] = payload.scheduler
|
|
|
|
if payload.sd3:
|
|
comfyui_prompt["5"]["class_type"] = "EmptySD3LatentImage"
|
|
|
|
if payload.steps:
|
|
comfyui_prompt["3"]["inputs"]["steps"] = payload.steps
|
|
|
|
comfyui_prompt["4"]["inputs"]["ckpt_name"] = model
|
|
comfyui_prompt["7"]["inputs"]["text"] = payload.negative_prompt
|
|
comfyui_prompt["3"]["inputs"]["seed"] = (
|
|
payload.seed if payload.seed else random.randint(0, 18446744073709551614)
|
|
)
|
|
|
|
# as Flux uses a completely different workflow, we must treat it specially
|
|
if payload.flux:
|
|
comfyui_prompt = json.loads(FLUX_DEFAULT_PROMPT)
|
|
comfyui_prompt["12"]["inputs"]["unet_name"] = model
|
|
comfyui_prompt["25"]["inputs"]["noise_seed"] = (
|
|
payload.seed if payload.seed else random.randint(0, 18446744073709551614)
|
|
)
|
|
|
|
if payload.sampler:
|
|
comfyui_prompt["16"]["inputs"]["sampler_name"] = payload.sampler
|
|
|
|
if payload.steps:
|
|
comfyui_prompt["17"]["inputs"]["steps"] = payload.steps
|
|
|
|
if payload.scheduler:
|
|
comfyui_prompt["17"]["inputs"]["scheduler"] = payload.scheduler
|
|
|
|
if payload.flux_weight_dtype:
|
|
comfyui_prompt["12"]["inputs"]["weight_dtype"] = payload.flux_weight_dtype
|
|
|
|
if payload.flux_fp8_clip:
|
|
comfyui_prompt["11"]["inputs"]["clip_name2"] = "t5xxl_fp8_e4m3fn.safetensors"
|
|
|
|
comfyui_prompt["5"]["inputs"]["batch_size"] = payload.n
|
|
comfyui_prompt["5"]["inputs"]["width"] = payload.width
|
|
comfyui_prompt["5"]["inputs"]["height"] = payload.height
|
|
|
|
# set the text prompt for our positive CLIPTextEncode
|
|
comfyui_prompt["6"]["inputs"]["text"] = payload.prompt
|
|
|
|
try:
|
|
ws = websocket.WebSocket()
|
|
ws.connect(f"{ws_url}/ws?clientId={client_id}")
|
|
log.info("WebSocket connection established.")
|
|
except Exception as e:
|
|
log.exception(f"Failed to connect to WebSocket server: {e}")
|
|
return None
|
|
|
|
try:
|
|
images = get_images(ws, comfyui_prompt, client_id, base_url)
|
|
except Exception as e:
|
|
log.exception(f"Error while receiving images: {e}")
|
|
images = None
|
|
|
|
ws.close()
|
|
|
|
return images
|