Improved Flux-Kontext nodes

Author: Aryan185

__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 .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']

flux_kontext_max_node.py flux_kontext.pyflux_kontext_max_node.py renamed to flux_kontext.py

@@ -6,31 +6,20 @@
 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 @@ def INPUT_TYPES(cls):
     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 @@ def generate_image(self, image, prompt, api_key, aspect_ratio, output_format, sa
             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 @@ def generate_image(self, image, prompt, api_key, aspect_ratio, output_format, sa
             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"
 }

flux_kontext_pro_node.py

@@ -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)