add: finetune
This commit is contained in:
@@ -3,6 +3,7 @@ import importlib
|
||||
|
||||
node_list = [
|
||||
"api_node",
|
||||
"finetune",
|
||||
]
|
||||
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
|
||||
+1
-141
@@ -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):
|
||||
|
||||
+122
@@ -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"
|
||||
@@ -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()
|
||||
@@ -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)"
|
||||
}
|
||||
@@ -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"
|
||||
Reference in New Issue
Block a user