extracted model loaders

FaceEnhancementProd.py

@@ -7,6 +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
+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.
 
@@ -125,23 +128,79 @@ from nodes import (
     ControlNetApplyAdvanced,
 )
 
+@torch.inference_mode()
+def load_models():
+    dualcliploader = DualCLIPLoader()
+    dualcliploader_94 = dualcliploader.load_clip(
+        clip_name1="t5xxl_fp16.safetensors",
+        clip_name2="clip_l.safetensors",
+        type="flux",
+        device="default",
+    )
+
+    vaeloader = VAELoader()
+    vaeloader_95 = vaeloader.load_vae(vae_name="ae.safetensors")
+
+    pulidfluxmodelloader = NODE_CLASS_MAPPINGS["PulidFluxModelLoader"]()
+    pulidfluxmodelloader_44 = pulidfluxmodelloader.load_model(
+        pulid_file="pulid_flux_v0.9.1.safetensors"
+    )
+
+    pulidfluxevacliploader = NODE_CLASS_MAPPINGS["PulidFluxEvaClipLoader"]()
+    pulidfluxevacliploader_45 = pulidfluxevacliploader.load_eva_clip()
+
+    pulidfluxinsightfaceloader = NODE_CLASS_MAPPINGS["PulidFluxInsightFaceLoader"]()
+    pulidfluxinsightfaceloader_46 = pulidfluxinsightfaceloader.load_insightface(
+        provider="CUDA"
+    )
+
+    controlnetloader = ControlNetLoader()
+    controlnetloader_49 = controlnetloader.load_controlnet(
+        control_net_name="Flux_Dev_ControlNet_Union_Pro_ShakkerLabs.safetensors"
+    )
+
+    unetloader = UNETLoader()
+    unetloader_93 = unetloader.load_unet(
+        unet_name="flux1-dev.safetensors", weight_dtype="default"
+    )
+
+    return {
+        "dualcliploader_94": dualcliploader_94,
+        "vaeloader_95": vaeloader_95,
+        "pulidfluxmodelloader_44": pulidfluxmodelloader_44,
+        "pulidfluxevacliploader_45": pulidfluxevacliploader_45,
+        "pulidfluxinsightfaceloader_46": pulidfluxinsightfaceloader_46,
+        "controlnetloader_49": controlnetloader_49,
+        "unetloader_93": unetloader_93
+    }
+
+def initialize_models():
+    global models
+    if models is None:
+        import_custom_nodes()  # Ensure NODE_CLASS_MAPPINGS is initialized
+        models = load_models()
+
+initialize_models()
 
 def main(
     face_image: str,
     input_image: str,
     output_image: str,
     dist_image: str = None,
-    positive_prompt: str = ""
+    positive_prompt: str = "",
+    # models: dict = None
 ):
-    import_custom_nodes()
+    global models
+    if models is None:
+        raise ValueError("Models must be initialized before calling main(). Call initialize_models() first.")
     with torch.inference_mode():
-        dualcliploader = DualCLIPLoader()
-        dualcliploader_94 = dualcliploader.load_clip(
-            clip_name1="t5xxl_fp16.safetensors",
-            clip_name2="clip_l.safetensors",
-            type="flux",
-            device="default",
-        )
+        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"]
 
         cliptextencode = CLIPTextEncode()
         cliptextencode_23 = cliptextencode.encode(
@@ -153,9 +212,6 @@ def main(
 
         loadimage_40 = loadimage.load_image(image=input_image)
 
-        vaeloader = VAELoader()
-        vaeloader_95 = vaeloader.load_vae(vae_name="ae.safetensors")
-
         vaeencode = VAEEncode()
         vaeencode_35 = vaeencode.encode(
             pixels=get_value_at_index(loadimage_40, 0),
@@ -169,37 +225,9 @@ def main(
             text=positive_prompt, clip=get_value_at_index(dualcliploader_94, 0)
         )
 
-        pulidfluxmodelloader = NODE_CLASS_MAPPINGS["PulidFluxModelLoader"]()
-        pulidfluxmodelloader_44 = pulidfluxmodelloader.load_model(
-            pulid_file="pulid_flux_v0.9.1.safetensors"
-        )
-
-        pulidfluxevacliploader = NODE_CLASS_MAPPINGS["PulidFluxEvaClipLoader"]()
-        pulidfluxevacliploader_45 = pulidfluxevacliploader.load_eva_clip()
-
-        pulidfluxinsightfaceloader = NODE_CLASS_MAPPINGS["PulidFluxInsightFaceLoader"]()
-        pulidfluxinsightfaceloader_46 = pulidfluxinsightfaceloader.load_insightface(
-            provider="CUDA"
-        )
-
-        controlnetloader = ControlNetLoader()
-        controlnetloader_49 = controlnetloader.load_controlnet(
-            control_net_name="Flux_Dev_ControlNet_Union_Pro_ShakkerLabs.safetensors"
-        )
-
         ksamplerselect = NODE_CLASS_MAPPINGS["KSamplerSelect"]()
         ksamplerselect_50 = ksamplerselect.get_sampler(sampler_name="euler")
 
-        unetloader = UNETLoader()
-        unetloader_93 = unetloader.load_unet(
-            unet_name="flux1-dev.safetensors", weight_dtype="default"
-        )
-
-        faceanalysismodels = NODE_CLASS_MAPPINGS["FaceAnalysisModels"]()
-        faceanalysismodels_118 = faceanalysismodels.load_models(
-            library="insightface", provider="CUDA"
-        )
-
         applypulidflux = NODE_CLASS_MAPPINGS["ApplyPulidFlux"]()
         setunioncontrolnettype = NODE_CLASS_MAPPINGS["SetUnionControlNetType"]()
         controlnetapplyadvanced = ControlNetApplyAdvanced()
@@ -269,24 +297,7 @@ def main(
             vae=get_value_at_index(vaeloader_95, 0),
         )
 
-        faceembeddistance_117 = faceembeddistance.analize(
-            similarity_metric="cosine",
-            filter_thresh=100,
-            filter_best=0,
-            generate_image_overlay=True,
-            analysis_models=get_value_at_index(faceanalysismodels_118, 0),
-            reference=get_value_at_index(loadimage_24, 0),
-            image=get_value_at_index(vaedecode_114, 0),
-        )
-
-        display_any_rgthree_121 = display_any_rgthree.main(
-            source=get_value_at_index(faceembeddistance_117, 1)
-        )
-
-        # 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):
@@ -307,26 +318,8 @@ def save_comfy_images(images, output_dirs):
 
 
 def enhance_face(face_image: str, input_image: str, output_image: str, dist_image: str = None, positive_prompt: str = ""):
+    initialize_models()  # Ensure models are loaded
     main(face_image, input_image, output_image, dist_image, positive_prompt)
 
 if __name__ == "__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)
+    pass

README.md

@@ -5,7 +5,7 @@ Enhancing faces in AI generated images.
 
 ### Prerequisites
 - Python 3.8 or higher
-- At least 50GB of free disk space for models and dependencies
+- At least 50GB of free disk space
 
 ### Setup
 
@@ -20,20 +20,16 @@ Enhancing faces in AI generated images.
    ```
    export HF_HOME=/path/to/your/huggingface_cache
    ```
-   This defines where models will be downloaded and then symlinked to the ComfyUI folder.
+   This is where models will be downloaded and then symlinked to ComfyUI.
 
-3. Create and activate a virtual environment:
+3. Create and activate a virtual environment with dependencies:
    ```
    python -m venv venv
    source venv/bin/activate
+   python -m pip install -r requirements.txt
    ```
 
-4. Install dependencies from requirements.txt:
-   ```
-   pip install -r requirements.txt
-   ```
-
-5. Run the installation script:
+4. Run the installation script:
    ```
    python install.py
    ```
@@ -56,30 +52,22 @@ These API keys are required for certain features of the application to work prop
 
 # Face Enhancement Gradio Demo
 
-A web interface for the face enhancement workflow using Gradio.
-
-## Features
-
-- Simple web interface for face enhancement
-- Upload input image and reference face image
-- Queue system to process jobs sequentially on a single GPU
-- Approximately 60 seconds processing time per image
+A web interface for the face enhancement workflow using Gradio. Between 45-60 seconds processing time per image.
 
 ## Setup
 
-1. Install dependencies:
+1.
 
 ```bash
-pip install -r requirements.txt
+python gradio_demo.py
 ```
 
-2. Run the Gradio demo:
-
+2. Port-forwarding 
 ```bash
-python gradio_demo.py
+ssh -L 7860:localhost:7860 root@[IP_ADDRESS] -p [RUNPOD_PORT] -i ~/.ssh/[PRIVATE_KEY_NAME]
 ```
 
-3. Open your browser and go to http://localhost:7860
+3. Go to http://localhost:7860
 
 ## Usage
 
@@ -91,6 +79,6 @@ python gradio_demo.py
 
 ## Notes
 
-- The demo uses a job queue to ensure only one job runs at a time
-- Processing takes approximately 60 seconds per image
+- The demo runs a ComfyUI server ephemerally
+- Processing takes approximately 45-60 seconds per image
 - Temporary files are created during processing and cleaned up afterward