diff --git a/__init__.py b/__init__.py index 621e772..2ff0ec1 100644 --- a/__init__.py +++ b/__init__.py @@ -11,6 +11,7 @@ from .fal_kling_video import FalAPI_kling_video from .fal_seedance_video import FalAPI_seedance_video from .nano_banana_pro import FalAPI_NanoBananaPro from .fal_qwen_edit import FalAPI_QwenEditPlus +from .fal_flux2_lora_edit import FalAPI_Flux2LoraEdit from .wavespeed_ai_image import WaveSpeedAI_Image from .string_lower import StringLower from .multi_alpha_composite import MultiAlphaComposite @@ -31,6 +32,7 @@ NODE_CLASS_MAPPINGS = { "FalAPI_seedance_video": FalAPI_seedance_video, "FalAPI_NanoBananaPro": FalAPI_NanoBananaPro, "FalAPI_QwenEditPlus": FalAPI_QwenEditPlus, + "FalAPI_Flux2LoraEdit": FalAPI_Flux2LoraEdit, "WaveSpeedAI_Image": WaveSpeedAI_Image, "StringLower": StringLower, "MultiAlphaComposite": MultiAlphaComposite, @@ -50,6 +52,7 @@ NODE_DISPLAY_NAME_MAPPINGS = { "FalAPI_seedance_video": "Fal API Seedance Video", "FalAPI_NanoBananaPro": "Nano Banana Pro (FAL)", "FalAPI_QwenEditPlus": "FAL Qwen Edit Plus", + "FalAPI_Flux2LoraEdit": "FAL Flux 2 LoRA Edit", "WaveSpeedAI_Image": "WaveSpeedAI Image Generation", "StringLower": "String to Lowercase", "MultiAlphaComposite": "Multi Alpha Composite", diff --git a/fal_flux2_lora_edit.py b/fal_flux2_lora_edit.py new file mode 100644 index 0000000..4fe7133 --- /dev/null +++ b/fal_flux2_lora_edit.py @@ -0,0 +1,302 @@ +import os +import time +import requests +import tempfile +import json +import random +import numpy as np +import torch +import asyncio +import folder_paths +import fal_client +from io import BytesIO +from PIL import Image +import base64 + +class FalAPI_Flux2LoraEdit: + @classmethod + def INPUT_TYPES(cls): + """ + An image generation node for `fal-ai/flux-2/lora/edit`. + """ + return { + "required": { + "api_key": ("STRING", { + "multiline": False, + "default": os.getenv("FAL_KEY", ""), + "display": "FAL API Key" + }), + "prompt": ("STRING", { + "multiline": True, + "default": "Make this donut realistic", + "display": "Prompt" + }), + "image1": ("IMAGE", {"display": "Image 1"}), + "width": ("INT", {"default": 1024, "min": 512, "max": 2048, "step": 8, "display": "Width"}), + "height": ("INT", {"default": 1024, "min": 512, "max": 2048, "step": 8, "display": "Height"}), + }, + "optional": { + "image2": ("IMAGE", {"display": "Image 2 (optional)"}), + "image3": ("IMAGE", {"display": "Image 3 (optional)"}), + "num_inference_steps": ("INT", { + "default": 28, + "min": 4, + "max": 50, + "step": 1, + "display": "Steps" + }), + "guidance_scale": ("FLOAT", { + "default": 2.5, + "min": 0.0, + "max": 20.0, + "step": 0.1, + "display": "Guidance Scale" + }), + "lora_url": ("STRING", { + "multiline": False, + "default": "", + "display": "LoRA URL (optional)" + }), + "lora_scale": ("FLOAT", { + "default": 1.0, + "min": 0.0, + "max": 2.0, + "step": 0.01, + "display": "LoRA Scale" + }), + "seed": ("INT", { + "default": -1, + "min": -1, + "max": 2147483647, + "display": "Seed (-1 for random)" + }), + "output_format": (["jpeg", "png", "webp"], { + "default": "png", + "display": "Output Format" + }), + "enable_safety_checker": ("BOOLEAN", { + "default": True, + "display": "Enable Safety Checker" + }), + "enable_prompt_expansion": ("BOOLEAN", { + "default": False, + "display": "Enable Prompt Expansion" + }), + "acceleration": (["none", "regular", "high"], { + "default": "regular", + "display": "Acceleration" + }), + } + } + + RETURN_TYPES = ("IMAGE", "STRING", "STRING",) + RETURN_NAMES = ("image", "image_path", "generation_info",) + FUNCTION = "generate_image" + CATEGORY = "FAL" + + def __init__(self): + self.temp_dir = folder_paths.get_temp_directory() + self.output_dir = os.path.join("output", "API", "FAL", "flux-2-lora-edit") + self.metadata_dir = os.path.join(self.output_dir, "metadata") + os.makedirs(self.output_dir, exist_ok=True) + os.makedirs(self.metadata_dir, exist_ok=True) + + def get_next_number(self): + valid_exts = {".png", ".jpg", ".webp", ".jpeg"} + files = [f for f in os.listdir(self.output_dir) + if os.path.splitext(f)[1].lower() in valid_exts] + numbers = [] + for file_name in files: + base, _ = os.path.splitext(file_name) + try: + numbers.append(int(base)) + except ValueError: + pass + return max(numbers) + 1 if numbers else 1 + + def save_image_and_metadata(self, pil_img, generation_info, number, ext="png"): + filename = f"{number:03d}.{ext}" + filepath = os.path.join(self.output_dir, filename) + pil_img.save(filepath) + + metadata_filename = f"{number:03d}_metadata.json" + metadata_filepath = os.path.join(self.metadata_dir, metadata_filename) + with open(metadata_filepath, "w", encoding="utf-8") as f: + json.dump(generation_info, f, indent=4, ensure_ascii=False) + return filepath, metadata_filepath + + def on_queue_update(self, update): + try: + if isinstance(update, fal_client.InProgress): + for log in update.logs: + print(f"[Flux2LoraEdit] {log['message']}") + except Exception as e: + print(f"[Flux2LoraEdit] Error in queue update: {e}") + + def tensor_to_tempfile(self, tensor): + pil_img = self.tensor_to_pil(tensor) + fd, filename = tempfile.mkstemp(suffix=".png") + os.close(fd) + pil_img.save(filename, format="PNG") + return filename + + def tensor_to_pil(self, tensor): + if len(tensor.shape) == 4: + tensor = tensor[0] + arr = tensor.cpu().numpy() + if arr.ndim == 3 and arr.shape[0] <= 4: + arr = np.transpose(arr, (1, 2, 0)) + arr = (arr * 255).clip(0, 255).astype("uint8") + return Image.fromarray(arr) + + def pil_to_tensor(self, pil_img): + img_array = np.array(pil_img).astype(np.float32) / 255.0 + return torch.from_numpy(img_array)[None,] + + async def generate_image(self, api_key, prompt, image1, width, height, + image2=None, image3=None, num_inference_steps=28, guidance_scale=2.5, + lora_url="", lora_scale=1.0, seed=-1, output_format="png", + enable_safety_checker=True, enable_prompt_expansion=False, acceleration="regular"): + + print("[Flux2LoraEdit] Starting image generation process...") + + if not api_key: + raise ValueError("A FAL API key is required.") + + os.environ["FAL_KEY"] = api_key + + temp_files = [] + + try: + image_urls = [] + all_images = [image1, image2, image3] + + # Upload images + for i, img_tensor in enumerate(all_images): + if img_tensor is not None: + print(f"[Flux2LoraEdit] Processing input image {i+1}...") + + # Convert to temp file and upload + temp_path = await asyncio.to_thread(self.tensor_to_tempfile, img_tensor) + temp_files.append(temp_path) + + print(f"[Flux2LoraEdit] Uploading image {i+1}...") + url = await asyncio.to_thread(fal_client.upload_file, temp_path) + image_urls.append(url) + + if not image_urls: + raise ValueError("At least one image input is required.") + + # Prepare arguments + arguments = { + "prompt": prompt, + "image_urls": image_urls, + "image_size": {"width": width, "height": height}, + "num_inference_steps": num_inference_steps, + "guidance_scale": guidance_scale, + "enable_safety_checker": enable_safety_checker, + "enable_prompt_expansion": enable_prompt_expansion, + "output_format": output_format, + "acceleration": acceleration, + "sync_mode": False, + "num_images": 1 + } + + # Add LoRA if provided + if lora_url.strip(): + arguments["loras"] = [{"path": lora_url, "scale": lora_scale}] + + if seed != -1: + arguments["seed"] = seed + + print(f"[Flux2LoraEdit] Arguments: {json.dumps(arguments, indent=2)}") + + endpoint = "fal-ai/flux-2/lora/edit" + print(f"[Flux2LoraEdit] Submitting to {endpoint}...") + + # Call API + result = await asyncio.to_thread( + fal_client.subscribe, + endpoint, + arguments=arguments, + with_logs=True, + on_queue_update=self.on_queue_update, + ) + + # Process results + if not result: + raise RuntimeError("No result returned from API") + + output_images = result.get("images", []) + if not output_images: + raise RuntimeError(f"No images in result: {result.keys()}") + + print(f"[Flux2LoraEdit] Downloaded {len(output_images)} images.") + + final_tensors = [] + saved_paths = [] + + for img_info in output_images: + img_url = img_info.get("url") + if not img_url: continue + + # Download + resp = await asyncio.to_thread(requests.get, img_url) + if resp.status_code == 200: + pil_img = Image.open(BytesIO(resp.content)).convert("RGB") + output_tensor = self.pil_to_tensor(pil_img) + final_tensors.append(output_tensor) + + # Save + number = self.get_next_number() + gen_info = { + "prompt": prompt, + "parameters": arguments, + "result": result + } + path, _ = await asyncio.to_thread(self.save_image_and_metadata, pil_img, gen_info, number, output_format) + saved_paths.append(path) + + if not final_tensors: + raise RuntimeError("Failed to process output images") + + if len(final_tensors) > 1: + batch_tensor = torch.cat(final_tensors, dim=0) + else: + batch_tensor = final_tensors[0] + + return (batch_tensor, ";".join(saved_paths), json.dumps(result, indent=2)) + + except Exception as e: + print(f"[Flux2LoraEdit] Error: {str(e)}") + # Clean up + for f in temp_files: + try: os.remove(f) + except: pass + + error_info = { + "error": f"Flux2LoraEdit generation failed: {str(e)}", + "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"), + } + return (torch.zeros((1, 64, 64, 3)), "", json.dumps(error_info, indent=2)) + + finally: + # Cleanup temp inputs + for f in temp_files: + try: os.remove(f) + except: pass + + @classmethod + def IS_CHANGED(cls, **kwargs): + return float("NaN") + +NODE_CLASS_MAPPINGS = { + "FalAPI_Flux2LoraEdit": FalAPI_Flux2LoraEdit +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "FalAPI_Flux2LoraEdit": "FAL Flux 2 LoRA Edit" +} + +__all__ = ["FalAPI_Flux2LoraEdit"] +