150 lines
5.5 KiB
Python
150 lines
5.5 KiB
Python
# nodes.py for ComfyUI-Kontext-API
|
|
"""
|
|
Custom ComfyUI nodes for the Fal Kontext API.
|
|
"""
|
|
|
|
import os
|
|
from typing import Any, Dict
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
from .api import call_kontext_api, FalKontextMaxMultiImageNode
|
|
|
|
# Debug flag - set to False to disable debug output
|
|
DEBUG = False
|
|
|
|
# ComfyUI imports - use proper error handling
|
|
def pil2tensor(image):
|
|
"""Convert PIL image to ComfyUI tensor format."""
|
|
# Convert PIL to numpy array
|
|
img_array = np.array(image).astype(np.float32) / 255.0
|
|
|
|
# Handle different image formats
|
|
if len(img_array.shape) == 2: # Grayscale
|
|
img_array = np.stack([img_array] * 3, axis=-1)
|
|
elif len(img_array.shape) == 3:
|
|
if img_array.shape[2] == 1: # Single channel
|
|
img_array = np.repeat(img_array, 3, axis=2)
|
|
elif img_array.shape[2] == 4: # RGBA
|
|
img_array = img_array[:, :, :3] # Drop alpha channel
|
|
|
|
# Convert to tensor in ComfyUI format (H, W, C)
|
|
tensor = torch.from_numpy(img_array)
|
|
|
|
if DEBUG:
|
|
print(f"[pil2tensor] Output shape: {tensor.shape}, dtype: {tensor.dtype}, min: {tensor.min():.3f}, max: {tensor.max():.3f}")
|
|
|
|
return tensor
|
|
|
|
class KontextAPINode:
|
|
"""
|
|
Main node for calling the Fal Kontext API.
|
|
Inputs: prompt, image, seed, aspect_ratio, disable prompt enhancement.
|
|
Outputs: single output image, info, passed_nsfw_filtering.
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"image": ("IMAGE", {}),
|
|
"seed": ("INT", {"default": -1, "min": -1, "max": 2147483647}),
|
|
"aspect_ratio": ([
|
|
"Match input image",
|
|
"21:9",
|
|
"16:9",
|
|
"4:3",
|
|
"3:2",
|
|
"1:1",
|
|
"2:3",
|
|
"3:4",
|
|
"9:16",
|
|
"9:21"
|
|
], {"default": "Match input image"}),
|
|
},
|
|
"optional": {
|
|
"disable_prompt_enhancement": ("BOOLEAN", {"default": False}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING", "BOOLEAN")
|
|
RETURN_NAMES = ("image", "info", "passed_nsfw_filtering")
|
|
|
|
FUNCTION = "execute"
|
|
|
|
CATEGORY = "image/generation"
|
|
|
|
def execute(self, prompt, image, seed, aspect_ratio, disable_prompt_enhancement=False):
|
|
if DEBUG:
|
|
print(f"[KontextAPINode] Input image shape: {image.shape}, dtype: {image.dtype}")
|
|
print(f"[KontextAPINode] Prompt: {prompt}")
|
|
print(f"[KontextAPINode] Seed: {seed}")
|
|
print(f"[KontextAPINode] Aspect Ratio: {aspect_ratio}")
|
|
|
|
# Validate inputs
|
|
if not prompt.strip():
|
|
return (image, "Error: Prompt cannot be empty", True) # Default to True when error
|
|
|
|
# Handle random seed
|
|
if seed == -1:
|
|
import random
|
|
seed = random.randint(0, 2147483647)
|
|
if DEBUG:
|
|
print(f"[KontextAPINode] Generated random seed: {seed}")
|
|
|
|
# Set values for API call
|
|
guidance_scale = 3.5 # Default
|
|
output_format = "jpeg" # Not used, but required by API
|
|
raw = disable_prompt_enhancement
|
|
image_prompt_strength = 0.1 # Default
|
|
num_inference_steps = 28 # Default
|
|
safety_tolerance = 6 # Max safety
|
|
num_images = 1
|
|
|
|
try:
|
|
output_images, info, passed_nsfw_filtering = call_kontext_api(
|
|
prompt, image, aspect_ratio, num_images, seed, guidance_scale,
|
|
output_format, raw, image_prompt_strength, num_inference_steps, safety_tolerance
|
|
)
|
|
|
|
# Convert PIL image to tensor for ComfyUI
|
|
if output_images and len(output_images) > 0:
|
|
# Take only the first image
|
|
output_image = output_images[0]
|
|
if DEBUG:
|
|
print(f"[KontextAPINode] Got {len(output_images)} images from API, using first one")
|
|
print(f"[KontextAPINode] Output PIL image size: {output_image.size}, mode: {output_image.mode}")
|
|
print(f"[KontextAPINode] Passed NSFW filtering: {passed_nsfw_filtering}")
|
|
|
|
# Convert to tensor
|
|
output_tensor = pil2tensor(output_image)
|
|
|
|
# Add batch dimension for ComfyUI
|
|
output_tensor = output_tensor.unsqueeze(0)
|
|
|
|
if DEBUG:
|
|
print(f"[KontextAPINode] Final output tensor shape: {output_tensor.shape}")
|
|
|
|
return (output_tensor, info, passed_nsfw_filtering)
|
|
else:
|
|
# Return original image if no output
|
|
print(f"[KontextAPINode] ERROR: No images returned from API")
|
|
return (image, f"Error: No images returned from API\n{info}", True)
|
|
|
|
except Exception as e:
|
|
import traceback
|
|
error_msg = f"Error: {str(e)}\n{traceback.format_exc()}"
|
|
print(f"[KontextAPINode] ERROR: {error_msg}")
|
|
return (image, error_msg, True)
|
|
|
|
# Add these for ComfyUI node discovery
|
|
NODE_CLASS_MAPPINGS = {
|
|
"KontextAPINode": KontextAPINode,
|
|
"FalKontextMaxMultiImageNode": FalKontextMaxMultiImageNode,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"KontextAPINode": "Fal Kontext API",
|
|
"FalKontextMaxMultiImageNode": "Fal Kontext Max (Multi-Image)",
|
|
}
|