diff --git a/__init__.py b/__init__.py index 5a5f445..b1139c2 100644 --- a/__init__.py +++ b/__init__.py @@ -93,6 +93,7 @@ class FlowMatchEulerSchedulerNode: }), "base_shift": ("FLOAT", { "default": 0.5, + "step": 0.01, "tooltip": "Stabilizes generation. Higher values = more consistent/predictable outputs. Z-Image-Turbo uses default 0.5." }), "invert_sigmas": (["disable", "enable"], { @@ -105,6 +106,7 @@ class FlowMatchEulerSchedulerNode: }), "max_shift": ("FLOAT", { "default": 1.15, + "step": 0.01, "tooltip": "Maximum variation allowed. Higher = more exaggerated/stylized results. Z-Image-Turbo uses default 1.15." }), "num_train_timesteps": ("INT", { @@ -113,10 +115,12 @@ class FlowMatchEulerSchedulerNode: }), "shift": ("FLOAT", { "default": 3.0, + "step": 0.01, "tooltip": "Global timestep schedule shift. Z-Image-Turbo uses 3.0 for optimal performance with the Turbo model." }), "shift_terminal": ("FLOAT", { "default": 0.0, + "step": 0.01, "tooltip": "End value for shifted schedule. Set to 0.0 to disable. Advanced parameter for timestep schedule control." }), "stochastic_sampling": (["disable", "enable"], { @@ -234,4 +238,19 @@ NODE_CLASS_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = { "FlowMatchEulerDiscreteScheduler (Custom)": "FlowMatch Euler Discrete Scheduler (Custom)", -} \ No newline at end of file +} + +from .extract_metadata_node import NODE_CLASS_MAPPINGS as METADATA_NODE_MAPPINGS +from .extract_metadata_node import NODE_DISPLAY_NAME_MAPPINGS as METADATA_DISPLAY_MAPPINGS + +NODE_CLASS_MAPPINGS.update(METADATA_NODE_MAPPINGS) +NODE_DISPLAY_NAME_MAPPINGS.update(METADATA_DISPLAY_MAPPINGS) + +# Import Nunchaku nodes +try: + from .nunchaku_compat import NODE_CLASS_MAPPINGS as NUNCHAKU_NODES + from .nunchaku_compat import NODE_DISPLAY_NAME_MAPPINGS as NUNCHAKU_NAMES + NODE_CLASS_MAPPINGS.update(NUNCHAKU_NODES) + NODE_DISPLAY_NAME_MAPPINGS.update(NUNCHAKU_NAMES) +except Exception as e: + print(f"[FlowMatch Scheduler] Could not load Nunchaku nodes: {e}") \ No newline at end of file diff --git a/batch_test_results.txt b/batch_test_results.txt new file mode 100644 index 0000000..784dcec --- /dev/null +++ b/batch_test_results.txt @@ -0,0 +1,122 @@ +================================================================================ +Image: ComfyUI-zit_00006_.png +Dimensions: 0x0 +Prompt: +胶片摄影，镜头语言，淡彩，暗调，氛围低光，个性视角，身穿朝鲜族民族服饰的少女，灵动俏皮，朦胧梦幻，高颜值，既视感，现场感，情绪氛围感拉满，透视感，暗朦，光朦，泛白，褪色，漏光，粒子高噪点，胶片颗粒质感，层次丰富，写意，朦胧美学，光的美学，lomo效果，超现实，高级感，杰作， +pornmaster bukkake, white cum on her face, white cum on her hair, White cum covered her hair, white cum covered her face. White cum covered her clothes +濃稠半透明精液從她的臉滴落，液體拉絲滴落, + +================================================================================ +Image: ComfyUI_00030_.png +Dimensions: 1792x1120 +Prompt: +A low-angle cinematic shot in a dense bamboo forest at golden hour, sunbeams filtering diagonally through the tall green bamboo canopy to cast sharp, dramatic shadows on the moss-covered ground. A young East Asian swordswoman in her mid-twenties, with long flowing black hair escaping her topknot, wears a traditional red silk martial arts robe edged with gold thread embroidery; her expression is intensely focused, muscles tensed in her arms and legs. She is captured mid-action during a horizontal sword slash: body twisted dynamically forward, left foot planted firmly on emerald moss, right arm fully extended holding a polished steel jian sword with a black lacquered hilt, the blade angled sharply downward. Bamboo stalks surround her, their smooth green surfaces textured with vertical grooves and dew drops catching the light; fallen bamboo leaves scatter the forest floor, partially covered in soft velvety moss. Mist drifts faintly near the ground, with background bamboo softly blurred to create depth. Color palette emphasizes vibrant jade greens of the bamboo, warm amber sunlight, and the rich crimson of the robe, accented by the metallic gleam of the sword. Highly detailed textures include the woven silk fabric rippling with movement, the sword's reflective edge, and the rough bark of nearby bamboo trunks. + +================================================================================ +Image: Comfy Edit ZimageVs_00004_.png +Dimensions: 0x0 +Prompt: + + +================================================================================ +Image: ComfyUI_00041_.png +Dimensions: 0x0 +Prompt: + + +================================================================================ +Image: ComfyUI-zit_00060_.png +Dimensions: 0x0 +Prompt: +masterpiece, best quality, photo realistic, 8k, a cybernetic woman with flowing nanotech ink tattoos animating across her skin, glossy black fluid moving like circuitry, sleek tech-editorial mood , hyperdetailed, dramatic lighting, cinematic shot, ultra detailed, intricate details, cinematic, photorealistic, masterpiece + +================================================================================ +Image: ComfyUI_00008_.png +Dimensions: 512x512 +Prompt: +['7', 0] + +================================================================================ +Image: ComfyUI-sqnd-multi-flux_00013_.png +Dimensions: 0x0 +Prompt: + + +================================================================================ +Image: ComfyUI-sqnd-multi-zimage_00013_.png +Dimensions: 0x0 +Prompt: + + +================================================================================ +Image: ComfyUI-zit_00028_.png +Dimensions: 0x0 +Prompt: +A close-up, explicit portrait of a 25-year-old beautiful woman in an ancient Egyptian royal sleep chamber at night. She has long black hair and a curvy figure with saggy, hanging breasts, visible nipples, and natural pubic hair. She wears elaborate golden body jewelry, including an underbra, armlets, thighlets, and a crotchless thong. The scene is captured from multiple angles—from behind, straight on, and from below—with dramatic foreshortening, focusing on her buttocks and visible labia and clitoral hood. The atmosphere is intimate, lit by candlelight with a warm, dim glow, casting soft shadows across her body and the silk bed she rests on. detailed labia, clitoral hood, visible pussy, from below, from behind + +================================================================================ +Image: ComfyUI_00050_.png +Dimensions: 1024x1024 +Prompt: +tranin passing, anime style, 4k ultra resolution, flat shading. + +================================================================================ +Image: ComfyUI-FlowmatchEuler-Simple_00004_.png +Dimensions: 1280x960 +Prompt: +A low-angle cinematic shot in a dense bamboo forest at golden hour, sunbeams filtering diagonally through the tall green bamboo canopy to cast sharp, dramatic shadows on the moss-covered ground. A young East Asian swordswoman in her mid-twenties, with long flowing black hair escaping her topknot, wears a traditional red silk martial arts robe edged with gold thread embroidery; her expression is intensely focused, muscles tensed in her arms and legs. She is captured mid-action during a horizontal sword slash: body twisted dynamically forward, left foot planted firmly on emerald moss, right arm fully extended holding a polished steel jian sword with a black lacquered hilt, the blade angled sharply downward. Bamboo stalks surround her, their smooth green surfaces textured with vertical grooves and dew drops catching the light; fallen bamboo leaves scatter the forest floor, partially covered in soft velvety moss. Mist drifts faintly near the ground, with background bamboo softly blurred to create depth. Color palette emphasizes vibrant jade greens of the bamboo, warm amber sunlight, and the rich crimson of the robe, accented by the metallic gleam of the sword. Highly detailed textures include the woven silk fabric rippling with movement, the sword's reflective edge, and the rough bark of nearby bamboo trunks. + +================================================================================ +Image: ComfyUI-zit_00042_.png +Dimensions: 0x0 +Prompt: +镜头从高处拍摄，在一片宁静花园的斑驳光线下，一位纤细的女子优雅地坐在一张磨损的石凳上，藤蔓和花朵环绕四周。她的身形纤细，但胸部巨大且圆润，胸部远远大于角色的头部，超巨乳，自然地向下垂。她微微前倾，双手叠放在膝上，嘴角带着淡淡的微笑，头微微倾向阳光。光线柔和地温暖地照在她裸露的肌肤上，勾勒出她身体的每一处曲线和轻柔的重力拉伸。镜头从女子侧面拍摄， + +================================================================================ +Image: ComfyUI-zit_00032_.png +Dimensions: 0x0 +Prompt: +a lomo photograph of a striking portrait of a naked woman with a detailed dragon tattoo on her back standing in front of a window, with her hand on the window sill, facing away from the camera but looking back, enveloped in the shadows of a dark room with an ethereal red glow cast from a neon light outside at night. her long wavy dark purple hair is tied back in a pony tail. her back is arched accentuating her equisite hour glass figure. the neon-lit sign and night time cityscape outside the window casts a red hue over the inside of the dimly lit apartment illuminating her back to reveal the dragon tattoo. the photograph has large dark vignetting and was shot on 35mm film with visible film grain and color splashing throughout the frame. while the woman is sharply in focus, the edges of the composition are soft, with a shallow depth of field, excellent bokeh. the film frame border can be seen in the image. the photograph was shot with a canon f1 using 800 iso film. dramatic cinematic lighting. the neon sign has chinese characters. light leaks and film borders visible. 4ft3rd4rk + +================================================================================ +Image: ComfyUI_00049_.png +Dimensions: 768x1280 +Prompt: +tranin passing, anime style, 4k ultra resolution, flat shading. + +================================================================================ +Image: ComfyUI-zit_00077_.png +Dimensions: 0x0 +Prompt: +masterpiece, best quality, photo realistic, 8k, a woman wearing a sculptural translucent mask carved from pure lightbeams, refracting prismatic colors across her face, futuristic beauty campaign energy , hyperdetailed, dramatic lighting, cinematic shot, ultra detailed, intricate details, cinematic, photorealistic, masterpiece + +================================================================================ +Image: ComfyUI_00046_.png +Dimensions: 0x0 +Prompt: +Gothic Glamour. "Back to the Future" Delorean car rushes through the mysterious night forest through the fog glowing in the moonlight.High detail, 10-bit color rendering, large-scale image.Half-turned to the viewer,action pose.Cinematic realism, high contrast, surround light, exceptional detail, 8k,. on the plate the text "F.M.E.D.S". a partly visible indication on a wooden signe reads "Eros Diffusion" with an arrow pointing backwards to wards the car. the driver is just a shadow inside the car and not well lit. + +================================================================================ +Image: ComfyUI-zit_00056_.png +Dimensions: 0x0 +Prompt: +masterpiece, best quality, photo realistic, 8k, a serene Japanese onsen scene with an otaku-styled woman relaxing in steaming mineral water, soft lantern light reflecting off wooden bath walls, subtle anime-inspired accessories, gentle mist rising around her, cherry blossoms drifting in the air, tranquil mountain backdrop, elegant editorial composition , hyperdetailed, dramatic lighting, cinematic shot, ultra detailed, intricate details, cinematic, photorealistic, masterpiece + +================================================================================ +Image: ComfyUI-zit_00026_.png +Dimensions: 0x0 +Prompt: +You are an assistant... Hyper-realistic cinematic shot of a nude bio-mechanical Asian woman m3tsumi1, her silicon skin is completely revealed. She sits leaning against a crumbling stucco wall, its texture rough and weathered. The wall is overtaken by nature, with creeping vines, vibrant wildflowers, and dense foliage bursting through cracks. The scene is set centuries after a devastating post-apocalyptic battle. The woman's exposed silicon parts show signs of wounds, and battle damage, with subtle LED lights flickering weakly in her circuitry. Dirt and grime coat her form, emphasizing the passage of time. Shafts of golden sunlight filter through the overgrown canopy above, casting dappled shadows across the scene. The atmosphere is one of eerie beauty and abandoned technology reclaimed by nature. Ultra-detailed textures, dramatic lighting, and a muted color palette dominated by earth tones and metallic hues. 8K resolution, photorealistic rendering, cinematic composition. + +================================================================================ +Image: \ +Dimensions: 768x1280 +Prompt: + + +================================================================================ +Image: ComfyUI_00052_.png +Dimensions: 1024x1024 +Prompt: +tranin passing, anime style, 4k ultra resolution, flat shading. + diff --git a/extract_metadata_node.py b/extract_metadata_node.py new file mode 100644 index 0000000..8df2c22 --- /dev/null +++ b/extract_metadata_node.py @@ -0,0 +1,106 @@ +import torch +import os +import json +from PIL import Image, ImageOps +import folder_paths +import numpy as np + +class ImageMetadataExtractor: + @classmethod + def INPUT_TYPES(s): + input_dir = folder_paths.get_input_directory() + files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))] + return {"required": + {"image": (sorted(files), {"image_upload": True})}, + } + + RETURN_TYPES = ("IMAGE", "STRING", "INT", "INT", "STRING") + RETURN_NAMES = ("image", "positive_prompt", "width", "height", "filename") + FUNCTION = "extract_metadata" + CATEGORY = "utils" + + def extract_metadata(self, image): + image_path = folder_paths.get_annotated_filepath(image) + img = Image.open(image_path) + + output_image = ImageOps.exif_transpose(img) + output_image = output_image.convert("RGB") + output_image = np.array(output_image).astype(np.float32) / 255.0 + output_image = torch.from_numpy(output_image)[None,] + + positive_prompt = "" + width = 0 + height = 0 + + # Extract from 'prompt' (API format) which is what ComfyUI uses for execution + if 'prompt' in img.info: + try: + prompt = json.loads(img.info['prompt']) + + # 1. Find Positive Prompt + # Strategy: Find KSampler -> positive input -> CLIPTextEncode -> text + ksampler_nodes = [] + for node_id, node in prompt.items(): + class_type = node.get('class_type', '') + if 'KSampler' in class_type or 'SamplerCustom' in class_type: + ksampler_nodes.append(node) + + for ksampler in ksampler_nodes: + inputs = ksampler.get('inputs', {}) + if 'positive' in inputs: + positive_link = inputs['positive'] + if isinstance(positive_link, list): # It's a link [node_id, slot_index] + positive_node_id = str(positive_link[0]) + if positive_node_id in prompt: + positive_node = prompt[positive_node_id] + if positive_node.get('class_type') == 'CLIPTextEncode': + positive_prompt = positive_node.get('inputs', {}).get('text', "") + break # Found it + + # Fallback: Look for any CLIPTextEncode with "positive" in title/meta if not found via KSampler + if not positive_prompt: + candidates = [] + for node_id, node in prompt.items(): + if node.get('class_type') == 'CLIPTextEncode': + title = node.get('_meta', {}).get('title', '').lower() + text = node.get('inputs', {}).get('text', "") + if 'positive' in title and 'negative' not in title: + candidates.append(text) + # Also consider just long text if no clear title match + elif len(text) > 50: + candidates.append(text) + + # Pick the longest candidate if any found + if candidates: + positive_prompt = max(candidates, key=len) + + # 2. Find Width/Height + # Strategy: Find EmptyLatentImage + for node_id, node in prompt.items(): + if node.get('class_type') == 'EmptyLatentImage': + width = node.get('inputs', {}).get('width', 0) + height = node.get('inputs', {}).get('height', 0) + break + + # Fallback: Look for width/height in any node if still 0 + if width == 0 or height == 0: + for node_id, node in prompt.items(): + inputs = node.get('inputs', {}) + if 'width' in inputs and 'height' in inputs: + width = inputs['width'] + height = inputs['height'] + break + + except Exception as e: + print(f"Error parsing metadata: {e}") + + return (output_image, positive_prompt, width, height, image) + +# Node registration +NODE_CLASS_MAPPINGS = { + "ImageMetadataExtractor": ImageMetadataExtractor +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "ImageMetadataExtractor": "Load Image ErosDiffusion" +} diff --git a/nunchaku_compat.py b/nunchaku_compat.py index a3fae47..f1b6557 100644 --- a/nunchaku_compat.py +++ b/nunchaku_compat.py @@ -8,7 +8,8 @@ import logging logger = logging.getLogger(__name__) -_original_apply_model = None +_original_module_call = None +_original_tiled_call = None _patch_applied = False def is_nunchaku_qwen_model(model): @@ -149,9 +150,13 @@ def patch_diffusion_model_forward(original_forward): return wrapper +# Global variables to store original functions +_original_module_call = None +_original_tiled_call = None + def apply_nunchaku_patches(): """Apply monkey patches to fix Nunchaku compatibility issues""" - global _patch_applied + global _patch_applied, _original_module_call, _original_tiled_call if _patch_applied: print("[Nunchaku Compat] Patches already applied") @@ -162,9 +167,12 @@ def apply_nunchaku_patches(): # The patch will be applied when models are loaded import torch.nn as nn + # Store original if not already stored + if _original_module_call is None: + _original_module_call = nn.Module.__call__ + # Patch torch.nn.Module's __call__ for modules that have txt_norm # This is tricky - we'll patch specific Nunchaku model classes when we detect them - original_module_call = nn.Module.__call__ def patched_module_call(self, *args, **kwargs): # ONLY patch Nunchaku diffusion models - very specific detection @@ -182,10 +190,10 @@ def apply_nunchaku_patches(): print(f"[Nunchaku Compat] Detected and patching Nunchaku diffusion model: {type(self).__name__}") self._nunchaku_patched = True - return patch_diffusion_model_forward(original_module_call)(self, *args, **kwargs) + return patch_diffusion_model_forward(_original_module_call)(self, *args, **kwargs) # For all other modules (VAE, etc.), use original __call__ without modification - return original_module_call(self, *args, **kwargs) + return _original_module_call(self, *args, **kwargs) nn.Module.__call__ = patched_module_call @@ -195,7 +203,8 @@ def apply_nunchaku_patches(): if 'ComfyUI-TiledDiffusion.tiled_diffusion' in sys.modules: tiled_diff = sys.modules['ComfyUI-TiledDiffusion.tiled_diffusion'] if hasattr(tiled_diff, 'TiledDiffusion'): - original_tiled_call = tiled_diff.TiledDiffusion.__call__ + if _original_tiled_call is None: + _original_tiled_call = tiled_diff.TiledDiffusion.__call__ def patched_tiled_call(self, model_function, kwargs): """Wrap TiledDiffusion to handle 5D tensors from Qwen Image models""" @@ -211,7 +220,7 @@ def apply_nunchaku_patches(): kwargs['input'] = x_in.squeeze(2) # Remove F dimension # Call original with 4D tensor - result = original_tiled_call(self, model_function, kwargs) + result = _original_tiled_call(self, model_function, kwargs) # Restore 5D shape if result is 4D if isinstance(result, torch.Tensor) and len(result.shape) == 4: @@ -222,7 +231,7 @@ def apply_nunchaku_patches(): else: print(f"[Nunchaku Compat] TiledDiffusion: Warning - 5D tensor with F={F} (not 1), cannot safely squeeze") - return original_tiled_call(self, model_function, kwargs) + return _original_tiled_call(self, model_function, kwargs) tiled_diff.TiledDiffusion.__call__ = patched_tiled_call print("[Nunchaku Compat] Successfully patched TiledDiffusion for 5D tensor support") @@ -238,5 +247,71 @@ def apply_nunchaku_patches(): traceback.print_exc() -# Auto-apply patches on import -apply_nunchaku_patches() +def remove_nunchaku_patches(): + """Remove Nunchaku compatibility patches""" + global _patch_applied, _original_module_call, _original_tiled_call + + if not _patch_applied: + print("[Nunchaku Compat] Patches not applied, nothing to remove") + return + + try: + import torch.nn as nn + + # Restore nn.Module.__call__ + if _original_module_call is not None: + nn.Module.__call__ = _original_module_call + print("[Nunchaku Compat] Restored original nn.Module.__call__") + + # Restore TiledDiffusion.__call__ + try: + import sys + if 'ComfyUI-TiledDiffusion.tiled_diffusion' in sys.modules: + tiled_diff = sys.modules['ComfyUI-TiledDiffusion.tiled_diffusion'] + if hasattr(tiled_diff, 'TiledDiffusion') and _original_tiled_call is not None: + tiled_diff.TiledDiffusion.__call__ = _original_tiled_call + print("[Nunchaku Compat] Restored original TiledDiffusion.__call__") + except Exception as e: + print(f"[Nunchaku Compat] Error restoring TiledDiffusion: {e}") + + _patch_applied = False + print("[Nunchaku Compat] Successfully removed Nunchaku compatibility patches") + + except Exception as e: + print(f"[Nunchaku Compat] Error removing patches: {e}") + import traceback + traceback.print_exc() + + +class NunchakuQwenPatches: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "mode": (["enable", "disable"], {"default": "enable"}), + }, + "optional": { + "model": ("MODEL",), + "image": ("IMAGE",), + } + } + + RETURN_TYPES = ("MODEL", "IMAGE",) + RETURN_NAMES = ("model", "image",) + FUNCTION = "execute" + CATEGORY = "utils" + + def execute(self, mode, model=None, image=None): + if mode == "enable": + apply_nunchaku_patches() + else: + remove_nunchaku_patches() + return (model, image) + +NODE_CLASS_MAPPINGS = { + "NunchakuQwenPatches": NunchakuQwenPatches +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "NunchakuQwenPatches": "Nunchaku Qwen Patches" +} diff --git a/output.txt b/output.txt new file mode 100644 index 0000000..c4f2092 Binary files /dev/null and b/output.txt differ diff --git a/test_extraction_batch.py b/test_extraction_batch.py new file mode 100644 index 0000000..0c5ede7 --- /dev/null +++ b/test_extraction_batch.py @@ -0,0 +1,58 @@ +import os +import random +import sys + +# Add custom node directory to path so we can import the class +sys.path.append(os.path.dirname(__file__)) + +# Mock folder_paths for the node to work +import folder_paths +folder_paths.get_annotated_filepath = lambda x: x # Just return the path as is + +from extract_metadata_node import ImageMetadataExtractor + +def main(): + output_dir = r"D:\ComfyUI7\ComfyUI\output" + output_file = "batch_test_results.txt" + + # Get all image files + all_files = [os.path.join(output_dir, f) for f in os.listdir(output_dir) + if f.lower().endswith(('.png', '.jpg', '.jpeg', '.webp'))] + + if not all_files: + print(f"No images found in {output_dir}") + return + + # Select 20 random images + num_samples = min(20, len(all_files)) + selected_files = random.sample(all_files, num_samples) + + extractor = ImageMetadataExtractor() + + print(f"Testing on {num_samples} images...") + + with open(output_file, "w", encoding="utf-8") as f: + for i, file_path in enumerate(selected_files): + try: + # We need to bypass the folder_paths.get_annotated_filepath call inside the node + # by mocking it, or just passing the absolute path if our mock above works. + # The node calls folder_paths.get_annotated_filepath(image) + # Our mock returns x, so we pass the full path. + + prompt, width, height = extractor.extract_metadata(file_path) + + separator = "=" * 80 + entry = f"{separator}\nImage: {os.path.basename(file_path)}\nDimensions: {width}x{height}\nPrompt:\n{prompt}\n" + + f.write(entry + "\n") + print(f"Processed {i+1}/{num_samples}: {os.path.basename(file_path)}") + + except Exception as e: + error_msg = f"Error processing {os.path.basename(file_path)}: {e}\n" + f.write(error_msg) + print(error_msg) + + print(f"Done. Results saved to {output_file}") + +if __name__ == "__main__": + main()