Buckets:

ACCC1380
/

molab-checkpoints

Files

xet

ACCC1380/molab-checkpoints / ComfyUI /comfy_api_nodes /nodes_meta.py

ACCC1380

1 day ago

download

raw

10.9 kB

import base64
import re
from io import BytesIO

import torch
from typing_extensions import override

from comfy_api.latest import IO, ComfyExtension
from comfy_api_nodes.apis.meta import (
    MuseImageEditRequest,
    MuseImageInput,
    MuseImageRequest,
    MuseImageResponse,
    MuseImageToolEnablement,
)
from comfy_api_nodes.util import (
    ApiEndpoint,
    bytesio_to_image_tensor,
    download_url_to_image_tensor,
    pad_images_to_common_channels,
    sync_op,
    upload_images_to_comfyapi,
    validate_string,
)

GENERATIONS_PATH = "/proxy/meta/v1/images/generations"
EDITS_PATH = "/proxy/meta/v1/images/edits"
MUSE_IMAGE_MODELS = ["muse-image-1.0"]
MAX_INPUT_IMAGES = 10
ASPECT_RATIOS = ["auto", "1:1", "3:2", "2:3", "4:3", "3:4", "5:4", "4:5", "16:9", "9:16", "21:9", "9:21", "2:1", "1:2"]
REASONING_STRENGTHS = ["high", "low"]

_IMAGE_REF_RE = re.compile(r"@image(?P<idx>\d*)(?!\w)", re.IGNORECASE | re.ASCII)


def _resolve_image_refs(prompt: str, total_images: int) -> str:
    parts = []
    pos = 0
    prev_end = -1
    for match in _IMAGE_REF_RE.finditer(prompt):
        start = match.start()
        if start > 0 and start != prev_end and (prompt[start - 1].isalnum() or prompt[start - 1] == "_"):
            continue
        idx = int(match.group("idx") or 1)
        if not 1 <= idx <= total_images:
            raise ValueError(
                f"The prompt references @Image{idx}, but only {total_images} reference images "
                f"are connected (a batched input counts once per image)."
            )
        parts.append(prompt[pos:start])
        parts.append(f"image {idx}")
        pos = match.end()
        prev_end = match.end()
    parts.append(prompt[pos:])
    return "".join(parts)


def _size(aspect_ratio: str) -> str | None:
    return None if aspect_ratio == "auto" else aspect_ratio.replace(":", "x")


async def _decode_images(cls: type[IO.ComfyNode], response: MuseImageResponse) -> torch.Tensor:
    images = []
    for item in response.data:
        if item.b64_json:
            images.append(bytesio_to_image_tensor(BytesIO(base64.b64decode(item.b64_json))))
        elif item.url:
            images.append(await download_url_to_image_tensor(item.url, cls=cls))
    if not images:
        raise Exception("The response contains no images.")
    return torch.cat(pad_images_to_common_channels(images))


def _reasoning_strength_input() -> IO.Combo.Input:
    return IO.Combo.Input(
        "reasoning_strength",
        options=REASONING_STRENGTHS,
        tooltip="How much the model thinks, plans and self-refines before rendering.",
    )


def _t2i_model_option(model_id: str) -> IO.DynamicCombo.Option:
    return IO.DynamicCombo.Option(
        model_id,
        [
            IO.String.Input(
                "prompt",
                multiline=True,
                default="",
                tooltip="Prompt describing the image. The model reasons about the prompt, and may use "
                "its built-in web and image search, before rendering.",
            ),
            IO.Combo.Input(
                "aspect_ratio",
                options=ASPECT_RATIOS,
                tooltip="Aspect ratio of the output. Images are rendered at about 2.5 megapixels "
                "(1:1 is 1600x1600, 16:9 is 2048x1152); 'auto' lets the model choose from the prompt.",
            ),
            _reasoning_strength_input(),
            *_tool_toggle_inputs(),
            _seed_input(),
        ],
    )


def _edit_model_option(model_id: str) -> IO.DynamicCombo.Option:
    return IO.DynamicCombo.Option(
        model_id,
        [
            IO.Autogrow.Input(
                "images",
                template=IO.Autogrow.TemplateNames(
                    IO.Image.Input("image"),
                    names=[f"image_{i}" for i in range(1, MAX_INPUT_IMAGES + 1)],
                    min=1,
                ),
                tooltip=f"1-{MAX_INPUT_IMAGES} reference images to edit or combine. Refer to them in the prompt "
                "as @Image1, @Image2, ..., numbered in input order; a batched input counts once per image.",
            ),
            IO.String.Input(
                "prompt",
                multiline=True,
                default="",
                tooltip="Editing instructions. Supports @Image1-style references to the input images.",
            ),
            IO.Combo.Input(
                "aspect_ratio",
                options=ASPECT_RATIOS,
                tooltip="Aspect ratio of the output. Images are rendered at about 2.5 megapixels "
                "(1:1 is 1600x1600, 16:9 is 2048x1152); 'auto' keeps the aspect ratio of the input.",
            ),
            _reasoning_strength_input(),
            *_tool_toggle_inputs(),
            _seed_input(),
        ],
    )


def _tool_toggle_inputs() -> list[IO.Boolean.Input]:
    return [
        IO.Boolean.Input(
            "enable_web_search",
            default=True,
            advanced=True,
            tooltip="Lets the model search the web for facts and live information while planning the image.",
        ),
        IO.Boolean.Input(
            "enable_image_search",
            default=True,
            advanced=True,
            tooltip="Lets the model search for reference images while planning the image.",
        ),
        IO.Boolean.Input(
            "enable_shell",
            default=True,
            advanced=True,
            tooltip="Lets the model run code while planning, for precise layouts, charts and diagrams; "
            "when off, quantities and alignment are approximated.",
        ),
    ]


def _tool_enablement(model: dict) -> MuseImageToolEnablement | None:
    if model["enable_web_search"] and model["enable_image_search"] and model["enable_shell"]:
        return None
    return MuseImageToolEnablement(
        enable_image_search=model["enable_image_search"],
        enable_web_search=model["enable_web_search"],
        enable_shell=model["enable_shell"],
    )


def _seed_input() -> IO.Int.Input:
    return IO.Int.Input(
        "seed",
        default=42,
        min=0,
        max=2147483647,
        step=1,
        display_mode=IO.NumberDisplay.number,
        control_after_generate=True,
        tooltip="Seed to determine if node should re-run; the API has no seed, "
        "so actual results are nondeterministic regardless of this value.",
    )


def _price_badge() -> IO.PriceBadge:
    return IO.PriceBadge(expr="""{"type":"usd","usd":0.0143}""")


class MetaMuseImageTextToImageApi(IO.ComfyNode):
    @classmethod
    def define_schema(cls):
        return IO.Schema(
            node_id="MetaMuseImageTextToImageApi",
            display_name="Meta Muse Image Text to Image",
            category="partner/image/Meta",
            description="Generates images from a text prompt using Meta's Muse Image model, "
            "which reasons about the prompt before rendering.",
            inputs=[
                IO.DynamicCombo.Input(
                    "model",
                    options=[_t2i_model_option(model_id) for model_id in MUSE_IMAGE_MODELS],
                    tooltip="Model to use.",
                ),
            ],
            outputs=[
                IO.Image.Output(),
            ],
            hidden=[
                IO.Hidden.auth_token_comfy_org,
                IO.Hidden.api_key_comfy_org,
                IO.Hidden.unique_id,
            ],
            is_api_node=True,
            price_badge=_price_badge(),
        )

    @classmethod
    async def execute(cls, model: dict):
        validate_string(model["prompt"], min_length=1)
        response = await sync_op(
            cls,
            ApiEndpoint(path=GENERATIONS_PATH, method="POST"),
            response_model=MuseImageResponse,
            asset_urls=True,
            data=MuseImageRequest(
                model=model["model"],
                prompt=model["prompt"],
                size=_size(model["aspect_ratio"]),
                reasoning_strength=model["reasoning_strength"],
                tool_enablement=_tool_enablement(model),
            ),
        )
        return IO.NodeOutput(await _decode_images(cls, response))


class MetaMuseImageEditApi(IO.ComfyNode):
    @classmethod
    def define_schema(cls):
        return IO.Schema(
            node_id="MetaMuseImageEditApi",
            display_name="Meta Muse Image Edit",
            category="partner/image/Meta",
            description=f"Edits or combines up to {MAX_INPUT_IMAGES} reference images guided by a text prompt "
            "using Meta's Muse Image model.",
            inputs=[
                IO.DynamicCombo.Input(
                    "model",
                    options=[_edit_model_option(model_id) for model_id in MUSE_IMAGE_MODELS],
                    tooltip="Model to use.",
                ),
            ],
            outputs=[
                IO.Image.Output(),
            ],
            hidden=[
                IO.Hidden.auth_token_comfy_org,
                IO.Hidden.api_key_comfy_org,
                IO.Hidden.unique_id,
            ],
            is_api_node=True,
            price_badge=_price_badge(),
        )

    @classmethod
    async def execute(cls, model: dict):
        validate_string(model["prompt"], min_length=1)
        reference_images = [image for key in model["images"] for image in model["images"][key]]
        if len(reference_images) > MAX_INPUT_IMAGES:
            raise ValueError(
                f"A maximum of {MAX_INPUT_IMAGES} reference images is supported; got {len(reference_images)} "
                f"(a batched input counts once per image)."
            )
        prompt = _resolve_image_refs(model["prompt"], len(reference_images))
        urls = await upload_images_to_comfyapi(
            cls,
            [image[..., :3] for image in reference_images],
            max_images=MAX_INPUT_IMAGES,
            mime_type="image/png",
            wait_label="Uploading reference images",
        )
        response = await sync_op(
            cls,
            ApiEndpoint(path=EDITS_PATH, method="POST"),
            response_model=MuseImageResponse,
            asset_urls=True,
            data=MuseImageEditRequest(
                model=model["model"],
                prompt=prompt,
                size=_size(model["aspect_ratio"]),
                reasoning_strength=model["reasoning_strength"],
                tool_enablement=_tool_enablement(model),
                images=[MuseImageInput(image_url=url) for url in urls],
            ),
        )
        return IO.NodeOutput(await _decode_images(cls, response))


class MetaApiExtension(ComfyExtension):
    @override
    async def get_node_list(self) -> list[type[IO.ComfyNode]]:
        return [
            MetaMuseImageTextToImageApi,
            MetaMuseImageEditApi,
        ]


async def comfy_entrypoint() -> MetaApiExtension:
    return MetaApiExtension()

Xet Storage Details

Size:: 10.9 kB
Xet hash:: 75cb0c27ab34db481523432ec3ad3d926de062a0e70608c3f4a113ca00d500a1

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	import base64
	import re
	from io import BytesIO

	import torch
	from typing_extensions import override

	from comfy_api.latest import IO, ComfyExtension
	from comfy_api_nodes.apis.meta import (
	MuseImageEditRequest,
	MuseImageInput,
	MuseImageRequest,
	MuseImageResponse,
	MuseImageToolEnablement,
	)
	from comfy_api_nodes.util import (
	ApiEndpoint,
	bytesio_to_image_tensor,
	download_url_to_image_tensor,
	pad_images_to_common_channels,
	sync_op,
	upload_images_to_comfyapi,
	validate_string,
	)

	GENERATIONS_PATH = "/proxy/meta/v1/images/generations"
	EDITS_PATH = "/proxy/meta/v1/images/edits"
	MUSE_IMAGE_MODELS = ["muse-image-1.0"]
	MAX_INPUT_IMAGES = 10
	ASPECT_RATIOS = ["auto", "1:1", "3:2", "2:3", "4:3", "3:4", "5:4", "4:5", "16:9", "9:16", "21:9", "9:21", "2:1", "1:2"]
	REASONING_STRENGTHS = ["high", "low"]

	_IMAGE_REF_RE = re.compile(r"@image(?P<idx>\d*)(?!\w)", re.IGNORECASE \| re.ASCII)


	def _resolve_image_refs(prompt: str, total_images: int) -> str:
	parts = []
	pos = 0
	prev_end = -1
	for match in _IMAGE_REF_RE.finditer(prompt):
	start = match.start()
	if start > 0 and start != prev_end and (prompt[start - 1].isalnum() or prompt[start - 1] == "_"):
	continue
	idx = int(match.group("idx") or 1)
	if not 1 <= idx <= total_images:
	raise ValueError(
	f"The prompt references @Image{idx}, but only {total_images} reference images "
	f"are connected (a batched input counts once per image)."
	)
	parts.append(prompt[pos:start])
	parts.append(f"image {idx}")
	pos = match.end()
	prev_end = match.end()
	parts.append(prompt[pos:])
	return "".join(parts)


	def _size(aspect_ratio: str) -> str \| None:
	return None if aspect_ratio == "auto" else aspect_ratio.replace(":", "x")


	async def _decode_images(cls: type[IO.ComfyNode], response: MuseImageResponse) -> torch.Tensor:
	images = []
	for item in response.data:
	if item.b64_json:
	images.append(bytesio_to_image_tensor(BytesIO(base64.b64decode(item.b64_json))))
	elif item.url:
	images.append(await download_url_to_image_tensor(item.url, cls=cls))
	if not images:
	raise Exception("The response contains no images.")
	return torch.cat(pad_images_to_common_channels(images))


	def _reasoning_strength_input() -> IO.Combo.Input:
	return IO.Combo.Input(
	"reasoning_strength",
	options=REASONING_STRENGTHS,
	tooltip="How much the model thinks, plans and self-refines before rendering.",
	)


	def _t2i_model_option(model_id: str) -> IO.DynamicCombo.Option:
	return IO.DynamicCombo.Option(
	model_id,
	[
	IO.String.Input(
	"prompt",
	multiline=True,
	default="",
	tooltip="Prompt describing the image. The model reasons about the prompt, and may use "
	"its built-in web and image search, before rendering.",
	),
	IO.Combo.Input(
	"aspect_ratio",
	options=ASPECT_RATIOS,
	tooltip="Aspect ratio of the output. Images are rendered at about 2.5 megapixels "
	"(1:1 is 1600x1600, 16:9 is 2048x1152); 'auto' lets the model choose from the prompt.",
	),
	_reasoning_strength_input(),
	*_tool_toggle_inputs(),
	_seed_input(),
	],
	)


	def _edit_model_option(model_id: str) -> IO.DynamicCombo.Option:
	return IO.DynamicCombo.Option(
	model_id,
	[
	IO.Autogrow.Input(
	"images",
	template=IO.Autogrow.TemplateNames(
	IO.Image.Input("image"),
	names=[f"image_{i}" for i in range(1, MAX_INPUT_IMAGES + 1)],
	min=1,
	),
	tooltip=f"1-{MAX_INPUT_IMAGES} reference images to edit or combine. Refer to them in the prompt "
	"as @Image1, @Image2, ..., numbered in input order; a batched input counts once per image.",
	),
	IO.String.Input(
	"prompt",
	multiline=True,
	default="",
	tooltip="Editing instructions. Supports @Image1-style references to the input images.",
	),
	IO.Combo.Input(
	"aspect_ratio",
	options=ASPECT_RATIOS,
	tooltip="Aspect ratio of the output. Images are rendered at about 2.5 megapixels "
	"(1:1 is 1600x1600, 16:9 is 2048x1152); 'auto' keeps the aspect ratio of the input.",
	),
	_reasoning_strength_input(),
	*_tool_toggle_inputs(),
	_seed_input(),
	],
	)


	def _tool_toggle_inputs() -> list[IO.Boolean.Input]:
	return [
	IO.Boolean.Input(
	"enable_web_search",
	default=True,
	advanced=True,
	tooltip="Lets the model search the web for facts and live information while planning the image.",
	),
	IO.Boolean.Input(
	"enable_image_search",
	default=True,
	advanced=True,
	tooltip="Lets the model search for reference images while planning the image.",
	),
	IO.Boolean.Input(
	"enable_shell",
	default=True,
	advanced=True,
	tooltip="Lets the model run code while planning, for precise layouts, charts and diagrams; "
	"when off, quantities and alignment are approximated.",
	),
	]


	def _tool_enablement(model: dict) -> MuseImageToolEnablement \| None:
	if model["enable_web_search"] and model["enable_image_search"] and model["enable_shell"]:
	return None
	return MuseImageToolEnablement(
	enable_image_search=model["enable_image_search"],
	enable_web_search=model["enable_web_search"],
	enable_shell=model["enable_shell"],
	)


	def _seed_input() -> IO.Int.Input:
	return IO.Int.Input(
	"seed",
	default=42,
	min=0,
	max=2147483647,
	step=1,
	display_mode=IO.NumberDisplay.number,
	control_after_generate=True,
	tooltip="Seed to determine if node should re-run; the API has no seed, "
	"so actual results are nondeterministic regardless of this value.",
	)


	def _price_badge() -> IO.PriceBadge:
	return IO.PriceBadge(expr="""{"type":"usd","usd":0.0143}""")


	class MetaMuseImageTextToImageApi(IO.ComfyNode):
	@classmethod
	def define_schema(cls):
	return IO.Schema(
	node_id="MetaMuseImageTextToImageApi",
	display_name="Meta Muse Image Text to Image",
	category="partner/image/Meta",
	description="Generates images from a text prompt using Meta's Muse Image model, "
	"which reasons about the prompt before rendering.",
	inputs=[
	IO.DynamicCombo.Input(
	"model",
	options=[_t2i_model_option(model_id) for model_id in MUSE_IMAGE_MODELS],
	tooltip="Model to use.",
	),
	],
	outputs=[
	IO.Image.Output(),
	],
	hidden=[
	IO.Hidden.auth_token_comfy_org,
	IO.Hidden.api_key_comfy_org,
	IO.Hidden.unique_id,
	],
	is_api_node=True,
	price_badge=_price_badge(),
	)

	@classmethod
	async def execute(cls, model: dict):
	validate_string(model["prompt"], min_length=1)
	response = await sync_op(
	cls,
	ApiEndpoint(path=GENERATIONS_PATH, method="POST"),
	response_model=MuseImageResponse,
	asset_urls=True,
	data=MuseImageRequest(
	model=model["model"],
	prompt=model["prompt"],
	size=_size(model["aspect_ratio"]),
	reasoning_strength=model["reasoning_strength"],
	tool_enablement=_tool_enablement(model),
	),
	)
	return IO.NodeOutput(await _decode_images(cls, response))


	class MetaMuseImageEditApi(IO.ComfyNode):
	@classmethod
	def define_schema(cls):
	return IO.Schema(
	node_id="MetaMuseImageEditApi",
	display_name="Meta Muse Image Edit",
	category="partner/image/Meta",
	description=f"Edits or combines up to {MAX_INPUT_IMAGES} reference images guided by a text prompt "
	"using Meta's Muse Image model.",
	inputs=[
	IO.DynamicCombo.Input(
	"model",
	options=[_edit_model_option(model_id) for model_id in MUSE_IMAGE_MODELS],
	tooltip="Model to use.",
	),
	],
	outputs=[
	IO.Image.Output(),
	],
	hidden=[
	IO.Hidden.auth_token_comfy_org,
	IO.Hidden.api_key_comfy_org,
	IO.Hidden.unique_id,
	],
	is_api_node=True,
	price_badge=_price_badge(),
	)

	@classmethod
	async def execute(cls, model: dict):
	validate_string(model["prompt"], min_length=1)
	reference_images = [image for key in model["images"] for image in model["images"][key]]
	if len(reference_images) > MAX_INPUT_IMAGES:
	raise ValueError(
	f"A maximum of {MAX_INPUT_IMAGES} reference images is supported; got {len(reference_images)} "
	f"(a batched input counts once per image)."
	)
	prompt = _resolve_image_refs(model["prompt"], len(reference_images))
	urls = await upload_images_to_comfyapi(
	cls,
	[image[..., :3] for image in reference_images],
	max_images=MAX_INPUT_IMAGES,
	mime_type="image/png",
	wait_label="Uploading reference images",
	)
	response = await sync_op(
	cls,
	ApiEndpoint(path=EDITS_PATH, method="POST"),
	response_model=MuseImageResponse,
	asset_urls=True,
	data=MuseImageEditRequest(
	model=model["model"],
	prompt=prompt,
	size=_size(model["aspect_ratio"]),
	reasoning_strength=model["reasoning_strength"],
	tool_enablement=_tool_enablement(model),
	images=[MuseImageInput(image_url=url) for url in urls],
	),
	)
	return IO.NodeOutput(await _decode_images(cls, response))


	class MetaApiExtension(ComfyExtension):
	@override
	async def get_node_list(self) -> list[type[IO.ComfyNode]]:
	return [
	MetaMuseImageTextToImageApi,
	MetaMuseImageEditApi,
	]


	async def comfy_entrypoint() -> MetaApiExtension:
	return MetaApiExtension()