mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-17 17:18:10 +00:00
103 lines
3.2 KiB
Python
103 lines
3.2 KiB
Python
from typing_extensions import override
|
|
|
|
import comfy.utils
|
|
import node_helpers
|
|
from comfy_api.latest import ComfyExtension, io
|
|
|
|
|
|
# fmt: off
|
|
BUCKETS_1024 = [
|
|
(512, 1792), (512, 1856), (512, 1920), (512, 1984), (512, 2048),
|
|
(576, 1600), (576, 1664), (576, 1728), (576, 1792),
|
|
(640, 1472), (640, 1536), (640, 1600),
|
|
(704, 1344), (704, 1408), (704, 1472),
|
|
(768, 1216), (768, 1280), (768, 1344),
|
|
(832, 1152), (832, 1216),
|
|
(896, 1088), (896, 1152),
|
|
(960, 1024), (960, 1088),
|
|
(1024, 960), (1024, 1024),
|
|
(1088, 896), (1088, 960),
|
|
(1152, 832), (1152, 896),
|
|
(1216, 768), (1216, 832),
|
|
(1280, 768),
|
|
(1344, 704), (1344, 768),
|
|
(1408, 704),
|
|
(1472, 640), (1472, 704),
|
|
(1536, 640),
|
|
(1600, 576), (1600, 640),
|
|
(1664, 576),
|
|
(1728, 576),
|
|
(1792, 512), (1792, 576),
|
|
(1856, 512),
|
|
(1920, 512),
|
|
(1984, 512),
|
|
(2048, 512),
|
|
]
|
|
# fmt: on
|
|
|
|
|
|
def _find_best_bucket(height: int, width: int) -> tuple[int, int]:
|
|
target_ratio = height / width
|
|
return min(BUCKETS_1024, key=lambda hw: abs(hw[0] / hw[1] - target_ratio))
|
|
|
|
|
|
def _resize_reference(image):
|
|
if image.shape[0] != 1:
|
|
raise ValueError("JoyImage reference inputs must contain one image each")
|
|
samples = image.movedim(-1, 1)
|
|
bucket_h, bucket_w = _find_best_bucket(samples.shape[2], samples.shape[3])
|
|
resized = comfy.utils.common_upscale(samples, bucket_w, bucket_h, "bilinear", "center")
|
|
return resized.movedim(1, -1)[:, :, :, :3]
|
|
|
|
|
|
def _encode(clip, prompt, vae, images):
|
|
resized_images = [_resize_reference(image) for image in images]
|
|
conditioning = clip.encode_from_tokens_scheduled(clip.tokenize(prompt, images=resized_images))
|
|
if vae is not None and resized_images:
|
|
ref_latents = [vae.encode(image) for image in resized_images]
|
|
conditioning = node_helpers.conditioning_set_values(
|
|
conditioning, {"reference_latents": ref_latents}, append=True,
|
|
)
|
|
return conditioning
|
|
|
|
|
|
class TextEncodeJoyImageEdit(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
image_template = io.Autogrow.TemplatePrefix(
|
|
io.Image.Input("image"),
|
|
prefix="image",
|
|
min=0,
|
|
max=6,
|
|
)
|
|
return io.Schema(
|
|
node_id="TextEncodeJoyImageEdit",
|
|
category="model/conditioning/joyimage",
|
|
inputs=[
|
|
io.Clip.Input("clip"),
|
|
io.String.Input("prompt", multiline=True, dynamic_prompts=True),
|
|
io.Vae.Input("vae", optional=True),
|
|
io.Autogrow.Input("images", template=image_template, optional=True),
|
|
],
|
|
outputs=[
|
|
io.Conditioning.Output(),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, clip, prompt, vae=None, images: io.Autogrow.Type = None) -> io.NodeOutput:
|
|
images = images or {}
|
|
return io.NodeOutput(_encode(clip, prompt, vae, list(images.values())))
|
|
|
|
|
|
class JoyImageExtension(ComfyExtension):
|
|
@override
|
|
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
|
return [
|
|
TextEncodeJoyImageEdit,
|
|
]
|
|
|
|
|
|
async def comfy_entrypoint() -> JoyImageExtension:
|
|
return JoyImageExtension()
|