Mister-Link
466e2b5734
Reproduces app.py's TeleStyleV2 style-transfer pipeline as native ComfyUI nodes (Qwen TeleStyle Loader + Qwen TeleStyle Apply), loading weights directly from ComfyUI's own model folders instead of bridging live MODEL/CLIP/VAE objects (which doubled memory and OOM'd on this box).
356 lines
14 KiB
Python
356 lines
14 KiB
Python
"""Qwen TeleStyle Loader: given just an architecture choice (2509/2511) and a device, resolves
|
|
(downloading if missing) the transformer/text_encoder/vae/lora weight files from ComfyUI's
|
|
own model folders, builds the diffusers pipeline once, and hands it to Qwen TeleStyle Apply
|
|
as a single TELESTYLE_MODELS socket - so Apply itself only needs to know about *generation*,
|
|
never about where any file lives.
|
|
|
|
Loading is a single, direct read of each flat safetensors file straight into this pipeline's
|
|
own modules - no ComfyUI MODEL/CLIP/VAE objects are involved, so nothing is ever resident in
|
|
memory twice (see the module docstring in the old comfy_bridge-based approach for why
|
|
bridging *live* comfy objects doubles memory and OOMs on boxes with much less system RAM
|
|
than VRAM, like this one: 62GB RAM vs 97GB VRAM).
|
|
"""
|
|
|
|
import gc
|
|
import json
|
|
import os
|
|
import threading
|
|
|
|
import torch
|
|
from safetensors.torch import load_file as load_safetensors, save_file as save_safetensors
|
|
|
|
import folder_paths
|
|
from transformers import AutoConfig, Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer, Qwen2VLProcessor
|
|
from diffusers.models import AutoencoderKLQwenImage, QwenImageTransformer2DModel
|
|
from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
|
|
|
|
try:
|
|
from . import vae_convert
|
|
except ImportError:
|
|
import vae_convert
|
|
|
|
try:
|
|
from .pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
|
|
except ImportError:
|
|
from pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
|
|
|
|
PIPELINE_DTYPE = torch.bfloat16
|
|
|
|
# text encoder + vae are shared across architectures (same base Qwen-Image-Edit family);
|
|
# only the transformer checkpoint and the two TeleStyleV2 LoRAs are architecture-specific.
|
|
TEXT_ENCODER_FILE = "qwen_2.5_vl_7b.safetensors"
|
|
VAE_FILE = "qwen_image_vae.safetensors"
|
|
|
|
# HF sources used only to auto-download a file ComfyUI doesn't have yet - once downloaded
|
|
# it lives permanently under ComfyUI's own model folders like any other checkpoint.
|
|
TEXT_ENCODER_REPO = "Comfy-Org/HunyuanVideo_1.5_repackaged"
|
|
TEXT_ENCODER_REPO_FILE = "split_files/text_encoders/qwen_2.5_vl_7b.safetensors"
|
|
VAE_REPO = "f5aiteam/VAE"
|
|
VAE_REPO_FILE = "qwen_image_vae.safetensors"
|
|
LORA_REPO = "Tele-AI/TeleStyleV2"
|
|
TRANSFORMER_SUBFOLDER = "transformer"
|
|
|
|
ARCHITECTURES = {
|
|
"2509": {
|
|
"transformer_file": "Qwen-Image-Edit-2509.safetensors",
|
|
# pre-merged single file (this pack's own repo) - no shard download/merge needed.
|
|
"transformer_single_file_repo": "troopa/TeleStyleV2-Qwen-Image-Edit-2509",
|
|
# fallback if the repo above is ever unavailable: merge from the original shards.
|
|
"transformer_sharded_repo": "Qwen/Qwen-Image-Edit-2509",
|
|
"style_lora": "diffusers-TeleStyleV2-QIE-2509-Lora-bf16.safetensors",
|
|
"speed_lora": "QIE-2509-Lightning-4steps-V1.0-bf16.safetensors",
|
|
},
|
|
"2511": {
|
|
"transformer_file": "Qwen-Image-Edit-2511.safetensors",
|
|
"transformer_single_file_repo": None,
|
|
"transformer_sharded_repo": "Qwen/Qwen-Image-Edit-2511",
|
|
"style_lora": "diffusers-TeleStyleV2-QIE-2511-Lora-bf16.safetensors",
|
|
"speed_lora": "QIE-2511-Lightning-4steps-V1.0-bf16.safetensors",
|
|
},
|
|
}
|
|
|
|
# config/tokenizer/processor/scheduler configs only (no weights) - bundled with this node,
|
|
# independent of app.py's project.
|
|
ASSET_MODEL_DIR = os.path.join(os.path.dirname(__file__), "assets", "Qwen-Image-Edit-2509")
|
|
|
|
_PIPELINE_CACHE = {"key": None, "pipe": None}
|
|
_PIPELINE_LOCK = threading.Lock()
|
|
|
|
|
|
def _existing_path(folder_name, filename):
|
|
# folder_paths.get_folder_paths() can list several directories for one folder_name
|
|
# (e.g. a legacy "unet" alias ahead of the real "diffusion_models" dir) - check every
|
|
# registered directory, not just the first, before concluding a file is missing.
|
|
for folder in folder_paths.get_folder_paths(folder_name):
|
|
candidate = os.path.join(folder, filename)
|
|
if os.path.isfile(candidate):
|
|
return candidate
|
|
return None
|
|
|
|
|
|
def _download_target_path(folder_name, filename):
|
|
return os.path.join(folder_paths.get_folder_paths(folder_name)[-1], filename)
|
|
|
|
|
|
def resolve_device(device):
|
|
if device == "auto":
|
|
return "cuda" if torch.cuda.is_available() else "cpu"
|
|
return device
|
|
|
|
|
|
def _ensure_single_file(folder_name, filename, repo_id, repo_file):
|
|
existing = _existing_path(folder_name, filename)
|
|
if existing is not None:
|
|
return existing
|
|
|
|
from huggingface_hub import hf_hub_download
|
|
|
|
target_path = _download_target_path(folder_name, filename)
|
|
os.makedirs(os.path.dirname(target_path), exist_ok=True)
|
|
downloaded = hf_hub_download(repo_id=repo_id, filename=repo_file)
|
|
if os.path.abspath(downloaded) != os.path.abspath(target_path):
|
|
os.symlink(downloaded, target_path)
|
|
return target_path
|
|
|
|
|
|
def _ensure_lora(filename, repo_file):
|
|
return _ensure_single_file("loras", filename, LORA_REPO, repo_file)
|
|
|
|
|
|
def _ensure_text_encoder(filename):
|
|
return _ensure_single_file("text_encoders", filename, TEXT_ENCODER_REPO, TEXT_ENCODER_REPO_FILE)
|
|
|
|
|
|
def _ensure_vae(filename):
|
|
return _ensure_single_file("vae", filename, VAE_REPO, VAE_REPO_FILE)
|
|
|
|
|
|
def _ensure_transformer(filename, single_file_repo, sharded_repo):
|
|
existing = _existing_path("diffusion_models", filename)
|
|
if existing is not None:
|
|
return existing
|
|
target_path = _download_target_path("diffusion_models", filename)
|
|
os.makedirs(os.path.dirname(target_path), exist_ok=True)
|
|
|
|
if single_file_repo:
|
|
# pre-merged single file - a plain download, same as any other model file.
|
|
from huggingface_hub import hf_hub_download
|
|
|
|
downloaded = hf_hub_download(repo_id=single_file_repo, filename=filename)
|
|
if os.path.abspath(downloaded) != os.path.abspath(target_path):
|
|
os.symlink(downloaded, target_path)
|
|
return target_path
|
|
|
|
# no pre-merged repo configured - fetch the original sharded release and merge once.
|
|
from huggingface_hub import snapshot_download
|
|
|
|
snapshot_dir = snapshot_download(repo_id=sharded_repo, allow_patterns=[f"{TRANSFORMER_SUBFOLDER}/*"])
|
|
shard_dir = os.path.join(snapshot_dir, TRANSFORMER_SUBFOLDER)
|
|
index_path = os.path.join(shard_dir, "diffusion_pytorch_model.safetensors.index.json")
|
|
with open(index_path) as f:
|
|
weight_map = json.load(f)["weight_map"]
|
|
shard_files = sorted(set(weight_map.values()))
|
|
|
|
# merges all shards into one file, so this needs ~the full transformer's size in RAM at
|
|
# once (~38GB) - only ever runs the first time this file is missing, never during
|
|
# regular inference.
|
|
merged = {}
|
|
for shard_file in shard_files:
|
|
merged.update(load_safetensors(os.path.join(shard_dir, shard_file)))
|
|
|
|
save_safetensors(merged, target_path)
|
|
del merged
|
|
gc.collect()
|
|
return target_path
|
|
|
|
|
|
def clear_pipeline_cache():
|
|
with _PIPELINE_LOCK:
|
|
pipe = _PIPELINE_CACHE.get("pipe")
|
|
_PIPELINE_CACHE["pipe"] = None
|
|
_PIPELINE_CACHE["key"] = None
|
|
if pipe is not None:
|
|
try:
|
|
pipe.to("cpu")
|
|
except Exception:
|
|
pass
|
|
del pipe
|
|
gc.collect()
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
|
|
def _load_state_dict(path, dtype):
|
|
sd = load_safetensors(path)
|
|
return {k: v.to(dtype=dtype) for k, v in sd.items()}
|
|
|
|
|
|
def _build_pipeline(transformer_path, text_encoder_path, vae_path, style_lora_path, speed_lora_path, device):
|
|
resolved_device = resolve_device(device)
|
|
|
|
vae_dir = os.path.join(ASSET_MODEL_DIR, "vae")
|
|
transformer_dir = os.path.join(ASSET_MODEL_DIR, "transformer")
|
|
text_encoder_dir = os.path.join(ASSET_MODEL_DIR, "text_encoder")
|
|
|
|
with torch.device("meta"):
|
|
vae_config = AutoencoderKLQwenImage.load_config(vae_dir)
|
|
vae_module = AutoencoderKLQwenImage.from_config(vae_config)
|
|
transformer_config = QwenImageTransformer2DModel.load_config(transformer_dir)
|
|
transformer_module = QwenImageTransformer2DModel.from_config(transformer_config)
|
|
text_encoder_config = AutoConfig.from_pretrained(text_encoder_dir)
|
|
|
|
# transformer weights are already in diffusers' own QwenImageTransformer2DModel key
|
|
# layout (this is literally the same file app.py's from_pretrained() reads), so this
|
|
# is a plain single load - no comfy involved, no second copy of anything.
|
|
transformer_sd = _load_state_dict(transformer_path, PIPELINE_DTYPE)
|
|
transformer_module.load_state_dict(transformer_sd, strict=True, assign=True)
|
|
del transformer_sd
|
|
|
|
# transformer_module.pos_embed (QwenEmbedRope/QwenEmbedLayer3DRope) precomputes its
|
|
# rotary frequency tables as plain python attributes in __init__ - deliberately not
|
|
# registered buffers (registering them as buffers loses the complex dtype's imaginary
|
|
# part) - so building under torch.device("meta") leaves them as empty meta tensors
|
|
# forever; load_state_dict can't fix this since they were never part of any checkpoint.
|
|
# Rebuilding just this small (config-only, no big weights) submodule for real replaces
|
|
# them with correctly computed data.
|
|
pos_embed = transformer_module.pos_embed
|
|
transformer_module.pos_embed = type(pos_embed)(
|
|
theta=pos_embed.theta, axes_dim=pos_embed.axes_dim, scale_rope=pos_embed.scale_rope
|
|
)
|
|
|
|
# the vae file on disk is in ComfyUI's own WanVAE key layout (top-level conv1/conv2,
|
|
# encoder.downsamples.N...), not diffusers' AutoencoderKLQwenImage layout - reuse the
|
|
# same structural remapping the old comfy-socket bridge used, just fed a plain state
|
|
# dict here instead of a live comfy VAE object.
|
|
vae_raw_sd = _load_state_dict(vae_path, PIPELINE_DTYPE)
|
|
vae_sd = vae_convert.convert_and_validate(
|
|
vae_raw_sd,
|
|
vae_module,
|
|
dim_mult=vae_config.get("dim_mult", [1, 2, 4, 4]),
|
|
num_res_blocks=vae_config.get("num_res_blocks", 2),
|
|
attn_scales=vae_config.get("attn_scales", []),
|
|
)
|
|
vae_module.load_state_dict(vae_sd, strict=True, assign=True)
|
|
del vae_raw_sd, vae_sd
|
|
|
|
# unlike a comfy CLIP object (conditioning-only, never carries lm_head), this flat file
|
|
# is the same checkpoint app.py loads directly and includes lm_head - so the
|
|
# prompt-enhancer (infer_with_content_prompt/infer_with_style_prompt) works here.
|
|
text_encoder_sd = _load_state_dict(text_encoder_path, PIPELINE_DTYPE)
|
|
text_encoder_module, te_loading_info = Qwen2_5_VLForConditionalGeneration.from_pretrained(
|
|
None,
|
|
config=text_encoder_config,
|
|
state_dict=text_encoder_sd,
|
|
dtype=PIPELINE_DTYPE,
|
|
low_cpu_mem_usage=True,
|
|
output_loading_info=True,
|
|
)
|
|
del text_encoder_sd
|
|
bad_loading = {
|
|
k: v
|
|
for k, v in te_loading_info.items()
|
|
if k in ("missing_keys", "unexpected_keys", "mismatched_keys") and v
|
|
}
|
|
if bad_loading:
|
|
raise RuntimeError(f"text encoder: from_pretrained reported loading problems: {bad_loading}")
|
|
has_lm_head = "lm_head.weight" not in (te_loading_info.get("missing_keys") or [])
|
|
|
|
tokenizer = Qwen2Tokenizer.from_pretrained(os.path.join(ASSET_MODEL_DIR, "tokenizer"))
|
|
processor = Qwen2VLProcessor.from_pretrained(os.path.join(ASSET_MODEL_DIR, "processor"))
|
|
scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(os.path.join(ASSET_MODEL_DIR, "scheduler"))
|
|
|
|
pipe = QwenImageEditPlusPipeline(
|
|
scheduler=scheduler,
|
|
vae=vae_module,
|
|
text_encoder=text_encoder_module,
|
|
tokenizer=tokenizer,
|
|
processor=processor,
|
|
transformer=transformer_module,
|
|
)
|
|
pipe.set_progress_bar_config(disable=True)
|
|
pipe.has_lm_head = has_lm_head
|
|
pipe.resolved_device = resolved_device
|
|
pipe.to(resolved_device)
|
|
|
|
pipe.load_lora_weights(style_lora_path, adapter_name="style")
|
|
pipe.load_lora_weights(speed_lora_path, adapter_name="dmd")
|
|
pipe.set_adapters(["style", "dmd"], adapter_weights=[1.0, 1.0])
|
|
pipe.fuse_lora(adapter_names=["style", "dmd"], lora_scale=1.0)
|
|
pipe.unload_lora_weights()
|
|
return pipe
|
|
|
|
|
|
class QwenTeleStyleLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"architecture": (list(ARCHITECTURES.keys()), {"default": "2509"}),
|
|
"device": (["cuda", "cpu", "auto"], {"default": "cuda" if torch.cuda.is_available() else "cpu"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("TELESTYLE_MODELS",)
|
|
RETURN_NAMES = ("models",)
|
|
FUNCTION = "load"
|
|
CATEGORY = "image/TeleStyleV2"
|
|
|
|
def load(self, architecture, device):
|
|
resolved_device = resolve_device(device)
|
|
arch = ARCHITECTURES[architecture]
|
|
|
|
transformer_path = _ensure_transformer(
|
|
arch["transformer_file"], arch["transformer_single_file_repo"], arch["transformer_sharded_repo"]
|
|
)
|
|
text_encoder_path = _ensure_text_encoder(TEXT_ENCODER_FILE)
|
|
vae_path = _ensure_vae(VAE_FILE)
|
|
style_lora_path = _ensure_lora(arch["style_lora"], arch["style_lora"])
|
|
speed_lora_path = _ensure_lora(arch["speed_lora"], arch["speed_lora"])
|
|
for label, path in (
|
|
("transformer", transformer_path),
|
|
("text_encoder", text_encoder_path),
|
|
("vae", vae_path),
|
|
("style_lora", style_lora_path),
|
|
("speed_lora", speed_lora_path),
|
|
):
|
|
if not path or not os.path.isfile(path):
|
|
raise FileNotFoundError(f"{label}: could not find or download {path}")
|
|
|
|
cache_key = (transformer_path, text_encoder_path, vae_path, style_lora_path, speed_lora_path, resolved_device)
|
|
with _PIPELINE_LOCK:
|
|
if _PIPELINE_CACHE["key"] == cache_key and _PIPELINE_CACHE["pipe"] is not None:
|
|
return (_PIPELINE_CACHE["pipe"],)
|
|
|
|
clear_pipeline_cache()
|
|
pipe = _build_pipeline(transformer_path, text_encoder_path, vae_path, style_lora_path, speed_lora_path, device)
|
|
with _PIPELINE_LOCK:
|
|
_PIPELINE_CACHE["key"] = cache_key
|
|
_PIPELINE_CACHE["pipe"] = pipe
|
|
return (pipe,)
|
|
|
|
|
|
class QwenTeleStyleUnload:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {}}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("status",)
|
|
FUNCTION = "unload"
|
|
CATEGORY = "image/TeleStyleV2"
|
|
|
|
def unload(self):
|
|
clear_pipeline_cache()
|
|
return ("Qwen TeleStyle pipeline cache cleared.",)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Qwen TeleStyle Loader": QwenTeleStyleLoader,
|
|
"Qwen TeleStyle Unload": QwenTeleStyleUnload,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Qwen TeleStyle Loader": "Qwen TeleStyle Loader",
|
|
"Qwen TeleStyle Unload": "Qwen TeleStyle Unload",
|
|
}
|