cleanup; readme

FaceEnhancementProd.py

@@ -4,13 +4,12 @@ import sys
 from typing import Sequence, Mapping, Any, Union
 import torch
 
-BASE_PATH = "./"
-COMFYUI_PATH = os.path.join(BASE_PATH, "ComfyUI")
+COMFYUI_PATH = "./ComfyUI"
 
 """
 To avoid loading the models each time, we store them in a global variable.
 """
-models = None
+COMFY_MODELS = None
 
 def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
     """Returns the value at the given index of a sequence or mapping.
@@ -74,7 +73,7 @@ def add_extra_model_paths() -> None:
     Parse the optional extra_model_paths.yaml file and add the parsed paths to the sys.path.
     """
     try:
-        from main import load_extra_path_config
+        from test import load_extra_path_config
     except ImportError:
         print(
             "Could not import load_extra_path_config from main.py. Looking in utils.extra_config instead."
@@ -177,10 +176,10 @@ def load_models():
     }
 
 def initialize_models():
-    global models
-    if models is None:
+    global COMFY_MODELS
+    if COMFY_MODELS is None:
         import_custom_nodes()  # Ensure NODE_CLASS_MAPPINGS is initialized
-        models = load_models()
+        COMFY_MODELS = load_models()
 
 initialize_models()
 
@@ -192,17 +191,17 @@ def main(
     positive_prompt: str = "",
     id_weight: float = 0.75,
 ):
-    global models
-    if models is None:
+    global COMFY_MODELS
+    if COMFY_MODELS is None:
         raise ValueError("Models must be initialized before calling main(). Call initialize_models() first.")
     with torch.inference_mode():
-        dualcliploader_94 = models["dualcliploader_94"]
-        vaeloader_95 = models["vaeloader_95"]
-        pulidfluxmodelloader_44 = models["pulidfluxmodelloader_44"]
-        pulidfluxevacliploader_45 = models["pulidfluxevacliploader_45"]
-        pulidfluxinsightfaceloader_46 = models["pulidfluxinsightfaceloader_46"]
-        controlnetloader_49 = models["controlnetloader_49"]
-        unetloader_93 = models["unetloader_93"]
+        dualcliploader_94 = COMFY_MODELS["dualcliploader_94"]
+        vaeloader_95 = COMFY_MODELS["vaeloader_95"]
+        pulidfluxmodelloader_44 = COMFY_MODELS["pulidfluxmodelloader_44"]
+        pulidfluxevacliploader_45 = COMFY_MODELS["pulidfluxevacliploader_45"]
+        pulidfluxinsightfaceloader_46 = COMFY_MODELS["pulidfluxinsightfaceloader_46"]
+        controlnetloader_49 = COMFY_MODELS["controlnetloader_49"]
+        unetloader_93 = COMFY_MODELS["unetloader_93"]
 
         cliptextencode = CLIPTextEncode()
         cliptextencode_23 = cliptextencode.encode(
@@ -237,10 +236,10 @@ def main(
         basicscheduler = NODE_CLASS_MAPPINGS["BasicScheduler"]()
         samplercustomadvanced = NODE_CLASS_MAPPINGS["SamplerCustomAdvanced"]()
         vaedecode = VAEDecode()
-        faceembeddistance = NODE_CLASS_MAPPINGS["FaceEmbedDistance"]()
-        display_any_rgthree = NODE_CLASS_MAPPINGS["Display Any (rgthree)"]()
-        image_comparer_rgthree = NODE_CLASS_MAPPINGS["Image Comparer (rgthree)"]()
-        saveimage = SaveImage()
+        # faceembeddistance = NODE_CLASS_MAPPINGS["FaceEmbedDistance"]()
+        # display_any_rgthree = NODE_CLASS_MAPPINGS["Display Any (rgthree)"]()
+        # image_comparer_rgthree = NODE_CLASS_MAPPINGS["Image Comparer (rgthree)"]()
+        # saveimage = SaveImage()
 
         applypulidflux_133 = applypulidflux.apply_pulid_flux(
             weight=id_weight,

README.md

@@ -27,7 +27,7 @@ A tool for improving facial consistency and quality in AI-generated images. Dram
 
 1. Set up your Hugging Face token:
    - Create a token at [Hugging Face](https://huggingface.co/settings/tokens)
-   - Log into Hugging Face and accept their terms of service to download Flux
+   - Log into Hugging Face and accept their terms of service to download [Flux.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev)
    - Set the following environment variables:
      ```
      export HUGGINGFACE_TOKEN=your_token_here
@@ -48,21 +48,26 @@ A tool for improving facial consistency and quality in AI-generated images. Dram
    ```
 
    This will
-   - Install ComfyUI, custom nodes, and required dependencies to your venv
+   - Install ComfyUI, custom nodes, and remaining dependencies to your venv
    - Download all required models (Flux.1-dev, ControlNet, text encoders, PuLID, and more)
 
 4. Run inference on one example:
 
    ```
-   python main.py --input examples/dany_gpt_1.png --ref examples/dany_face.jpg --out examples/dany_enhanced.png
+   python test.py --input examples/dany_gpt_1.png --ref examples/dany_face.jpg --out examples/dany_enhanced.png
    ```
 
-## Running on ComfyUI
-
-Using the ComfyUI workflows is the fastest way to get started. Run `python run_comfy.py`
-- `./workflows/FaceEnhancementProd.json` for face enhancement
-- `./workflows/FaceEmbedDist.json` for computing the face embedding distance
+   <details>
+   <summary>Arguments</summary>
 
+   - `--input` (str): Path to the input image.
+   - `--ref` (str): Path to the reference face image.
+   - `--output` (str): Path to save the output image.
+   - `--crop` (store_true): Flag to face crop the reference image. Default: False.
+   - `--upscale` (store_true): Flag to upscale the (cropped) reference image. Default: False.
+   - `--caption` (store_true): Flag to caption the input image. Default: False.
+   - `--id_weight` (float): Face ID weight. Default: 0.75.
+   </details>
 
 ## Gradio Demo
 
@@ -72,15 +77,22 @@ A simple web interface for the face enhancement workflow.
 
 2. Go to http://localhost:7860. You may need to enable port forwarding.
 
+## Running on ComfyUI
+
+Using the ComfyUI workflows is the fastest way to get started. Run `python run_comfy.py`. There are two workflows:
+- `./workflows/FaceEnhancementProd.json` for face enhancement
+- `./workflows/FaceEmbedDist.json` for computing the [face embedding distance](https://github.com/cubiq/ComfyUI_FaceAnalysis)
+
+
 ### Notes
 - The script and demo run a ComfyUI server ephemerally
-- Gradio demo is faster than the script because models remain loaded in memory
+- Gradio demo is faster than the script because the models remain loaded in memory and ComfyUI server is booted up.
 - All images are saved in `./ComfyUI/input/scratch/`
-- Temporary files are created during processing and cleaned up afterward
-- Face cropping and upscaling are not applied to the reference image; this will be added in an update.
+- `FaceEnhancementProd.py` was created with the [ComfyUI-to-Python-Extension](https://github.com/pydn/ComfyUI-to-Python-Extension) and re-engineered for efficiency and function.
+- Face cropping, upscaling, and captioning are unavailable; these will be added in an update.
 
 ### Troubleshooting
 
-- **Out of memory errors**: If your GPU has less than 48 GB VRAM, install [Flux.1-dev at float8 precision](https://huggingface.co/Comfy-Org/flux1-dev).
+- **Out of memory errors**: If your GPU has less than 48 GB VRAM, install [Flux.1-dev at fp8 precision](https://huggingface.co/Comfy-Org/flux1-dev).
 - **Face detection issues**: This method works for photorealistic images of people. It may not work on cartoons, anime characters, or non-human subjects.
 - **Downloading models fails**: Check your Hugging Face token has proper permissions.

demo.py

@@ -1,37 +1,27 @@
+import os
 from install import install
 
-# Global variable to track if install() has been run
-INSTALLED = False
-if not INSTALLED:
-    install()
-    INSTALLED = True
+if "HF_DEMO" in os.environ:
+    # Global variable to track if install() has been run; only for deploying on HF space
+    INSTALLED = False
+    if not INSTALLED:
+        install()
+        INSTALLED = True
 
 import gradio as gr
-import os
 import tempfile
 import hashlib
 import io
 import pickle
-import pathlib
 import sys
-from main import process_face
+from test import process_face
 from PIL import Image
 
-CACHE_DIR = "./cache"
-
-# Ensure cache directory exists
-os.makedirs(CACHE_DIR, exist_ok=True)
+INPUT_CACHE_DIR = "./cache"
+os.makedirs(INPUT_CACHE_DIR, exist_ok=True)
 
 def get_image_hash(img):
-    """
-    Generate a hash of the image content.
-    
-    Args:
-        img: PIL Image
-        
-    Returns:
-        str: Hash of the image
-    """
+    """Generate a hash of the image content."""
     img_bytes = io.BytesIO()
     img.save(img_bytes, format='PNG')
     return hashlib.md5(img_bytes.getvalue()).hexdigest()
@@ -51,7 +41,7 @@ def enhance_face_gradio(input_image, ref_image):
     input_hash = get_image_hash(input_image)
     ref_hash = get_image_hash(ref_image)
     combined_hash = f"{input_hash}_{ref_hash}"
-    cache_path = os.path.join(CACHE_DIR, f"{combined_hash}.pkl")
+    cache_path = os.path.join(INPUT_CACHE_DIR, f"{combined_hash}.pkl")
     
     # Check if result exists in cache
     if os.path.exists(cache_path):
@@ -78,7 +68,6 @@ def enhance_face_gradio(input_image, ref_image):
     ref_image.save(ref_path)
     
     try:
-        # Process the face
         process_face(
             input_path=input_path,
             ref_path=ref_path,
@@ -109,9 +98,7 @@ def enhance_face_gradio(input_image, ref_image):
     return result_img
 
 def create_gradio_interface():
-    # Create the Gradio interface
     with gr.Blocks(title="Face Enhancement Demo") as demo:
-        # Add instructions at the top
         gr.Markdown("""
         # Face Enhancement
         ### Instructions
@@ -124,10 +111,8 @@ def create_gradio_interface():
         For more information, check out my [blog post](https://rishidesai.github.io/posts/face-enhancement-techniques/).
         """, elem_id="instructions")
 
-        # Add a horizontal line for separation
         gr.Markdown("---")
 
-        # Main interface layout
         with gr.Row():
             with gr.Column():
                 input_image = gr.Image(label="Target Image", type="pil")
@@ -144,7 +129,6 @@ def create_gradio_interface():
             queue=True  # Enable queue for sequential processing
         )
 
-        # Add examples using gr.Examples
         gr.Markdown("## Examples\nClick on an example to load the images into the interface.")
         example_inps = [
             ["examples/dany_gpt_1.png", "examples/dany_face.jpg"],

install.py

@@ -1,62 +1,47 @@
 import os
 
-# Define paths
-BASE_PATH = "./"
-COMFYUI_PATH = os.path.join(BASE_PATH, "ComfyUI")
-MODEL_PATH = os.path.join(COMFYUI_PATH, "models")
+COMFYUI_PATH = "./ComfyUI"
 
 CACHE_PATH = os.getenv('HF_HOME')
 os.makedirs(CACHE_PATH, exist_ok=True)
 
 
-def run_command(command):
-    """Run a shell command using os.system() (simpler than subprocess)."""
+def run_cmd(command):
+    """Run a shell command"""
     print(f"🔄 Running: {command}")
     exit_code = os.system(command)
-
     if exit_code != 0:
         print(f"❌ Command failed: {command} (Exit Code: {exit_code})")
-        exit(1)  # Exit on failure
-
+        exit(1)
 
-def manage_git_repo(repo_url, install_path, requirements=False, submodules=False):
-    """Clone or update a git repository and handle its dependencies.
 
-    Args:
-        repo_url: URL of the git repository
-        install_path: Where to install/update the repository
-        requirements: Whether to install requirements.txt
-        submodules: Whether to update git submodules
-    """
-    # Save the original directory
+def install_git_repo(repo_url, install_path, requirements=False, submodules=False):
+    """Clone or update a git repository and handle its dependencies"""
     original_dir = os.getcwd()
 
     if not os.path.exists(install_path) or not os.path.isdir(install_path) or not os.path.exists(
             os.path.join(install_path, ".git")):
         print(f"📂 Cloning {os.path.basename(install_path)}...")
-        run_command(f"git clone {repo_url} {install_path}")
+        run_cmd(f"git clone {repo_url} {install_path}")
     else:
         print(f"🔄 {os.path.basename(install_path)} exists. Checking for updates...")
 
     # Change to repo directory and update
     os.chdir(install_path)
-    run_command("git pull")
+    run_cmd("git pull")
 
     if submodules:
-        run_command("git submodule update --init --recursive")
-
+        run_cmd("git submodule update --init --recursive")
     if requirements:
-        run_command("python -m pip install -r requirements.txt")
+        run_cmd("python -m pip install -r requirements.txt")
 
     print(f"✅ {os.path.basename(install_path)} installed and updated.")
-
-    # Change back to the original directory
     os.chdir(original_dir)
 
 
 def install_comfyui():
     """Clone and set up ComfyUI if not already installed."""
-    manage_git_repo(
+    install_git_repo(
         "https://github.com/comfyanonymous/ComfyUI.git",
         COMFYUI_PATH,
         requirements=True
@@ -77,10 +62,7 @@ def download_huggingface_models():
         {"repo_id": "comfyanonymous/flux_text_encoders", "filename": "clip_l.safetensors", "folder": "text_encoders"},
     ]
 
-    # Get the Hugging Face token from the environment variable
-    huggingface_token = os.getenv('HUGGINGFACE_TOKEN')
-
-    # Dictionary mapping repo_ids to specific filenames
+    # More specific filenames
     filename_mappings = {
         "Shakker-Labs/FLUX.1-dev-ControlNet-Union-Pro": "Flux_Dev_ControlNet_Union_Pro_ShakkerLabs.safetensors",
     }
@@ -92,9 +74,9 @@ def download_huggingface_models():
                 filename=model["filename"],
                 cache_dir=CACHE_PATH,
                 repo_type=model.get("repo_type", "model"),
-                token=huggingface_token
+                token=os.getenv('HUGGINGFACE_TOKEN')
             )
-            target_dir = os.path.join(MODEL_PATH, model["folder"])
+            target_dir = os.path.join(COMFYUI_PATH, "models", model["folder"])
             os.makedirs(target_dir, exist_ok=True)
 
             # Use mapping if it exists, otherwise use original filename
@@ -114,7 +96,7 @@ def download_and_extract_antelopev2():
     """Download and extract AntelopeV2 model for insightface."""
     import zipfile, requests, shutil
 
-    base_path = os.path.join(MODEL_PATH, "insightface/models")
+    base_path = os.path.join(COMFYUI_PATH, "models", "insightface/models")
     model_target_path = os.path.join(base_path, "antelopev2")
     download_url = "https://huggingface.co/MonsterMMORPG/tools/resolve/main/antelopev2.zip"
     zip_path = os.path.join(base_path, "antelopev2.zip")
@@ -174,12 +156,7 @@ def download_and_extract_antelopev2():
 def install_custom_nodes():
     """Install all custom nodes for ComfyUI."""
 
-    custom_nodes_git = [
-        {
-            "repo": "https://github.com/ltdrdata/ComfyUI-Manager",
-            "name": "ComfyUI-Manager",
-            "requirements": True
-        },
+    custom_nodes = [
         {
             "repo": "https://github.com/sipie800/ComfyUI-PuLID-Flux-Enhanced",
             "name": "ComfyUI-PuLID-Flux-Enhanced",
@@ -190,7 +167,7 @@ def install_custom_nodes():
             "name": "rgthree-comfy",
             "requirements": True
         },
-        {  # we already have insightface so don't need requirements (for dlib)
+        {  # we already have insightface so don't need requirements (no dlib)
             "repo": "https://github.com/cubiq/ComfyUI_FaceAnalysis",
             "name": "ComfyUI_FaceAnalysis",
             "requirements": False
@@ -202,10 +179,10 @@ def install_custom_nodes():
         # },
     ]
 
-    for node in custom_nodes_git:
+    for node in custom_nodes:
         repo_name = node["name"]
         repo_path = os.path.join(COMFYUI_PATH, "custom_nodes", repo_name)
-        manage_git_repo(
+        install_git_repo(
             node["repo"],
             repo_path,
             requirements=node.get("requirements", False),
@@ -221,7 +198,7 @@ def install():
     install_custom_nodes()
     download_huggingface_models()
     download_and_extract_antelopev2()
-    print("🎉 Setup Complete! Run `run.py` to start ComfyUI.")  
+    print("🎉 Setup Complete!")
 
 if __name__ == "__main__":
     install()

run_comfy.py

@@ -1,18 +1,19 @@
 import os
 import subprocess
+import argparse
 
 COMFYUI_PATH = "./ComfyUI"
-PORT = 8000
 
-
-def run_comfyui():
+def run_comfyui(port):
     """Launch ComfyUI with external access."""
     os.chdir(COMFYUI_PATH)
-    print(f"🚀 Launching ComfyUI on port {PORT}...")
-
-    subprocess.run(f"python main.py --listen 0.0.0.0 --port {PORT} --disable-auto-launch", shell=True)
+    print(f"🚀 Launching ComfyUI on port {port}...")
 
+    subprocess.run(f"python main.py --listen 0.0.0.0 --port {port} --disable-auto-launch", shell=True)
 
 if __name__ == "__main__":
-    run_comfyui()
-    print(f"Enable port-forwarding\nssh -L {PORT}:localhost:{PORT} [NAME]@[IP_ADDRESS] -p [SERVER_PORT] -i [PRIVATE_KEY]")
+    parser = argparse.ArgumentParser(description="Run ComfyUI")
+    parser.add_argument('--port', type=int, default=8000, help='Port to run ComfyUI on')
+    args = parser.parse_args()
+
+    run_comfyui(args.port)

main.py → test.py

@@ -12,88 +12,82 @@ def parse_args():
     parser.add_argument('--output', type=str, required=True, help='Path to save the output image')
     parser.add_argument('--id_weight', type=float, default=0.75, help='face ID weight')
     args = parser.parse_args()
-    
-    # Validate input file exists
+
     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):
         parser.error(f"Output directory does not exist: {output_dir}")
-    
     return args
 
 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
 
 def process_face(input_path, ref_path, crop=False, upscale=False, output_path=None, id_weight=0.75):
     """
     Process a face image using the given parameters.
-        
+
     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", id_weight=id_weight)
-    
+
     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, 
+        input_path=args.input,
+        ref_path=args.ref,
+        crop=args.crop,
+        upscale=args.upscale,
         output_path=args.output,
         id_weight=args.id_weight
     )
 
 if __name__ == "__main__":
-    main() 
+    main()