feat: clip text encode

__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",
 ]

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,
+    )

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",
 ]

src/domain/model/_conditioning.py

@@ -0,0 +1,9 @@
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass
+class Conditioning:
+    conditioning: torch.Tensor
+    path: str

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"]

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

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",
+]

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

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",
 ]

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,)

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",
+]

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)

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"]

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)

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"]

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