diff --git a/__init__.py b/__init__.py index 41c9b55..6f8ca3c 100644 --- a/__init__.py +++ b/__init__.py @@ -1,5 +1,4 @@ -from .flux_kontext_pro_node import NODE_CLASS_MAPPINGS as PRO_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as PRO_DISPLAY -from .flux_kontext_max_node import NODE_CLASS_MAPPINGS as MAX_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as MAX_DISPLAY +from .flux_kontext import NODE_CLASS_MAPPINGS as PRO_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as PRO_DISPLAY from .gemini_node import NODE_CLASS_MAPPINGS as GEMINI_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as GEMINI_DISPLAY from .gemini_diarisation import NODE_CLASS_MAPPINGS as GEMINI_DIAR_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as GEMINI_DIAR_DISPLAY from .gpt_image1 import NODE_CLASS_MAPPINGS as GPT_IMAGE_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as GPT_IMAGE_DISPLAY @@ -13,7 +12,7 @@ from .gemini_tts import NODE_CLASS_MAPPINGS as GEMINI_TTS_MAPPINGS, NODE_DISPLAY from .elevenlabs_tts import NODE_CLASS_MAPPINGS as ELEVENLABS_TTS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as ELEVENLABS_TTS_DISPLAY from .cleanup import NODE_CLASS_MAPPINGS as CLEANUP_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS as CLEANUP_DISPLAY -NODE_CLASS_MAPPINGS = {**PRO_MAPPINGS, **MAX_MAPPINGS, **GEMINI_MAPPINGS, **GEMINI_DIAR_MAPPINGS, **GPT_IMAGE_MAPPINGS, **IMAGEN_IMAGE_MAPPINGS, **IMAGEN_EDIT_MAPPINGS, **VEO_MAPPINGS, **VEO_GEMINI_MAPPINGS, **GEMINI_SEGMENT_MAPPINGS, **NANO_BANANA_MAPPINGS, **GEMINI_TTS_MAPPINGS, **ELEVENLABS_TTS_MAPPINGS, **CLEANUP_MAPPINGS} -NODE_DISPLAY_NAME_MAPPINGS = {**PRO_DISPLAY, **MAX_DISPLAY, **GEMINI_DISPLAY, **GEMINI_DIAR_DISPLAY, **GPT_IMAGE_DISPLAY, **IMAGEN_IMAGE_DISPLAY, **IMAGEN_EDIT_DISPLAY, **VEO_DISPLAY, **VEO_GEMINI_DISPLAY, **GEMINI_SEGMENT_DISPLAY, **NANO_BANANA_DISPLAY, **GEMINI_TTS_DISPLAY, **ELEVENLABS_TTS_DISPLAY, **CLEANUP_DISPLAY} +NODE_CLASS_MAPPINGS = {**PRO_MAPPINGS, **GEMINI_MAPPINGS, **GEMINI_DIAR_MAPPINGS, **GPT_IMAGE_MAPPINGS, **IMAGEN_IMAGE_MAPPINGS, **IMAGEN_EDIT_MAPPINGS, **VEO_MAPPINGS, **VEO_GEMINI_MAPPINGS, **GEMINI_SEGMENT_MAPPINGS, **NANO_BANANA_MAPPINGS, **GEMINI_TTS_MAPPINGS, **ELEVENLABS_TTS_MAPPINGS, **CLEANUP_MAPPINGS} +NODE_DISPLAY_NAME_MAPPINGS = {**PRO_DISPLAY, **GEMINI_DISPLAY, **GEMINI_DIAR_DISPLAY, **GPT_IMAGE_DISPLAY, **IMAGEN_IMAGE_DISPLAY, **IMAGEN_EDIT_DISPLAY, **VEO_DISPLAY, **VEO_GEMINI_DISPLAY, **GEMINI_SEGMENT_DISPLAY, **NANO_BANANA_DISPLAY, **GEMINI_TTS_DISPLAY, **ELEVENLABS_TTS_DISPLAY, **CLEANUP_DISPLAY} __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/flux_kontext_max_node.py b/flux_kontext.py similarity index 57% rename from flux_kontext_max_node.py rename to flux_kontext.py index 76d1a65..f7de13b 100644 --- a/flux_kontext_max_node.py +++ b/flux_kontext.py @@ -6,31 +6,20 @@ import numpy as np from PIL import Image import io -class FluxKontextMaxNode: +class FluxKontextNode: @classmethod def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE",), - "prompt": ("STRING", { - "multiline": True, - "default": "Make this a 90s cartoon" - }), - "api_key": ("STRING", { - "default": "" - }), - "aspect_ratio": (["1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3", "5:4", "4:5", "21:9", "9:21", "2:1", "1:2", "match_input_image"], { - "default": "match_input_image" - }), - "output_format": (["jpg", "png"], { - "default": "jpg" - }), - "safety_tolerance": ("INT", { - "default": 2, - "min": 0, - "max": 6, - "step": 1 - }), + "prompt": ("STRING", {"multiline": True, "default": "Make this a 90s cartoon"}), + "api_key": ("STRING", {"default": ""}), + "model": (["flux-kontext-dev", "flux-kontext-max", "flux-kontext-pro"], {"default": "flux-kontext-dev"}), + "aspect_ratio": (["1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3", "5:4", "4:5", "21:9", "9:21", "2:1", "1:2", "match_input_image"], {"default": "match_input_image"}), + "output_format": (["jpg", "png"], {"default": "jpg"}), + "safety_tolerance": ("INT", {"default": 2, "min": 0, "max": 6, "step": 1}), + "seed": ("INT", {"default": 69, "min": 1, "max": 2147483646, "step": 1}), + "prompt_upsampling": ("BOOLEAN", {"default": False}) } } @@ -39,7 +28,7 @@ class FluxKontextMaxNode: FUNCTION = "generate_image" CATEGORY = "image/edit" - def generate_image(self, image, prompt, api_key, aspect_ratio, output_format, safety_tolerance): + def generate_image(self, image, prompt, api_key, model, aspect_ratio, output_format, safety_tolerance, seed, prompt_upsampling): try: os.environ["REPLICATE_API_TOKEN"] = api_key @@ -53,16 +42,21 @@ class FluxKontextMaxNode: pil_image.save(img_buffer, format='PNG') img_buffer.seek(0) - # Run Replicate model + # Build input dict with all parameters for both models + replicate_input = { + "prompt": prompt, + "input_image": img_buffer, + "aspect_ratio": aspect_ratio, + "output_format": output_format, + "safety_tolerance": safety_tolerance, + "seed": seed, + "prompt_upsampling": prompt_upsampling + } + + # Run Replicate model with selected model output = replicate.run( - "black-forest-labs/flux-kontext-max", - input={ - "prompt": prompt, - "input_image": img_buffer, - "aspect_ratio": aspect_ratio, - "output_format": output_format, - "safety_tolerance": safety_tolerance - } + f"black-forest-labs/{model}", + input=replicate_input ) # Get URL from output @@ -85,9 +79,9 @@ class FluxKontextMaxNode: return (torch.zeros((1, 512, 512, 3)),) NODE_CLASS_MAPPINGS = { - "FluxKontextMaxNode": FluxKontextMaxNode + "FluxKontextNode": FluxKontextNode } NODE_DISPLAY_NAME_MAPPINGS = { - "FluxKontextMaxNode": "Flux Kontext Max" + "FluxKontextNode": "Flux Kontext" } \ No newline at end of file diff --git a/flux_kontext_pro_node.py b/flux_kontext_pro_node.py deleted file mode 100644 index faaaee3..0000000 --- a/flux_kontext_pro_node.py +++ /dev/null @@ -1,93 +0,0 @@ -import replicate -import os -import requests -import torch -import numpy as np -from PIL import Image -import io - -class FluxKontextProNode: - @classmethod - def INPUT_TYPES(cls): - return { - "required": { - "image": ("IMAGE",), - "prompt": ("STRING", { - "multiline": True, - "default": "Make this a 90s cartoon" - }), - "api_key": ("STRING", { - "default": "" - }), - "aspect_ratio": (["1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3", "5:4", "4:5", "21:9", "9:21", "2:1", "1:2", "match_input_image"], { - "default": "match_input_image" - }), - "output_format": (["jpg", "png"], { - "default": "jpg" - }), - "safety_tolerance": ("INT", { - "default": 2, - "min": 0, - "max": 6, - "step": 1 - }), - } - } - - RETURN_TYPES = ("IMAGE",) - RETURN_NAMES = ("image",) - FUNCTION = "generate_image" - CATEGORY = "image/edit" - - def generate_image(self, image, prompt, api_key, aspect_ratio, output_format, safety_tolerance): - try: - os.environ["REPLICATE_API_TOKEN"] = api_key - - # Convert tensor to PIL and save to buffer - tensor = image.squeeze(0) if len(image.shape) == 4 else image - if tensor.max() <= 1.0: - tensor = (tensor * 255).clamp(0, 255).byte() - pil_image = Image.fromarray(tensor.cpu().numpy(), 'RGB') - - img_buffer = io.BytesIO() - pil_image.save(img_buffer, format='PNG') - img_buffer.seek(0) - - # Run Replicate model - output = replicate.run( - "black-forest-labs/flux-kontext-pro", - input={ - "prompt": prompt, - "input_image": img_buffer, - "aspect_ratio": aspect_ratio, - "output_format": output_format, - "safety_tolerance": safety_tolerance - } - ) - - # Get URL from output - output_url = output if isinstance(output, str) else (output[0] if isinstance(output, list) and output else str(output)) - - # Download and convert back to tensor - response = requests.get(output_url, timeout=30) - response.raise_for_status() - - downloaded_image = Image.open(io.BytesIO(response.content)) - if downloaded_image.mode != 'RGB': - downloaded_image = downloaded_image.convert('RGB') - - np_image = np.array(downloaded_image).astype(np.float32) / 255.0 - output_tensor = torch.from_numpy(np_image).unsqueeze(0) - - return (output_tensor,) - - except Exception as e: - return (torch.zeros((1, 512, 512, 3)),) - -NODE_CLASS_MAPPINGS = { - "FluxKontextProNode": FluxKontextProNode -} - -NODE_DISPLAY_NAME_MAPPINGS = { - "FluxKontextProNode": "Flux Kontext Pro" -} \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 1136cfa..8ebb8f0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "externalapi-helpers" description = "Various ComfyUI nodes for Gemini, Replicate and OpenAI" -version = "1.0.3" +version = "1.0.4" license = {file = "LICENSE"} # classifiers = [ # # For OS-independent nodes (works on all operating systems)