470 lines
16 KiB
Python
470 lines
16 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
import glob
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
|
|
try:
|
|
import folder_paths
|
|
except Exception:
|
|
folder_paths = None
|
|
|
|
|
|
DIFFSYNTH_INSTALL_HINT = (
|
|
"DiffSynth-Studio is required for Qwen-Image-Layered-Control-V2.\n"
|
|
"Install it in the ComfyUI Python environment:\n"
|
|
" git clone https://github.com/modelscope/DiffSynth-Studio.git\n"
|
|
" cd DiffSynth-Studio\n"
|
|
" pip install -e ."
|
|
)
|
|
|
|
DEFAULT_V1_MODEL = "qwen_image_layered_control_bf16.safetensors"
|
|
DEFAULT_V2_LORA = "Qwen-Image-Layered-Control-V2.safetensors"
|
|
QWEN_IMAGE_MODEL_ID = "Qwen/Qwen-Image"
|
|
QWEN_LAYERED_MODEL_ID = "Qwen/Qwen-Image-Layered"
|
|
QWEN_TEXT_ENCODER_PATTERN = "text_encoder/model*.safetensors"
|
|
QWEN_TOKENIZER_PATTERN = "tokenizer/"
|
|
QWEN_LAYERED_VAE_PATTERN = "vae/diffusion_pytorch_model.safetensors"
|
|
|
|
|
|
def _load_diffsynth():
|
|
try:
|
|
from diffsynth.pipelines.qwen_image import ModelConfig, QwenImagePipeline
|
|
except Exception as exc:
|
|
raise RuntimeError(DIFFSYNTH_INSTALL_HINT) from exc
|
|
return ModelConfig, QwenImagePipeline
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class QwenImageLayeredV2Pipeline:
|
|
pipe: Any
|
|
v2_lora: str
|
|
|
|
|
|
def _float8_or_bfloat16():
|
|
return getattr(torch, "float8_e4m3fn", torch.bfloat16)
|
|
|
|
|
|
def _vram_config(mode: str) -> dict[str, Any]:
|
|
fp8 = _float8_or_bfloat16()
|
|
if mode == "low":
|
|
return {
|
|
"offload_dtype": "disk",
|
|
"offload_device": "disk",
|
|
"onload_dtype": fp8,
|
|
"onload_device": "cpu",
|
|
"preparing_dtype": fp8,
|
|
"preparing_device": "cuda",
|
|
"computation_dtype": torch.bfloat16,
|
|
"computation_device": "cuda",
|
|
}
|
|
if mode == "high":
|
|
return {
|
|
"offload_dtype": torch.bfloat16,
|
|
"offload_device": "cpu",
|
|
"onload_dtype": torch.bfloat16,
|
|
"onload_device": "cuda",
|
|
"preparing_dtype": torch.bfloat16,
|
|
"preparing_device": "cuda",
|
|
"computation_dtype": torch.bfloat16,
|
|
"computation_device": "cuda",
|
|
}
|
|
return {
|
|
"offload_dtype": fp8,
|
|
"offload_device": "cpu",
|
|
"onload_dtype": fp8,
|
|
"onload_device": "cpu",
|
|
"preparing_dtype": fp8,
|
|
"preparing_device": "cuda",
|
|
"computation_dtype": torch.bfloat16,
|
|
"computation_device": "cuda",
|
|
}
|
|
|
|
|
|
def _vram_limit(device: str, margin_gb: float) -> float | None:
|
|
if device.startswith("cuda") and torch.cuda.is_available():
|
|
total_gb = torch.cuda.mem_get_info(device)[1] / (1024**3)
|
|
return max(total_gb - margin_gb, 1.0)
|
|
return None
|
|
|
|
|
|
def _filename_list(folder_name: str, preferred: str) -> list[str]:
|
|
if folder_paths is None:
|
|
return [preferred]
|
|
|
|
names = list(folder_paths.get_filename_list(folder_name))
|
|
if preferred and preferred not in names:
|
|
names.insert(0, preferred)
|
|
return names or [preferred]
|
|
|
|
|
|
def _resolve_comfy_model_path(folder_name: str, name_or_path: str) -> str:
|
|
path = Path(name_or_path)
|
|
if path.exists():
|
|
return str(path)
|
|
|
|
if folder_paths is not None:
|
|
full_path = folder_paths.get_full_path(folder_name, name_or_path)
|
|
if full_path:
|
|
return full_path
|
|
|
|
return name_or_path
|
|
|
|
|
|
def _diffsynth_model_base_path() -> Path:
|
|
if os.environ.get("DIFFSYNTH_MODEL_BASE_PATH"):
|
|
return Path(os.environ["DIFFSYNTH_MODEL_BASE_PATH"])
|
|
if folder_paths is not None:
|
|
return Path(folder_paths.models_dir)
|
|
return Path("./models")
|
|
|
|
|
|
def _cached_repo_files(model_id: str, file_pattern: str) -> list[str]:
|
|
pattern = file_pattern + "*" if file_pattern.endswith("/") else file_pattern
|
|
return glob.glob(str(_diffsynth_model_base_path() / model_id / pattern))
|
|
|
|
|
|
def _require_cached_repo_files(model_id: str, file_pattern: str, label: str):
|
|
if _cached_repo_files(model_id, file_pattern):
|
|
return
|
|
|
|
expected = _diffsynth_model_base_path() / model_id / file_pattern
|
|
raise RuntimeError(
|
|
f"{label} is not available locally for DiffSynth.\n"
|
|
f"Expected cache: {expected}\n"
|
|
"The ComfyUI FP8 text encoder qwen_2.5_vl_7b_fp8_scaled.safetensors is not compatible with "
|
|
"DiffSynth QwenImagePipeline. Enable official download in the Loader, or pre-download the "
|
|
"official Qwen/Qwen-Image text_encoder and tokenizer files into the DiffSynth cache."
|
|
)
|
|
|
|
|
|
def _local_model_config(ModelConfig, folder_name: str, name_or_path: str, vram: dict[str, Any] | None = None):
|
|
kwargs = dict(vram or {})
|
|
path = Path(_resolve_comfy_model_path(folder_name, name_or_path))
|
|
if path.exists():
|
|
return ModelConfig(path=str(path), **kwargs)
|
|
raise FileNotFoundError(f"Model file not found in ComfyUI models/{folder_name}: {name_or_path}")
|
|
|
|
|
|
def _repo_model_config(
|
|
ModelConfig,
|
|
model_id: str,
|
|
file_pattern: str,
|
|
vram: dict[str, Any] | None = None,
|
|
allow_download: bool = False,
|
|
download_source: str = "modelscope",
|
|
):
|
|
kwargs = dict(vram or {})
|
|
return ModelConfig(
|
|
model_id=model_id,
|
|
origin_file_pattern=file_pattern,
|
|
local_model_path=str(_diffsynth_model_base_path()),
|
|
skip_download=not allow_download,
|
|
download_source=download_source,
|
|
**kwargs,
|
|
)
|
|
|
|
|
|
def _lora_path(lora: str):
|
|
lora_path = Path(_resolve_comfy_model_path("loras", lora))
|
|
if lora_path.exists():
|
|
return str(lora_path)
|
|
raise FileNotFoundError(f"LoRA file not found in ComfyUI models/loras: {lora}")
|
|
|
|
|
|
def _tensor_to_pil(image: torch.Tensor, mode: str = "RGBA") -> Image.Image:
|
|
if image.dim() == 4:
|
|
image = image[0]
|
|
arr = image.detach().cpu().float().numpy()
|
|
arr = (arr * 255.0).clip(0, 255).astype(np.uint8)
|
|
|
|
if arr.ndim == 2:
|
|
arr = arr[:, :, None]
|
|
channels = arr.shape[-1]
|
|
if channels == 1:
|
|
arr = np.repeat(arr, 3, axis=-1)
|
|
elif channels == 2:
|
|
arr = np.concatenate([arr[:, :, :1]] * 3 + [arr[:, :, 1:2]], axis=-1)
|
|
elif channels > 4:
|
|
arr = arr[:, :, :4]
|
|
|
|
if mode == "RGBA" and arr.shape[-1] == 3:
|
|
alpha = np.full((*arr.shape[:2], 1), 255, dtype=np.uint8)
|
|
arr = np.concatenate([arr, alpha], axis=-1)
|
|
if mode == "RGB" and arr.shape[-1] == 4:
|
|
arr = arr[:, :, :3]
|
|
return Image.fromarray(arr, mode)
|
|
|
|
|
|
def _pil_to_tensor(image: Image.Image) -> torch.Tensor:
|
|
arr = np.asarray(image).astype(np.float32) / 255.0
|
|
if arr.ndim == 2:
|
|
arr = arr[:, :, None]
|
|
return torch.from_numpy(arr).unsqueeze(0)
|
|
|
|
|
|
def _round_dimension(value: int, multiple: int = 16) -> int:
|
|
return max(256, int(round(value / multiple) * multiple))
|
|
|
|
|
|
def _mask_to_array(mask: torch.Tensor, width: int, height: int, invert: bool = False) -> np.ndarray:
|
|
if mask.dim() == 4:
|
|
if mask.shape[-1] <= 4:
|
|
mask = mask[0, :, :, 0]
|
|
else:
|
|
mask = mask[0, 0]
|
|
elif mask.dim() == 3:
|
|
mask = mask[0]
|
|
|
|
arr = mask.detach().cpu().float().numpy()
|
|
arr = np.clip(arr, 0.0, 1.0)
|
|
if invert:
|
|
arr = 1.0 - arr
|
|
|
|
img = Image.fromarray((arr * 255.0).astype(np.uint8), "L")
|
|
img = img.resize((width, height), Image.Resampling.BILINEAR)
|
|
return np.asarray(img, dtype=np.uint8)
|
|
|
|
|
|
def _context_from_masks(
|
|
target_mask: torch.Tensor,
|
|
remove_mask: torch.Tensor | None,
|
|
width: int,
|
|
height: int,
|
|
invert_target: bool = False,
|
|
invert_remove: bool = False,
|
|
) -> Image.Image:
|
|
red = _mask_to_array(target_mask, width, height, invert_target)
|
|
green = np.zeros_like(red)
|
|
if remove_mask is not None:
|
|
green = _mask_to_array(remove_mask, width, height, invert_remove)
|
|
|
|
rgba = np.zeros((height, width, 4), dtype=np.uint8)
|
|
rgba[:, :, 0] = red
|
|
rgba[:, :, 1] = green
|
|
rgba[:, :, 3] = np.maximum(red, green)
|
|
return Image.fromarray(rgba, "RGBA")
|
|
|
|
|
|
def _first_pil_image(result: Any) -> Image.Image:
|
|
if isinstance(result, Image.Image):
|
|
return result
|
|
if hasattr(result, "images"):
|
|
return _first_pil_image(result.images)
|
|
if isinstance(result, (list, tuple)):
|
|
if not result:
|
|
raise RuntimeError("DiffSynth returned an empty image list.")
|
|
return _first_pil_image(result[0])
|
|
raise TypeError(f"Unsupported DiffSynth output type: {type(result)!r}")
|
|
|
|
|
|
class QwenImageLayeredV2Loader:
|
|
CATEGORY = "QwenImageLayered/V2"
|
|
RETURN_TYPES = ("QWEN_LAYERED_V2_PIPELINE",)
|
|
RETURN_NAMES = ("pipeline",)
|
|
FUNCTION = "load"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"v1_model": (_filename_list("diffusion_models", DEFAULT_V1_MODEL), {"default": DEFAULT_V1_MODEL}),
|
|
"v2_lora": (_filename_list("loras", DEFAULT_V2_LORA), {"default": DEFAULT_V2_LORA}),
|
|
"allow_official_download": ("BOOLEAN", {"default": False}),
|
|
"download_source": (["modelscope", "huggingface"], {"default": "modelscope"}),
|
|
"device": ("STRING", {"default": "cuda"}),
|
|
"vram_mode": (["low", "medium", "high"], {"default": "medium"}),
|
|
"vram_margin_gb": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 16.0, "step": 0.25}),
|
|
},
|
|
}
|
|
|
|
def load(
|
|
self,
|
|
v1_model: str,
|
|
v2_lora: str,
|
|
allow_official_download: bool,
|
|
download_source: str,
|
|
device: str,
|
|
vram_mode: str,
|
|
vram_margin_gb: float,
|
|
):
|
|
ModelConfig, QwenImagePipeline = _load_diffsynth()
|
|
vram = _vram_config(vram_mode)
|
|
if not allow_official_download:
|
|
_require_cached_repo_files(QWEN_IMAGE_MODEL_ID, QWEN_TEXT_ENCODER_PATTERN, "Official Qwen text encoder")
|
|
_require_cached_repo_files(QWEN_IMAGE_MODEL_ID, QWEN_TOKENIZER_PATTERN, "Official Qwen tokenizer")
|
|
_require_cached_repo_files(QWEN_LAYERED_MODEL_ID, QWEN_LAYERED_VAE_PATTERN, "Official Qwen layered VAE")
|
|
|
|
pipe = QwenImagePipeline.from_pretrained(
|
|
torch_dtype=torch.bfloat16,
|
|
device=device,
|
|
model_configs=[
|
|
_local_model_config(ModelConfig, "diffusion_models", v1_model, vram),
|
|
_repo_model_config(
|
|
ModelConfig,
|
|
QWEN_IMAGE_MODEL_ID,
|
|
QWEN_TEXT_ENCODER_PATTERN,
|
|
vram,
|
|
allow_official_download,
|
|
download_source,
|
|
),
|
|
_repo_model_config(
|
|
ModelConfig,
|
|
QWEN_LAYERED_MODEL_ID,
|
|
QWEN_LAYERED_VAE_PATTERN,
|
|
vram,
|
|
allow_official_download,
|
|
download_source,
|
|
),
|
|
],
|
|
tokenizer_config=_repo_model_config(
|
|
ModelConfig,
|
|
QWEN_IMAGE_MODEL_ID,
|
|
QWEN_TOKENIZER_PATTERN,
|
|
allow_download=allow_official_download,
|
|
download_source=download_source,
|
|
),
|
|
vram_limit=_vram_limit(device, vram_margin_gb),
|
|
)
|
|
pipe.load_lora(pipe.dit, _lora_path(v2_lora))
|
|
return (QwenImageLayeredV2Pipeline(pipe=pipe, v2_lora=v2_lora),)
|
|
|
|
|
|
class QwenImageLayeredV2BrushContext:
|
|
CATEGORY = "QwenImageLayered/V2"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("context_image",)
|
|
FUNCTION = "build"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"target_mask": ("MASK",),
|
|
"invert_target": ("BOOLEAN", {"default": False}),
|
|
},
|
|
"optional": {
|
|
"remove_mask": ("MASK",),
|
|
"invert_remove": ("BOOLEAN", {"default": False}),
|
|
},
|
|
}
|
|
|
|
def build(
|
|
self,
|
|
image: torch.Tensor,
|
|
target_mask: torch.Tensor,
|
|
invert_target: bool,
|
|
remove_mask: torch.Tensor | None = None,
|
|
invert_remove: bool = False,
|
|
):
|
|
_, height, width, _ = image.shape
|
|
context = _context_from_masks(
|
|
target_mask=target_mask,
|
|
remove_mask=remove_mask,
|
|
width=width,
|
|
height=height,
|
|
invert_target=invert_target,
|
|
invert_remove=invert_remove,
|
|
)
|
|
return (_pil_to_tensor(context),)
|
|
|
|
|
|
class QwenImageLayeredV2Decompose:
|
|
CATEGORY = "QwenImageLayered/V2"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("layer_image",)
|
|
FUNCTION = "decompose"
|
|
OUTPUT_NODE = False
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("QWEN_LAYERED_V2_PIPELINE",),
|
|
"image": ("IMAGE",),
|
|
"prompt": ("STRING", {"default": "", "multiline": True}),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
|
|
"steps": ("INT", {"default": 10, "min": 1, "max": 100}),
|
|
"cfg_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 20.0, "step": 0.1}),
|
|
"width": ("INT", {"default": 1024, "min": 256, "max": 4096, "step": 16}),
|
|
"height": ("INT", {"default": 1024, "min": 256, "max": 4096, "step": 16}),
|
|
"layer_num": ("INT", {"default": 0, "min": 0, "max": 32}),
|
|
"use_input_size": ("BOOLEAN", {"default": True}),
|
|
},
|
|
"optional": {
|
|
"negative_prompt": ("STRING", {"default": "", "multiline": True}),
|
|
"context_image": ("IMAGE",),
|
|
"target_mask": ("MASK",),
|
|
"remove_mask": ("MASK",),
|
|
},
|
|
}
|
|
|
|
def decompose(
|
|
self,
|
|
pipeline: QwenImageLayeredV2Pipeline,
|
|
image: torch.Tensor,
|
|
prompt: str,
|
|
seed: int,
|
|
steps: int,
|
|
cfg_scale: float,
|
|
width: int,
|
|
height: int,
|
|
layer_num: int,
|
|
use_input_size: bool = True,
|
|
negative_prompt: str = "",
|
|
context_image: torch.Tensor | None = None,
|
|
target_mask: torch.Tensor | None = None,
|
|
remove_mask: torch.Tensor | None = None,
|
|
):
|
|
if use_input_size:
|
|
_, input_height, input_width, _ = image.shape
|
|
width = _round_dimension(input_width)
|
|
height = _round_dimension(input_height)
|
|
|
|
input_image = _tensor_to_pil(image, "RGBA").resize((width, height), Image.Resampling.LANCZOS)
|
|
|
|
context = None
|
|
if context_image is not None:
|
|
context = _tensor_to_pil(context_image, "RGBA").resize((width, height), Image.Resampling.LANCZOS)
|
|
elif target_mask is not None:
|
|
context = _context_from_masks(target_mask, remove_mask, width, height)
|
|
|
|
result = pipeline.pipe(
|
|
prompt=prompt,
|
|
negative_prompt=negative_prompt,
|
|
seed=seed,
|
|
height=height,
|
|
width=width,
|
|
layer_input_image=input_image,
|
|
layer_num=layer_num,
|
|
context_image=context,
|
|
num_inference_steps=steps,
|
|
cfg_scale=cfg_scale,
|
|
)
|
|
return (_pil_to_tensor(_first_pil_image(result).convert("RGBA")),)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"QwenImageLayeredV2Loader": QwenImageLayeredV2Loader,
|
|
"QwenImageLayeredV2BrushContext": QwenImageLayeredV2BrushContext,
|
|
"QwenImageLayeredV2Decompose": QwenImageLayeredV2Decompose,
|
|
# Compatibility aliases for early local workflows.
|
|
"QwenImageLayeredLoader": QwenImageLayeredV2Loader,
|
|
"QwenImageLayeredDecompose": QwenImageLayeredV2Decompose,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"QwenImageLayeredV2Loader": "QwenImage Layered V2 Loader",
|
|
"QwenImageLayeredV2BrushContext": "QwenImage Layered V2 Brush Context",
|
|
"QwenImageLayeredV2Decompose": "QwenImage Layered V2 Decompose",
|
|
"QwenImageLayeredLoader": "QwenImage Layered V2 Loader (compat)",
|
|
"QwenImageLayeredDecompose": "QwenImage Layered V2 Decompose (compat)",
|
|
}
|