From 1998249c46e29eb51dd367d97ce029c0c5e56a88 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?C=CC=A7ag=CC=86an=20Kiraz?= Date: Sun, 8 Jun 2025 02:09:00 +0300 Subject: [PATCH] add: finetune --- __init__.py | 1 + nodes/api_node.py | 142 +---------------- nodes/base.py | 122 ++++++++++++++ nodes/config.py | 59 +++++++ nodes/finetune.py | 397 ++++++++++++++++++++++++++++++++++++++++++++++ nodes/status.py | 9 ++ 6 files changed, 589 insertions(+), 141 deletions(-) create mode 100644 nodes/base.py create mode 100644 nodes/config.py create mode 100644 nodes/finetune.py create mode 100644 nodes/status.py diff --git a/__init__.py b/__init__.py index 0befe34..7b890ae 100644 --- a/__init__.py +++ b/__init__.py @@ -3,6 +3,7 @@ import importlib node_list = [ "api_node", +"finetune", ] NODE_CLASS_MAPPINGS = {} diff --git a/nodes/api_node.py b/nodes/api_node.py index 62718c6..54e6a11 100644 --- a/nodes/api_node.py +++ b/nodes/api_node.py @@ -1,144 +1,4 @@ -import requests -from PIL import Image -import io -import numpy as np -import torch -import os -import configparser -import time -from enum import Enum -from urllib.parse import urljoin - -class Status(Enum): - TASK_NOT_FOUND = "Task not found" - PENDING = "Pending" - REQUEST_MODERATED = "Request Moderated" - CONTENT_MODERATED = "Content Moderated" - READY = "Ready" - ERROR = "Error" - -class ConfigLoader: - def __init__(self): - current_dir = os.path.dirname(os.path.abspath(__file__)) - parent_dir = os.path.dirname(current_dir) - config_path = os.path.join(parent_dir, "config.ini") - self.config = configparser.ConfigParser() - self.config.read(config_path) - self.set_x_key() - - def get_key(self, section, key): - try: - return self.config[section][key] - except KeyError: - raise KeyError(f"{key} not found in section {section} of config file.") - - def create_url(self, path): - try: - base_url = self.get_key('API', 'BASE_URL') - return urljoin(base_url, path) - except KeyError as e: - raise KeyError(f"Error constructing URL: {str(e)}") - - def set_x_key(self): - try: - x_key = self.get_key('API', 'X_KEY') - os.environ["X_KEY"] = x_key - except KeyError as e: - print(f"Error: {str(e)}") - - -config_loader = ConfigLoader() - -class BaseFlux: - RETURN_TYPES = ("IMAGE",) - FUNCTION = "generate_image" - CATEGORY = "BFL" - - def process_result(self, result, output_format="jpeg"): - try: - sample_url = result['result']['sample'] - img_response = requests.get(sample_url) - img = Image.open(io.BytesIO(img_response.content)) - - with io.BytesIO() as output: - img.save(output, format=output_format.upper()) - output.seek(0) - img_converted = Image.open(output) - - img_array = np.array(img_converted).astype(np.float32) / 255.0 - img_tensor = torch.from_numpy(img_array)[None,] - return (img_tensor,) - except KeyError as e: - print(f"KeyError: Missing expected key {e}") - return self.create_blank_image() - except Exception as e: - print(f"Error processing image result: {str(e)}") - return self.create_blank_image() - - def create_blank_image(self): - blank_img = Image.new('RGB', (512, 512), color='black') - img_array = np.array(blank_img).astype(np.float32) / 255.0 - img_tensor = torch.from_numpy(img_array)[None,] - return (img_tensor,) - - def check_multiple_of_32(self, width, height): - if width % 32 != 0 or height % 32 != 0: - raise ValueError(f"Width {width} and height {height} must be multiples of 32.") - - def post_request(self, url_path, arguments): - post_url = config_loader.create_url(url_path) - headers = {"x-key": os.environ["X_KEY"]} - response = requests.post(post_url, json=arguments, headers=headers) - - if response.status_code == 200: - return response.json().get("id") - else: - print(f"Error initiating request: {response.status_code}, {response.text}") - return None - - def get_result(self, task_id, output_format="jpeg", attempt=1, max_attempts=20): - if attempt > max_attempts: - print(f"Max attempts reached for task_id {task_id}. Image not ready.") - return self.create_blank_image() - - get_url = config_loader.create_url(f"get_result?id={task_id}") - headers = {"x-key": os.environ["X_KEY"]} - result_response = requests.get(get_url, headers=headers) - - if result_response.status_code == 200: - try: - result = result_response.json() - status = result.get("status") - if Status(status) == Status.READY: - return self.process_result(result, output_format=output_format) - elif Status(status) == Status.PENDING: - print(f"Attempt {attempt}: Image not ready, status is '{status}'. Retrying in 5 seconds...") - time.sleep(5) - return self.get_result(task_id, output_format, attempt + 1, max_attempts) - else: - print(f"Error: Unexpected status '{status}'") - return self.create_blank_image() - except ValueError as e: - print(f"Error parsing JSON response: {str(e)}") - print(f"Response content: {result_response.text}") - return self.create_blank_image() - else: - print(f"Error fetching result: {result_response.status_code}, {result_response.text}") - return self.create_blank_image() - - def generate_image(self, url_path, arguments): - if "width" in arguments and "height" in arguments: - self.check_multiple_of_32(arguments["width"], arguments["height"]) - - try: - task_id = self.post_request(url_path, arguments) - if task_id: - print(f"Task ID '{task_id}'") - return self.get_result(task_id, output_format=arguments.get("output_format", "jpeg")) - return self.create_blank_image() - except Exception as e: - print(f"Error generating image: {str(e)}") - return self.create_blank_image() +from .base import BaseFlux class FluxPro11(BaseFlux): diff --git a/nodes/base.py b/nodes/base.py new file mode 100644 index 0000000..80b8d8a --- /dev/null +++ b/nodes/base.py @@ -0,0 +1,122 @@ +import requests +from PIL import Image +import io +import numpy as np +import torch +import os +import time +from .status import Status +from .config import config_loader + +class BaseFlux: + RETURN_TYPES = ("IMAGE",) + FUNCTION = "generate_image" + CATEGORY = "BFL" + + def process_result(self, result, output_format="jpeg"): + try: + sample_url = result['result']['sample'] + img_response = requests.get(sample_url) + img = Image.open(io.BytesIO(img_response.content)) + + with io.BytesIO() as output: + img.save(output, format=output_format.upper()) + output.seek(0) + img_converted = Image.open(output) + + img_array = np.array(img_converted).astype(np.float32) / 255.0 + img_tensor = torch.from_numpy(img_array)[None,] + return (img_tensor,) + except KeyError as e: + print(f"KeyError: Missing expected key {e}") + return self.create_blank_image() + except Exception as e: + print(f"Error processing image result: {str(e)}") + return self.create_blank_image() + + def create_blank_image(self): + blank_img = Image.new('RGB', (512, 512), color='black') + img_array = np.array(blank_img).astype(np.float32) / 255.0 + img_tensor = torch.from_numpy(img_array)[None,] + return (img_tensor,) + + def check_multiple_of_32(self, width, height): + if width % 32 != 0 or height % 32 != 0: + raise ValueError(f"Width {width} and height {height} must be multiples of 32.") + + def post_request(self, url_path, arguments): + post_url = config_loader.create_url(url_path) + headers = {"x-key": os.environ["X_KEY"]} + response = requests.post(post_url, json=arguments, headers=headers) + + if response.status_code == 200: + return response.json().get("id") + else: + print(f"Error initiating request: {response.status_code}, {response.text}") + return None + + def get_result(self, task_id, output_format="jpeg", max_attempts=20): + headers = {"x-key": os.environ["X_KEY"]} + get_url = config_loader.create_url(f"get_result?id={task_id}") + attempt = 1 + + while attempt <= max_attempts: + try: + result_response = requests.get(get_url, headers=headers) + + if result_response.status_code != 200: + print(f"HTTP Error on attempt {attempt}/{max_attempts}: {result_response.status_code}, {result_response.text}") + attempt += 1 + if attempt <= max_attempts: + time.sleep(5) + continue + + result = result_response.json() + status = result.get("status") + + if Status(status) == Status.READY: + return self.process_result(result, output_format=output_format) + elif Status(status) == Status.PENDING: + print(f"Attempt {attempt}/{max_attempts}: Image not ready, status is '{status}'. Retrying in 5 seconds...") + attempt += 1 + if attempt <= max_attempts: + time.sleep(5) + elif Status(status) in [Status.ERROR, Status.CONTENT_MODERATED, Status.REQUEST_MODERATED]: + print(f"Terminal status '{status}' - stopping retries") + break + else: + print(f"Unknown status '{status}' on attempt {attempt}/{max_attempts}") + attempt += 1 + if attempt <= max_attempts: + time.sleep(5) + + except ValueError as e: + print(f"JSON parsing error on attempt {attempt}/{max_attempts}: {str(e)}") + attempt += 1 + if attempt <= max_attempts: + time.sleep(5) + except Exception as e: + print(f"Unexpected error on attempt {attempt}/{max_attempts}: {str(e)}") + attempt += 1 + if attempt <= max_attempts: + time.sleep(5) + + print(f"All attempts exhausted for task_id {task_id}. Returning blank image.") + return self.create_blank_image() + + def generate_image(self, url_path, arguments): + if "width" in arguments and "height" in arguments: + self.check_multiple_of_32(arguments["width"], arguments["height"]) + + try: + task_id = self.post_request(url_path, arguments) + if task_id: + print(f"Task ID '{task_id}'") + return self.get_result(task_id, output_format=arguments.get("output_format", "jpeg")) + return self.create_blank_image() + except Exception as e: + print(f"Error generating image: {str(e)}") + return self.create_blank_image() + +class BaseFinetuneFlux(BaseFlux): + CATEGORY = "BFL/Finetune" \ No newline at end of file diff --git a/nodes/config.py b/nodes/config.py new file mode 100644 index 0000000..0711cf5 --- /dev/null +++ b/nodes/config.py @@ -0,0 +1,59 @@ +import os +import configparser +from urllib.parse import urljoin + +class ConfigLoader: + def __init__(self): + current_dir = os.path.dirname(os.path.abspath(__file__)) + parent_dir = os.path.dirname(current_dir) + config_path = os.path.join(parent_dir, "config.ini") + self.config = configparser.ConfigParser() + self.config.read(config_path) + self.set_x_key() + + # Regional endpoints for finetuning (required by BFL API) + self.regional_endpoints = { + "us": "https://api.us1.bfl.ai", + "eu": "https://api.eu1.bfl.ai" + } + + def get_key(self, section, key): + try: + return self.config[section][key] + except KeyError: + raise KeyError(f"{key} not found in section {section} of config file.") + + def create_url(self, path, region=None): + """ + Create URL for API endpoints. + + Args: + path: API endpoint path + region: Optional region for finetuning operations ("us" or "eu") + If provided, uses regional endpoint instead of global + """ + try: + if region and region in self.regional_endpoints: + base_url = self.regional_endpoints[region] + return urljoin(base_url, f"/v1/{path}") + else: + base_url = self.get_key('API', 'BASE_URL') + return urljoin(base_url, path) + except KeyError as e: + raise KeyError(f"Error constructing URL: {str(e)}") + + def get_regional_endpoint(self, region): + """Get the full regional endpoint URL for a given region.""" + if region not in self.regional_endpoints: + raise ValueError(f"Invalid region '{region}'. Must be one of: {list(self.regional_endpoints.keys())}") + return self.regional_endpoints[region] + + def set_x_key(self): + try: + x_key = self.get_key('API', 'X_KEY') + os.environ["X_KEY"] = x_key + except KeyError as e: + print(f"Error: {str(e)}") + +# Create a singleton instance to be shared across modules +config_loader = ConfigLoader() \ No newline at end of file diff --git a/nodes/finetune.py b/nodes/finetune.py new file mode 100644 index 0000000..fb38d84 --- /dev/null +++ b/nodes/finetune.py @@ -0,0 +1,397 @@ +import requests +import os +import json +from .base import BaseFinetuneFlux +from .config import config_loader + +class FluxFinetune: + RETURN_TYPES = ("STRING",) + FUNCTION = "create_finetune" + CATEGORY = "BFL/Finetune" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "file_data": ("STRING", {"default": "", "multiline": True}), + "finetune_comment": ("STRING", {"default": "my-finetune"}), + "trigger_word": ("STRING", {"default": "TOK"}), + "mode": (["character", "product", "style", "general"], {"default": "general"}), + "region": (["us", "eu"], {"default": "us"}), + "iterations": ("INT", {"default": 300, "min": 100, "max": 2000}), + "learning_rate": ("FLOAT", {"default": 0.00001, "min": 0.000001, "max": 0.01}), + "captioning": ("BOOLEAN", {"default": True}), + "priority": (["speed", "quality", "high_res_only"], {"default": "quality"}), + "finetune_type": (["full", "lora"], {"default": "full"}), + "lora_rank": ("INT", {"default": 32, "min": 8, "max": 128}) + }, + "optional": { + "webhook_url": ("STRING", {"default": ""}), + "webhook_secret": ("STRING", {"default": ""}) + } + } + + def create_finetune(self, file_data, finetune_comment, trigger_word, mode, region, iterations, learning_rate, + captioning, priority, finetune_type, lora_rank, webhook_url="", webhook_secret=""): + + # Auto-adjust learning rate based on finetune_type if using default + if learning_rate == 0.00001: # Default value + if finetune_type == "lora": + learning_rate = 0.0001 # 10x higher for LoRA as per docs + print(f"Auto-adjusted learning rate to {learning_rate} for LoRA finetune") + + arguments = { + "file_data": file_data, + "finetune_comment": finetune_comment, + "trigger_word": trigger_word, + "mode": mode, + "iterations": iterations, + "learning_rate": learning_rate, + "captioning": captioning, + "priority": priority, + "finetune_type": finetune_type, + "lora_rank": lora_rank + } + + if webhook_url: + arguments["webhook_url"] = webhook_url + if webhook_secret: + arguments["webhook_secret"] = webhook_secret + + try: + post_url = config_loader.create_url("finetune", region=region) + headers = {"x-key": os.environ["X_KEY"]} + response = requests.post(post_url, json=arguments, headers=headers) + + # Log the full response details + print("=" * 80) + print("FINETUNE RESPONSE:") + print("=" * 80) + print(f"Status Code: {response.status_code}") + print(f"Response Body:") + try: + response_json = response.json() + print(json.dumps(response_json, indent=2)) + except: + print(f"Raw Response Text: {response.text}") + print("=" * 80) + + if response.status_code == 200: + result = response.json() + finetune_id = result.get("id", "Unknown") + regional_endpoint = config_loader.get_regional_endpoint(region) + print(f"✅ Finetune created successfully with ID: {finetune_id}") + print(f"Region: {region.upper()} ({regional_endpoint})") + print(f"⚠️ Remember: Use the same region ({region}) for inference!") + return (str(finetune_id),) + else: + print(f"❌ Error creating finetune: {response.status_code}") + return ("Error",) + except Exception as e: + print("=" * 80) + print("FINETUNE EXCEPTION:") + print("=" * 80) + print(f"Exception Type: {type(e).__name__}") + print(f"Exception Message: {str(e)}") + import traceback + print(f"Full Traceback:") + traceback.print_exc() + print("=" * 80) + return ("Error",) + +class FluxProFinetune(BaseFinetuneFlux): + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "finetune_id": ("STRING", {"default": "my-finetune"}), + "region": (["us", "eu"], {"default": "us"}), + "prompt": ("STRING", {"default": "", "multiline": True}), + "finetune_strength": ("FLOAT", {"default": 1.2, "min": 0.0, "max": 2.0}), + "steps": ("INT", {"default": 40, "min": 1, "max": 50}), + "guidance": ("FLOAT", {"default": 2.5, "min": 1.5, "max": 5}), + "width": ("INT", {"default": 1024, "min": 256, "max": 1440}), + "height": ("INT", {"default": 768, "min": 256, "max": 1440}), + "prompt_upsampling": ("BOOLEAN", {"default": False}), + "safety_tolerance": ("INT", {"default": 2, "min": 1, "max": 6}), + "output_format": (["jpeg", "png"], {"default": "jpeg"}) + }, + "optional": { + "image_prompt": ("STRING", {"default": ""}), + "seed": ("INT", {"default": -1}), + "webhook_url": ("STRING", {"default": ""}), + "webhook_secret": ("STRING", {"default": ""}) + } + } + + def generate_image(self, finetune_id, region, prompt, finetune_strength, steps, guidance, width, height, + prompt_upsampling, safety_tolerance, output_format, image_prompt="", seed=-1, + webhook_url="", webhook_secret=""): + arguments = { + "finetune_id": finetune_id, + "prompt": prompt, + "finetune_strength": finetune_strength, + "steps": steps, + "guidance": guidance, + "width": width, + "height": height, + "prompt_upsampling": prompt_upsampling, + "safety_tolerance": safety_tolerance, + "output_format": output_format + } + + if image_prompt: + arguments["image_prompt"] = image_prompt + if seed != -1: + arguments["seed"] = seed + if webhook_url: + arguments["webhook_url"] = webhook_url + if webhook_secret: + arguments["webhook_secret"] = webhook_secret + + return self.generate_regional_image("flux-pro-finetuned", arguments, region) + + def generate_regional_image(self, endpoint, arguments, region): + if "width" in arguments and "height" in arguments: + self.check_multiple_of_32(arguments["width"], arguments["height"]) + + try: + post_url = config_loader.create_url(endpoint, region=region) + headers = {"x-key": os.environ["X_KEY"]} + response = requests.post(post_url, json=arguments, headers=headers) + + if response.status_code == 200: + task_id = response.json().get("id") + if task_id: + regional_endpoint = config_loader.get_regional_endpoint(region) + print(f"Finetune Task ID '{task_id}' (Region: {region.upper()} - {regional_endpoint})") + return self.get_result(task_id, output_format=arguments.get("output_format", "jpeg")) + else: + print(f"Error initiating regional finetune request: {response.status_code}, {response.text}") + return self.create_blank_image() + except Exception as e: + print(f"Error generating regional finetune image: {str(e)}") + return self.create_blank_image() + +class FluxProDepthFinetune(BaseFinetuneFlux): + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "finetune_id": ("STRING", {"default": "my-finetune"}), + "prompt": ("STRING", {"default": "", "multiline": True}), + "control_image": ("STRING", {"default": None}), + "finetune_strength": ("FLOAT", {"default": 1.1, "min": 0.1, "max": 2.0}), + "prompt_upsampling": ("BOOLEAN", {"default": False}), + "steps": ("INT", {"default": 50, "min": 15, "max": 50}), + "output_format": (["jpeg", "png"], {"default": "jpeg"}), + "guidance": ("FLOAT", {"default": 15, "min": 1.5, "max": 100.0}), + "safety_tolerance": ("INT", {"default": 2, "min": 0, "max": 6}) + }, + "optional": { + "seed": ("INT", {"default": -1}), + "webhook_url": ("STRING", {"default": ""}), + "webhook_secret": ("STRING", {"default": ""}) + } + } + + def generate_image(self, finetune_id, prompt, control_image, finetune_strength, prompt_upsampling, + steps, output_format, guidance, safety_tolerance, seed=-1, webhook_url="", webhook_secret=""): + arguments = { + "finetune_id": finetune_id, + "prompt": prompt, + "control_image": control_image, + "finetune_strength": finetune_strength, + "prompt_upsampling": prompt_upsampling, + "steps": steps, + "output_format": output_format, + "guidance": guidance, + "safety_tolerance": safety_tolerance + } + + if seed != -1: + arguments["seed"] = seed + if webhook_url: + arguments["webhook_url"] = webhook_url + if webhook_secret: + arguments["webhook_secret"] = webhook_secret + + return super().generate_image("flux-pro-1.0-depth-finetuned", arguments) + +class FluxProCannyFinetune(BaseFinetuneFlux): + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "finetune_id": ("STRING", {"default": "my-finetune"}), + "prompt": ("STRING", {"default": "", "multiline": True}), + "control_image": ("STRING", {"default": None}), + "finetune_strength": ("FLOAT", {"default": 1.1, "min": 0.1, "max": 2.0}), + "canny_low_threshold": ("INT", {"default": 250, "min": 0, "max": 255}), + "canny_high_threshold": ("INT", {"default": 250, "min": 0, "max": 255}), + "prompt_upsampling": ("BOOLEAN", {"default": False}), + "steps": ("INT", {"default": 50, "min": 15, "max": 50}), + "output_format": (["jpeg", "png"], {"default": "jpeg"}), + "guidance": ("FLOAT", {"default": 30, "min": 1.5, "max": 100.0}), + "safety_tolerance": ("INT", {"default": 2, "min": 0, "max": 6}) + }, + "optional": { + "seed": ("INT", {"default": -1}), + "webhook_url": ("STRING", {"default": ""}), + "webhook_secret": ("STRING", {"default": ""}) + } + } + + def generate_image(self, finetune_id, prompt, control_image, finetune_strength, canny_low_threshold, + canny_high_threshold, prompt_upsampling, steps, output_format, guidance, safety_tolerance, + seed=-1, webhook_url="", webhook_secret=""): + arguments = { + "finetune_id": finetune_id, + "prompt": prompt, + "control_image": control_image, + "finetune_strength": finetune_strength, + "canny_low_threshold": canny_low_threshold, + "canny_high_threshold": canny_high_threshold, + "prompt_upsampling": prompt_upsampling, + "steps": steps, + "output_format": output_format, + "guidance": guidance, + "safety_tolerance": safety_tolerance + } + + if seed != -1: + arguments["seed"] = seed + if webhook_url: + arguments["webhook_url"] = webhook_url + if webhook_secret: + arguments["webhook_secret"] = webhook_secret + + return super().generate_image("flux-pro-1.0-canny-finetuned", arguments) + +class FluxProFillFinetune(BaseFinetuneFlux): + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "finetune_id": ("STRING", {"default": "my-finetune"}), + "image": ("STRING", {"default": None}), + "finetune_strength": ("FLOAT", {"default": 1.1, "min": 0.1, "max": 2.0}), + "steps": ("INT", {"default": 28, "min": 15, "max": 50}), + "prompt_upsampling": ("BOOLEAN", {"default": False}), + "guidance": ("FLOAT", {"default": 60.0, "min": 1.5, "max": 100.0}), + "safety_tolerance": ("INT", {"default": 2, "min": 0, "max": 6}), + "output_format": (["jpeg", "png"], {"default": "jpeg"}) + }, + "optional": { + "mask": ("STRING", {"default": ""}), + "prompt": ("STRING", {"default": "", "multiline": True}), + "seed": ("INT", {"default": -1}), + "webhook_url": ("STRING", {"default": ""}), + "webhook_secret": ("STRING", {"default": ""}) + } + } + + def generate_image(self, finetune_id, image, finetune_strength, steps, prompt_upsampling, guidance, + safety_tolerance, output_format, mask="", prompt="", seed=-1, webhook_url="", webhook_secret=""): + arguments = { + "finetune_id": finetune_id, + "image": image, + "finetune_strength": finetune_strength, + "steps": steps, + "prompt_upsampling": prompt_upsampling, + "guidance": guidance, + "safety_tolerance": safety_tolerance, + "output_format": output_format + } + + if mask: + arguments["mask"] = mask + if prompt: + arguments["prompt"] = prompt + if seed != -1: + arguments["seed"] = seed + if webhook_url: + arguments["webhook_url"] = webhook_url + if webhook_secret: + arguments["webhook_secret"] = webhook_secret + + return super().generate_image("flux-pro-1.0-fill-finetuned", arguments) + +class FluxPro11UltraFinetune(BaseFinetuneFlux): + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "finetune_id": ("STRING", {"default": "my-finetune"}), + "region": (["us", "eu"], {"default": "us"}), + "prompt": ("STRING", {"default": "", "multiline": True}), + "finetune_strength": ("FLOAT", {"default": 1.2, "min": 0.0, "max": 2.0}), + "aspect_ratio": (["16:9", "4:3", "1:1", "3:2", "21:9", "9:16", "3:4", "2:3", "9:21"], {"default": "16:9"}), + "safety_tolerance": ("INT", {"default": 2, "min": 1, "max": 6}), + "output_format": (["jpeg", "png"], {"default": "jpeg"}), + "raw": ("BOOLEAN", {"default": False}) + }, + "optional": { + "seed": ("INT", {"default": -1}), + "webhook_url": ("STRING", {"default": ""}), + "webhook_secret": ("STRING", {"default": ""}) + } + } + + def generate_image(self, finetune_id, region, prompt, finetune_strength, aspect_ratio, safety_tolerance, + output_format, raw, seed=-1, webhook_url="", webhook_secret=""): + arguments = { + "finetune_id": finetune_id, + "prompt": prompt, + "finetune_strength": finetune_strength, + "aspect_ratio": aspect_ratio, + "safety_tolerance": safety_tolerance, + "output_format": output_format, + "raw": raw + } + + if seed != -1: + arguments["seed"] = seed + if webhook_url: + arguments["webhook_url"] = webhook_url + if webhook_secret: + arguments["webhook_secret"] = webhook_secret + + return self.generate_regional_image("flux-pro-1.1-ultra-finetuned", arguments, region) + + def generate_regional_image(self, endpoint, arguments, region): + try: + post_url = config_loader.create_url(endpoint, region=region) + headers = {"x-key": os.environ["X_KEY"]} + response = requests.post(post_url, json=arguments, headers=headers) + + if response.status_code == 200: + task_id = response.json().get("id") + if task_id: + regional_endpoint = config_loader.get_regional_endpoint(region) + print(f"Finetune Task ID '{task_id}' (Region: {region.upper()} - {regional_endpoint})") + return self.get_result(task_id, output_format=arguments.get("output_format", "jpeg")) + else: + print(f"Error initiating regional finetune request: {response.status_code}, {response.text}") + return self.create_blank_image() + except Exception as e: + print(f"Error generating regional finetune image: {str(e)}") + return self.create_blank_image() + +NODE_CLASS_MAPPINGS = { + "FluxFinetune_BFL": FluxFinetune, + "FluxProFinetune_BFL": FluxProFinetune, + "FluxProDepthFinetune_BFL": FluxProDepthFinetune, + "FluxProCannyFinetune_BFL": FluxProCannyFinetune, + "FluxProFillFinetune_BFL": FluxProFillFinetune, + "FluxPro11UltraFinetune_BFL": FluxPro11UltraFinetune +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "FluxFinetune_BFL": "Flux Finetune Creator (BFL)", + "FluxProFinetune_BFL": "Flux Pro Finetune (BFL)", + "FluxProDepthFinetune_BFL": "Flux Pro Depth Finetune (BFL)", + "FluxProCannyFinetune_BFL": "Flux Pro Canny Finetune (BFL)", + "FluxProFillFinetune_BFL": "Flux Pro Fill Finetune (BFL)", + "FluxPro11UltraFinetune_BFL": "Flux Pro 1.1 Ultra Finetune (BFL)" +} \ No newline at end of file diff --git a/nodes/status.py b/nodes/status.py new file mode 100644 index 0000000..ba4efe5 --- /dev/null +++ b/nodes/status.py @@ -0,0 +1,9 @@ +from enum import Enum + +class Status(Enum): + TASK_NOT_FOUND = "Task not found" + PENDING = "Pending" + REQUEST_MODERATED = "Request Moderated" + CONTENT_MODERATED = "Content Moderated" + READY = "Ready" + ERROR = "Error" \ No newline at end of file