flux 2 lora
This commit is contained in:
@@ -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",
|
||||
|
||||
@@ -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"]
|
||||
|
||||
Reference in New Issue
Block a user