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'[ZImageI2L] Unloading {len(loaded_models)} models for pipeline loading') 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'[ZImageI2L] Unloading {len(loaded_models)} models for LoRA generation') 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, }