demo.py

import os
from install import install

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 tempfile
import hashlib
import io
import pickle
import sys
from test import process_face
from PIL import Image

INPUT_CACHE_DIR = "./cache"
os.makedirs(INPUT_CACHE_DIR, exist_ok=True)

def get_image_hash(img):
    """Generate a hash of the image content."""
    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):
    """
    Wrapper function for process_face that works with Gradio.
    
    Args:
        input_image: Input image from Gradio
        ref_image: Reference face image from Gradio
        
    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(INPUT_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, \
         tempfile.NamedTemporaryFile(suffix=".png", delete=False) as output_file:
        
        input_path = input_file.name
        ref_path = ref_file.name
        output_path = output_file.name
    
    # Save uploaded images to temporary files
    input_image.save(input_path)
    ref_image.save(ref_path)
    
    try:
        process_face(
            input_path=input_path,
            ref_path=ref_path,
            crop=False,
            upscale=False,
            output_path=output_path
        )
    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)
    
    # 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():
    with gr.Blocks(title="Face Enhancement Demo") as demo:
        gr.Markdown("""
        # Face Enhancement
        ### Instructions
        1. Upload the target image you want to enhance
        2. Upload a high-quality reference face image
        3. Click 'Enhance Face'

        Processing takes about 60 seconds. Due to the constraints of this demo, face cropping and upscaling are not applied to the reference image.

        For more information, check out my [blog post](https://rishidesai.github.io/posts/face-enhancement-techniques/).
        """, elem_id="instructions")

        gr.Markdown("---")

        with gr.Row():
            with gr.Column():
                input_image = gr.Image(label="Target Image", type="pil")
                ref_image = gr.Image(label="Reference Face", type="pil")
                enhance_button = gr.Button("Enhance Face")
            
            with gr.Column():
                output_image = gr.Image(label="Enhanced Result")
        
        enhance_button.click(
            fn=enhance_face_gradio,
            inputs=[input_image, ref_image],
            outputs=output_image,
            queue=True  # Enable queue for sequential processing
        )

        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"],
            ["examples/dany_gpt_2.png", "examples/dany_face.jpg"],
            ["examples/tim_gpt_1.png", "examples/tim_face.jpg"],
            ["examples/tim_gpt_2.png", "examples/tim_face.jpg"],
            ["examples/elon_gpt.png", "examples/elon_face.png"],
        ]
        gr.Examples(examples=example_inps, inputs=[input_image, ref_image], outputs=output_image)

    # Launch the Gradio app with queue
    demo.queue(max_size=99)
    
    try:
        demo.launch()
    except OSError as e:
        print(f"Error starting server: {e}")
        sys.exit(1)


if __name__ == "__main__":
    create_gradio_interface()