From 955fa8a7fa666a678e48582be584543b5f2f388d Mon Sep 17 00:00:00 2001 From: Limitex <76650151+Limitex@users.noreply.github.com> Date: Mon, 3 Nov 2025 03:51:43 +0900 Subject: [PATCH] feat: clip text encode --- __init__.py | 2 + src/di/di.py | 18 +++++- src/domain/model/__init__.py | 2 + src/domain/model/_conditioning.py | 9 +++ src/domain/repositories/__init__.py | 3 +- .../repositories/_text_encoder_repository.py | 14 +++++ src/infra/diffusers/__init__.py | 7 ++- .../diffusers/_text_encoder_repository.py | 46 +++++++++++++++ src/nodes/__init__.py | 6 +- src/nodes/_diffuser_clip_text_encode.py | 34 +++++++++++ src/nodes/dto/__init__.py | 8 ++- src/nodes/dto/_conditioning_dto.py | 57 +++++++++++++++++++ src/service/__init__.py | 3 +- src/service/_clip_text_encode_service.py | 19 +++++++ src/ui/__init__.py | 3 +- src/ui/_clip_text_encode_handler.py | 27 +++++++++ 16 files changed, 250 insertions(+), 8 deletions(-) create mode 100644 src/domain/model/_conditioning.py create mode 100644 src/domain/repositories/_text_encoder_repository.py create mode 100644 src/infra/diffusers/_text_encoder_repository.py create mode 100644 src/nodes/_diffuser_clip_text_encode.py create mode 100644 src/nodes/dto/_conditioning_dto.py create mode 100644 src/service/_clip_text_encode_service.py create mode 100644 src/ui/_clip_text_encode_handler.py diff --git a/__init__.py b/__init__.py index a89779d..a149547 100644 --- a/__init__.py +++ b/__init__.py @@ -3,6 +3,7 @@ from .src.nodes import ( NODE_DISPLAY_NAME_MAPPINGS, ComfyUIAutoencoderDTO, ComfyUIClipDTO, + ComfyUIConditioningDTO, ComfyUIPipelineDTO, ) @@ -12,4 +13,5 @@ __all__ = [ "ComfyUIPipelineDTO", "ComfyUIAutoencoderDTO", "ComfyUIClipDTO", + "ComfyUIConditioningDTO", ] diff --git a/src/di/di.py b/src/di/di.py index 915143a..c63b5a6 100644 --- a/src/di/di.py +++ b/src/di/di.py @@ -2,9 +2,14 @@ from dependency_injector import containers, providers -from ..infra.diffusers import DiffusersAutoencoderRepository, DiffusersPipelineRepository +from ..infra.diffusers import ( + DiffusersAutoencoderRepository, + DiffusersPipelineRepository, +) +from ..infra.diffusers._text_encoder_repository import DiffusersTextEncoderRepository from ..service import AutoencoderService, PipelineService -from ..ui import AutoencoderHandler, PipelineHandler +from ..service._clip_text_encode_service import ClipTextEncodeService +from ..ui import AutoencoderHandler, ClipTextEncodeHandler, PipelineHandler class Container(containers.DeclarativeContainer): @@ -13,6 +18,7 @@ class Container(containers.DeclarativeContainer): # 1. Repositories pipeline_repository = providers.Factory(DiffusersPipelineRepository) autoencoder_repository = providers.Factory(DiffusersAutoencoderRepository) + text_encoder_repository = providers.Factory(DiffusersTextEncoderRepository) # 2. Services pipeline_service = providers.Factory( @@ -23,6 +29,10 @@ class Container(containers.DeclarativeContainer): AutoencoderService, autoencoder_repo=autoencoder_repository, ) + clip_text_encode_service = providers.Factory( + ClipTextEncodeService, + text_encoder_repo=text_encoder_repository, + ) # 3. Handler pipeline_handler = providers.Factory( @@ -33,3 +43,7 @@ class Container(containers.DeclarativeContainer): AutoencoderHandler, autoencoder_service=autoencoder_service, ) + clip_text_encode_handler = providers.Factory( + ClipTextEncodeHandler, + clip_text_encode_service=clip_text_encode_service, + ) diff --git a/src/domain/model/__init__.py b/src/domain/model/__init__.py index 8cf9670..4f60a9d 100644 --- a/src/domain/model/__init__.py +++ b/src/domain/model/__init__.py @@ -1,9 +1,11 @@ from ._autoencoder import Autoencoder from ._clip import Clip +from ._conditioning import Conditioning from ._pipeline import Pipeline __all__ = [ "Pipeline", "Autoencoder", "Clip", + "Conditioning", ] diff --git a/src/domain/model/_conditioning.py b/src/domain/model/_conditioning.py new file mode 100644 index 0000000..f827afe --- /dev/null +++ b/src/domain/model/_conditioning.py @@ -0,0 +1,9 @@ +from dataclasses import dataclass + +import torch + + +@dataclass +class Conditioning: + conditioning: torch.Tensor + path: str diff --git a/src/domain/repositories/__init__.py b/src/domain/repositories/__init__.py index 8e067bd..3ab409d 100644 --- a/src/domain/repositories/__init__.py +++ b/src/domain/repositories/__init__.py @@ -1,4 +1,5 @@ from ._autoencoder_repository import AutoencoderRepository from ._pipeline_repository import PipelineRepository +from ._text_encoder_repository import TextEncoderRepository -__all__ = ["PipelineRepository", "AutoencoderRepository"] +__all__ = ["PipelineRepository", "AutoencoderRepository", "TextEncoderRepository"] diff --git a/src/domain/repositories/_text_encoder_repository.py b/src/domain/repositories/_text_encoder_repository.py new file mode 100644 index 0000000..278557e --- /dev/null +++ b/src/domain/repositories/_text_encoder_repository.py @@ -0,0 +1,14 @@ +from __future__ import annotations + +from abc import ABC, abstractmethod + +import torch +from transformers import CLIPTextModel, CLIPTokenizer + + +class TextEncoderRepository(ABC): + @abstractmethod + def encode( + self, tokenizer: CLIPTokenizer, text_encoder: CLIPTextModel, text: str + ) -> torch.Tensor: + pass diff --git a/src/infra/diffusers/__init__.py b/src/infra/diffusers/__init__.py index 71049f3..766ce98 100644 --- a/src/infra/diffusers/__init__.py +++ b/src/infra/diffusers/__init__.py @@ -1,4 +1,9 @@ from ._autoencoder_repository import DiffusersAutoencoderRepository from ._pipeline_repository import DiffusersPipelineRepository +from ._text_encoder_repository import DiffusersTextEncoderRepository -__all__ = ["DiffusersPipelineRepository", "DiffusersAutoencoderRepository"] +__all__ = [ + "DiffusersPipelineRepository", + "DiffusersAutoencoderRepository", + "DiffusersTextEncoderRepository", +] diff --git a/src/infra/diffusers/_text_encoder_repository.py b/src/infra/diffusers/_text_encoder_repository.py new file mode 100644 index 0000000..0457d6d --- /dev/null +++ b/src/infra/diffusers/_text_encoder_repository.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +from typing import Any + +import torch +from comfy.model_management import get_torch_device # pyright: ignore[reportMissingImports] +from transformers import CLIPTextModel, CLIPTokenizer + +from ...domain.repositories._text_encoder_repository import TextEncoderRepository + + +class DiffusersTextEncoderRepository(TextEncoderRepository): + def __init__(self) -> None: + self.device = get_torch_device() + + def encode( + self, tokenizer: CLIPTokenizer, text_encoder: CLIPTextModel, text: str + ) -> torch.Tensor: + max_length = tokenizer.model_max_length + input_ids = tokenizer(text, return_tensors="pt").input_ids + text_length = input_ids.shape[-1] + + if max_length < text_length: + text_ids = input_ids.to(self.device) + else: + text_ids = tokenizer( + text, + truncation=False, + padding="max_length", + max_length=max_length, + return_tensors="pt", + ).input_ids.to(self.device) + + concat_embeds: list[torch.Tensor] = [] + with torch.no_grad(): + for i in range(0, text_ids.shape[-1], max_length): + segment_ids = text_ids[:, i : i + max_length] + outputs: Any = text_encoder(segment_ids) + if hasattr(outputs, "last_hidden_state"): + embeds = outputs.last_hidden_state + else: + embeds = outputs[0] + concat_embeds.append(embeds) + + text_embeds: torch.Tensor = torch.cat(concat_embeds, dim=1) + return text_embeds diff --git a/src/nodes/__init__.py b/src/nodes/__init__.py index 6803e59..866ba12 100644 --- a/src/nodes/__init__.py +++ b/src/nodes/__init__.py @@ -1,7 +1,8 @@ from ..di import Container +from ._diffuser_clip_text_encode import DiffusersClipTextEncode from ._diffusers_pipeline_loader import DiffusersPipelineLoader from ._diffusers_vae_loader import DiffusersVaeLoader -from .dto import ComfyUIAutoencoderDTO, ComfyUIClipDTO, ComfyUIPipelineDTO +from .dto import ComfyUIAutoencoderDTO, ComfyUIClipDTO, ComfyUIConditioningDTO, ComfyUIPipelineDTO container = Container() container.wire(modules=[__name__]) @@ -9,10 +10,12 @@ container.wire(modules=[__name__]) NODE_CLASS_MAPPINGS = { DiffusersPipelineLoader.__name__: DiffusersPipelineLoader, DiffusersVaeLoader.__name__: DiffusersVaeLoader, + DiffusersClipTextEncode.__name__: DiffusersClipTextEncode, } NODE_DISPLAY_NAME_MAPPINGS = { DiffusersPipelineLoader.__name__: "Diffusers Pipeline Loader", DiffusersVaeLoader.__name__: "Diffusers VAE Loader", + DiffusersClipTextEncode.__name__: "Diffusers CLIP Text Encode", } __all__ = [ @@ -21,4 +24,5 @@ __all__ = [ "ComfyUIPipelineDTO", "ComfyUIAutoencoderDTO", "ComfyUIClipDTO", + "ComfyUIConditioningDTO", ] diff --git a/src/nodes/_diffuser_clip_text_encode.py b/src/nodes/_diffuser_clip_text_encode.py new file mode 100644 index 0000000..b21074d --- /dev/null +++ b/src/nodes/_diffuser_clip_text_encode.py @@ -0,0 +1,34 @@ +from dependency_injector.wiring import Provide, inject + +from ..di import Container +from ..ui import ClipTextEncodeHandler +from .dto import ComfyUIClipDTO, ComfyUIConditioningDTO + + +class DiffusersClipTextEncode: + def __init__(self) -> None: + pass + + @classmethod + def INPUT_TYPES(cls) -> dict[str, dict[str, tuple[str, ...] | tuple[str, dict[str, bool]]]]: + return { + "required": { + "clip": (ComfyUIClipDTO.COMFY_TYPE,), + "text": ("STRING", {"multiline": True}), + }, + } + + RETURN_TYPES = (ComfyUIConditioningDTO.COMFY_TYPE,) + FUNCTION = "execute" + CATEGORY = "Diffusers" + + @inject + def execute( + self, + clip: ComfyUIClipDTO, + text: str, + handler: ClipTextEncodeHandler = Provide[Container.clip_text_encode_handler], + ) -> tuple[ComfyUIConditioningDTO]: + conditioning_model = handler.encode(clip.tokenizer, clip.text_encoder, clip.path, text) + conditioning_dto = ComfyUIConditioningDTO.from_domain(conditioning_model) + return (conditioning_dto,) diff --git a/src/nodes/dto/__init__.py b/src/nodes/dto/__init__.py index 4853e89..00cca12 100644 --- a/src/nodes/dto/__init__.py +++ b/src/nodes/dto/__init__.py @@ -1,5 +1,11 @@ from ._autoencoder_dto import ComfyUIAutoencoderDTO from ._clip import ComfyUIClipDTO +from ._conditioning_dto import ComfyUIConditioningDTO from ._pipeline_dto import ComfyUIPipelineDTO -__all__ = ["ComfyUIPipelineDTO", "ComfyUIAutoencoderDTO", "ComfyUIClipDTO"] +__all__ = [ + "ComfyUIPipelineDTO", + "ComfyUIAutoencoderDTO", + "ComfyUIClipDTO", + "ComfyUIConditioningDTO", +] diff --git a/src/nodes/dto/_conditioning_dto.py b/src/nodes/dto/_conditioning_dto.py new file mode 100644 index 0000000..c9f113f --- /dev/null +++ b/src/nodes/dto/_conditioning_dto.py @@ -0,0 +1,57 @@ +from dataclasses import dataclass +from typing import ClassVar + +import torch + +from ...domain.model import Conditioning + + +@dataclass +class ComfyUIConditioningDTO: + """Data Transfer Object for Diffusers Conditioning in ComfyUI. + + This DTO is designed to be used across different custom nodes. + Other node developers can import and use this type for conditioning operations. + + Attributes: + conditioning: Tensor representing text embeddings/conditioning + path: Path to the source model directory + + Class Attributes: + COMFY_TYPE: Type name used in ComfyUI's RETURN_TYPES and INPUT_TYPES + + Example: + ```python + from your_node_package.nodes.dto import ComfyUIConditioningDTO + + class YourCustomNode: + @classmethod + def INPUT_TYPES(cls): + return {"required": {"conditioning": (ComfyUIConditioningDTO.COMFY_TYPE,)}} + + RETURN_TYPES = (ComfyUIConditioningDTO.COMFY_TYPE,) + + def execute(self, conditioning: ComfyUIConditioningDTO): + # Use the conditioning tensor for generation + result = process_with_conditioning(conditioning.conditioning) + return (result,) + ``` + """ + + # ClassVar to avoid being treated as a dataclass field + COMFY_TYPE: ClassVar[str] = "DIFFUSERS_CONDITIONING" + + conditioning: torch.Tensor + path: str + + @classmethod + def from_domain(cls, conditioning: Conditioning) -> "ComfyUIConditioningDTO": + """Create DTO from domain model. + + Args: + conditioning: Conditioning from domain layer + + Returns: + ComfyUIConditioningDTO instance + """ + return cls(conditioning=conditioning.conditioning, path=conditioning.path) diff --git a/src/service/__init__.py b/src/service/__init__.py index d35cd52..cd9ca0c 100644 --- a/src/service/__init__.py +++ b/src/service/__init__.py @@ -1,4 +1,5 @@ from ._autoencoder_service import AutoencoderService +from ._clip_text_encode_service import ClipTextEncodeService from ._pipeline_service import PipelineService -__all__ = ["PipelineService", "AutoencoderService"] +__all__ = ["PipelineService", "AutoencoderService", "ClipTextEncodeService"] diff --git a/src/service/_clip_text_encode_service.py b/src/service/_clip_text_encode_service.py new file mode 100644 index 0000000..71bd446 --- /dev/null +++ b/src/service/_clip_text_encode_service.py @@ -0,0 +1,19 @@ +from __future__ import annotations + +from transformers import CLIPTextModel, CLIPTokenizer + +from ..domain.model import Conditioning +from ..domain.repositories._text_encoder_repository import TextEncoderRepository + + +class ClipTextEncodeService: + def __init__(self, text_encoder_repo: TextEncoderRepository) -> None: + self.text_encoder_repo = text_encoder_repo + + def encode( + self, tokenizer: CLIPTokenizer, text_encoder: CLIPTextModel, path: str, text: str + ) -> Conditioning: + embeddings = self.text_encoder_repo.encode(tokenizer, text_encoder, text) + if embeddings is None: + raise RuntimeError("Failed to generate embeddings.") + return Conditioning(conditioning=embeddings, path=path) diff --git a/src/ui/__init__.py b/src/ui/__init__.py index 32cba9a..91f6061 100644 --- a/src/ui/__init__.py +++ b/src/ui/__init__.py @@ -1,4 +1,5 @@ from ._autoencoder_handler import AutoencoderHandler +from ._clip_text_encode_handler import ClipTextEncodeHandler from ._pipeline_handler import PipelineHandler -__all__ = ["PipelineHandler", "AutoencoderHandler"] +__all__ = ["PipelineHandler", "AutoencoderHandler", "ClipTextEncodeHandler"] diff --git a/src/ui/_clip_text_encode_handler.py b/src/ui/_clip_text_encode_handler.py new file mode 100644 index 0000000..bd83181 --- /dev/null +++ b/src/ui/_clip_text_encode_handler.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +from transformers import CLIPTextModel, CLIPTokenizer + +from ..domain.model import Conditioning +from ..service._clip_text_encode_service import ClipTextEncodeService + + +class ClipTextEncodeHandler: + def __init__(self, clip_text_encode_service: ClipTextEncodeService) -> None: + self.clip_text_encode_service = clip_text_encode_service + + def encode( + self, tokenizer: CLIPTokenizer, text_encoder: CLIPTextModel, path: str, text: str + ) -> Conditioning: + if not text: + raise ValueError("Text input is empty.") + if path is None: + raise ValueError("Clip path is None.") + if tokenizer is None or text_encoder is None: + raise ValueError("Tokenizer or Text Encoder is None.") + if not isinstance(text, str): + raise TypeError("Text input must be a string.") + conditioning = self.clip_text_encode_service.encode(tokenizer, text_encoder, path, text) + if conditioning is None: + raise RuntimeError("Failed to encode text into conditioning.") + return conditioning