153 lines
5.0 KiB
Python
153 lines
5.0 KiB
Python
from .modelverse_api.utils import imageurl2tensor
|
|
from .modelverse_api.client import ModelverseClient
|
|
from .modelverse_api.requests.flux_dev import FluxDev
|
|
import torch
|
|
from comfy.comfy_types.node_typing import IO
|
|
|
|
|
|
class FluxDevNode:
|
|
"""
|
|
Flux Image Generator Node (Dev)
|
|
|
|
This node uses Modelverse's Flux model to generate high-quality images.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"client": ("MODELVERSE_API_CLIENT",),
|
|
"prompt": (IO.STRING, {"multiline": True, "default": "", "tooltip": "Text description of the image to generate"}),
|
|
"width": (IO.INT, {
|
|
"default": 1024,
|
|
"min": 512,
|
|
"max": 1536,
|
|
"step": 8,
|
|
"display": "number",
|
|
"tooltip": "Image width (512 to 1536)"
|
|
}),
|
|
"height": (IO.INT, {
|
|
"default": 1024,
|
|
"min": 512,
|
|
"max": 1536,
|
|
"step": 8,
|
|
"display": "number",
|
|
"tooltip": "Image height (512 to 1536)"
|
|
}),
|
|
"strength": (IO.FLOAT, {
|
|
"default": 0.8,
|
|
"min": 0.01,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
"display": "number",
|
|
"tooltip": "Strength of the image-to-image transformation (0.01 to 1.0)"
|
|
}),
|
|
"seed": (IO.INT, {
|
|
"default": -1,
|
|
"min": -1,
|
|
"max": 0xffffffffffffffff,
|
|
"control_after_generate": True,
|
|
"tooltip": "Random seed for reproducible results. -1 for random seed"
|
|
}),
|
|
"num_images": (IO.INT, {
|
|
"default": 1,
|
|
"min": 1,
|
|
"max": 4,
|
|
"step": 1,
|
|
"display": "number",
|
|
"tooltip": "Number of images to generate in a single request (1 to 4)"
|
|
}),
|
|
"num_requests": (IO.INT, {
|
|
"default": 1,
|
|
"min": 1,
|
|
"max": 10,
|
|
"step": 1,
|
|
"display": "number",
|
|
"tooltip": "Number of request to make (1 to 10)"
|
|
}),
|
|
"num_inference_steps": (IO.INT, {
|
|
"default": 28,
|
|
"min": 1,
|
|
"max": 50,
|
|
"step": 1,
|
|
"display": "number",
|
|
"tooltip": "Number of inference steps (1 to 50)"
|
|
}),
|
|
"guidance_scale": (IO.FLOAT, {
|
|
"default": 3.5,
|
|
"min": 0.0,
|
|
"max": 10.0,
|
|
"step": 0.1,
|
|
"display": "number",
|
|
"tooltip": "Guidance scale for generation (0.0 to 10.0)"
|
|
}),
|
|
},
|
|
"optional": {
|
|
"image": (IO.IMAGE, {
|
|
"tooltip": "The image for reference.",
|
|
"forceInput": False,
|
|
"default": None
|
|
})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", )
|
|
RETURN_NAMES = ("image",)
|
|
|
|
CATEGORY = "UCLOUD_MODELVERSE/Flux"
|
|
FUNCTION = "execute"
|
|
|
|
async def execute(self,
|
|
client,
|
|
prompt,
|
|
width=864,
|
|
height=1536,
|
|
strength=0.8,
|
|
seed=-1,
|
|
num_images=1,
|
|
num_requests=1,
|
|
num_inference_steps=28,
|
|
guidance_scale=3.5,
|
|
image=None):
|
|
|
|
if prompt is None or prompt == "":
|
|
raise ValueError("Prompt is required")
|
|
print("INFO:", "Running Flux Dev.")
|
|
|
|
client = ModelverseClient(client["api_key"])
|
|
|
|
tasks = [client.async_send_request(FluxDev(
|
|
prompt=prompt,
|
|
image=image,
|
|
strength=strength,
|
|
guidance_scale=guidance_scale,
|
|
num_images=num_images,
|
|
num_inference_steps=num_inference_steps,
|
|
seed=seed+i,
|
|
width=width,
|
|
height=height,
|
|
))
|
|
for i in range(num_requests)]
|
|
|
|
image_urls = await client.run_tasks(tasks)
|
|
|
|
output_images_list = []
|
|
for image_url in image_urls:
|
|
if not image_url:
|
|
print(
|
|
"WARN:", "No image URLs in the generated result in current request. Skipping...")
|
|
output_images = imageurl2tensor(image_url)
|
|
output_images_list.append(output_images)
|
|
print(
|
|
"INFO:", f"{len(output_images_list)}/{num_requests} request made successfully.")
|
|
return (torch.cat(output_images_list, dim=0),)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Modelverse FluxDevNode": FluxDevNode
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Modelverse FluxDevNode": "Modelverse Flux Dev"
|
|
}
|