476 lines
16 KiB
Python
476 lines
16 KiB
Python
"""
|
|
ComfyUI custom nodes for LongCat-Image
|
|
https://github.com/meituan-longcat/LongCat-Image
|
|
"""
|
|
|
|
from typing_extensions import override
|
|
from comfy_api.latest import ComfyExtension, io
|
|
from comfy_api.latest._io import comfytype, ComfyTypeIO
|
|
import torch
|
|
import folder_paths
|
|
import os
|
|
|
|
|
|
# Define custom type for LongCat pipeline
|
|
@comfytype(io_type="LONGCAT_PIPE")
|
|
class LongCatPipe(ComfyTypeIO):
|
|
"""Custom type for LongCat-Image pipeline"""
|
|
Type = dict
|
|
|
|
|
|
class LongCatImageModelLoader(io.ComfyNode):
|
|
"""
|
|
Load LongCat-Image models for text-to-image or image editing.
|
|
"""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="LongCatImageModelLoader",
|
|
display_name="LongCat-Image Model Loader",
|
|
category="loaders",
|
|
inputs=[
|
|
io.String.Input(
|
|
"model_path",
|
|
default="",
|
|
multiline=False,
|
|
tooltip="Path to LongCat-Image model directory (e.g., 'LongCat-Image', 'LongCat-Image-Edit')"
|
|
),
|
|
io.Combo.Input(
|
|
"dtype",
|
|
options=["bfloat16", "float16", "float32"],
|
|
default="bfloat16",
|
|
tooltip="Data type for model weights"
|
|
),
|
|
io.Combo.Input(
|
|
"enable_cpu_offload",
|
|
options=["false", "true"],
|
|
default="true",
|
|
tooltip="Enable CPU offload to save VRAM (~17-19GB required). Slower but prevents OOM on low VRAM GPUs."
|
|
),
|
|
io.Combo.Input(
|
|
"attention_backend",
|
|
options=["default", "sage"],
|
|
default="default",
|
|
tooltip="Choose attention backend. 'sage' uses SageAttention (requires CUDA and the sageattention package)."
|
|
),
|
|
],
|
|
outputs=[
|
|
LongCatPipe.Output(
|
|
display_name="LongCat Pipeline",
|
|
tooltip="LongCat-Image pipeline for generation"
|
|
),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, model_path, dtype, enable_cpu_offload, attention_backend) -> io.NodeOutput:
|
|
try:
|
|
from transformers import AutoProcessor
|
|
from diffusers.models.transformers.transformer_longcat_image import LongCatImageTransformer2DModel
|
|
from diffusers import LongCatImagePipeline, LongCatImageEditPipeline
|
|
except ImportError as e:
|
|
raise ImportError(
|
|
f"Failed to import LongCat-Image dependencies. "
|
|
f"Please install: pip install -r custom_nodes/comfyui_longcat_image/requirements.txt\n"
|
|
f"Error: {e}"
|
|
)
|
|
|
|
# Find model path
|
|
if not model_path:
|
|
raise ValueError("model_path is required")
|
|
|
|
# Check if it's an absolute path
|
|
if os.path.isabs(model_path):
|
|
checkpoint_dir = model_path
|
|
else:
|
|
# Try to find in standard ComfyUI model directories
|
|
base_paths = [
|
|
os.path.join(folder_paths.models_dir, "diffusion_models"),
|
|
os.path.join(folder_paths.models_dir, "checkpoints"),
|
|
folder_paths.models_dir,
|
|
]
|
|
|
|
checkpoint_dir = None
|
|
for base_path in base_paths:
|
|
potential_path = os.path.join(base_path, model_path)
|
|
if os.path.exists(potential_path):
|
|
checkpoint_dir = potential_path
|
|
break
|
|
|
|
if checkpoint_dir is None:
|
|
raise ValueError(
|
|
f"Model not found at '{model_path}'. "
|
|
f"Please provide a valid path or place the model in: {base_paths[0]}"
|
|
)
|
|
|
|
# Map dtype string to torch dtype
|
|
dtype_map = {
|
|
"bfloat16": torch.bfloat16,
|
|
"float16": torch.float16,
|
|
"float32": torch.float32,
|
|
}
|
|
torch_dtype = dtype_map[dtype]
|
|
|
|
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
|
|
# Optional SageAttention backend
|
|
if attention_backend == "sage":
|
|
if device.type != "cuda":
|
|
raise ValueError("SageAttention requires a CUDA-capable GPU.")
|
|
try:
|
|
import torch.nn.functional as F
|
|
from sageattention import sageattn
|
|
except ImportError as e:
|
|
raise ImportError(
|
|
"SageAttention selected but the package is not installed. "
|
|
"Install it with: pip install sageattention"
|
|
) from e
|
|
# Replace PyTorch SDPA with SageAttention globally, but fall back to
|
|
# the original implementation when a head_dim is unsupported or a
|
|
# mask/dropout is requested.
|
|
if not getattr(F.scaled_dot_product_attention, "_longcat_sage_wrapped", False):
|
|
orig_sdpa = F.scaled_dot_product_attention
|
|
|
|
def _longcat_sage_sdpa(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None, **kwargs):
|
|
# SageAttention currently supports limited head dimensions (typically <=256).
|
|
head_dim = query.shape[-1]
|
|
if attn_mask is not None or dropout_p not in (0, 0.0) or head_dim > 256:
|
|
return orig_sdpa(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale, **kwargs)
|
|
try:
|
|
return sageattn(query, key, value, tensor_layout="HND", is_causal=is_causal, sm_scale=scale)
|
|
except Exception:
|
|
# Fallback for any unsupported shapes/dtypes
|
|
return orig_sdpa(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale, **kwargs)
|
|
|
|
_longcat_sage_sdpa._longcat_sage_wrapped = True
|
|
F.scaled_dot_product_attention = _longcat_sage_sdpa
|
|
|
|
# Load text processor
|
|
text_processor = AutoProcessor.from_pretrained(
|
|
checkpoint_dir,
|
|
subfolder='tokenizer'
|
|
)
|
|
|
|
# Determine if this is an edit model by checking the path
|
|
is_edit_model = "edit" in model_path.lower()
|
|
|
|
# Determine CPU offload setting
|
|
cpu_offload = enable_cpu_offload == "true"
|
|
|
|
# Load transformer
|
|
transformer = LongCatImageTransformer2DModel.from_pretrained(
|
|
checkpoint_dir,
|
|
subfolder='transformer',
|
|
torch_dtype=torch_dtype,
|
|
use_safetensors=True
|
|
)
|
|
|
|
# Only move transformer to device if CPU offload is disabled
|
|
if not cpu_offload:
|
|
transformer = transformer.to(device)
|
|
|
|
# Load appropriate pipeline
|
|
if is_edit_model:
|
|
pipe = LongCatImageEditPipeline.from_pretrained(
|
|
checkpoint_dir,
|
|
transformer=transformer,
|
|
text_processor=text_processor,
|
|
torch_dtype=torch_dtype,
|
|
)
|
|
else:
|
|
pipe = LongCatImagePipeline.from_pretrained(
|
|
checkpoint_dir,
|
|
transformer=transformer,
|
|
text_processor=text_processor,
|
|
torch_dtype=torch_dtype,
|
|
)
|
|
|
|
# Apply CPU offload or move to device based on user preference
|
|
if cpu_offload:
|
|
# Enable CPU offload to save VRAM (Requires ~17-19 GB); slower but prevents OOM
|
|
pipe.enable_model_cpu_offload()
|
|
|
|
else:
|
|
# Load all models to GPU at once (Faster inference but requires more VRAM)
|
|
pipe.to(device, torch_dtype)
|
|
|
|
pipeline_data = {
|
|
"pipe": pipe,
|
|
"device": device,
|
|
"dtype": torch_dtype,
|
|
"is_edit": is_edit_model,
|
|
}
|
|
|
|
return io.NodeOutput(pipeline_data)
|
|
|
|
|
|
class LongCatImageTextToImage(io.ComfyNode):
|
|
"""
|
|
Generate images using LongCat-Image text-to-image pipeline.
|
|
"""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="LongCatImageTextToImage",
|
|
display_name="LongCat-Image Text to Image",
|
|
category="sampling",
|
|
inputs=[
|
|
LongCatPipe.Input(
|
|
"longcat_pipeline",
|
|
display_name="LongCat Pipeline",
|
|
tooltip="LongCat-Image pipeline from loader"
|
|
),
|
|
io.String.Input(
|
|
"prompt",
|
|
default="",
|
|
multiline=True,
|
|
tooltip="Text prompt for image generation"
|
|
),
|
|
io.String.Input(
|
|
"negative_prompt",
|
|
default="",
|
|
multiline=True,
|
|
tooltip="Negative prompt"
|
|
),
|
|
io.Int.Input(
|
|
"width",
|
|
default=1344,
|
|
min=64,
|
|
max=4096,
|
|
step=64,
|
|
tooltip="Image width"
|
|
),
|
|
io.Int.Input(
|
|
"height",
|
|
default=768,
|
|
min=64,
|
|
max=4096,
|
|
step=64,
|
|
tooltip="Image height"
|
|
),
|
|
io.Int.Input(
|
|
"steps",
|
|
default=50,
|
|
min=1,
|
|
max=200,
|
|
step=1,
|
|
display_mode=io.NumberDisplay.number,
|
|
tooltip="Number of inference steps"
|
|
),
|
|
io.Float.Input(
|
|
"guidance_scale",
|
|
default=4.5,
|
|
min=0.0,
|
|
max=20.0,
|
|
step=0.1,
|
|
tooltip="Guidance scale for classifier-free guidance"
|
|
),
|
|
io.Int.Input(
|
|
"seed",
|
|
default=43,
|
|
min=0,
|
|
max=0xffffffffffffffff,
|
|
tooltip="Random seed for generation"
|
|
),
|
|
io.Combo.Input(
|
|
"enable_cfg_renorm",
|
|
options=["true", "false"],
|
|
default="true",
|
|
tooltip="Enable CFG renormalization"
|
|
),
|
|
io.Combo.Input(
|
|
"enable_prompt_rewrite",
|
|
options=["true", "false"],
|
|
default="true",
|
|
tooltip="Enable built-in prompt rewriting using text encoder"
|
|
),
|
|
],
|
|
outputs=[
|
|
io.Image.Output(
|
|
tooltip="Generated image"
|
|
),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(
|
|
cls,
|
|
longcat_pipeline,
|
|
prompt,
|
|
negative_prompt,
|
|
width,
|
|
height,
|
|
steps,
|
|
guidance_scale,
|
|
seed,
|
|
enable_cfg_renorm,
|
|
enable_prompt_rewrite,
|
|
) -> io.NodeOutput:
|
|
if not longcat_pipeline:
|
|
raise ValueError("longcat_pipeline input is required")
|
|
|
|
pipeline = longcat_pipeline["pipe"]
|
|
|
|
if longcat_pipeline.get("is_edit", False):
|
|
raise ValueError("This is an edit pipeline. Use LongCatImageEdit node instead.")
|
|
|
|
# Convert string bools to actual bools
|
|
cfg_renorm = enable_cfg_renorm == "true"
|
|
prompt_rewrite = enable_prompt_rewrite == "true"
|
|
|
|
# Generate image
|
|
generator = torch.Generator("cpu").manual_seed(seed)
|
|
|
|
result = pipeline(
|
|
prompt,
|
|
negative_prompt=negative_prompt,
|
|
height=height,
|
|
width=width,
|
|
guidance_scale=guidance_scale,
|
|
num_inference_steps=steps,
|
|
num_images_per_prompt=1,
|
|
generator=generator,
|
|
enable_cfg_renorm=cfg_renorm,
|
|
enable_prompt_rewrite=prompt_rewrite,
|
|
)
|
|
|
|
image = result.images[0]
|
|
|
|
# Convert PIL image to tensor format expected by ComfyUI
|
|
import numpy as np
|
|
image_np = np.array(image).astype(np.float32) / 255.0
|
|
image_tensor = torch.from_numpy(image_np)[None,]
|
|
|
|
return io.NodeOutput(image_tensor)
|
|
|
|
|
|
class LongCatImageEdit(io.ComfyNode):
|
|
"""
|
|
Edit images using LongCat-Image-Edit pipeline.
|
|
"""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="LongCatImageEdit",
|
|
display_name="LongCat-Image Edit",
|
|
category="sampling",
|
|
inputs=[
|
|
LongCatPipe.Input(
|
|
"longcat_pipeline",
|
|
display_name="LongCat Pipeline",
|
|
tooltip="LongCat-Image-Edit pipeline from loader"
|
|
),
|
|
io.Image.Input(
|
|
"image",
|
|
tooltip="Input image to edit"
|
|
),
|
|
io.String.Input(
|
|
"prompt",
|
|
default="",
|
|
multiline=True,
|
|
tooltip="Edit instruction prompt"
|
|
),
|
|
io.String.Input(
|
|
"negative_prompt",
|
|
default="",
|
|
multiline=True,
|
|
tooltip="Negative prompt"
|
|
),
|
|
io.Int.Input(
|
|
"steps",
|
|
default=50,
|
|
min=1,
|
|
max=200,
|
|
step=1,
|
|
display_mode=io.NumberDisplay.number,
|
|
tooltip="Number of inference steps"
|
|
),
|
|
io.Float.Input(
|
|
"guidance_scale",
|
|
default=4.5,
|
|
min=0.0,
|
|
max=20.0,
|
|
step=0.1,
|
|
tooltip="Guidance scale for classifier-free guidance"
|
|
),
|
|
io.Int.Input(
|
|
"seed",
|
|
default=43,
|
|
min=0,
|
|
max=0xffffffffffffffff,
|
|
tooltip="Random seed for generation"
|
|
),
|
|
],
|
|
outputs=[
|
|
io.Image.Output(
|
|
tooltip="Edited image"
|
|
),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(
|
|
cls,
|
|
longcat_pipeline,
|
|
image,
|
|
prompt,
|
|
negative_prompt,
|
|
steps,
|
|
guidance_scale,
|
|
seed,
|
|
) -> io.NodeOutput:
|
|
if not longcat_pipeline:
|
|
raise ValueError("longcat_pipeline input is required")
|
|
|
|
pipeline = longcat_pipeline["pipe"]
|
|
|
|
if not longcat_pipeline.get("is_edit", False):
|
|
raise ValueError("This is not an edit pipeline. Use LongCatImageTextToImage node instead.")
|
|
|
|
# Convert tensor to PIL Image
|
|
from PIL import Image
|
|
import numpy as np
|
|
|
|
# ComfyUI images are in format [B, H, W, C] in range [0, 1]
|
|
image_np = (image[0].cpu().numpy() * 255).astype(np.uint8)
|
|
pil_image = Image.fromarray(image_np).convert('RGB')
|
|
|
|
# Generate edited image
|
|
generator = torch.Generator("cpu").manual_seed(seed)
|
|
|
|
result = pipeline(
|
|
pil_image,
|
|
prompt,
|
|
negative_prompt=negative_prompt,
|
|
guidance_scale=guidance_scale,
|
|
num_inference_steps=steps,
|
|
num_images_per_prompt=1,
|
|
generator=generator,
|
|
)
|
|
|
|
edited_image = result.images[0]
|
|
|
|
# Convert PIL image back to tensor
|
|
image_np = np.array(edited_image).astype(np.float32) / 255.0
|
|
image_tensor = torch.from_numpy(image_np)[None,]
|
|
|
|
return io.NodeOutput(image_tensor)
|
|
|
|
|
|
class LongCatImageExtension(ComfyExtension):
|
|
@override
|
|
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
|
return [
|
|
LongCatImageModelLoader,
|
|
LongCatImageTextToImage,
|
|
LongCatImageEdit,
|
|
]
|
|
|
|
|
|
async def comfy_entrypoint() -> LongCatImageExtension:
|
|
"""ComfyUI calls this to load the extension and its nodes."""
|
|
return LongCatImageExtension()
|