clean up, readme

FaceEnhancementProd.py

@@ -7,7 +7,9 @@ import torch
 BASE_PATH = "./"
 COMFYUI_PATH = os.path.join(BASE_PATH, "ComfyUI")
 
-# Declare models as a global variable at the top of the script
+"""
+To avoid loading the models each time, we store them in a global variable.
+"""
 models = None
 
 def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:

README.md

@@ -40,14 +40,20 @@ This will
 - Install ComfyUI, custom nodes, and required 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
+```
+
 ## 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 embed distance
+- `./workflows/FaceEmbedDist.json` for computing the face embedding distance
 
 
-## Configuration
+<!-- ## Configuration
 
 Create a .env file in the project root directory with your API keys:
 ```
@@ -55,9 +61,9 @@ touch .env
 echo "FAL_API_KEY=your_fal_api_key_here" >> .env
 ```
 
-The FAL API key is used for face upscaling during preprocessing. You can get one at [fal.ai](https://fal.ai/).
+The FAL API key is used for face upscaling during preprocessing. You can get one at [fal.ai](https://fal.ai/). -->
 
-# Gradio Demo
+## Gradio Demo
 
 A simple web interface for the face enhancement workflow. 
 
@@ -70,6 +76,7 @@ python gradio_demo.py
 
 ### Notes
 - The script and demo run a ComfyUI server ephemerally
+- Gradio demo faster than the script since models remain loaded in memory
 - All images are saved in ./ComfyUI/input/scratch/
 - Temporary files are created during processing and cleaned up afterward
 

demo.py

@@ -1,10 +1,33 @@
 import gradio as gr
 import os
 import tempfile
+import hashlib
+import io
+import pickle
+import pathlib
+import sys
 from main import process_face
 from PIL import Image
 
 PORT = 7860
+CACHE_DIR = "./cache"
+
+# Ensure cache directory exists
+os.makedirs(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
+    """
+    img_bytes = io.BytesIO()
+    img.save(img_bytes, format='PNG')
+    return hashlib.md5(img_bytes.getvalue()).hexdigest()
 
 def enhance_face_gradio(input_image, ref_image):
     """
@@ -17,6 +40,23 @@ def enhance_face_gradio(input_image, ref_image):
     Returns:
         PIL Image: Enhanced image
     """
+    # Generate hashes for both images
+    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")
+    
+    # Check if result exists in cache
+    if os.path.exists(cache_path):
+        try:
+            with open(cache_path, 'rb') as f:
+                result_img = pickle.load(f)
+                print(f"Returning cached result for images with hash {combined_hash}")
+                return result_img
+        except (pickle.PickleError, IOError) as e:
+            print(f"Error loading from cache: {e}")
+            # Continue to processing if cache load fails
+    
     # Create temporary files for input, reference, and output
     with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as input_file, \
          tempfile.NamedTemporaryFile(suffix=".png", delete=False) as ref_file, \
@@ -39,18 +79,27 @@ def enhance_face_gradio(input_image, ref_image):
             upscale=False,
             output_path=output_path
         )
-        pass
     except Exception as e:
         # Handle the error, log it, and return an error message
         print(f"Error processing face: {e}")
         return "An error occurred while processing the face. Please try again."
-
     finally:
         # Clean up temporary input and reference files
         os.unlink(input_path)
         os.unlink(ref_path)
     
-    return Image.open(output_path)
+    # Load the output image
+    result_img = Image.open(output_path)
+    
+    # Cache the result
+    try:
+        with open(cache_path, 'wb') as f:
+            pickle.dump(result_img, f)
+            print(f"Cached result for images with hash {combined_hash}")
+    except (pickle.PickleError, IOError) as e:
+        print(f"Error caching result: {e}")
+    
+    return result_img
 
 def create_gradio_interface():
     # Create the Gradio interface

main.py

@@ -10,6 +10,7 @@ def parse_args():
     parser.add_argument('--crop', action='store_true', help='Whether to crop the image')
     parser.add_argument('--upscale', action='store_true', help='Whether to upscale the image')
     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
@@ -46,7 +47,7 @@ def create_scratch_dir():
     
     return new_dir
 
-def process_face(input_path, ref_path, crop=False, upscale=False, output_path=None):
+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.
     
@@ -83,7 +84,7 @@ def process_face(input_path, ref_path, crop=False, upscale=False, output_path=No
     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=0.75)
+    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}")
@@ -97,7 +98,8 @@ def main():
         ref_path=args.ref, 
         crop=args.crop, 
         upscale=args.upscale, 
-        output_path=args.output
+        output_path=args.output,
+        id_weight=args.id_weight
     )
 
 if __name__ == "__main__":

requirements.txt

@@ -3,7 +3,6 @@ huggingface_hub[hf_transfer]
 comfy-cli
 python-dotenv
 requests
-openai
 fal-client
 gradio>=3.50.2
 pillow>=10.0.0

workflows/FaceDistanceProd.json

@@ -0,0 +1,404 @@
+{
+  "last_node_id": 8,
+  "last_link_id": 8,
+  "nodes": [
+    {
+      "id": 1,
+      "type": "FaceEmbedDistance",
+      "pos": [
+        4701.9052734375,
+        1165.7012939453125
+      ],
+      "size": [
+        315,
+        170
+      ],
+      "flags": {},
+      "order": 4,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "analysis_models",
+          "type": "ANALYSIS_MODELS",
+          "link": 1
+        },
+        {
+          "name": "reference",
+          "type": "IMAGE",
+          "link": 5
+        },
+        {
+          "name": "image",
+          "type": "IMAGE",
+          "link": 7
+        }
+      ],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            2
+          ],
+          "slot_index": 0
+        },
+        {
+          "name": "distance",
+          "type": "FLOAT",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "FaceEmbedDistance"
+      },
+      "widgets_values": [
+        "cosine",
+        100,
+        0,
+        true
+      ]
+    },
+    {
+      "id": 2,
+      "type": "PreviewImage",
+      "pos": [
+        5130.59228515625,
+        1011.3296508789062
+      ],
+      "size": [
+        375.2000732421875,
+        279.4446716308594
+      ],
+      "flags": {},
+      "order": 6,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 2
+        }
+      ],
+      "outputs": [],
+      "properties": {
+        "Node name for S&R": "PreviewImage"
+      },
+      "widgets_values": []
+    },
+    {
+      "id": 3,
+      "type": "FaceEmbedDistance",
+      "pos": [
+        4699.32275390625,
+        1407.7418212890625
+      ],
+      "size": [
+        315,
+        170
+      ],
+      "flags": {},
+      "order": 5,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "analysis_models",
+          "type": "ANALYSIS_MODELS",
+          "link": 3
+        },
+        {
+          "name": "reference",
+          "type": "IMAGE",
+          "link": 6
+        },
+        {
+          "name": "image",
+          "type": "IMAGE",
+          "link": 8
+        }
+      ],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            4
+          ],
+          "slot_index": 0
+        },
+        {
+          "name": "distance",
+          "type": "FLOAT",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "FaceEmbedDistance"
+      },
+      "widgets_values": [
+        "cosine",
+        100,
+        0,
+        true
+      ]
+    },
+    {
+      "id": 5,
+      "type": "PreviewImage",
+      "pos": [
+        5118.4765625,
+        1408.8367919921875
+      ],
+      "size": [
+        375.2000732421875,
+        279.4446716308594
+      ],
+      "flags": {},
+      "order": 7,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 4
+        }
+      ],
+      "outputs": [],
+      "properties": {
+        "Node name for S&R": "PreviewImage"
+      },
+      "widgets_values": []
+    },
+    {
+      "id": 4,
+      "type": "FaceAnalysisModels",
+      "pos": [
+        4041.2080078125,
+        1888.7574462890625
+      ],
+      "size": [
+        315,
+        82
+      ],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "ANALYSIS_MODELS",
+          "type": "ANALYSIS_MODELS",
+          "links": [
+            1,
+            3
+          ],
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "FaceAnalysisModels"
+      },
+      "widgets_values": [
+        "insightface",
+        "CUDA"
+      ]
+    },
+    {
+      "id": 6,
+      "type": "LoadImage",
+      "pos": [
+        4040.829345703125,
+        790.5192260742188
+      ],
+      "size": [
+        315,
+        314
+      ],
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            5,
+            6
+          ],
+          "slot_index": 0
+        },
+        {
+          "name": "MASK",
+          "type": "MASK",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "LoadImage"
+      },
+      "widgets_values": [
+        "dany_face.png",
+        "image"
+      ]
+    },
+    {
+      "id": 7,
+      "type": "LoadImage",
+      "pos": [
+        4036.8935546875,
+        1155.6903076171875
+      ],
+      "size": [
+        315,
+        314
+      ],
+      "flags": {},
+      "order": 2,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            7
+          ],
+          "slot_index": 0
+        },
+        {
+          "name": "MASK",
+          "type": "MASK",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "LoadImage"
+      },
+      "widgets_values": [
+        "chatgpt_dany_1.png",
+        "image"
+      ]
+    },
+    {
+      "id": 8,
+      "type": "LoadImage",
+      "pos": [
+        4037.78955078125,
+        1522.9161376953125
+      ],
+      "size": [
+        315,
+        314
+      ],
+      "flags": {},
+      "order": 3,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "IMAGE",
+          "type": "IMAGE",
+          "links": [
+            8
+          ],
+          "slot_index": 0
+        },
+        {
+          "name": "MASK",
+          "type": "MASK",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "LoadImage"
+      },
+      "widgets_values": [
+        "chatgpt_dany_2.png",
+        "image"
+      ]
+    }
+  ],
+  "links": [
+    [
+      1,
+      4,
+      0,
+      1,
+      0,
+      "ANALYSIS_MODELS"
+    ],
+    [
+      2,
+      1,
+      0,
+      2,
+      0,
+      "IMAGE"
+    ],
+    [
+      3,
+      4,
+      0,
+      3,
+      0,
+      "ANALYSIS_MODELS"
+    ],
+    [
+      4,
+      3,
+      0,
+      5,
+      0,
+      "IMAGE"
+    ],
+    [
+      5,
+      6,
+      0,
+      1,
+      1,
+      "IMAGE"
+    ],
+    [
+      6,
+      6,
+      0,
+      3,
+      1,
+      "IMAGE"
+    ],
+    [
+      7,
+      7,
+      0,
+      1,
+      2,
+      "IMAGE"
+    ],
+    [
+      8,
+      8,
+      0,
+      3,
+      2,
+      "IMAGE"
+    ]
+  ],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 0.8390545288824041,
+      "offset": [
+        -4016.1357592321388,
+        -971.9468925845375
+      ]
+    },
+    "node_versions": {
+      "comfyui_faceanalysis": "4919e4e931db0edb219ba5086b3c10b8af750631",
+      "comfy-core": "0.3.26"
+    },
+    "ue_links": []
+  },
+  "version": 0.4
+}