From fcb48ec62822bd89179c39db90dea2a62740d8ca Mon Sep 17 00:00:00 2001 From: niknah Date: Thu, 16 Jul 2026 16:25:44 +1000 Subject: [PATCH] Added loader nodes so that the models use ComfyUI VRAM management. Faster generations. --- MiniT2I.py | 103 ++++++++------ __init__.py | 4 +- example_workflows/MiniT2I_example.json | 177 +++++++++++++------------ pipeline.py | 37 ++++-- 4 files changed, 188 insertions(+), 133 deletions(-) diff --git a/MiniT2I.py b/MiniT2I.py index fcaf928..e090813 100644 --- a/MiniT2I.py +++ b/MiniT2I.py @@ -1,15 +1,66 @@ -from pathlib import Path -from typing import List, Union -import numpy as np from comfy_api.latest import io -from PIL import Image import torch -from diffusers import DiffusionPipeline from .pipeline import MiniT2IPipeline -# import os +from comfy.model_patcher import ModelPatcher +import comfy.model_management as mm +from transformers import T5EncoderModel # AutoTokenizer, + model_downloaded = False + +class MiniT2ITextEncoder(io.ComfyNode): + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="MiniT2ITextEncoder", + display_name="MiniT2I Text Encoder Loader", + category="MiniT2I", + inputs=[ + ], + outputs=[ + io.Model.Output(), + ], + ) + + @classmethod + def execute(cls) -> io.NodeOutput: + load_device = mm.get_torch_device() + offload_device = mm.intermediate_device() + + text_encoder = T5EncoderModel.from_pretrained( + "google/flan-t5-large", + torch_dtype=torch.float32, + local_files_only=model_downloaded, + ) + + return io.NodeOutput(ModelPatcher(text_encoder, load_device, offload_device),) + + +class MiniT2ILoader(io.ComfyNode): + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="MiniT2ILoader", + display_name="MiniT2I Model Loader", + category="MiniT2I", + inputs=[ + io.Combo.Input("model_type", options=["b16","l16"], tooltip="l16=large, b16=normal"), + ], + outputs=[ + io.Model.Output(), + ], + ) + + @classmethod + def execute(cls, model_type) -> io.NodeOutput: + transformer = MiniT2IPipeline.load_transformer(model_type) + load_device = mm.get_torch_device() + offload_device = mm.intermediate_device() + + return io.NodeOutput(ModelPatcher(transformer, load_device, offload_device),) + + class MiniT2ISampler(io.ComfyNode): """ An example node @@ -58,7 +109,8 @@ class MiniT2ISampler(io.ComfyNode): display_mode=io.NumberDisplay.number, lazy=True, ), - io.Combo.Input("model_type", options=["b16","l16"]), + io.Model.Input("model", tooltip="Use MiniT2I Loader", optional=True), + io.Model.Input("text_encoder", tooltip="Use MiniT2I Text Encoder Loader", optional=True), io.Int.Input( "seed", default=1, @@ -90,43 +142,14 @@ class MiniT2ISampler(io.ComfyNode): # return [] -# @classmethod -# def pil2tensor(cls, image: Union[Image.Image, List[Image.Image]]) -> torch.Tensor: -# """ -# Convert PIL image(s) to tensor, matching ComfyUI's implementation. -# -# Args: -# image: Single PIL Image or list of PIL Images -# -# Returns: -# torch.Tensor: Image tensor with values normalized to [0, 1] -# """ -# if isinstance(image, list): -# if len(image) == 0: -# return torch.empty(0) -# return torch.cat([cls.pil2tensor(img) for img in image], dim=0) -# -# # Convert PIL image to RGB if needed -# if image.mode == 'RGBA': -# image = image.convert('RGB') -# elif image.mode != 'RGB': -# image = image.convert('RGB') -# -# # Convert to numpy array and normalize to [0, 1] -# img_array = np.array(image).astype(np.float32) / 255.0 -# -# # Return tensor with shape [1, H, W, 3] -# return torch.from_numpy(img_array)[None,] @classmethod - def execute(cls, prompt, steps, guidance, model_type, seed) -> io.NodeOutput: + def execute(cls, prompt, steps, guidance, model, text_encoder, seed) -> io.NodeOutput: global model_downloaded torch.manual_seed(seed) + # transformer = model -# Gets the absolute directory of the running script - script_dir = Path(__file__).resolve().parent - HUB_MODEL_ID = "MiniT2I/MiniT2I" pipe = MiniT2IPipeline.from_pretrained( @@ -141,7 +164,9 @@ class MiniT2ISampler(io.ComfyNode): output = pipe( prompt, - model_type=model_type, + model.model, text_encoder.model, +# model_type=model_type, + model_dir=model.model.config._name_or_path.parent if model else None, guidance_scale=guidance, num_inference_steps=steps, torch_dtype=torch.bfloat16, diff --git a/__init__.py b/__init__.py index 60e8222..e07d19d 100644 --- a/__init__.py +++ b/__init__.py @@ -1,6 +1,6 @@ from typing_extensions import override from comfy_api.latest import ComfyExtension, io -from .MiniT2I import MiniT2ISampler +from .MiniT2I import MiniT2ISampler, MiniT2ILoader, MiniT2ITextEncoder class MiniT2IExtension(ComfyExtension): @@ -8,6 +8,8 @@ class MiniT2IExtension(ComfyExtension): async def get_node_list(self) -> list[type[io.ComfyNode]]: return [ MiniT2ISampler, + MiniT2ILoader, + MiniT2ITextEncoder, ] diff --git a/example_workflows/MiniT2I_example.json b/example_workflows/MiniT2I_example.json index e59bcdf..914ce8b 100644 --- a/example_workflows/MiniT2I_example.json +++ b/example_workflows/MiniT2I_example.json @@ -1,50 +1,15 @@ { - "id": "e572e121-bf51-413f-bba1-ca525f710414", + "id": "03998d2d-877c-420e-92b6-4358d2ed18a9", "revision": 0, - "last_node_id": 4, - "last_link_id": 2, + "last_node_id": 11, + "last_link_id": 5, "nodes": [ - { - "id": 1, - "type": "PreviewImage", - "pos": [ - 1397.0351491908762, - -950.3260055448459 - ], - "size": [ - 140, - 246 - ], - "flags": {}, - "order": 2, - "mode": 0, - "inputs": [ - { - "name": "images", - "type": "IMAGE", - "link": 1 - } - ], - "outputs": [ - { - "name": "images", - "type": "IMAGE", - "links": null - } - ], - "properties": { - "cnr_id": "comfy-core", - "ver": "0.26.0", - "Node name for S&R": "PreviewImage" - }, - "widgets_values": [] - }, { "id": 4, "type": "PreviewImage", "pos": [ - 1416.1947412443299, - -549.015734080902 + 1402.7689200070347, + 828.7114550770376 ], "size": [ 140, @@ -57,7 +22,7 @@ { "name": "images", "type": "IMAGE", - "link": 2 + "link": 3 } ], "outputs": [ @@ -69,21 +34,21 @@ ], "properties": { "cnr_id": "comfy-core", - "ver": "0.26.0", + "ver": "0.27.0", "Node name for S&R": "PreviewImage" }, "widgets_values": [] }, { - "id": 2, - "type": "MiniT2I", + "id": 11, + "type": "MiniT2ITextEncoder", "pos": [ - 850.8583116297305, - -918.896001453913 + 503.65560995074725, + 973.0844817026747 ], "size": [ - 400, - 208 + 230.43125, + 26 ], "flags": {}, "order": 0, @@ -91,35 +56,30 @@ "inputs": [], "outputs": [ { - "name": "IMAGE", - "type": "IMAGE", + "name": "MODEL", + "type": "MODEL", "links": [ - 1 + 5 ] } ], "properties": { - "Node name for S&R": "MiniT2I" + "cnr_id": "MiniT2I-ComfyUI", + "ver": "1190fbb080b090dd158ec32996e9def1844ed15a", + "Node name for S&R": "MiniT2ITextEncoder" }, - "widgets_values": [ - "Purple monkey dishwasher", - 100, - 2.5, - "b16", - 1, - "increment" - ] + "widgets_values": [] }, { - "id": 3, - "type": "MiniT2I", + "id": 2, + "type": "MiniT2ILoader", "pos": [ - 820.5800443091753, - -530.5854576396309 + 545.8728693240846, + 753.3909621095928 ], "size": [ - 400, - 208 + 270, + 58 ], "flags": {}, "order": 1, @@ -127,23 +87,68 @@ "inputs": [], "outputs": [ { - "name": "IMAGE", - "type": "IMAGE", + "name": "MODEL", + "type": "MODEL", "links": [ - 2 + 1 ] } ], "properties": { + "cnr_id": "MiniT2I-ComfyUI", + "ver": "1190fbb080b090dd158ec32996e9def1844ed15a", + "Node name for S&R": "MiniT2ILoader" + }, + "widgets_values": [ + "b16" + ] + }, + { + "id": 1, + "type": "MiniT2I", + "pos": [ + 905.9168723189537, + 828.3846739799442 + ], + "size": [ + 400, + 204 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 1 + }, + { + "name": "text_encoder", + "type": "MODEL", + "link": 5 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 3 + ] + } + ], + "properties": { + "cnr_id": "MiniT2I-ComfyUI", + "ver": "1190fbb080b090dd158ec32996e9def1844ed15a", "Node name for S&R": "MiniT2I" }, "widgets_values": [ - "Purple monkey dishwasher", + "test", 100, - 6, - "l16", - 1, - "increment" + 2.5, + 1400, + "randomize" ] } ], @@ -154,32 +159,36 @@ 0, 1, 0, - "IMAGE" + "MODEL" ], [ - 2, 3, + 1, 0, 4, 0, "IMAGE" + ], + [ + 5, + 11, + 0, + 1, + 1, + "MODEL" ] ], "groups": [], "config": {}, "extra": { "ds": { - "scale": 1.588548164722016, + "scale": 1.3513057093103993, "offset": [ - -504.72890778875325, - 878.5672286853679 + -350.91547909735857, + -626.8584323803433 ] }, - "frontendVersion": "1.45.19", - "VHS_latentpreview": false, - "VHS_latentpreviewrate": 0, - "VHS_MetadataImage": true, - "VHS_KeepIntermediate": true + "frontendVersion": "1.45.20" }, "version": 0.4 } \ No newline at end of file diff --git a/pipeline.py b/pipeline.py index 377d75e..012142c 100644 --- a/pipeline.py +++ b/pipeline.py @@ -611,7 +611,6 @@ class MiniT2ITextToImagePipeline(nn.Module): finally: self.transformer.model.cfg.n_T = old_steps - if output_type == "pt": images = (images.clamp(-1, 1) * 0.5 + 0.5).permute(0, 2, 3, 1).to(torch.float16).cpu() pass @@ -679,11 +678,27 @@ class MiniT2IPipeline(DiffusionPipeline): ) ) + @classmethod + def load_transformer(cls, + model_type: str = "b16", + repo_id_or_path: Union[str, os.PathLike] = "MiniT2I/MiniT2I", + revision: Optional[str] = None, + local_files_only: bool = False, + torch_dtype: Optional[torch.dtype] = torch.bfloat16, + cache_dir: Optional[Union[str, os.PathLike]] = None, + ): + model_dir = cls._resolve_model_type(model_type) + root = cls._resolve_root(repo_id_or_path, model_dir, revision, cache_dir, local_files_only) + model_root = root / model_dir + return MiniT2IMMJiTModel.from_pretrained(model_root / "transformer", torch_dtype=torch_dtype) + @torch.no_grad() def __call__( self, prompt: Union[str, List[str]], - model_type: str = "b16", + transformer, text_encoder, + #model_type: str = "b16", + model_dir: Path, repo_id_or_path: Union[str, os.PathLike] = "MiniT2I/MiniT2I", torch_dtype: Optional[torch.dtype] = torch.bfloat16, text_encoder_dtype: torch.dtype = torch.float32, @@ -693,18 +708,22 @@ class MiniT2IPipeline(DiffusionPipeline): cache_dir: Optional[Union[str, os.PathLike]] = None, **kwargs, ): - model_dir = self._resolve_model_type(model_type) + if model_dir is None: + model_dir = self._resolve_model_type("b16") root = self._resolve_root(repo_id_or_path, model_dir, revision, cache_dir, local_files_only) model_root = root / model_dir - transformer = MiniT2IMMJiTModel.from_pretrained(model_root / "transformer", torch_dtype=torch_dtype) + if transformer is None: + transformer = MiniT2IMMJiTModel.from_pretrained(model_root / "transformer", torch_dtype=torch_dtype) scheduler = MiniT2IFlowMatchScheduler.from_pretrained(model_root / "scheduler") text_encoder_name = transformer.mmjit_config.llm tokenizer = AutoTokenizer.from_pretrained(text_encoder_name, local_files_only=local_files_only) - text_encoder = T5EncoderModel.from_pretrained( - text_encoder_name, - torch_dtype=text_encoder_dtype, - local_files_only=local_files_only, - ) + + if text_encoder is None: + text_encoder = T5EncoderModel.from_pretrained( + text_encoder_name, + torch_dtype=text_encoder_dtype, + local_files_only=local_files_only, + ) pipe = MiniT2ITextToImagePipeline( transformer=transformer, scheduler=scheduler,