245 lines
10 KiB
Python
245 lines
10 KiB
Python
import uuid
|
|
import os
|
|
import folder_paths
|
|
from diffsynth.pipelines.qwen_image import (
|
|
QwenImagePipeline, ModelConfig,
|
|
QwenImageUnit_Image2LoRAEncode, QwenImageUnit_Image2LoRADecode
|
|
)
|
|
from modelscope import snapshot_download
|
|
from safetensors.torch import save_file
|
|
import torch
|
|
from PIL import Image
|
|
import numpy as np
|
|
from diffsynth.utils.lora import merge_lora
|
|
from diffsynth import load_state_dict
|
|
|
|
def setup_model_download_path():
|
|
if 'DIFFSYNTH_MODEL_BASE_PATH' not in os.environ:
|
|
os.environ['DIFFSYNTH_MODEL_BASE_PATH'] = folder_paths.models_dir
|
|
print(f"[RH_QwenImageI2L] Set DIFFSYNTH_MODEL_BASE_PATH to: {folder_paths.models_dir}")
|
|
|
|
class AnyComboList(list):
|
|
"""
|
|
A JSON-serializable list subtype used as a ComfyUI socket type.
|
|
|
|
ComfyUI validates linked socket types by calling `received_type != input_type`.
|
|
For combo types, `input_type` is a plain Python list generated by
|
|
`folder_paths.get_filename_list(...)`, which can change when new files appear.
|
|
That makes two otherwise-compatible combo types fail validation.
|
|
|
|
By overriding `__ne__` to treat any list as compatible, we keep the UI behavior
|
|
(still a list/COMBO type) while making validation stable across list changes.
|
|
"""
|
|
|
|
def __ne__(self, other):
|
|
if isinstance(other, list):
|
|
return False
|
|
return super().__ne__(other)
|
|
|
|
class RunningHub_ImageQwenI2L_Loader_Style:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ('RH_QwenImageI2LPipeline', )
|
|
RETURN_NAMES = ('QwenImageI2LPipeline', )
|
|
FUNCTION = "load"
|
|
CATEGORY = "RunningHub/ImageQwenI2L"
|
|
|
|
def __init__(self):
|
|
self.vram_config_disk_offload = {
|
|
"offload_dtype": "disk",
|
|
"offload_device": "disk",
|
|
"onload_dtype": "disk",
|
|
"onload_device": "disk",
|
|
"preparing_dtype": torch.bfloat16,
|
|
"preparing_device": "cuda",
|
|
"computation_dtype": torch.bfloat16,
|
|
"computation_device": "cuda",
|
|
}
|
|
# self.encoder_path = os.path.join(folder_paths.models_dir, 'DiffSynth-Studio', 'General-Image-Encoders')
|
|
# self.i2l_path = os.path.join(folder_paths.models_dir, 'DiffSynth-Studio', 'Qwen-Image-i2L')
|
|
# self.processor_path = os.path.join(folder_paths.models_dir, 'DiffSynth-Studio', 'Qwen-Image-Edit')
|
|
|
|
def load(self):
|
|
setup_model_download_path()
|
|
|
|
model_configs = [
|
|
ModelConfig(model_id='DiffSynth-Studio/General-Image-Encoders', origin_file_pattern="SigLIP2-G384/model.safetensors", local_model_path=folder_paths.models_dir, **self.vram_config_disk_offload),
|
|
ModelConfig(model_id='DiffSynth-Studio/General-Image-Encoders', origin_file_pattern="DINOv3-7B/model.safetensors", local_model_path=folder_paths.models_dir, **self.vram_config_disk_offload),
|
|
ModelConfig(model_id='DiffSynth-Studio/Qwen-Image-i2L', origin_file_pattern="Qwen-Image-i2L-Style.safetensors", local_model_path=folder_paths.models_dir, **self.vram_config_disk_offload),
|
|
]
|
|
processor_config = ModelConfig(model_id='Qwen/Qwen-Image-Edit', origin_file_pattern="processor/", local_model_path=folder_paths.models_dir)
|
|
|
|
pipe = QwenImagePipeline.from_pretrained(
|
|
torch_dtype=torch.bfloat16,
|
|
device="cuda",
|
|
model_configs=model_configs,
|
|
processor_config=processor_config,
|
|
vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 2,
|
|
)
|
|
return (pipe, )
|
|
|
|
class RunningHub_ImageQwenI2L_Loader_CFB:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ('RH_QwenImageI2LPipeline', )
|
|
RETURN_NAMES = ('QwenImageI2LPipeline', )
|
|
FUNCTION = "load"
|
|
CATEGORY = "RunningHub/ImageQwenI2L"
|
|
|
|
def __init__(self):
|
|
self.vram_config_disk_offload = {
|
|
"offload_dtype": "disk",
|
|
"offload_device": "disk",
|
|
"onload_dtype": "disk",
|
|
"onload_device": "disk",
|
|
"preparing_dtype": torch.bfloat16,
|
|
"preparing_device": "cuda",
|
|
"computation_dtype": torch.bfloat16,
|
|
"computation_device": "cuda",
|
|
}
|
|
|
|
def load(self):
|
|
setup_model_download_path()
|
|
|
|
model_configs = [
|
|
ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="text_encoder/model*.safetensors", local_model_path=folder_paths.models_dir, **self.vram_config_disk_offload),
|
|
ModelConfig(model_id="DiffSynth-Studio/General-Image-Encoders", origin_file_pattern="SigLIP2-G384/model.safetensors", local_model_path=folder_paths.models_dir, **self.vram_config_disk_offload),
|
|
ModelConfig(model_id="DiffSynth-Studio/General-Image-Encoders", origin_file_pattern="DINOv3-7B/model.safetensors", local_model_path=folder_paths.models_dir, **self.vram_config_disk_offload),
|
|
ModelConfig(model_id="DiffSynth-Studio/Qwen-Image-i2L", origin_file_pattern="Qwen-Image-i2L-Coarse.safetensors", local_model_path=folder_paths.models_dir, **self.vram_config_disk_offload),
|
|
ModelConfig(model_id="DiffSynth-Studio/Qwen-Image-i2L", origin_file_pattern="Qwen-Image-i2L-Fine.safetensors", local_model_path=folder_paths.models_dir, **self.vram_config_disk_offload),
|
|
]
|
|
processor_config = ModelConfig(model_id="Qwen/Qwen-Image-Edit", origin_file_pattern="processor/", local_model_path=folder_paths.models_dir)
|
|
|
|
pipe = QwenImagePipeline.from_pretrained(
|
|
torch_dtype=torch.bfloat16,
|
|
device="cuda",
|
|
model_configs=model_configs,
|
|
processor_config=processor_config,
|
|
vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 2,
|
|
)
|
|
pipe.is_cfb = True
|
|
return (pipe, )
|
|
|
|
class RunningHub_ImageQwenI2L_LoraGenerator:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("RH_QwenImageI2LPipeline", ),
|
|
"training_images": ("IMAGE", ),
|
|
"seed": ("INT", {"default": 42, "min": 0, "max": 0xffffffffffffffff}),
|
|
"custom_lora_name": ("STRING", {"default": "", "multiline": False}),
|
|
}
|
|
}
|
|
|
|
# Match ComfyUI's LoRA dropdown input type (combo list from folder_paths),
|
|
# but keep it validation-stable even if the LoRA file list changes at runtime.
|
|
RETURN_TYPES = (AnyComboList(folder_paths.get_filename_list("loras")), 'STRING')
|
|
RETURN_NAMES = ('lora_name', 'lora_path')
|
|
FUNCTION = "generate"
|
|
CATEGORY = "RunningHub/ImageQwenI2L"
|
|
|
|
OUTPUT_NODE = True
|
|
|
|
def tensor_2_pil(self, img_tensor):
|
|
i = 255. * img_tensor.squeeze().cpu().numpy()
|
|
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
|
|
return img
|
|
|
|
def __init__(self):
|
|
# Note: Do NOT pre-generate a file name here.
|
|
# A single node instance can be executed multiple times; the LoRA file name
|
|
# should be unique per execution to avoid overwriting previous outputs.
|
|
pass
|
|
|
|
def _make_unique_lora_name(self) -> str:
|
|
return f"i2l_style_lora_{str(uuid.uuid4())}.safetensors"
|
|
|
|
def _normalize_user_lora_name(self, name):
|
|
"""
|
|
Normalize user-provided LoRA file name.
|
|
- Accepts "file" or "file.safetensors"
|
|
- Ensures the final name ends with exactly one ".safetensors"
|
|
- Prevents path traversal by stripping directory components
|
|
"""
|
|
if name is None:
|
|
return None
|
|
name = str(name).strip()
|
|
if not name:
|
|
return None
|
|
|
|
# Prevent path traversal / accidental directories: collapse separators then take basename.
|
|
name = name.replace("/", "_").replace("\\", "_")
|
|
name = os.path.basename(name)
|
|
if not name or name in {".", ".."}:
|
|
return None
|
|
|
|
ext = ".safetensors"
|
|
lowered = name.lower()
|
|
# Strip repeated ".safetensors" suffixes, then append exactly one.
|
|
while lowered.endswith(ext):
|
|
name = name[: -len(ext)]
|
|
lowered = name.lower()
|
|
name = name.strip()
|
|
if not name:
|
|
return None
|
|
return f"{name}{ext}"
|
|
|
|
def _get_lora_save_dir(self) -> str:
|
|
"""
|
|
Resolve the actual LoRA directory in a cross-platform way.
|
|
|
|
Prefer ComfyUI's folder registry (supports custom LoRA paths and avoids
|
|
cwd-related issues), then fall back to the default models/loras folder.
|
|
"""
|
|
try:
|
|
lora_dirs = folder_paths.get_folder_paths("loras")
|
|
if isinstance(lora_dirs, (list, tuple)) and len(lora_dirs) > 0 and lora_dirs[0]:
|
|
return os.path.abspath(lora_dirs[0])
|
|
except Exception:
|
|
# Keep a safe fallback for older ComfyUI versions / unexpected environments.
|
|
pass
|
|
return os.path.abspath(os.path.join(folder_paths.models_dir, "loras"))
|
|
|
|
def generate(self, pipeline, training_images, **kwargs):
|
|
training_images = [self.tensor_2_pil(image) for image in training_images]
|
|
training_images = [image.convert("RGB") for image in training_images]
|
|
lora_save_dir = self._get_lora_save_dir()
|
|
os.makedirs(lora_save_dir, exist_ok=True)
|
|
custom_lora_name = kwargs.get("custom_lora_name", "")
|
|
lora_name = self._normalize_user_lora_name(custom_lora_name) or self._make_unique_lora_name()
|
|
lora_path = os.path.join(lora_save_dir, lora_name)
|
|
with torch.no_grad():
|
|
embs = QwenImageUnit_Image2LoRAEncode().process(pipeline, image2lora_images=training_images)
|
|
lora = QwenImageUnit_Image2LoRADecode().process(pipeline, **embs)["lora"]
|
|
|
|
if hasattr(pipeline, 'is_cfb') and pipeline.is_cfb:
|
|
print('[kiki] is_cfb:', pipeline.is_cfb)
|
|
setup_model_download_path()
|
|
lora_bias = ModelConfig(model_id="DiffSynth-Studio/Qwen-Image-i2L", origin_file_pattern="Qwen-Image-i2L-Bias.safetensors", local_model_path=folder_paths.models_dir)
|
|
lora_bias.download_if_necessary()
|
|
lora_bias = load_state_dict(lora_bias.path, torch_dtype=torch.bfloat16, device="cuda")
|
|
lora = merge_lora([lora, lora_bias])
|
|
save_file(lora, lora_path)
|
|
# lora_name is a filename under models/loras (e.g. *.safetensors)
|
|
return (lora_name, os.path.normpath(lora_path))
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"RunningHub_ImageQwenI2L_Loader(Style)": RunningHub_ImageQwenI2L_Loader_Style,
|
|
"RunningHub_ImageQwenI2L_Loader(CFB)": RunningHub_ImageQwenI2L_Loader_CFB,
|
|
"RunningHub_ImageQwenI2L_LoraGenerator": RunningHub_ImageQwenI2L_LoraGenerator,
|
|
}
|
|
|