Skip to content

Commit

Permalink
chore: make comfy workflow can generate image with a random seed (lan…
Browse files Browse the repository at this point in the history
  • Loading branch information
hjlarry authored and JunXu01 committed Nov 9, 2024
1 parent 33518e8 commit 72624e6
Show file tree
Hide file tree
Showing 3 changed files with 24 additions and 1 deletion.
13 changes: 12 additions & 1 deletion api/core/tools/provider/builtin/comfyui/tools/comfyui_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,6 @@ def set_prompt_by_ksampler(self, origin_prompt: dict, positive_prompt: str, nega
prompt = origin_prompt.copy()
id_to_class_type = {id: details["class_type"] for id, details in prompt.items()}
k_sampler = [key for key, value in id_to_class_type.items() if value == "KSampler"][0]
prompt.get(k_sampler)["inputs"]["seed"] = random.randint(10**14, 10**15 - 1)
positive_input_id = prompt.get(k_sampler)["inputs"]["positive"][0]
prompt.get(positive_input_id)["inputs"]["text"] = positive_prompt

Expand All @@ -72,6 +71,18 @@ def set_prompt_images_by_default(self, origin_prompt: dict, image_names: list[st
prompt.get(load_image)["inputs"]["image"] = image_name
return prompt

def set_prompt_seed_by_id(self, origin_prompt: dict, seed_id: str) -> dict:
prompt = origin_prompt.copy()
if seed_id not in prompt:
raise Exception("Not a valid seed node")
if "seed" in prompt[seed_id]["inputs"]:
prompt[seed_id]["inputs"]["seed"] = random.randint(10**14, 10**15 - 1)
elif "noise_seed" in prompt[seed_id]["inputs"]:
prompt[seed_id]["inputs"]["noise_seed"] = random.randint(10**14, 10**15 - 1)
else:
raise Exception("Not a valid seed node")
return prompt

def track_progress(self, prompt: dict, ws: WebSocket, prompt_id: str):
node_ids = list(prompt.keys())
finished_nodes = []
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,9 @@ def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMe
else:
prompt = comfyui.set_prompt_images_by_default(prompt, image_names)

if seed_id := tool_parameters.get("seed_id"):
prompt = comfyui.set_prompt_seed_by_id(prompt, seed_id)

images = comfyui.generate_image_by_prompt(prompt)
result = []
for img in images:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -52,3 +52,12 @@ parameters:
en_US: When the workflow has multiple image nodes, enter the ID list of these nodes, and the images will be passed to ComfyUI in the order of the list.
zh_Hans: 当工作流有多个图片节点时,输入这些节点的ID列表,图片将按列表顺序传给ComfyUI
form: form
- name: seed_id
type: string
label:
en_US: Seed Node Id
zh_Hans: 种子节点ID
human_description:
en_US: If you need to generate different images each time, you need to enter the ID of the seed node.
zh_Hans: 如果需要每次生成时使用不同的种子,需要输入包含种子的节点的ID
form: form

0 comments on commit 72624e6

Please sign in to comment.