diff --git a/__init__.py b/__init__.py index bd50b46..ca9fa73 100644 --- a/__init__.py +++ b/__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", ] diff --git a/src/di/di.py b/src/di/di.py index 45a71b9..777efdc 100644 --- a/src/di/di.py +++ b/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, + ) diff --git a/src/domain/model/__init__.py b/src/domain/model/__init__.py index 46bea47..d51212d 100644 --- a/src/domain/model/__init__.py +++ b/src/domain/model/__init__.py @@ -1,3 +1,7 @@ +from .autoencoder import AutoencoderModel from .pipeline import PipelineModel -__all__ = ["PipelineModel"] +__all__ = [ + "PipelineModel", + "AutoencoderModel", +] diff --git a/src/domain/model/autoencoder.py b/src/domain/model/autoencoder.py new file mode 100644 index 0000000..cdaba2a --- /dev/null +++ b/src/domain/model/autoencoder.py @@ -0,0 +1,9 @@ +from dataclasses import dataclass + +from diffusers import AutoencoderKL + + +@dataclass +class AutoencoderModel: + autoencoder: AutoencoderKL + path: str diff --git a/src/domain/repositories/__init__.py b/src/domain/repositories/__init__.py index 6a974d6..6355550 100644 --- a/src/domain/repositories/__init__.py +++ b/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"] diff --git a/src/domain/repositories/i_autoencoder_repository.py b/src/domain/repositories/i_autoencoder_repository.py new file mode 100644 index 0000000..37830da --- /dev/null +++ b/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 diff --git a/src/infra/diffusers/__init__.py b/src/infra/diffusers/__init__.py index c6fbc41..9b19ace 100644 --- a/src/infra/diffusers/__init__.py +++ b/src/infra/diffusers/__init__.py @@ -1,3 +1,4 @@ +from .autoencoder_repository import DiffusersAutoencoderRepository from .pipeline_repository import DiffusersPipelineRepository -__all__ = ["DiffusersPipelineRepository"] +__all__ = ["DiffusersPipelineRepositoryDiffusersAutoencoderRepository"] diff --git a/src/infra/diffusers/autoencoder_repository.py b/src/infra/diffusers/autoencoder_repository.py new file mode 100644 index 0000000..51ff593 --- /dev/null +++ b/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 diff --git a/src/nodes/__init__.py b/src/nodes/__init__.py index b954c04..46f4cdf 100644 --- a/src/nodes/__init__.py +++ b/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", ] diff --git a/src/nodes/diffusers_vae_loader.py b/src/nodes/diffusers_vae_loader.py new file mode 100644 index 0000000..02a1f2a --- /dev/null +++ b/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,) diff --git a/src/nodes/dto/__init__.py b/src/nodes/dto/__init__.py index 7325eea..917341c 100644 --- a/src/nodes/dto/__init__.py +++ b/src/nodes/dto/__init__.py @@ -1,3 +1,4 @@ +from .autoencoder_dto import ComfyUIAutoencoderDTO from .pipeline_dto import ComfyUIPipelineDTO -__all__ = ["ComfyUIPipelineDTO"] +__all__ = ["ComfyUIPipelineDTO", "ComfyUIAutoencoderDTO"] diff --git a/src/nodes/dto/autoencoder_dto.py b/src/nodes/dto/autoencoder_dto.py new file mode 100644 index 0000000..cb88439 --- /dev/null +++ b/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) diff --git a/src/service/__init__.py b/src/service/__init__.py index 75e4881..054dcf9 100644 --- a/src/service/__init__.py +++ b/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"] diff --git a/src/service/create_autoencoder_service.py b/src/service/create_autoencoder_service.py new file mode 100644 index 0000000..98f9375 --- /dev/null +++ b/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) diff --git a/src/ui/__init__.py b/src/ui/__init__.py index dc68e44..3aea75c 100644 --- a/src/ui/__init__.py +++ b/src/ui/__init__.py @@ -1,3 +1,4 @@ +from .autoencoder_handler import AutoencoderHandler from .pipeline_handler import PipelineHandler -__all__ = ["PipelineHandler"] +__all__ = ["PipelineHandler", "AutoencoderHandler"] diff --git a/src/ui/autoencoder_handler.py b/src/ui/autoencoder_handler.py new file mode 100644 index 0000000..e044701 --- /dev/null +++ b/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)