nodes/niknah-MiniT2I-ComfyUI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Added loader nodes so that the models use ComfyUI VRAM management.

Browse Source 
Faster generations.
This commit is contained in:
niknah
2026-07-16 16:25:44 +10:00
parent 1190fbb080
commit fcb48ec628
4 changed files with 188 additions and 133 deletions
Download Patch File Download Diff File Expand all files Collapse all files
MiniT2I.py
+64 -39
View File
@@ -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,
__init__.py
+3 -1
View File
@@ -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,
]
example_workflows/MiniT2I_example.json
+93 -84
View File
@@ -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
}
pipeline.py
+28 -9
View File
@@ -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,