Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
74 changes: 46 additions & 28 deletions comfy_extras/nodes_photomaker.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
import comfy.clip_model
import comfy.clip_vision
import comfy.ops
from typing_extensions import override
from comfy_api.latest import ComfyExtension, io

# code for model from: https://github.com/TencentARC/PhotoMaker/blob/main/photomaker/model.py under Apache License Version 2.0
VISION_CONFIG_DICT = {
Expand Down Expand Up @@ -116,41 +118,52 @@ def forward(self, id_pixel_values, prompt_embeds, class_tokens_mask):
return updated_prompt_embeds


class PhotoMakerLoader:
class PhotoMakerLoader(io.ComfyNode):
@classmethod
def INPUT_TYPES(s):
return {"required": { "photomaker_model_name": (folder_paths.get_filename_list("photomaker"), )}}

RETURN_TYPES = ("PHOTOMAKER",)
FUNCTION = "load_photomaker_model"

CATEGORY = "_for_testing/photomaker"
def define_schema(cls):
return io.Schema(
node_id="PhotoMakerLoader",
category="_for_testing/photomaker",
inputs=[
io.Combo.Input("photomaker_model_name", options=folder_paths.get_filename_list("photomaker")),
],
outputs=[
io.Photomaker.Output(),
],
is_experimental=True,
)

def load_photomaker_model(self, photomaker_model_name):
@classmethod
def execute(cls, photomaker_model_name):
photomaker_model_path = folder_paths.get_full_path_or_raise("photomaker", photomaker_model_name)
photomaker_model = PhotoMakerIDEncoder()
data = comfy.utils.load_torch_file(photomaker_model_path, safe_load=True)
if "id_encoder" in data:
data = data["id_encoder"]
photomaker_model.load_state_dict(data)
return (photomaker_model,)
return io.NodeOutput(photomaker_model)


class PhotoMakerEncode:
class PhotoMakerEncode(io.ComfyNode):
@classmethod
def INPUT_TYPES(s):
return {"required": { "photomaker": ("PHOTOMAKER",),
"image": ("IMAGE",),
"clip": ("CLIP", ),
"text": ("STRING", {"multiline": True, "dynamicPrompts": True, "default": "photograph of photomaker"}),
}}

RETURN_TYPES = ("CONDITIONING",)
FUNCTION = "apply_photomaker"

CATEGORY = "_for_testing/photomaker"
def define_schema(cls):
return io.Schema(
node_id="PhotoMakerEncode",
category="_for_testing/photomaker",
inputs=[
io.Photomaker.Input("photomaker"),
io.Image.Input("image"),
io.Clip.Input("clip"),
io.String.Input("text", multiline=True, dynamic_prompts=True, default="photograph of photomaker"),
],
outputs=[
io.Conditioning.Output(),
],
is_experimental=True,
)

def apply_photomaker(self, photomaker, image, clip, text):
@classmethod
def execute(cls, photomaker, image, clip, text):
special_token = "photomaker"
pixel_values = comfy.clip_vision.clip_preprocess(image.to(photomaker.load_device)).float()
try:
Expand Down Expand Up @@ -178,11 +191,16 @@ def apply_photomaker(self, photomaker, image, clip, text):
else:
out = cond

return ([[out, {"pooled_output": pooled}]], )
return io.NodeOutput([[out, {"pooled_output": pooled}]])


NODE_CLASS_MAPPINGS = {
"PhotoMakerLoader": PhotoMakerLoader,
"PhotoMakerEncode": PhotoMakerEncode,
}
class PhotomakerExtension(ComfyExtension):
@override
async def get_node_list(self) -> list[type[io.ComfyNode]]:
return [
PhotoMakerLoader,
PhotoMakerEncode,
]

async def comfy_entrypoint() -> PhotomakerExtension:
return PhotomakerExtension()
Loading