feat: create vae loader

__init__.py

@@ -1,6 +1,7 @@
 from .src.nodes import (
     NODE_CLASS_MAPPINGS,
     NODE_DISPLAY_NAME_MAPPINGS,
+    ComfyUIAutoencoderDTO,
     ComfyUIPipelineDTO,
 )
 
@@ -8,4 +9,5 @@ __all__ = [
     "NODE_CLASS_MAPPINGS",
     "NODE_DISPLAY_NAME_MAPPINGS",
     "ComfyUIPipelineDTO",
+    "ComfyUIAutoencoderDTO",
 ]

src/di/di.py

@@ -1,22 +1,34 @@
 from dependency_injector import containers, providers
 
+from ..infra.diffusers.autoencoder_repository import DiffusersAutoencoderRepository
 from ..infra.diffusers.pipeline_repository import DiffusersPipelineRepository
+from ..service.create_autoencoder_service import CreateAutoencoderService
 from ..service.create_pipeline_service import CreatePipelineService
+from ..ui.autoencoder_handler import AutoencoderHandler
 from ..ui.pipeline_handler import PipelineHandler
 
 
 class Container(containers.DeclarativeContainer):
     # 1. Repositories
     pipeline_repository = providers.Factory(DiffusersPipelineRepository)
+    autoencoder_repository = providers.Factory(DiffusersAutoencoderRepository)
 
     # 2. Services
     create_pipeline_service = providers.Factory(
         CreatePipelineService,
         pipeline_repo=pipeline_repository,
     )
+    create_autoencoder_service = providers.Factory(
+        CreateAutoencoderService,
+        autoencoder_repo=autoencoder_repository,
+    )
 
     # 3. Handler
     pipeline_handler = providers.Factory(
         PipelineHandler,
         create_pipeline_service=create_pipeline_service,
     )
+    autoencoder_handler = providers.Factory(
+        AutoencoderHandler,
+        create_autoencoder_service=create_autoencoder_service,
+    )

src/domain/model/__init__.py

@@ -1,3 +1,7 @@
+from .autoencoder import AutoencoderModel
 from .pipeline import PipelineModel
 
-__all__ = ["PipelineModel"]
+__all__ = [
+    "PipelineModel",
+    "AutoencoderModel",
+]

src/domain/model/autoencoder.py

@@ -0,0 +1,9 @@
+from dataclasses import dataclass
+
+from diffusers import AutoencoderKL
+
+
+@dataclass
+class AutoencoderModel:
+    autoencoder: AutoencoderKL
+    path: str

src/domain/repositories/__init__.py

@@ -1,3 +1,4 @@
+from .i_autoencoder_repository import IAutoencoderRepository
 from .i_pipeline_repository import IPipelineRepository
 
-__all__ = ["IPipelineRepository"]
+__all__ = ["IPipelineRepository", "IAutoencoderRepository"]

src/domain/repositories/i_autoencoder_repository.py

@@ -0,0 +1,14 @@
+from abc import ABC, abstractmethod
+
+import torch
+from diffusers import AutoencoderKL
+
+
+class IAutoencoderRepository(ABC):
+    @abstractmethod
+    def convert_and_save_from_single_file(self, checkpoint_path: str) -> str:
+        pass
+
+    @abstractmethod
+    def load_autoencoder_from_path(self, model_path: str, dtype: torch.dtype) -> AutoencoderKL:
+        pass

src/infra/diffusers/__init__.py

@@ -1,3 +1,4 @@
+from .autoencoder_repository import DiffusersAutoencoderRepository
 from .pipeline_repository import DiffusersPipelineRepository
 
-__all__ = ["DiffusersPipelineRepository"]
+__all__ = ["DiffusersPipelineRepositoryDiffusersAutoencoderRepository"]

src/infra/diffusers/autoencoder_repository.py

@@ -0,0 +1,225 @@
+import io
+import os
+from typing import Any
+
+import folder_paths  # pyright: ignore[reportMissingImports]
+import requests
+import torch
+from diffusers import AutoencoderKL
+from diffusers.pipelines.stable_diffusion.convert_from_ckpt import (
+    assign_to_checkpoint,
+    conv_attn_to_linear,
+    create_vae_diffusers_config,
+    renew_vae_attention_paths,
+    renew_vae_resnet_paths,
+)
+from omegaconf import OmegaConf
+
+from ...domain.repositories import IAutoencoderRepository
+
+
+class DiffusersAutoencoderRepository(IAutoencoderRepository):
+    def __init__(self) -> None:
+        self.cache_dir = folder_paths.get_temp_directory()
+
+    # Reference from : https://github.com/huggingface/diffusers/blob/main/scripts/convert_vae_pt_to_diffusers.py
+    def convert_and_save_from_single_file(self, checkpoint_path: str) -> str:
+        checkpoint_name = os.path.basename(checkpoint_path)
+        ckpt_cache_path = os.path.join(self.cache_dir, checkpoint_name)
+
+        self._vae_pt_to_vae_diffuser(checkpoint_path, ckpt_cache_path)
+        return ckpt_cache_path
+
+    def load_autoencoder_from_path(self, model_path: str, dtype: torch.dtype) -> AutoencoderKL:
+        vae: AutoencoderKL = AutoencoderKL.from_pretrained(  # type: ignore[no-untyped-call]
+            pretrained_model_name_or_path=model_path,
+            torch_dtype=dtype,
+            cache_dir=self.cache_dir,
+        )
+        return vae
+
+    # Reference from : https://github.com/huggingface/diffusers/blob/main/scripts/convert_vae_pt_to_diffusers.py
+    def _custom_convert_ldm_vae_checkpoint(
+        self, checkpoint: dict[str, Any], config: dict[str, Any]
+    ) -> dict[str, Any]:
+        vae_state_dict = checkpoint
+
+        new_checkpoint: dict[str, Any] = {}
+
+        new_checkpoint["encoder.conv_in.weight"] = vae_state_dict["encoder.conv_in.weight"]
+        new_checkpoint["encoder.conv_in.bias"] = vae_state_dict["encoder.conv_in.bias"]
+        new_checkpoint["encoder.conv_out.weight"] = vae_state_dict["encoder.conv_out.weight"]
+        new_checkpoint["encoder.conv_out.bias"] = vae_state_dict["encoder.conv_out.bias"]
+        new_checkpoint["encoder.conv_norm_out.weight"] = vae_state_dict["encoder.norm_out.weight"]
+        new_checkpoint["encoder.conv_norm_out.bias"] = vae_state_dict["encoder.norm_out.bias"]
+
+        new_checkpoint["decoder.conv_in.weight"] = vae_state_dict["decoder.conv_in.weight"]
+        new_checkpoint["decoder.conv_in.bias"] = vae_state_dict["decoder.conv_in.bias"]
+        new_checkpoint["decoder.conv_out.weight"] = vae_state_dict["decoder.conv_out.weight"]
+        new_checkpoint["decoder.conv_out.bias"] = vae_state_dict["decoder.conv_out.bias"]
+        new_checkpoint["decoder.conv_norm_out.weight"] = vae_state_dict["decoder.norm_out.weight"]
+        new_checkpoint["decoder.conv_norm_out.bias"] = vae_state_dict["decoder.norm_out.bias"]
+
+        new_checkpoint["quant_conv.weight"] = vae_state_dict["quant_conv.weight"]
+        new_checkpoint["quant_conv.bias"] = vae_state_dict["quant_conv.bias"]
+        new_checkpoint["post_quant_conv.weight"] = vae_state_dict["post_quant_conv.weight"]
+        new_checkpoint["post_quant_conv.bias"] = vae_state_dict["post_quant_conv.bias"]
+
+        # Retrieves the keys for the encoder down blocks only
+        num_down_blocks = len(
+            {".".join(layer.split(".")[:3]) for layer in vae_state_dict if "encoder.down" in layer}
+        )
+        down_blocks = {
+            layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key]
+            for layer_id in range(num_down_blocks)
+        }
+
+        # Retrieves the keys for the decoder up blocks only
+        num_up_blocks = len(
+            {".".join(layer.split(".")[:3]) for layer in vae_state_dict if "decoder.up" in layer}
+        )
+        up_blocks = {
+            layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key]
+            for layer_id in range(num_up_blocks)
+        }
+
+        for i in range(num_down_blocks):
+            resnets = [
+                key
+                for key in down_blocks[i]
+                if f"down.{i}" in key and f"down.{i}.downsample" not in key
+            ]
+
+            if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict:
+                new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = (
+                    vae_state_dict.pop(f"encoder.down.{i}.downsample.conv.weight")
+                )
+                new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = (
+                    vae_state_dict.pop(f"encoder.down.{i}.downsample.conv.bias")
+                )
+
+            paths = renew_vae_resnet_paths(resnets)  # type: ignore[no-untyped-call]
+            meta_path = {"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"}
+            assign_to_checkpoint(  # type: ignore[no-untyped-call]
+                paths,
+                new_checkpoint,
+                vae_state_dict,
+                additional_replacements=[meta_path],
+                config=config,
+            )
+
+        mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key]
+        num_mid_res_blocks = 2
+        for i in range(1, num_mid_res_blocks + 1):
+            resnets = [key for key in mid_resnets if f"encoder.mid.block_{i}" in key]
+
+            paths = renew_vae_resnet_paths(resnets)  # type: ignore[no-untyped-call]
+            meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}
+            assign_to_checkpoint(  # type: ignore[no-untyped-call]
+                paths,
+                new_checkpoint,
+                vae_state_dict,
+                additional_replacements=[meta_path],
+                config=config,
+            )
+
+        mid_attentions = [key for key in vae_state_dict if "encoder.mid.attn" in key]
+        paths = renew_vae_attention_paths(mid_attentions)  # type: ignore[no-untyped-call]
+        meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"}
+        assign_to_checkpoint(  # type: ignore[no-untyped-call]
+            paths,
+            new_checkpoint,
+            vae_state_dict,
+            additional_replacements=[meta_path],
+            config=config,
+        )
+        conv_attn_to_linear(new_checkpoint)  # type: ignore[no-untyped-call]
+
+        for i in range(num_up_blocks):
+            block_id = num_up_blocks - 1 - i
+            resnets = [
+                key
+                for key in up_blocks[block_id]
+                if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key
+            ]
+
+            if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict:
+                new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"] = vae_state_dict[
+                    f"decoder.up.{block_id}.upsample.conv.weight"
+                ]
+                new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"] = vae_state_dict[
+                    f"decoder.up.{block_id}.upsample.conv.bias"
+                ]
+
+            paths = renew_vae_resnet_paths(resnets)  # type: ignore[no-untyped-call]
+            meta_path = {"old": f"up.{block_id}.block", "new": f"up_blocks.{i}.resnets"}
+            assign_to_checkpoint(  # type: ignore[no-untyped-call]
+                paths,
+                new_checkpoint,
+                vae_state_dict,
+                additional_replacements=[meta_path],
+                config=config,
+            )
+
+        mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key]
+        num_mid_res_blocks = 2
+        for i in range(1, num_mid_res_blocks + 1):
+            resnets = [key for key in mid_resnets if f"decoder.mid.block_{i}" in key]
+
+            paths = renew_vae_resnet_paths(resnets)  # type: ignore[no-untyped-call]
+            meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}
+            assign_to_checkpoint(  # type: ignore[no-untyped-call]
+                paths,
+                new_checkpoint,
+                vae_state_dict,
+                additional_replacements=[meta_path],
+                config=config,
+            )
+
+        mid_attentions = [key for key in vae_state_dict if "decoder.mid.attn" in key]
+        paths = renew_vae_attention_paths(mid_attentions)  # type: ignore[no-untyped-call]
+        meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"}
+        assign_to_checkpoint(  # type: ignore[no-untyped-call]
+            paths,
+            new_checkpoint,
+            vae_state_dict,
+            additional_replacements=[meta_path],
+            config=config,
+        )
+        conv_attn_to_linear(new_checkpoint)  # type: ignore[no-untyped-call]
+        return new_checkpoint
+
+    # Reference from : https://github.com/huggingface/diffusers/blob/main/scripts/convert_vae_pt_to_diffusers.py
+    def _vae_pt_to_vae_diffuser(
+        self,
+        checkpoint_path: str,
+        output_path: str,
+    ) -> str:
+        # Only support V1
+        r = requests.get(
+            "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml"
+        )
+        io_obj = io.BytesIO(r.content)
+
+        original_config = OmegaConf.load(io_obj)
+        image_size = 512
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        if checkpoint_path.endswith("safetensors"):
+            from safetensors import safe_open
+
+            checkpoint: dict[str, Any] = {}
+            with safe_open(checkpoint_path, framework="pt", device="cpu") as f:  # type: ignore[no-untyped-call]
+                for key in f.keys():
+                    checkpoint[key] = f.get_tensor(key)
+        else:
+            checkpoint = torch.load(checkpoint_path, map_location=device)["state_dict"]
+
+        # Convert the VAE model.
+        vae_config = create_vae_diffusers_config(original_config, image_size=image_size)
+        converted_vae_checkpoint = self._custom_convert_ldm_vae_checkpoint(checkpoint, vae_config)
+
+        vae = AutoencoderKL(**vae_config)  # type: ignore[no-untyped-call]
+        vae.load_state_dict(converted_vae_checkpoint)  # type: ignore[attr-defined]
+        vae.save_pretrained(output_path)  # type: ignore[attr-defined]
+
+        return output_path

src/nodes/__init__.py

@@ -1,15 +1,23 @@
 from ..di import Container
 from .diffusers_pipeline_loader import DiffusersPipelineLoader
-from .dto import ComfyUIPipelineDTO
+from .diffusers_vae_loader import DiffusersVaeLoader
+from .dto import ComfyUIAutoencoderDTO, ComfyUIPipelineDTO
 
 container = Container()
 container.wire(modules=[__name__])
 
-NODE_CLASS_MAPPINGS = {DiffusersPipelineLoader.__name__: DiffusersPipelineLoader}
-NODE_DISPLAY_NAME_MAPPINGS = {DiffusersPipelineLoader.__name__: "Diffusers Pipeline Loader"}
+NODE_CLASS_MAPPINGS = {
+    DiffusersPipelineLoader.__name__: DiffusersPipelineLoader,
+    DiffusersVaeLoader.__name__: DiffusersVaeLoader,
+}
+NODE_DISPLAY_NAME_MAPPINGS = {
+    DiffusersPipelineLoader.__name__: "Diffusers Pipeline Loader",
+    DiffusersVaeLoader.__name__: "Diffusers VAE Loader",
+}
 
 __all__ = [
     "NODE_CLASS_MAPPINGS",
     "NODE_DISPLAY_NAME_MAPPINGS",
     "ComfyUIPipelineDTO",
+    "ComfyUIAutoencoderDTO",
 ]

src/nodes/diffusers_vae_loader.py

@@ -0,0 +1,33 @@
+import folder_paths  # pyright: ignore[reportMissingImports]
+from dependency_injector.wiring import Provide, inject
+
+from ..di import Container
+from ..ui import AutoencoderHandler
+from .dto import ComfyUIAutoencoderDTO
+
+
+class DiffusersVaeLoader:
+    def __init__(self) -> None:
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s) -> dict[str, dict[str, tuple[tuple[str, ...], ...]]]:
+        return {
+            "required": {
+                "vae_name": (folder_paths.get_filename_list("vae"),),
+            },
+        }
+
+    RETURN_TYPES = (ComfyUIAutoencoderDTO.COMFY_TYPE,)
+    FUNCTION = "execute"
+    CATEGORY = "Diffusers"
+
+    @inject
+    def execute(
+        self,
+        vae_name: str,
+        handler: AutoencoderHandler = Provide[Container.autoencoder_handler],
+    ) -> tuple[ComfyUIAutoencoderDTO]:
+        autoencoder_model = handler.create(vae_name)
+        autoencoder_dto = ComfyUIAutoencoderDTO.from_domain(autoencoder_model)
+        return (autoencoder_dto,)

src/nodes/dto/__init__.py

@@ -1,3 +1,4 @@
+from .autoencoder_dto import ComfyUIAutoencoderDTO
 from .pipeline_dto import ComfyUIPipelineDTO
 
-__all__ = ["ComfyUIPipelineDTO"]
+__all__ = ["ComfyUIPipelineDTO", "ComfyUIAutoencoderDTO"]

src/nodes/dto/autoencoder_dto.py

@@ -0,0 +1,59 @@
+from dataclasses import dataclass
+from typing import ClassVar
+
+from diffusers import AutoencoderKL
+
+from ...domain.model import AutoencoderModel
+
+
+@dataclass
+class ComfyUIAutoencoderDTO:
+    """
+    Data Transfer Object for Diffusers Autoencoder in ComfyUI.
+
+    This DTO is designed to be used across different custom nodes.
+    Other node developers can import and use this type for autoencoder operations.
+
+    Attributes:
+        autoencoder: AutoencoderKL instance from diffusers library
+        path: Path to the loaded 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 ComfyUIAutoencoderDTO
+
+        class YourCustomNode:
+            @classmethod
+            def INPUT_TYPES(cls):
+                return {"required": {"autoencoder": (ComfyUIAutoencoderDTO.COMFY_TYPE,)}}
+
+            RETURN_TYPES = (ComfyUIAutoencoderDTO.COMFY_TYPE,)
+
+            def execute(self, autoencoder: ComfyUIAutoencoderDTO):
+                # Use the autoencoder for encoding/decoding
+                encoded = autoencoder.autoencoder.encode(image)
+                return (encoded,)
+        ```
+    """
+
+    # ClassVar to avoid being treated as a dataclass field
+    COMFY_TYPE: ClassVar[str] = "DIFFUSERS_VAE"
+
+    autoencoder: AutoencoderKL
+    path: str
+
+    @classmethod
+    def from_domain(cls, model: AutoencoderModel) -> "ComfyUIAutoencoderDTO":
+        """
+        Create DTO from domain model.
+
+        Args:
+            model: AutoencoderModel from domain layer
+
+        Returns:
+            ComfyUIAutoencoderDTO instance
+        """
+        return cls(autoencoder=model.autoencoder, path=model.path)

src/service/__init__.py

@@ -1,3 +1,4 @@
+from .create_autoencoder_service import CreateAutoencoderService
 from .create_pipeline_service import CreatePipelineService
 
-__all__ = ["CreatePipelineService"]
+__all__ = ["CreatePipelineService", "CreateAutoencoderService"]

src/service/create_autoencoder_service.py

@@ -0,0 +1,25 @@
+import os
+
+import torch
+
+from ..domain.model import AutoencoderModel
+from ..domain.repositories import IAutoencoderRepository
+
+
+class CreateAutoencoderService:
+    def __init__(self, autoencoder_repo: IAutoencoderRepository) -> None:
+        self.autoencoder_repo = autoencoder_repo
+        self.dtype = torch.float32
+
+    def create_from_checkpoint(self, checkpoint_full_path: str) -> AutoencoderModel:
+        path = self.autoencoder_repo.convert_and_save_from_single_file(checkpoint_full_path)
+        if not os.path.exists(path):
+            raise RuntimeError(f"Failed to convert VAE checkpoint: {checkpoint_full_path}")
+
+        vae = self.autoencoder_repo.load_autoencoder_from_path(path, self.dtype)
+        if vae is None:
+            raise RuntimeError(
+                f"Failed to create autoencoder from checkpoint: {checkpoint_full_path}"
+            )
+
+        return AutoencoderModel(autoencoder=vae, path=path)

src/ui/__init__.py

@@ -1,3 +1,4 @@
+from .autoencoder_handler import AutoencoderHandler
 from .pipeline_handler import PipelineHandler
 
-__all__ = ["PipelineHandler"]
+__all__ = ["PipelineHandler", "AutoencoderHandler"]

src/ui/autoencoder_handler.py

@@ -0,0 +1,17 @@
+import os
+
+import folder_paths  # pyright: ignore[reportMissingImports]
+
+from ..domain.model import AutoencoderModel
+from ..service.create_autoencoder_service import CreateAutoencoderService
+
+
+class AutoencoderHandler:
+    def __init__(self, create_autoencoder_service: CreateAutoencoderService) -> None:
+        self.create_autoencoder_service = create_autoencoder_service
+
+    def create(self, checkpoint_name: str) -> AutoencoderModel:
+        checkpoint_path = folder_paths.get_full_path("vae", checkpoint_name)
+        if not os.path.exists(checkpoint_path):
+            raise FileNotFoundError(f"VAE file not found: {checkpoint_name}")
+        return self.create_autoencoder_service.create_from_checkpoint(checkpoint_path)