204 lines
7.7 KiB
Python
204 lines
7.7 KiB
Python
import uuid
|
|
import os
|
|
import folder_paths
|
|
from diffsynth.pipelines.z_image import (
|
|
ZImagePipeline, ModelConfig,
|
|
ZImageUnit_Image2LoRAEncode, ZImageUnit_Image2LoRADecode
|
|
)
|
|
|
|
from safetensors.torch import save_file
|
|
import torch
|
|
from PIL import Image
|
|
import numpy as np
|
|
|
|
import gc
|
|
from comfy import model_management as mm
|
|
import shutil
|
|
|
|
CHECKING_ZIMAGE_LORA_PREFIX = 'zimage_i2l_lora'
|
|
|
|
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_ZImageI2L_Loader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ('RH_ZImageI2LPipeline', )
|
|
RETURN_NAMES = ('ZImageI2LPipeline', )
|
|
FUNCTION = "load"
|
|
CATEGORY = "RunningHub/ZImageI2L"
|
|
|
|
def __init__(self):
|
|
self.vram_config_disk_offload = {
|
|
"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",
|
|
}
|
|
# 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):
|
|
|
|
loaded_models = mm.current_loaded_models
|
|
print(f'[kiki] loaded_models:', len(loaded_models))
|
|
if len(loaded_models) > 0:
|
|
mm.unload_all_models()
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
|
|
pipe = ZImagePipeline.from_pretrained(
|
|
torch_dtype=torch.bfloat16,
|
|
device="cuda",
|
|
model_configs=[
|
|
ModelConfig(model_id="Tongyi-MAI/Z-Image", origin_file_pattern="transformer/*.safetensors", **self.vram_config_disk_offload),
|
|
ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="text_encoder/*.safetensors", **self.vram_config_disk_offload),
|
|
ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="vae/diffusion_pytorch_model.safetensors", **self.vram_config_disk_offload),
|
|
ModelConfig(model_id="DiffSynth-Studio/General-Image-Encoders", origin_file_pattern="SigLIP2-G384/model.safetensors", **self.vram_config_disk_offload),
|
|
ModelConfig(model_id="DiffSynth-Studio/General-Image-Encoders", origin_file_pattern="DINOv3-7B/model.safetensors", **self.vram_config_disk_offload),
|
|
ModelConfig(model_id="DiffSynth-Studio/Z-Image-i2L", origin_file_pattern="model.safetensors", **self.vram_config_disk_offload),
|
|
],
|
|
tokenizer_config=ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="tokenizer/"),
|
|
vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 2,
|
|
)
|
|
return (pipe, )
|
|
|
|
class RunningHub_ZImageI2L_LoraGenerator:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"pipeline": ("RH_ZImageI2LPipeline", ),
|
|
"training_images": ("IMAGE", ),
|
|
"seed": ("INT", {"default": 42, "min": 0, "max": 0xffffffffffffffff}),
|
|
}
|
|
}
|
|
|
|
# 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")), 'LORA_PATH')
|
|
RETURN_NAMES = ('lora_name', 'lora_path')
|
|
FUNCTION = "generate"
|
|
CATEGORY = "RunningHub/ZImageI2L"
|
|
|
|
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):
|
|
self.lora_name = f"{CHECKING_ZIMAGE_LORA_PREFIX}_{str(uuid.uuid4())}.safetensors"
|
|
|
|
def generate(self, pipeline, training_images, **kwargs):
|
|
|
|
loaded_models = mm.current_loaded_models
|
|
print(f'[kiki] loaded_models before generate:', len(loaded_models))
|
|
if len(loaded_models) > 0:
|
|
mm.unload_all_models()
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
|
|
training_images = [self.tensor_2_pil(image) for image in training_images]
|
|
training_images = [image.convert("RGB") for image in training_images]
|
|
with torch.no_grad():
|
|
embs = ZImageUnit_Image2LoRAEncode().process(pipeline, image2lora_images=training_images)
|
|
lora = ZImageUnit_Image2LoRADecode().process(pipeline, **embs)["lora"]
|
|
lora_path = os.path.join(folder_paths.models_dir, 'loras', self.lora_name)
|
|
save_file(lora, lora_path)
|
|
# lora_name is a filename under models/loras (e.g. *.safetensors)
|
|
return (self.lora_name, lora_path)
|
|
|
|
|
|
class RunningHub_ZImageI2L_Saver:
|
|
"""
|
|
RH platform compatible output node for saving LoRA files.
|
|
Follows ComfyUI native output specification for RunningHub integration.
|
|
"""
|
|
def __init__(self):
|
|
self.type = "output" # Required for RH platform
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"lora_path": ("LORA_PATH", {"forceInput": True}),
|
|
"filename_prefix": ("STRING", {"default": "zimage_lora"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = () # Output node returns via UI dict
|
|
FUNCTION = "save"
|
|
CATEGORY = "RunningHub/ZImageI2L"
|
|
OUTPUT_NODE = True # Required for RH platform
|
|
|
|
def save(self, lora_path, filename_prefix="lora"):
|
|
|
|
lora_path = str(lora_path)
|
|
|
|
# Check if source file exists
|
|
if not os.path.exists(lora_path):
|
|
raise ValueError(f"illegal lora path")
|
|
|
|
if CHECKING_ZIMAGE_LORA_PREFIX not in lora_path:
|
|
raise ValueError(f"illegal lora path")
|
|
|
|
# Use folder_paths to get proper output directory (RH platform requirement)
|
|
output_dir = folder_paths.get_directory_by_type("output")
|
|
full_output_folder, filename, counter, subfolder, filename_prefix = \
|
|
folder_paths.get_save_image_path(filename_prefix, output_dir)
|
|
|
|
# Build final filename with counter
|
|
file_name_with_ext = f"{filename}_{counter:05}_.safetensors"
|
|
full_path = os.path.join(full_output_folder, file_name_with_ext)
|
|
|
|
# Copy the LoRA file to output directory
|
|
shutil.copy2(lora_path, full_path)
|
|
print(f'[RH] LoRA saved to output: {full_path}')
|
|
|
|
# Return UI dict for RH platform to capture the file
|
|
return {
|
|
"ui": {
|
|
"images": [
|
|
{
|
|
"filename": file_name_with_ext,
|
|
"subfolder": subfolder,
|
|
"type": self.type
|
|
}
|
|
]
|
|
}
|
|
}
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"RunningHub_ZImageI2L_Loader": RunningHub_ZImageI2L_Loader,
|
|
"RunningHub_ZImageI2L_LoraGenerator": RunningHub_ZImageI2L_LoraGenerator,
|
|
"RunningHub_ZImageI2L_Saver": RunningHub_ZImageI2L_Saver,
|
|
}
|
|
|