From bb9b1d575512d8aaa716ff9d9a0eec0505605c09 Mon Sep 17 00:00:00 2001 From: Tair Date: Sun, 19 Jan 2025 14:23:35 +0000 Subject: [PATCH] reimagine --- __init__.py | 5 ++- nodes/__init__.py | 3 +- nodes/reimagine_node.py | 72 +++++++++++++++++++++++++++++++++ nodes/text_2_image_base_node.py | 6 +-- nodes/text_2_image_fast_node.py | 6 +-- nodes/text_2_image_hd_node.py | 2 +- 6 files changed, 85 insertions(+), 9 deletions(-) create mode 100644 nodes/reimagine_node.py diff --git a/__init__.py b/__init__.py index 50284f2..55ecf2b 100644 --- a/__init__.py +++ b/__init__.py @@ -1,5 +1,6 @@ from .nodes import (EraserNode, GenFillNode, ShotByTextNode, ShotByImageNode, TailoredGenNode, - TailoredModelInfoNode, Text2ImageBaseNode, Text2ImageFastNode, Text2ImageHDNode) + TailoredModelInfoNode, Text2ImageBaseNode, Text2ImageFastNode, Text2ImageHDNode, + ReimagineNode) # Map the node class to a name used internally by ComfyUI NODE_CLASS_MAPPINGS = { "BriaEraser": EraserNode, # Return the class, not an instance @@ -11,6 +12,7 @@ NODE_CLASS_MAPPINGS = { "Text2ImageBaseNode": Text2ImageBaseNode, "Text2ImageFastNode": Text2ImageFastNode, "Text2ImageHDNode": Text2ImageHDNode, + "ReimagineNode": ReimagineNode, } # Map the node display name to the one shown in the ComfyUI node interface NODE_DISPLAY_NAME_MAPPINGS = { @@ -23,4 +25,5 @@ NODE_DISPLAY_NAME_MAPPINGS = { "Text2ImageBaseNode": "Bria Text2Image Base", "Text2ImageFastNode": "Bria Text2Image Fast", "Text2ImageHDNode": "Bria Text2Image HD", + "ReimagineNode": "Bria Reimagine", } diff --git a/nodes/__init__.py b/nodes/__init__.py index b0f09aa..e09abad 100644 --- a/nodes/__init__.py +++ b/nodes/__init__.py @@ -6,4 +6,5 @@ from .tailored_gen_node import TailoredGenNode from .tailored_model_info_node import TailoredModelInfoNode from .text_2_image_base_node import Text2ImageBaseNode from .text_2_image_fast_node import Text2ImageFastNode -from .text_2_image_hd_node import Text2ImageHDNode \ No newline at end of file +from .text_2_image_hd_node import Text2ImageHDNode +from .reimagine_node import ReimagineNode \ No newline at end of file diff --git a/nodes/reimagine_node.py b/nodes/reimagine_node.py new file mode 100644 index 0000000..91385bc --- /dev/null +++ b/nodes/reimagine_node.py @@ -0,0 +1,72 @@ +import requests + +from .common import postprocess_image, preprocess_image, image_to_base64 + + +class ReimagineNode(): + @classmethod + def INPUT_TYPES(self): + return { + "required": { + "api_key": ("STRING", ), + }, + "optional": { + "prompt": ("STRING",), + "seed": ("INT", {"default": -1}), + "steps_num": ("INT", {"default": 12}), # if used with tailored, possibly get this from the tailored model info node + "structure_ref_influence": ("FLOAT", {"default": 0.75}), + "fast": ("INT", {"default": 0}), # if used with tailored, possibly get this from the tailored model info node + "structure_image": ("IMAGE", ), + "tailored_model_id": ("STRING", ), + "tailored_model_influence": ("FLOAT", {"default": 0.5}), + "tailored_generation_prefix": ("STRING",), # if used with tailored, possibly get this from the tailored model info node + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("output_image",) + CATEGORY = "API Nodes" + FUNCTION = "execute" # This is the method that will be executed + + def __init__(self): + self.api_url = "https://engine.prod.bria-api.com/v1/reimagine" #"http://0.0.0.0:5000/v1/reimagine" + + def execute( + self, api_key, prompt, seed, + steps_num, fast, structure_ref_influence, structure_image=None, + tailored_model_id=None, tailored_model_influence=None, tailored_generation_prefix=None, + ): + fast = bool(fast) + payload = { + "prompt": tailored_generation_prefix + prompt, + "num_results": 1, + "sync": True, + "seed": seed, + "steps_num": steps_num, + "include_generation_prefix": False, + } + if structure_image is not None: + structure_image = preprocess_image(structure_image) + structure_image = image_to_base64(structure_image) + payload["structure_image_file"] = structure_image + payload["structure_ref_influence"] = structure_ref_influence + if tailored_model_id is not None and tailored_model_id != "": + payload["tailored_model_id"] = tailored_model_id + payload["tailored_model_influence"] = tailored_model_influence + response = requests.post( + self.api_url, + json=payload, + headers={"api_token": api_key} + ) + print("api_key" + api_key) + print("api_url" + self.api_url) + print("payload" + str(payload)) + + print("now") + if response.status_code == 200: + response_dict = response.json() + image_response = requests.get(response_dict['result'][0]["urls"][0]) + result_image = postprocess_image(image_response.content) + return (result_image,) + else: + raise Exception(f"Error: API request failed with status code {response.status_code} and text {response.text}") diff --git a/nodes/text_2_image_base_node.py b/nodes/text_2_image_base_node.py index 78a28e8..2463845 100644 --- a/nodes/text_2_image_base_node.py +++ b/nodes/text_2_image_base_node.py @@ -9,9 +9,9 @@ class Text2ImageBaseNode(): return { "required": { "api_key": ("STRING", ), - "prompt": ("STRING",), }, "optional": { + "prompt": ("STRING",), "aspect_ratio": (["1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9"], {"default": "4:3"}), "seed": ("INT", {"default": -1}), "negative_prompt": ("STRING", {"default": ""}), @@ -25,7 +25,7 @@ class Text2ImageBaseNode(): "guidance_method_2": (["controlnet_canny", "controlnet_depth", "controlnet_recoloring", "controlnet_color_grid"], {"default": "controlnet_canny"}), "guidance_method_2_scale": ("FLOAT", {"default": 1.0}), "guidance_method_2_image": ("IMAGE", ), - "image_prompt_mode": (["none", "regular", "style_only"], {"default": "none"}), + "image_prompt_mode": (["regular", "style_only"], {"default": "regular"}), "image_prompt_image": ("IMAGE", ), "image_prompt_scale": ("FLOAT", {"default": 1.0}), } @@ -72,7 +72,7 @@ class Text2ImageBaseNode(): payload["guidance_method_2"] = guidance_method_2 payload["guidance_method_2_scale"] = guidance_method_2_scale payload["guidance_method_2_image_file"] = guidance_method_2_image - if image_prompt_mode != "none": + if image_prompt_image != "none": image_prompt_image = preprocess_image(image_prompt_image) image_prompt_image = image_to_base64(image_prompt_image) payload["image_prompt_mode"] = image_prompt_mode diff --git a/nodes/text_2_image_fast_node.py b/nodes/text_2_image_fast_node.py index 402ab0a..442d3da 100644 --- a/nodes/text_2_image_fast_node.py +++ b/nodes/text_2_image_fast_node.py @@ -9,9 +9,9 @@ class Text2ImageFastNode(): return { "required": { "api_key": ("STRING", ), - "prompt": ("STRING",), }, "optional": { + "prompt": ("STRING",), "aspect_ratio": (["1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9"], {"default": "4:3"}), "seed": ("INT", {"default": -1}), "steps_num": ("INT", {"default": 8}), @@ -22,7 +22,7 @@ class Text2ImageFastNode(): "guidance_method_2": (["controlnet_canny", "controlnet_depth", "controlnet_recoloring", "controlnet_color_grid"], {"default": "controlnet_canny"}), "guidance_method_2_scale": ("FLOAT", {"default": 1.0}), "guidance_method_2_image": ("IMAGE", ), - "image_prompt_mode": (["none", "regular", "style_only"], {"default": "none"}), + "image_prompt_mode": (["regular", "style_only"], {"default": "regular"}), "image_prompt_image": ("IMAGE", ), "image_prompt_scale": ("FLOAT", {"default": 1.0}), } @@ -65,7 +65,7 @@ class Text2ImageFastNode(): payload["guidance_method_2"] = guidance_method_2 payload["guidance_method_2_scale"] = guidance_method_2_scale payload["guidance_method_2_image_file"] = guidance_method_2_image - if image_prompt_mode != "none": + if image_prompt_image != "none": image_prompt_image = preprocess_image(image_prompt_image) image_prompt_image = image_to_base64(image_prompt_image) payload["image_prompt_mode"] = image_prompt_mode diff --git a/nodes/text_2_image_hd_node.py b/nodes/text_2_image_hd_node.py index fafba48..a8df502 100644 --- a/nodes/text_2_image_hd_node.py +++ b/nodes/text_2_image_hd_node.py @@ -9,9 +9,9 @@ class Text2ImageHDNode(): return { "required": { "api_key": ("STRING", ), - "prompt": ("STRING",), }, "optional": { + "prompt": ("STRING",), "aspect_ratio": (["1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9"], {"default": "4:3"}), "seed": ("INT", {"default": -1}), "negative_prompt": ("STRING", {"default": ""}),