app.py

# app.py (Complete Version for Hugging Face Spaces)

import os
import gradio as gr
from refacer import Refacer
import imageio
from PIL import Image
import tempfile
import base64
import shutil
import time

# --- Configuration (Hardcoded for Hugging Face Spaces) ---
NUM_FACES = 8
FORCE_CPU = True

# --- Initialize Refacer ---
refacer = Refacer(force_cpu=FORCE_CPU)

# --- Core Functions ---
def run_image(*vars):
    image_path = vars[0]
    origins = vars[1:(NUM_FACES + 1)]
    destinations = vars[(NUM_FACES + 1):(NUM_FACES * 2) + 1]
    thresholds = vars[(NUM_FACES * 2) + 1:-2]
    face_mode = vars[-2]
    partial_reface_ratio = vars[-1]
    disable_similarity = (face_mode in ["Single Face", "Multiple Faces"])
    multiple_faces_mode = (face_mode == "Multiple Faces")
    faces = [{'origin': origins[k] if not multiple_faces_mode else None, 'destination': destinations[k], 'threshold': thresholds[k] if not multiple_faces_mode else 0.0} for k in range(NUM_FACES) if destinations[k] is not None]
    return refacer.reface_image(image_path, faces, disable_similarity=disable_similarity, multiple_faces_mode=multiple_faces_mode, partial_reface_ratio=partial_reface_ratio)

def run_video(*vars):
    video_path = vars[0]
    origins = vars[1:(NUM_FACES + 1)]
    destinations = vars[(NUM_FACES + 1):(NUM_FACES * 2) + 1]
    thresholds = vars[(NUM_FACES * 2) + 1:-3]
    preview = vars[-3]
    face_mode = vars[-2]
    partial_reface_ratio = vars[-1]
    disable_similarity = (face_mode in ["Single Face", "Multiple Faces"])
    multiple_faces_mode = (face_mode == "Multiple Faces")
    faces = [{'origin': origins[k] if not multiple_faces_mode else None, 'destination': destinations[k], 'threshold': thresholds[k] if not multiple_faces_mode else 0.0} for k in range(NUM_FACES) if destinations[k] is not None]
    mp4_path, gif_path = refacer.reface(video_path, faces, preview=preview, disable_similarity=disable_similarity, multiple_faces_mode=multiple_faces_mode, partial_reface_ratio=partial_reface_ratio)
    return mp4_path, gif_path

def load_first_frame(filepath):
    if filepath is None: return None
    with imageio.get_reader(filepath) as reader:
        return reader.get_data(0)

def extract_faces_auto(filepath, max_faces=5, isvideo=False):
    if filepath is None: return [None] * max_faces
    if isvideo and os.path.getsize(filepath) > 10 * 1024 * 1024:
        print("Video too large for auto-extract.")
        return [None] * max_faces
    try:
        frame = load_first_frame(filepath)
        faces = refacer.extract_faces_from_image(frame, max_faces=max_faces)
        return faces + [None] * (max_faces - len(faces))
    except Exception as e:
        print(f"Could not extract faces: {e}")
        return [None] * max_faces

def toggle_tabs_and_faces(mode):
    if mode == "Single Face":
        return [gr.update(visible=(i == 0)) for i in range(NUM_FACES)] + [gr.update(visible=False)] * NUM_FACES
    elif mode == "Multiple Faces":
        return [gr.update(visible=True)] * NUM_FACES + [gr.update(visible=False)] * NUM_FACES
    else: # Faces By Match
        return [gr.update(visible=True)] * NUM_FACES + [gr.update(visible=True)] * NUM_FACES

def handle_tif_preview(filepath):
    if filepath is None: return None
    with Image.open(filepath) as img, tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as temp_file:
        img.convert('RGB').save(temp_file.name)
        return temp_file.name

# --- UI ---
theme = gr.themes.Base(primary_hue="blue", secondary_hue="cyan")

with gr.Blocks(theme=theme, title="NeoRefacer - AI Refacer") as demo:
    with open("icon.png", "rb") as f:
        icon_data = base64.b64encode(f.read()).decode()
    gr.Markdown(f'<div style="display: flex; align-items: center;"><img src="data:image/png;base64,{icon_data}" style="width:40px;height:40px;margin-right:10px;"><span style="font-size: 2em; font-weight: bold; color:#2563eb;">NeoRefacer</span></div>')

    for mode in ["Image", "GIF", "TIFF", "Video"]:
        with gr.Tab(f"{mode} Mode"):
            is_video_or_gif = mode in ["Video", "GIF"]
            
            with gr.Row():
                if is_video_or_gif:
                    input_component = gr.Video(label=f"Original {mode}")
                    if mode == "GIF":
                        output_main = gr.Video(label="Refaced GIF (MP4)", interactive=False, format="mp4")
                        output_secondary = gr.Image(label="Refaced GIF (GIF)", type="filepath")
                    else: # Video
                        output_main = gr.Video(label="Refaced Video", interactive=False, format="mp4")
                        output_secondary = gr.File(visible=False) # Dummy component
                elif mode == "TIFF":
                    input_component = gr.File(label="Original TIF", file_types=[".tif", ".tiff"])
                    output_main = gr.Image(label="Refaced TIF Preview", type="filepath")
                    output_secondary = gr.File(label="Refaced TIF (Download)", interactive=False)
                else: # Image
                    input_component = gr.Image(label="Original image", type="filepath")
                    output_main = gr.Image(label="Refaced image", interactive=False, type="filepath")
            
            with gr.Row():
                face_mode_radio = gr.Radio(["Single Face", "Multiple Faces", "Faces By Match"], value="Single Face", label="Replacement Mode")
                partial_reface_slider = gr.Slider(label="Reface Ratio (0=Full, 0.5=Half)", minimum=0.0, maximum=0.5, value=0.0, step=0.1)
                reface_btn = gr.Button(f"Reface {mode}", variant="primary")
                if is_video_or_gif:
                    preview_checkbox = gr.Checkbox(label="Preview (fast)", value=False)
            
            origins, destinations, thresholds, face_tabs = [], [], [], []
            for i in range(NUM_FACES):
                with gr.Tab(f"Face #{i+1}", visible=(i==0)) as tab:
                    with gr.Row():
                        origin_img = gr.Image(label="Face to replace (Match)", visible=False)
                        dest_img = gr.Image(label="Destination face (Target)")
                    thresh_slider = gr.Slider(label="Threshold (for Match mode)", minimum=0.0, maximum=1.0, value=0.2)
                origins.append(origin_img); destinations.append(dest_img); thresholds.append(thresh_slider); face_tabs.append(tab)

            # Event Handlers
            face_mode_radio.change(toggle_tabs_and_faces, inputs=face_mode_radio, outputs=face_tabs + origins)
            
            if mode == "TIFF":
                tif_preview = gr.Image(label="TIF Preview", type="filepath") # Specific for TIFF
                input_component.change(handle_tif_preview, inputs=input_component, outputs=tif_preview)
                input_component.change(lambda fp: extract_faces_auto(fp, max_faces=NUM_FACES), inputs=input_component, outputs=destinations)
                reface_btn.click(lambda fp, *args: (handle_tif_preview(run_image(fp, *args)), run_image(fp, *args)), inputs=[input_component] + origins + destinations + thresholds + [face_mode_radio, partial_reface_slider], outputs=[output_main, output_secondary])
            elif is_video_or_gif:
                input_component.change(lambda fp: extract_faces_auto(fp, max_faces=NUM_FACES, isvideo=True), inputs=input_component, outputs=destinations)
                reface_btn.click(run_video, inputs=[input_component] + origins + destinations + thresholds + [preview_checkbox, face_mode_radio, partial_reface_slider], outputs=[output_main, output_secondary])
            else: # Image
                input_component.change(lambda fp: extract_faces_auto(fp, max_faces=NUM_FACES), inputs=input_component, outputs=destinations)
                reface_btn.click(run_image, inputs=[input_component] + origins + destinations + thresholds + [face_mode_radio, partial_reface_slider], outputs=[output_main])
    
    # Load initial state for all tabs
    demo.load(lambda: toggle_tabs_and_faces("Single Face"), outputs=face_tabs + origins)

# --- Launch app ---
demo.queue().launch()