import os import sys import json import torch import numpy as np from PIL import Image from typing import List, Union, Dict, Any, Optional import math from transformers import AutoTokenizer, AutoModelForCausalLM, AutoProcessor, AutoConfig, CLIPImageProcessor, CLIPVisionModelWithProjection from diffusers.models import AutoencoderKL from diffusers.schedulers import FlowMatchEulerDiscreteScheduler from safetensors.torch import load_file as load_safetensors import folder_paths import comfy.model_management import comfy.utils import comfy.sd import comfy.model_patcher # Add current directory to sys.path to ensure imports work current_dir = os.path.dirname(os.path.abspath(__file__)) if current_dir not in sys.path: sys.path.insert(0, current_dir) from longcat_image.pipelines.pipeline_longcat_image_edit import LongCatImageEditPipeline from longcat_image.pipelines.pipeline_longcat_image import ( LongCatImagePipeline, SYSTEM_PROMPT_EN, SYSTEM_PROMPT_ZH, get_prompt_language, ) from longcat_image.models.longcat_image_dit import LongCatImageTransformer2DModel from longcat_image.dataset.data_utils import MULTI_ASPECT_RATIO_1024, MULTI_ASPECT_RATIO_512, MULTI_ASPECT_RATIO_256 from longcat_image.utils.model_utils import calculate_shift, retrieve_timesteps, prepare_pos_ids, split_quotation def tensor2pil(image: torch.Tensor) -> List[Image.Image]: # image: [B, H, W, C] batch_count = image.size(0) pil_images = [] for i in range(batch_count): img = image[i].cpu().numpy() img = (img * 255.0).clip(0, 255).astype(np.uint8) pil_images.append(Image.fromarray(img)) return pil_images def pil2tensor(images: List[Image.Image]) -> torch.Tensor: tensors = [] for img in images: # Convert to numpy, normalize to 0-1 np_img = np.array(img).astype(np.float32) / 255.0 tensors.append(torch.from_numpy(np_img)) # Stack to [B, H, W, C] if len(tensors) > 0: return torch.stack(tensors, dim=0) else: return torch.empty(0) DTYPE_MAP = { "bf16": torch.bfloat16, "fp16": torch.float16, "float16": torch.float16, "fp32": torch.float32, "float32": torch.float32, } RESOURCE_DIR = os.path.join(current_dir, "longcat_image", "resources") BUNDLED_TOKENIZER_DIR = os.path.join(RESOURCE_DIR, "tokenizer") BUNDLED_TEXT_ENCODER_DIR = os.path.join(RESOURCE_DIR, "text_encoder") BUNDLED_TEXT_ENCODER_CONFIG = os.path.join(RESOURCE_DIR, "text_encoder", "config.json") BUNDLED_VAE_CONFIG = os.path.join(RESOURCE_DIR, "vae", "config.json") def _resolve_device(device_str: str): if not device_str or device_str == "auto": return comfy.model_management.get_torch_device() lowered = device_str.lower() if lowered.startswith("cuda"): if not torch.cuda.is_available(): print("Requested CUDA device but CUDA is not available; falling back to default device.") return comfy.model_management.get_torch_device() return torch.device(device_str) if lowered == "mps": if not torch.backends.mps.is_available(): print("Requested MPS device but MPS is not available; falling back to default device.") return comfy.model_management.get_torch_device() return torch.device("mps") if lowered == "cpu": return torch.device("cpu") # Fallback for any other torch-recognized device string return torch.device(device_str) def list_available_devices(): devices = ["auto"] if torch.cuda.is_available(): devices.append("cuda") for idx in range(torch.cuda.device_count()): devices.append(f"cuda:{idx}") if torch.backends.mps.is_available(): devices.append("mps") devices.append("cpu") # Deduplicate while preserving order deduped = [] for dev in devices: if dev not in deduped: deduped.append(dev) return deduped def _load_config_from_metadata(safetensors_path: str) -> Optional[Dict[str, Any]]: try: from safetensors import safe_open with safe_open(safetensors_path, framework="pt", device="cpu") as f: meta = f.metadata() or {} for key in ["config", "model_config", "config_json"]: if key in meta: try: return json.loads(meta[key]) except Exception: continue except Exception: pass return None def _load_bundled_config(rel_path: str) -> Optional[Dict[str, Any]]: path = os.path.join(RESOURCE_DIR, rel_path) if os.path.exists(path): try: with open(path, "r", encoding="utf-8") as f: return json.load(f) except Exception: return None return None class LongCatCLIPWrapper: def __init__( self, text_encoder, tokenizer, text_processor, image_encoder=None, feature_extractor=None, ): self.text_encoder = text_encoder self.tokenizer = tokenizer self.text_processor = text_processor self.image_encoder = image_encoder self.feature_extractor = feature_extractor self.image_processor_vl = getattr(text_processor, "image_processor", None) self.prompt_template_encode_prefix = '<|im_start|>system\nAs an image captioning expert, generate a descriptive text prompt based on an image content, suitable for input to a text-to-image model.<|im_end|>\n<|im_start|>user\n' self.prompt_template_encode_suffix = '<|im_end|>\n<|im_start|>assistant\n' self.prompt_template_encode_start_idx = 36 self.prompt_template_encode_end_idx = 5 self.default_sample_size = 128 self.max_tokenizer_len = 512 self.image_token = "<|image_pad|>" self.edit_prompt_template_encode_prefix = "<|im_start|>system\nAs an image editing expert, first analyze the content and attributes of the input image(s). Then, based on the user's editing instructions, clearly and precisely determine how to modify the given image(s), ensuring that only the specified parts are altered and all other aspects remain consistent with the original(s).<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>" self.edit_prompt_template_encode_suffix = '<|im_end|>\n<|im_start|>assistant\n' self.edit_prompt_template_encode_start_idx = 67 self.edit_prompt_template_encode_end_idx = 5 @property def device(self): return next(self.text_encoder.parameters()).device @property def dtype(self): return next(self.text_encoder.parameters()).dtype def rewrite_prompt(self, prompt: str) -> str: language = get_prompt_language(prompt) if language == 'zh': question = SYSTEM_PROMPT_ZH + f"\n用户输入为：{prompt}\n改写后的prompt为：" else: question = SYSTEM_PROMPT_EN + f"\nUser Input: {prompt}\nRewritten prompt:" messages = [ { "role": "user", "content": [ {"type": "text", "text": question}, ], } ] text = self.text_processor.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) inputs = self.text_processor(text=[text], padding=True, return_tensors="pt") inputs = inputs.to(self.device) generated_ids = self.text_encoder.generate( **inputs, max_new_tokens=self.max_tokenizer_len ) generated_ids_trimmed = [ out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids) ] output_text = self.text_processor.batch_decode( generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False, )[0] return output_text def _encode_tokens(self, prompt: str, prefix: str, suffix: str): prompt = prompt.strip('"') if prompt.startswith('"') and prompt.endswith('"') else prompt all_tokens = [] for clean_prompt_sub, matched in split_quotation(prompt): if matched: for sub_word in clean_prompt_sub: tokens = self.tokenizer(sub_word, add_special_tokens=False)['input_ids'] all_tokens.extend(tokens) else: tokens = self.tokenizer(clean_prompt_sub, add_special_tokens=False)['input_ids'] all_tokens.extend(tokens) all_tokens = all_tokens[: self.max_tokenizer_len] text_tokens_and_mask = self.tokenizer.pad( {'input_ids': [all_tokens]}, max_length=self.max_tokenizer_len, padding='max_length', return_attention_mask=True, return_tensors='pt', ) prefix_tokens = self.tokenizer(prefix, add_special_tokens=False)['input_ids'] suffix_tokens = self.tokenizer(suffix, add_special_tokens=False)['input_ids'] prefix_tokens_mask = torch.tensor( [1] * len(prefix_tokens), dtype=text_tokens_and_mask.attention_mask[0].dtype ) suffix_tokens_mask = torch.tensor( [1] * len(suffix_tokens), dtype=text_tokens_and_mask.attention_mask[0].dtype ) prefix_tokens = torch.tensor(prefix_tokens, dtype=text_tokens_and_mask.input_ids.dtype) suffix_tokens = torch.tensor(suffix_tokens, dtype=text_tokens_and_mask.input_ids.dtype) input_ids = torch.cat((prefix_tokens, text_tokens_and_mask.input_ids[0], suffix_tokens), dim=-1) attention_mask = torch.cat( (prefix_tokens_mask, text_tokens_and_mask.attention_mask[0], suffix_tokens_mask), dim=-1 ) return input_ids.unsqueeze(0), attention_mask.unsqueeze(0) @torch.inference_mode() def encode_text(self, prompts: Union[str, List[str]], rewrite: bool = False): if isinstance(prompts, str): prompts = [prompts] prompt = prompts[0] if rewrite: prompt = self.rewrite_prompt(prompt) input_ids, attention_mask = self._encode_tokens( prompt, self.prompt_template_encode_prefix, self.prompt_template_encode_suffix, ) input_ids = input_ids.to(self.device) attention_mask = attention_mask.to(self.device) text_output = self.text_encoder( input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True, ) prompt_embeds = text_output.hidden_states[-1].detach() prompt_embeds = prompt_embeds[ :, self.prompt_template_encode_start_idx: -self.prompt_template_encode_end_idx, : ] text_ids = prepare_pos_ids( modality_id=0, type='text', start=(0, 0), num_token=prompt_embeds.shape[1], ).to(self.device, dtype=torch.float64) return prompt_embeds, text_ids @torch.inference_mode() def encode_edit(self, images: List[Image.Image], prompts: Union[str, List[str]]): if self.image_processor_vl is None: raise ValueError("Loaded LongCat CLIP does not contain a vision-language processor; please use the edit checkpoint.") if isinstance(prompts, str): prompts = [prompts] prompt = prompts[0] raw_vl_input = self.image_processor_vl(images=images, return_tensors="pt") pixel_values = raw_vl_input['pixel_values'] image_grid_thw = raw_vl_input['image_grid_thw'] prompt = prompt.strip('"') if prompt.startswith('"') and prompt.endswith('"') else prompt text = self.edit_prompt_template_encode_prefix num_images = len(images) if num_images > 1: single_image_placeholder = "<|vision_start|><|image_pad|><|vision_end|>" if single_image_placeholder in text: multi_image_placeholder = single_image_placeholder * num_images text = text.replace(single_image_placeholder, multi_image_placeholder) merge_length = self.image_processor_vl.merge_size ** 2 image_idx = 0 while self.image_token in text: if image_idx < len(image_grid_thw): grid = image_grid_thw[image_idx] num_image_tokens = grid.prod() // merge_length text = text.replace(self.image_token, "<|placeholder|>" * int(num_image_tokens), 1) image_idx += 1 else: break text = text.replace("<|placeholder|>", self.image_token) input_ids, attention_mask = self._encode_tokens( prompt, text, self.edit_prompt_template_encode_suffix, ) pixel_values = pixel_values.to(self.device) image_grid_thw = image_grid_thw.to(self.device) input_ids = input_ids.to(self.device) attention_mask = attention_mask.to(self.device) text_output = self.text_encoder( input_ids=input_ids, attention_mask=attention_mask, pixel_values=pixel_values, image_grid_thw=image_grid_thw, output_hidden_states=True, ) prompt_embeds = text_output.hidden_states[-1].detach() prompt_embeds = prompt_embeds[ :, self.edit_prompt_template_encode_start_idx: -self.edit_prompt_template_encode_end_idx, : ] text_ids = prepare_pos_ids( modality_id=0, type='text', start=(0, 0), num_token=prompt_embeds.shape[1], ).to(self.device, dtype=torch.float64) return prompt_embeds, text_ids class LongCatVAEWrapper: def __init__(self, vae: AutoencoderKL): self.vae = vae self.scaling_factor = getattr(vae.config, "scaling_factor", 1.0) self.shift_factor = getattr(vae.config, "shift_factor", 0.0) self.latent_channels = getattr(vae.config, "latent_channels", 16) self.vae_scale_factor = 2 ** (len(getattr(vae.config, "block_out_channels", [1])) - 1) def to(self, device, dtype=None): target_dtype = dtype if dtype is not None else None if target_dtype is None: self.vae = self.vae.to(device) else: self.vae = self.vae.to(device=device, dtype=target_dtype) return self def encode(self, image_tensor: torch.Tensor) -> torch.Tensor: # Expect image_tensor shape [B, H, W, C] in range 0..1 pixel_values = image_tensor.permute(0, 3, 1, 2) posterior = self.vae.encode(pixel_values).latent_dist latents = posterior.mode() latents = (latents - self.shift_factor) * self.scaling_factor return latents def decode(self, latents: torch.Tensor) -> torch.Tensor: latents = (latents / self.scaling_factor) + self.shift_factor decoded = self.vae.decode(latents, return_dict=False)[0] decoded = decoded.permute(0, 2, 3, 1) return decoded def get_dtype(self): return self.vae.dtype class LongCatCLIPLoader: @classmethod def INPUT_TYPES(cls): return { "required": { "model_name": (folder_paths.get_filename_list("text_encoders"),), "dtype": (["bf16", "fp16", "fp32"], {"default": "bf16"}), "device": (list_available_devices(),), } } RETURN_TYPES = ("CLIP",) RETURN_NAMES = ("clip",) FUNCTION = "load_clip" CATEGORY = "LongCat" def load_clip(self, model_name: str, dtype: str, device: str): model_path = folder_paths.get_full_path("text_encoders", model_name) base_dir = model_path if os.path.isdir(model_path) else os.path.dirname(model_path) torch_dtype = DTYPE_MAP.get(dtype, torch.bfloat16) target_device = _resolve_device(device) tokenizer_dir = BUNDLED_TOKENIZER_DIR tokenizer = AutoTokenizer.from_pretrained( tokenizer_dir, trust_remote_code=True ) text_processor = AutoProcessor.from_pretrained( tokenizer_dir, trust_remote_code=True ) image_encoder = None feature_extractor = None if os.path.isdir(model_path): text_encoder = AutoModelForCausalLM.from_pretrained( model_path, subfolder="text_encoder", torch_dtype=torch_dtype, trust_remote_code=True ).to(target_device) # Optional vision encoder if present vision_dir = os.path.join(model_path, "image_encoder") if os.path.isdir(vision_dir): try: image_encoder = CLIPVisionModelWithProjection.from_pretrained( model_path, subfolder="image_encoder", torch_dtype=torch_dtype, trust_remote_code=True ).to(target_device) feature_extractor = CLIPImageProcessor.from_pretrained( model_path, subfolder="feature_extractor", trust_remote_code=True ) except Exception: image_encoder = None feature_extractor = None else: # Single-file weights; use bundled config/tokenizer config_dict = _load_config_from_metadata(model_path) if config_dict is not None: text_config = AutoConfig.from_dict(config_dict) else: text_config = AutoConfig.from_pretrained(BUNDLED_TEXT_ENCODER_DIR, trust_remote_code=True) text_encoder = AutoModelForCausalLM.from_config(text_config, trust_remote_code=True).to( target_device, dtype=torch_dtype ) state_dict = load_safetensors(model_path, device="cpu") missing, unexpected = text_encoder.load_state_dict(state_dict, strict=False) if missing or unexpected: print(f"Text encoder load: missing={len(missing)}, unexpected={len(unexpected)}") clip = LongCatCLIPWrapper( text_encoder=text_encoder, tokenizer=tokenizer, text_processor=text_processor, image_encoder=image_encoder, feature_extractor=feature_extractor, ) return (clip,) class LongCatVAELoader: @classmethod def INPUT_TYPES(cls): return { "required": { "model_name": (folder_paths.get_filename_list("vae"),), "dtype": (["bf16", "fp16", "fp32"], {"default": "bf16"}), "device": (list_available_devices(),), } } RETURN_TYPES = ("VAE",) RETURN_NAMES = ("vae",) FUNCTION = "load_vae" CATEGORY = "LongCat" def load_vae(self, model_name: str, dtype: str, device: str): model_path = folder_paths.get_full_path("vae", model_name) base_dir = model_path if os.path.isdir(model_path) else os.path.dirname(model_path) torch_dtype = DTYPE_MAP.get(dtype, torch.bfloat16) target_device = _resolve_device(device) if os.path.isdir(model_path): vae = AutoencoderKL.from_pretrained( model_path, subfolder="vae", torch_dtype=torch_dtype ).to(target_device) else: # Single-file safetensors config_dict = _load_config_from_metadata(model_path) if config_dict is None: config_dict = _load_bundled_config(os.path.join("vae", "config.json")) if config_dict is None: raise ValueError("Bundled VAE config not found.") vae = AutoencoderKL.from_config(config_dict).to(target_device, dtype=torch_dtype) state_dict = load_safetensors(model_path, device="cpu") missing, unexpected = vae.load_state_dict(state_dict, strict=False) if missing or unexpected: print(f"VAE load: missing={len(missing)}, unexpected={len(unexpected)}") wrapper = LongCatVAEWrapper(vae) return (wrapper,) class TextEncodeLongCatImage: @classmethod def INPUT_TYPES(s): return { "required": { "clip": ("CLIP",), "prompt": ("STRING", {"multiline": True, "default": ""}), "rewrite_prompt": ("BOOLEAN", {"default": True}), } } RETURN_TYPES = ("CONDITIONING",) RETURN_NAMES = ("conditioning",) FUNCTION = "encode" CATEGORY = "LongCat" def encode(self, clip: LongCatCLIPWrapper, prompt: str, rewrite_prompt: bool = True): prompt_embeds, text_ids = clip.encode_text(prompt, rewrite=rewrite_prompt) conditioning = { "prompt_embeds": prompt_embeds, "text_ids": text_ids, } return ([[prompt_embeds, conditioning]],) class TextEncodeLongCatImageEdit: @classmethod def INPUT_TYPES(s): return { "required": { "clip": ("CLIP",), "prompt": ("STRING", {"multiline": True, "default": ""}), }, "optional": { "vae": ("VAE",), "image1": ("IMAGE",), "image2": ("IMAGE",), "image3": ("IMAGE",), } } RETURN_TYPES = ("CONDITIONING",) RETURN_NAMES = ("conditioning",) FUNCTION = "encode" CATEGORY = "LongCat" def encode(self, clip: LongCatCLIPWrapper, prompt, vae=None, image1=None, image2=None, image3=None): images = [] if image1 is not None: images.append(image1) if image2 is not None: images.append(image2) if image3 is not None: images.append(image3) if images and vae is None: raise ValueError("VAE is required when images are provided.") if images: pil_images: List[Image.Image] = [] for img in images: pil_images.extend(tensor2pil(img)) prompt_embeds, text_ids = clip.encode_edit(pil_images, prompt) else: prompt_embeds, text_ids = clip.encode_text(prompt, rewrite=False) conditioning = { "prompt_embeds": prompt_embeds, "text_ids": text_ids, "images": images if len(images) > 0 else None, "vae": vae, } return ([[prompt_embeds, conditioning]],) class LoadLongCatModel: @classmethod def INPUT_TYPES(s): return { "required": { "model_path": ("STRING", {"default": ""}), "dtype": (["bf16", "fp16", "fp32"], {"default": "bf16"}), "device": (list_available_devices(),), "gpu_only": ("BOOLEAN", {"default": False}), } } RETURN_TYPES = ("MODEL",) RETURN_NAMES = ("model",) FUNCTION = "load_model" CATEGORY = "LongCat" def load_model(self, model_path: str, dtype: str, device: str, gpu_only: bool): if not model_path: raise ValueError("Please provide a LongCat checkpoint directory or transformer weight file.") torch_dtype = DTYPE_MAP.get(dtype, torch.bfloat16) target_device = _resolve_device(device) scheduler = None if os.path.isdir(model_path): transformer = LongCatImageTransformer2DModel.from_pretrained( model_path, subfolder="transformer", torch_dtype=torch_dtype, use_safetensors=True, ).to(target_device) try: scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( model_path, subfolder="scheduler" ) except Exception: scheduler = FlowMatchEulerDiscreteScheduler( num_train_timesteps=1000, shift=3.0, use_dynamic_shifting=True, base_shift=0.5, max_shift=1.15, base_image_seq_len=256, max_image_seq_len=4096, ) else: config = { "patch_size": 1, "in_channels": 64, "num_layers": 10, "num_single_layers": 20, "attention_head_dim": 128, "num_attention_heads": 24, "joint_attention_dim": 3584, "pooled_projection_dim": 3584, "axes_dims_rope": [16, 56, 56], } transformer = LongCatImageTransformer2DModel(**config).to(target_device) sd = comfy.utils.load_torch_file(model_path) missing, unexpected = transformer.load_state_dict(sd, strict=False) if len(missing) > 0 or len(unexpected) > 0: print( f"LongCat Model Load: Missing {len(missing)} keys, Unexpected {len(unexpected)} keys" ) scheduler = FlowMatchEulerDiscreteScheduler( num_train_timesteps=1000, shift=3.0, use_dynamic_shifting=True, base_shift=0.5, max_shift=1.15, base_image_seq_len=256, max_image_seq_len=4096, ) class LongCatInnerModel: def __init__(self, diffusion_model, scheduler): self.diffusion_model = diffusion_model self.latent_format = None self.manual_cast_dtype = None self.longcat_scheduler = scheduler def get_dtype(self): return self.diffusion_model.dtype def is_adm(self): return False def to(self, device): self.diffusion_model.to(device) return self def state_dict(self, *args, **kwargs): return self.diffusion_model.state_dict(*args, **kwargs) def load_state_dict(self, *args, **kwargs): return self.diffusion_model.load_state_dict(*args, **kwargs) def parameters(self, *args, **kwargs): return self.diffusion_model.parameters(*args, **kwargs) def modules(self, *args, **kwargs): return self.diffusion_model.modules(*args, **kwargs) def named_modules(self, *args, **kwargs): return self.diffusion_model.named_modules(*args, **kwargs) inner_model = LongCatInnerModel(transformer, scheduler) load_device = target_device if gpu_only: offload_device = load_device else: offload_device = comfy.model_management.unet_offload_device() patcher = comfy.model_patcher.ModelPatcher( inner_model, load_device=load_device, offload_device=offload_device ) patcher.longcat_scheduler = scheduler return (patcher,) class LongCatSizePicker: @classmethod def INPUT_TYPES(s): # Flatten the dictionaries to a list of strings "resolution (aspect_ratio)" # or just use the keys/values. # The user said "size: a list of avaliable size(from data_utils.py)" # Let's combine all sizes sizes = [] for name, mapping in [("1024", MULTI_ASPECT_RATIO_1024), ("512", MULTI_ASPECT_RATIO_512), ("256", MULTI_ASPECT_RATIO_256)]: for ratio, (h, w) in mapping.items(): sizes.append(f"{name} - {ratio} ({int(h)}x{int(w)})") return { "required": { "size": (sizes,), } } RETURN_TYPES = ("LATENT", "INT", "INT") RETURN_NAMES = ("latent", "width", "height") FUNCTION = "pick_size" CATEGORY = "LongCat" def pick_size(self, size): # Format: "1024 - 1.0 (1024x1024)" try: part = size.split("(")[1].split(")")[0] h_str, w_str = part.split("x") height = int(h_str) width = int(w_str) except: height = 1024 width = 1024 # Create empty latent # Standard ComfyUI latent is [1, 4, H//8, W//8] latent = torch.zeros([1, 4, height // 8, width // 8]) return ({"samples": latent}, width, height) class LongCatImageResizer: @classmethod def INPUT_TYPES(s): return { "required": { "image": ("IMAGE",), "method": (["keep proportion", "stretch", "fill/crop", "pad"],), "interpolation": (["nearest", "bilinear", "bicubic", "lanczos"],), } } RETURN_TYPES = ("IMAGE", "INT", "INT") RETURN_NAMES = ("image", "width", "height") FUNCTION = "resize" CATEGORY = "LongCat" def resize(self, image, method, interpolation): # image is [B, H, W, C] pil_images = tensor2pil(image) out_images = [] # Determine interpolation interp_map = { "nearest": Image.NEAREST, "bilinear": Image.BILINEAR, "bicubic": Image.BICUBIC, "lanczos": Image.LANCZOS } interp = interp_map.get(interpolation, Image.BICUBIC) final_w, final_h = 0, 0 for img in pil_images: w, h = img.size # Find nearest supported resolution all_sizes = [] for mapping in [MULTI_ASPECT_RATIO_1024, MULTI_ASPECT_RATIO_512, MULTI_ASPECT_RATIO_256]: for _, (th, tw) in mapping.items(): all_sizes.append((int(tw), int(th))) # Find nearest size best_size = None min_dist = float('inf') for tw, th in all_sizes: dist = (tw - w)**2 + (th - h)**2 if dist < min_dist: min_dist = dist best_size = (tw, th) target_w, target_h = best_size final_w, final_h = target_w, target_h if method == "stretch": resized = img.resize((target_w, target_h), interp) out_images.append(resized) elif method == "keep proportion": # Find size with closest aspect ratio best_ar_size = None min_ar_diff = float('inf') current_ar = w / h for tw, th in all_sizes: ar = tw / th diff = abs(ar - current_ar) if diff < min_ar_diff: min_ar_diff = diff best_ar_size = (tw, th) target_w, target_h = best_ar_size final_w, final_h = target_w, target_h resized = img.resize((target_w, target_h), interp) out_images.append(resized) elif method == "pad": ratio = min(target_w/w, target_h/h) new_w = int(w * ratio) new_h = int(h * ratio) resized = img.resize((new_w, new_h), interp) new_img = Image.new("RGB", (target_w, target_h), (0, 0, 0)) new_img.paste(resized, ((target_w - new_w)//2, (target_h - new_h)//2)) out_images.append(new_img) elif method == "fill/crop": ratio = max(target_w/w, target_h/h) new_w = int(w * ratio) new_h = int(h * ratio) resized = img.resize((new_w, new_h), interp) left = (new_w - target_w)//2 top = (new_h - target_h)//2 cropped = resized.crop((left, top, left + target_w, top + target_h)) out_images.append(cropped) return (pil2tensor(out_images), final_w, final_h) class LongCatSampler: @classmethod def INPUT_TYPES(s): return { "required": { "model": ("MODEL",), "positive": ("CONDITIONING",), "negative": ("CONDITIONING",), "latent_image": ("LATENT",), "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), "steps": ("INT", {"default": 50, "min": 1, "max": 1000}), "cfg": ("FLOAT", {"default": 4.5, "min": 0.0, "max": 20.0, "step": 0.1}), "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ), "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ), "cfg_norm": ("BOOLEAN", {"default": True}), "cfg_renorm_min": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01}), "control_after_generate": (["fixed", "increment", "decrement", "random"], {"default": "fixed"}), } } RETURN_TYPES = ("LATENT",) RETURN_NAMES = ("latent",) FUNCTION = "sample" CATEGORY = "LongCat" def sample(self, model, positive, negative, latent_image, seed, steps, cfg, sampler_name, scheduler, cfg_norm, cfg_renorm_min, control_after_generate): comfy.model_management.load_model_gpu(model) transformer = model.model.diffusion_model device = comfy.model_management.get_torch_device() sched = getattr(model, "longcat_scheduler", None) if sched is None and hasattr(model, "model"): sched = getattr(model.model, "longcat_scheduler", None) if sched is None: sched = FlowMatchEulerDiscreteScheduler( num_train_timesteps=1000, shift=3.0, use_dynamic_shifting=True, base_shift=0.5, max_shift=1.15, base_image_seq_len=256, max_image_seq_len=4096, ) samples = latent_image["samples"] batch_size, _, h, w = samples.shape height = h * 8 width = w * 8 generator = torch.Generator(device=device).manual_seed(seed) def _pack_latents(latents, batch, num_channels_latents, h_px, w_px): latents = latents.view(batch, num_channels_latents, h_px // 2, 2, w_px // 2, 2) latents = latents.permute(0, 2, 4, 1, 3, 5) latents = latents.reshape(batch, (h_px // 2) * (w_px // 2), num_channels_latents * 4) return latents def _unpack_latents(latents, h_px, w_px, vae_scale_factor): batch, num_patches, channels = latents.shape h_px = 2 * (int(h_px) // (vae_scale_factor * 2)) w_px = 2 * (int(w_px) // (vae_scale_factor * 2)) latents = latents.view(batch, h_px // 2, w_px // 2, channels // 4, 2, 2) latents = latents.permute(0, 3, 1, 4, 2, 5) latents = latents.reshape(batch, channels // (2 * 2), h_px, w_px) return latents # Unpack conditioning pos_cond = positive[0][0] pos_dict = positive[0][1] neg_cond = negative[0][0] neg_dict = negative[0][1] prompt_embeds = pos_cond.to(device, dtype=transformer.dtype) neg_prompt_embeds = neg_cond.to(device, dtype=transformer.dtype) text_ids = pos_dict["text_ids"].to(device, dtype=torch.float64) neg_text_ids = neg_dict.get("text_ids", text_ids).to(device, dtype=torch.float64) text_len = prompt_embeds.shape[1] neg_text_len = neg_prompt_embeds.shape[1] target_len = max(text_len, neg_text_len) if prompt_embeds.shape[1] != target_len: pad = torch.zeros( prompt_embeds.shape[0], target_len - prompt_embeds.shape[1], prompt_embeds.shape[2], device=device, dtype=prompt_embeds.dtype, ) prompt_embeds = torch.cat([prompt_embeds, pad], dim=1) if neg_prompt_embeds.shape[1] != target_len: pad = torch.zeros( neg_prompt_embeds.shape[0], target_len - neg_prompt_embeds.shape[1], neg_prompt_embeds.shape[2], device=device, dtype=neg_prompt_embeds.dtype, ) neg_prompt_embeds = torch.cat([neg_prompt_embeds, pad], dim=1) images = pos_dict.get("images") or neg_dict.get("images") vae = pos_dict.get("vae") or neg_dict.get("vae") if images is not None and not isinstance(vae, LongCatVAEWrapper): raise ValueError("LongCat VAE Loader output is required when conditioning on images.") num_channels_latents = 16 vae_scale_factor = 8 if isinstance(vae, LongCatVAEWrapper): num_channels_latents = vae.latent_channels vae_scale_factor = vae.vae_scale_factor latent_height = 2 * (int(height) // (vae_scale_factor * 2)) latent_width = 2 * (int(width) // (vae_scale_factor * 2)) latents = torch.randn( (batch_size, num_channels_latents, latent_height, latent_width), generator=generator, device=device, dtype=transformer.dtype, ) latents = _pack_latents(latents, batch_size, num_channels_latents, latent_height, latent_width) start_pos = max(target_len, 512) latent_image_ids = prepare_pos_ids( modality_id=1, type='image', start=(start_pos, start_pos), height=latent_height // 2, width=latent_width // 2, ).to(device, dtype=torch.float64) image_latents = None image_latents_ids = None if images is not None: stacked_images = torch.cat(images, dim=0) if len(images) > 1 else images[0] stacked_images = stacked_images.to(device=device, dtype=transformer.dtype) stacked_images_chw = stacked_images.permute(0, 3, 1, 2) stacked_images_chw = torch.nn.functional.interpolate( stacked_images_chw, size=(height, width), mode="bicubic", align_corners=False ) stacked_images = stacked_images_chw.permute(0, 2, 3, 1) encoded = vae.encode(stacked_images) encoded = encoded.to(device=device, dtype=transformer.dtype) image_latents = _pack_latents( encoded, encoded.shape[0], num_channels_latents, latent_height, latent_width ) image_latents_ids = prepare_pos_ids( modality_id=2, type="image", start=(target_len, target_len), height=latent_height // 2, width=latent_width // 2, ).to(device, dtype=torch.float64) if cfg > 1.0: prompt_embeds = torch.cat([neg_prompt_embeds, prompt_embeds], dim=0) if image_latents is not None: latent_image_ids = torch.cat([latent_image_ids, image_latents_ids], dim=0) sigmas = np.linspace(1.0, 1.0 / steps, steps) image_seq_len = latents.shape[1] mu = calculate_shift( image_seq_len, sched.config.get("base_image_seq_len", 256) if hasattr(sched, "config") else 256, sched.config.get("max_image_seq_len", 4096) if hasattr(sched, "config") else 4096, sched.config.get("base_shift", 0.5) if hasattr(sched, "config") else 0.5, sched.config.get("max_shift", 1.15) if hasattr(sched, "config") else 1.15, ) timesteps, num_inference_steps = retrieve_timesteps( sched, steps, device, sigmas=sigmas, mu=mu, ) with torch.inference_mode(): for _, t in enumerate(timesteps): latent_model_input = latents if image_latents is not None: latent_model_input = torch.cat([latents, image_latents], dim=1) latent_model_input = torch.cat([latent_model_input] * 2) if cfg > 1.0 else latent_model_input timestep = t.expand(latent_model_input.shape[0]).to(latents.dtype) noise_pred = transformer( hidden_states=latent_model_input, timestep=timestep / 1000, guidance=None, encoder_hidden_states=prompt_embeds, txt_ids=text_ids, img_ids=latent_image_ids, return_dict=False, )[0] if image_latents is not None: noise_pred = noise_pred[:, :image_seq_len] if cfg > 1.0: noise_pred_uncond, noise_pred_text = noise_pred.chunk(2, dim=0) noise_pred = noise_pred_uncond + cfg * (noise_pred_text - noise_pred_uncond) if cfg_norm: cond_norm = torch.norm(noise_pred_text, dim=-1, keepdim=True) noise_norm = torch.norm(noise_pred, dim=-1, keepdim=True) scale = (cond_norm / (noise_norm + 1e-8)).clamp(min=cfg_renorm_min, max=1.0) noise_pred = noise_pred * scale latents = sched.step(noise_pred, t, latents, return_dict=False)[0] unpacked = _unpack_latents(latents, latent_height, latent_width, vae_scale_factor) return ({"samples": unpacked},) NODE_CLASS_MAPPINGS = { "LongCatCLIPLoader": LongCatCLIPLoader, "LongCatVAELoader": LongCatVAELoader, "TextEncodeLongCatImage": TextEncodeLongCatImage, "TextEncodeLongCatImageEdit": TextEncodeLongCatImageEdit, "LongCatSizePicker": LongCatSizePicker, "LongCatImageResizer": LongCatImageResizer, "LongCatSampler": LongCatSampler, "LoadLongCatModel": LoadLongCatModel, } NODE_DISPLAY_NAME_MAPPINGS = { "LongCatCLIPLoader": "LongCat CLIP Loader", "LongCatVAELoader": "LongCat VAE Loader", "TextEncodeLongCatImage": "Text Encode LongCat Image", "TextEncodeLongCatImageEdit": "Text Encode LongCat Image Edit", "LongCatSizePicker": "LongCat Size Picker", "LongCatImageResizer": "LongCat Image Resizer", "LongCatSampler": "LongCat Sampler", "LoadLongCatModel": "Load LongCat Model", }