addig params

FaceEnhancementProd.py

@@ -61,7 +61,7 @@ def add_comfyui_directory_to_sys_path() -> None:
     """
     Add 'ComfyUI' to the sys.path
     """
-    return COMFYUI_PATH
+    sys.path.append(COMFYUI_PATH)
 
 
 def add_extra_model_paths() -> None:
@@ -126,7 +126,13 @@ from nodes import (
 )
 
 
-def main():
+def main(
+    face_image: str,
+    input_image: str,
+    output_image: str,
+    dist_image: str = None,
+    positive_prompt: str = ""
+):
     import_custom_nodes()
     with torch.inference_mode():
         dualcliploader = DualCLIPLoader()
@@ -143,9 +149,9 @@ def main():
         )
 
         loadimage = LoadImage()
-        loadimage_24 = loadimage.load_image(image="woman_face.jpg")
+        loadimage_24 = loadimage.load_image(image=face_image)
 
-        loadimage_40 = loadimage.load_image(image="woman1_gpt_2.png")
+        loadimage_40 = loadimage.load_image(image=input_image)
 
         vaeloader = VAELoader()
         vaeloader_95 = vaeloader.load_vae(vae_name="ae.safetensors")
@@ -160,7 +166,7 @@ def main():
         randomnoise_39 = randomnoise.get_noise(noise_seed=random.randint(1, 2**64))
 
         cliptextencode_42 = cliptextencode.encode(
-            text="", clip=get_value_at_index(dualcliploader_94, 0)
+            text=positive_prompt, clip=get_value_at_index(dualcliploader_94, 0)
         )
 
         pulidfluxmodelloader = NODE_CLASS_MAPPINGS["PulidFluxModelLoader"]()
@@ -277,16 +283,50 @@ def main():
             source=get_value_at_index(faceembeddistance_117, 1)
         )
 
-        saveimage_128 = saveimage.save_images(
-            filename_prefix="FaceEnhanced",
-            images=get_value_at_index(vaedecode_114, 0),
-        )
+        # Save using direct image saving
+        save_comfy_images(get_value_at_index(vaedecode_114, 0), [output_image])
+        if dist_image:
+            save_comfy_images(get_value_at_index(faceembeddistance_117, 0), [dist_image])
+
+
+def save_comfy_images(images, output_dirs):
+    # images is a PyTorch tensor with shape [batch_size, height, width, channels]
+    import numpy as np
+    from PIL import Image
+    
+    for idx, image in enumerate(images):
+        # Create the output directory if it doesn't exist
+        output_dir = os.path.dirname(output_dirs[idx])
+        if output_dir and not os.path.exists(output_dir):
+            os.makedirs(output_dir, exist_ok=True)
+            
+        numpy_image = 255. * image.cpu().numpy()
+        numpy_image = np.clip(numpy_image, 0, 255).astype(np.uint8)
+        pil_image = Image.fromarray(numpy_image)
+        pil_image.save(output_dirs[idx])
 
-        saveimage_129 = saveimage.save_images(
-            filename_prefix="FaceEmbedDist",
-            images=get_value_at_index(faceembeddistance_117, 0),
-        )
 
+def enhance_face(face_image: str, input_image: str, output_image: str, dist_image: str = None, positive_prompt: str = ""):
+    main(face_image, input_image, output_image, dist_image, positive_prompt)
 
 if __name__ == "__main__":
-    main()
+    import sys
+    
+    if len(sys.argv) < 4:
+        print("Usage: python FaceEnhancementProd.py face_image input_image output_image [dist_image] [positive_prompt]")
+        sys.exit(1)
+    
+    face_image = sys.argv[1]
+    input_image = sys.argv[2]
+    output_image = sys.argv[3]
+    
+    dist_image = None
+    positive_prompt = ""
+    
+    if len(sys.argv) > 4:
+        dist_image = sys.argv[4]
+    
+    if len(sys.argv) > 5:
+        positive_prompt = sys.argv[5]
+    
+    main(face_image, input_image, output_image, dist_image, positive_prompt)

main.py

@@ -1,6 +1,7 @@
 import argparse
 import os
-from utils import crop_face, upscale_image
+import shutil
+from FaceEnhancementProd import enhance_face
 
 def parse_args():
     parser = argparse.ArgumentParser(description='Face Enhancement Tool')
@@ -15,6 +16,10 @@ def parse_args():
     if not os.path.exists(args.input):
         parser.error(f"Input file does not exist: {args.input}")
     
+    # Validate reference file exists
+    if not os.path.exists(args.ref):
+        parser.error(f"Reference file does not exist: {args.ref}")
+    
     # Validate output directory exists
     output_dir = os.path.dirname(args.output)
     if output_dir and not os.path.exists(output_dir):
@@ -22,22 +27,78 @@ def parse_args():
     
     return args
 
-def main():
-    args = parse_args()
-    print(f"Processing image: {args.input}")
-    print(f"Crop enabled: {args.crop}")
-    print(f"Upscale enabled: {args.upscale}")
-    print(f"Output will be saved to: {args.output}")
-    
-    face_image = args.ref
-    if args.crop:
-        crop_face(face_image, "./scratch/cropped_face.png")
-        face_image = "./scratch/cropped_face.png"
+def create_scratch_dir():
+    """Create a new numbered directory in ./ComfyUI/input/scratch"""
+    base_dir = "./ComfyUI/input/scratch"
+    
+    # Create base directory if it doesn't exist
+    os.makedirs(base_dir, exist_ok=True)
+    
+    # Get existing directories and find the next number
+    existing_dirs = [d for d in os.listdir(base_dir) if os.path.isdir(os.path.join(base_dir, d)) and d.isdigit()]
+    next_dir_num = 1
+    if existing_dirs:
+        next_dir_num = max([int(d) for d in existing_dirs]) + 1
+    
+    # Create the new directory
+    new_dir = os.path.join(base_dir, str(next_dir_num))
+    os.makedirs(new_dir, exist_ok=True)
+    
+    return new_dir
 
-    if args.upscale:
-        upscale_image(face_image, "./scratch/upscaled_face.png")
-        face_image = "./scratch/upscaled_face.png"
+def process_face(input_path, ref_path, crop=False, upscale=False, output_path=None):
+    """
+    Process a face image using the given parameters.
+    
+    Args:
+        input_path (str): Path to the input image
+        ref_path (str): Path to the reference image
+        crop (bool): Whether to crop the image
+        upscale (bool): Whether to upscale the image
+        output_path (str): Path to save the output image
+        
+    Returns:
+        str: Path to the scratch directory used for processing
+    """
+    print(f"Processing image: {input_path}")
+    print(f"Reference image: {ref_path}")
+    print(f"Output will be saved to: {output_path}")
+    
+    # Create a new scratch directory for this run
+    scratch_dir = create_scratch_dir()
+    print(f"Created scratch directory: {scratch_dir}")
     
+    # Copy input and reference images to scratch directory
+    input_filename = os.path.basename(input_path)
+    ref_filename = os.path.basename(ref_path)
+    
+    scratch_input = os.path.join(scratch_dir, input_filename)
+    scratch_ref = os.path.join(scratch_dir, ref_filename)
+    
+    shutil.copy(input_path, scratch_input)
+    shutil.copy(ref_path, scratch_ref)
+    
+    # Convert paths to ComfyUI format (relative to ComfyUI/input/)
+    # For example: "./ComfyUI/input/scratch/1/image.png" becomes "scratch/1/image.png"
+    comfy_ref_path = os.path.relpath(scratch_ref, "./ComfyUI/input")
+    comfy_input_path = os.path.relpath(scratch_input, "./ComfyUI/input")
+    
+    enhance_face(comfy_ref_path, comfy_input_path, output_path, dist_image=f"{output_path}_dist.png")
+    
+    print(f"Enhanced image saved to: {output_path}")
+    print(f"Working files are in: {scratch_dir}")
+    
+    return scratch_dir
+
+def main():
+    args = parse_args()
+    return process_face(
+        input_path=args.input, 
+        ref_path=args.ref, 
+        crop=args.crop, 
+        upscale=args.upscale, 
+        output_path=args.output
+    )
 
 if __name__ == "__main__":
     main()