|
1 | 1 | import node_helpers
|
| 2 | +from typing_extensions import override |
| 3 | +from comfy_api.latest import ComfyExtension, io |
2 | 4 |
|
3 | 5 |
|
4 |
| -class ReferenceLatent: |
| 6 | +class ReferenceLatent(io.ComfyNode): |
5 | 7 | @classmethod
|
6 |
| - def INPUT_TYPES(s): |
7 |
| - return {"required": {"conditioning": ("CONDITIONING", ), |
8 |
| - }, |
9 |
| - "optional": {"latent": ("LATENT", ),} |
10 |
| - } |
| 8 | + def define_schema(cls): |
| 9 | + return io.Schema( |
| 10 | + node_id="ReferenceLatent", |
| 11 | + category="advanced/conditioning/edit_models", |
| 12 | + description="This node sets the guiding latent for an edit model. If the model supports it you can chain multiple to set multiple reference images.", |
| 13 | + inputs=[ |
| 14 | + io.Conditioning.Input("conditioning"), |
| 15 | + io.Latent.Input("latent", optional=True), |
| 16 | + ], |
| 17 | + outputs=[ |
| 18 | + io.Conditioning.Output(), |
| 19 | + ] |
| 20 | + ) |
11 | 21 |
|
12 |
| - RETURN_TYPES = ("CONDITIONING",) |
13 |
| - FUNCTION = "append" |
14 |
| - |
15 |
| - CATEGORY = "advanced/conditioning/edit_models" |
16 |
| - DESCRIPTION = "This node sets the guiding latent for an edit model. If the model supports it you can chain multiple to set multiple reference images." |
17 |
| - |
18 |
| - def append(self, conditioning, latent=None): |
| 22 | + @classmethod |
| 23 | + def execute(cls, conditioning, latent=None) -> io.NodeOutput: |
19 | 24 | if latent is not None:
|
20 | 25 | conditioning = node_helpers.conditioning_set_values(conditioning, {"reference_latents": [latent["samples"]]}, append=True)
|
21 |
| - return (conditioning, ) |
| 26 | + return io.NodeOutput(conditioning) |
| 27 | + |
| 28 | + |
| 29 | +class EditModelExtension(ComfyExtension): |
| 30 | + @override |
| 31 | + async def get_node_list(self) -> list[type[io.ComfyNode]]: |
| 32 | + return [ |
| 33 | + ReferenceLatent, |
| 34 | + ] |
22 | 35 |
|
23 | 36 |
|
24 |
| -NODE_CLASS_MAPPINGS = { |
25 |
| - "ReferenceLatent": ReferenceLatent, |
26 |
| -} |
| 37 | +def comfy_entrypoint() -> EditModelExtension: |
| 38 | + return EditModelExtension() |
0 commit comments