357 lines
11 KiB
Python
357 lines
11 KiB
Python
"""
|
|
Runtime wrappers for official ComfyUI Qwen image pipelines.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import importlib
|
|
from typing import Any
|
|
|
|
try:
|
|
import folder_paths
|
|
except Exception:
|
|
folder_paths = None
|
|
|
|
try:
|
|
comfy_core_nodes = importlib.import_module("nodes")
|
|
except Exception as exc:
|
|
comfy_core_nodes = None
|
|
_CORE_NODES_ERR = exc
|
|
else:
|
|
_CORE_NODES_ERR = None
|
|
|
|
try:
|
|
nodes_sd3 = importlib.import_module("comfy_extras.nodes_sd3")
|
|
except Exception as exc:
|
|
nodes_sd3 = None
|
|
_SD3_ERR = exc
|
|
else:
|
|
_SD3_ERR = None
|
|
|
|
try:
|
|
nodes_model_advanced = importlib.import_module("comfy_extras.nodes_model_advanced")
|
|
except Exception as exc:
|
|
nodes_model_advanced = None
|
|
_MODEL_ADVANCED_ERR = exc
|
|
else:
|
|
_MODEL_ADVANCED_ERR = None
|
|
|
|
try:
|
|
nodes_qwen = importlib.import_module("comfy_extras.nodes_qwen")
|
|
except Exception as exc:
|
|
nodes_qwen = None
|
|
_QWEN_ERR = exc
|
|
else:
|
|
_QWEN_ERR = None
|
|
|
|
try:
|
|
nodes_flux = importlib.import_module("comfy_extras.nodes_flux")
|
|
except Exception as exc:
|
|
nodes_flux = None
|
|
_FLUX_ERR = exc
|
|
else:
|
|
_FLUX_ERR = None
|
|
|
|
try:
|
|
nodes_cfg = importlib.import_module("comfy_extras.nodes_cfg")
|
|
except Exception as exc:
|
|
nodes_cfg = None
|
|
_CFG_ERR = exc
|
|
else:
|
|
_CFG_ERR = None
|
|
|
|
from . import discovery
|
|
|
|
|
|
_QWEN_CLIP_TYPE = "qwen_image"
|
|
|
|
|
|
def _unwrap_first(result: Any):
|
|
if hasattr(result, "result"):
|
|
values = result.result
|
|
if values is None:
|
|
return None
|
|
return values[0]
|
|
if isinstance(result, tuple):
|
|
return result[0]
|
|
if isinstance(result, list):
|
|
return result[0]
|
|
return result
|
|
|
|
|
|
def _require_module(module: Any, label: str, import_error: Exception | None = None):
|
|
if module is None:
|
|
raise RuntimeError(
|
|
f"[LLS] Required ComfyUI runtime component '{label}' is unavailable. "
|
|
"Use a ComfyUI build that includes official Qwen image support."
|
|
) from import_error
|
|
return module
|
|
|
|
|
|
def _require_class(module: Any, class_name: str, import_error: Exception | None = None):
|
|
_require_module(module, class_name, import_error)
|
|
cls = getattr(module, class_name, None)
|
|
if cls is None:
|
|
raise RuntimeError(
|
|
f"[LLS] Required ComfyUI runtime component '{class_name}' is unavailable. "
|
|
"Use a ComfyUI build that includes official Qwen image support."
|
|
) from import_error
|
|
return cls
|
|
|
|
|
|
def _load_qwen_clip(clip_name: str):
|
|
clip_loader_cls = _require_class(comfy_core_nodes, "CLIPLoader", _CORE_NODES_ERR)
|
|
loader = clip_loader_cls()
|
|
try:
|
|
return loader.load_clip(clip_name, type=_QWEN_CLIP_TYPE, device="default")[0]
|
|
except TypeError:
|
|
return loader.load_clip(clip_name, type=_QWEN_CLIP_TYPE)[0]
|
|
|
|
|
|
def _load_qwen_vae(vae_name: str):
|
|
vae_loader_cls = _require_class(comfy_core_nodes, "VAELoader", _CORE_NODES_ERR)
|
|
return vae_loader_cls().load_vae(vae_name)[0]
|
|
|
|
|
|
def _load_qwen_model(model_name: str):
|
|
unet_loader_cls = _require_class(comfy_core_nodes, "UNETLoader", _CORE_NODES_ERR)
|
|
return unet_loader_cls().load_unet(model_name, "default")[0]
|
|
|
|
|
|
def _load_model_only_lora(model, lora_name: str, strength_model: float):
|
|
lora_loader_cls = _require_class(comfy_core_nodes, "LoraLoaderModelOnly", _CORE_NODES_ERR)
|
|
return lora_loader_cls().load_lora_model_only(model, lora_name, float(strength_model))[0]
|
|
|
|
|
|
def _resolve_generation_model(model_name: str, model):
|
|
if model is not None:
|
|
return model
|
|
return _load_qwen_model(model_name)
|
|
|
|
|
|
def _prepare_turbo_request(resolver, model_name: str, steps: int, cfg: float, enable_turbo_mode: bool, turbo_lora_name: str):
|
|
if not enable_turbo_mode:
|
|
return None, int(steps), float(cfg)
|
|
|
|
resolved_lora = resolver(model_name, turbo_lora_name)
|
|
profile = discovery.get_qwen_turbo_profile(model_name)
|
|
if profile is None:
|
|
raise RuntimeError(f"[LLS] No supported turbo preset exists for '{model_name}'.")
|
|
|
|
return resolved_lora, int(profile["steps"]), float(profile["cfg"])
|
|
|
|
|
|
def _patch_qwen_model_sampling(model, shift: float):
|
|
patcher_cls = _require_class(nodes_model_advanced, "ModelSamplingAuraFlow", _MODEL_ADVANCED_ERR)
|
|
return patcher_cls().patch_aura(model, float(shift))[0]
|
|
|
|
|
|
def _apply_cfg_norm(model, strength: float):
|
|
cfg_cls = _require_class(nodes_cfg, "CFGNorm", _CFG_ERR)
|
|
return _unwrap_first(cfg_cls.execute(model, float(strength)))
|
|
|
|
|
|
def _encode_clip_text(clip, prompt: str):
|
|
encoder_cls = _require_class(comfy_core_nodes, "CLIPTextEncode", _CORE_NODES_ERR)
|
|
return encoder_cls().encode(clip, prompt)[0]
|
|
|
|
|
|
def _create_empty_qwen_latent(width: int, height: int, batch_size: int):
|
|
latent_cls = _require_class(nodes_sd3, "EmptySD3LatentImage", _SD3_ERR)
|
|
latent_node = latent_cls()
|
|
if hasattr(latent_node, "generate"):
|
|
return latent_node.generate(width, height, batch_size)[0]
|
|
return latent_node.execute(width, height, batch_size)[0]
|
|
|
|
|
|
def _encode_qwen_edit_conditioning(clip, prompt: str, vae, image1, image2=None, image3=None):
|
|
encoder_cls = _require_class(nodes_qwen, "TextEncodeQwenImageEditPlus", _QWEN_ERR)
|
|
return _unwrap_first(
|
|
encoder_cls.execute(
|
|
clip,
|
|
prompt,
|
|
vae=vae,
|
|
image1=image1,
|
|
image2=image2,
|
|
image3=image3,
|
|
)
|
|
)
|
|
|
|
|
|
def _scale_qwen_edit_image(image):
|
|
scaler_cls = _require_class(nodes_flux, "FluxKontextImageScale", _FLUX_ERR)
|
|
return _unwrap_first(scaler_cls.execute(image))
|
|
|
|
|
|
def _apply_reference_latents_method(conditioning, reference_latents_method: str):
|
|
method_cls = _require_class(nodes_flux, "FluxKontextMultiReferenceLatentMethod", _FLUX_ERR)
|
|
return _unwrap_first(method_cls.execute(conditioning, reference_latents_method))
|
|
|
|
|
|
def _encode_image_to_latent(vae, image):
|
|
encoder_cls = _require_class(comfy_core_nodes, "VAEEncode", _CORE_NODES_ERR)
|
|
return encoder_cls().encode(vae, image)[0]
|
|
|
|
|
|
def _sample_qwen(model, positive, negative, latent, steps: int, seed: int, cfg: float, sampler_name: str, scheduler: str):
|
|
sampler_cls = _require_class(comfy_core_nodes, "KSampler", _CORE_NODES_ERR)
|
|
return sampler_cls().sample(
|
|
model,
|
|
int(seed),
|
|
int(steps),
|
|
float(cfg),
|
|
sampler_name,
|
|
scheduler,
|
|
positive,
|
|
negative,
|
|
latent,
|
|
denoise=1.0,
|
|
)[0]
|
|
|
|
|
|
def _decode_qwen_latent(vae, latent):
|
|
decoder_cls = _require_class(comfy_core_nodes, "VAEDecode", _CORE_NODES_ERR)
|
|
return decoder_cls().decode(vae, latent)[0]
|
|
|
|
|
|
def _resolve_qwen_resources(model_name: str, validator):
|
|
validator(model_name)
|
|
clip_name = discovery.resolve_qwen_text_encoder_name()
|
|
vae_name = discovery.resolve_qwen_vae_name()
|
|
|
|
if clip_name is None:
|
|
raise RuntimeError(
|
|
"[LLS] Missing Qwen text encoder. Place a compatible Qwen VL text encoder in "
|
|
"ComfyUI/models/text_encoders/."
|
|
)
|
|
if vae_name is None:
|
|
raise RuntimeError(
|
|
"[LLS] Missing Qwen VAE. Place qwen_image_vae.safetensors in ComfyUI/models/vae/."
|
|
)
|
|
|
|
return model_name, clip_name, vae_name
|
|
|
|
|
|
def run_qwen_text_to_image(
|
|
model_name: str,
|
|
prompt: str,
|
|
negative_prompt: str,
|
|
width: int,
|
|
height: int,
|
|
steps: int,
|
|
seed: int,
|
|
batch_size: int,
|
|
cfg: float,
|
|
sampler_name: str,
|
|
scheduler: str,
|
|
shift: float,
|
|
enable_turbo_mode: bool,
|
|
turbo_lora_name: str,
|
|
turbo_strength: float,
|
|
model=None,
|
|
):
|
|
try:
|
|
model_name, clip_name, vae_name = _resolve_qwen_resources(
|
|
model_name,
|
|
discovery.validate_qwen_text_model_name,
|
|
)
|
|
resolved_turbo_lora, effective_steps, effective_cfg = _prepare_turbo_request(
|
|
discovery.resolve_qwen_text_turbo_lora,
|
|
model_name,
|
|
steps,
|
|
cfg,
|
|
enable_turbo_mode,
|
|
turbo_lora_name,
|
|
)
|
|
model = _resolve_generation_model(model_name, model)
|
|
if resolved_turbo_lora is not None:
|
|
model = _load_model_only_lora(model, resolved_turbo_lora, float(turbo_strength))
|
|
model = _patch_qwen_model_sampling(model, shift)
|
|
clip = _load_qwen_clip(clip_name)
|
|
vae = _load_qwen_vae(vae_name)
|
|
positive = _encode_clip_text(clip, prompt)
|
|
negative = _encode_clip_text(clip, negative_prompt)
|
|
latent = _create_empty_qwen_latent(int(width), int(height), int(batch_size))
|
|
sampled = _sample_qwen(
|
|
model,
|
|
positive,
|
|
negative,
|
|
latent,
|
|
effective_steps,
|
|
int(seed),
|
|
effective_cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
)
|
|
return _decode_qwen_latent(vae, sampled)
|
|
except Exception as exc:
|
|
if str(exc).startswith("[LLS]"):
|
|
raise
|
|
raise RuntimeError(f"[LLS] Qwen text-to-image failed: {exc}") from exc
|
|
|
|
|
|
def run_qwen_image_edit(
|
|
model_name: str,
|
|
image,
|
|
image2,
|
|
image3,
|
|
prompt: str,
|
|
negative_prompt: str,
|
|
steps: int,
|
|
seed: int,
|
|
cfg: float,
|
|
sampler_name: str,
|
|
scheduler: str,
|
|
shift: float,
|
|
cfg_norm_strength: float,
|
|
reference_latents_method: str,
|
|
enable_turbo_mode: bool,
|
|
turbo_lora_name: str,
|
|
turbo_strength: float,
|
|
model=None,
|
|
):
|
|
try:
|
|
if image is None:
|
|
raise RuntimeError("[LLS] Missing source image. Connect an IMAGE input.")
|
|
|
|
model_name, clip_name, vae_name = _resolve_qwen_resources(
|
|
model_name,
|
|
discovery.validate_qwen_edit_model_name,
|
|
)
|
|
resolved_turbo_lora, effective_steps, effective_cfg = _prepare_turbo_request(
|
|
discovery.resolve_qwen_edit_turbo_lora,
|
|
model_name,
|
|
steps,
|
|
cfg,
|
|
enable_turbo_mode,
|
|
turbo_lora_name,
|
|
)
|
|
model = _resolve_generation_model(model_name, model)
|
|
if resolved_turbo_lora is not None:
|
|
model = _load_model_only_lora(model, resolved_turbo_lora, float(turbo_strength))
|
|
model = _patch_qwen_model_sampling(model, shift)
|
|
model = _apply_cfg_norm(model, cfg_norm_strength)
|
|
clip = _load_qwen_clip(clip_name)
|
|
vae = _load_qwen_vae(vae_name)
|
|
scaled = _scale_qwen_edit_image(image)
|
|
positive = _encode_qwen_edit_conditioning(clip, prompt, vae, scaled, image2, image3)
|
|
negative = _encode_qwen_edit_conditioning(clip, negative_prompt, vae, scaled, image2, image3)
|
|
positive = _apply_reference_latents_method(positive, reference_latents_method)
|
|
negative = _apply_reference_latents_method(negative, reference_latents_method)
|
|
latent = _encode_image_to_latent(vae, scaled)
|
|
sampled = _sample_qwen(
|
|
model,
|
|
positive,
|
|
negative,
|
|
latent,
|
|
effective_steps,
|
|
int(seed),
|
|
effective_cfg,
|
|
sampler_name,
|
|
scheduler,
|
|
)
|
|
return _decode_qwen_latent(vae, sampled)
|
|
except Exception as exc:
|
|
if str(exc).startswith("[LLS]"):
|
|
raise
|
|
raise RuntimeError(f"[LLS] Qwen image edit failed: {exc}") from exc
|