app.py

# -*- coding:UTF-8 -*-
#!/usr/bin/env python
import numpy as np
import gradio as gr
import roop.globals
from roop.core import (
    start,
    decode_execution_providers,
    suggest_max_memory,
    suggest_execution_threads,
)
from roop.processors.frame.core import get_frame_processors_modules
from roop.utilities import normalize_output_path
import os
from PIL import Image
from datetime import datetime
from huggingface_hub import HfApi, login
from datasets import load_dataset, Dataset
import json
import shutil
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

class FaceIntegrDataset:
    def __init__(self, repo_id="Arrcttacsrks/face_integrData"):
        # Get token from environment variable
        self.token = os.getenv('hf_token')
        if not self.token:
            raise ValueError("HF_TOKEN environment variable is not set")
            
        self.repo_id = repo_id
        self.api = HfApi()
        
        # Login to Hugging Face
        login(self.token)
        
        # Create local temp directory for organizing files
        self.temp_dir = "temp_dataset"
        os.makedirs(self.temp_dir, exist_ok=True)

    def create_date_folder(self):
        """Create folder structure based on current date"""
        current_date = datetime.now().strftime("%Y-%m-%d")
        folder_path = os.path.join(self.temp_dir, current_date)
        os.makedirs(folder_path, exist_ok=True)
        return folder_path, current_date

    def save_metadata(self, source_path, target_path, output_path, timestamp):
        """Save metadata for the face swap operation"""
        metadata = {
            "timestamp": timestamp,
            "source_image": source_path,
            "target_image": target_path,
            "output_image": output_path,
            "date_created": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        }
        return metadata

    def upload_to_hf(self, local_folder, date_folder):
        """Upload files to Hugging Face dataset"""
        try:
            # Upload the files
            self.api.upload_folder(
                folder_path=local_folder,
                repo_id=self.repo_id,
                repo_type="dataset",
                path_in_repo=date_folder
            )
            return True
        except Exception as e:
            print(f"Error uploading to Hugging Face: {str(e)}")
            return False

def swap_face(source_file, target_file, doFaceEnhancer):
    folder_path = None
    try:
        # Initialize dataset handler
        dataset_handler = FaceIntegrDataset()
        
        # Create date-based folder
        folder_path, date_folder = dataset_handler.create_date_folder()
        
        # Generate timestamp for unique identification
        timestamp = datetime.now().strftime("%S-%M-%H-%d-%m-%Y")
        
        # Save input images with timestamp in folder
        source_path = os.path.join(folder_path, f"source_{timestamp}.jpg")
        target_path = os.path.join(folder_path, f"target_{timestamp}.jpg")
        output_path = os.path.join(folder_path, f"OutputImage{timestamp}.jpg")
        
        # Save the input images
        if source_file is None or target_file is None:
            raise ValueError("Source and target images are required")
            
        source_image = Image.fromarray(source_file)
        source_image.save(source_path)
        target_image = Image.fromarray(target_file)
        target_image.save(target_path)
        
        print("source_path: ", source_path)
        print("target_path: ", target_path)
        
        # Set global paths
        roop.globals.source_path = source_path
        roop.globals.target_path = target_path
        roop.globals.output_path = normalize_output_path(
            roop.globals.source_path, 
            roop.globals.target_path, 
            output_path
        )
        
        # Configure face processing options
        if doFaceEnhancer:
            roop.globals.frame_processors = ["face_swapper", "face_enhancer"]
        else:
            roop.globals.frame_processors = ["face_swapper"]
        
        # Set global parameters
        roop.globals.headless = True
        roop.globals.keep_fps = True
        roop.globals.keep_audio = True
        roop.globals.keep_frames = False
        roop.globals.many_faces = False
        roop.globals.video_encoder = "libx264"
        roop.globals.video_quality = 18
        roop.globals.max_memory = suggest_max_memory()
        roop.globals.execution_providers = decode_execution_providers(["cuda"])
        roop.globals.execution_threads = suggest_execution_threads()
        
        print(
            "start process",
            roop.globals.source_path,
            roop.globals.target_path,
            roop.globals.output_path,
        )
        
        # Check frame processors
        for frame_processor in get_frame_processors_modules(roop.globals.frame_processors):
            if not frame_processor.pre_check():
                return None
        
        # Process the face swap
        start()
        
        # Save metadata
        metadata = dataset_handler.save_metadata(
            f"source_{timestamp}.jpg",
            f"target_{timestamp}.jpg",
            f"OutputImage{timestamp}.jpg",
            timestamp
        )
        
        # Save metadata to JSON file in the same folder
        metadata_path = os.path.join(folder_path, f"metadata_{timestamp}.json")
        with open(metadata_path, 'w') as f:
            json.dump(metadata, f, indent=4)
        
        # Upload to Hugging Face
        upload_success = dataset_handler.upload_to_hf(folder_path, date_folder)
        
        if upload_success:
            print(f"Successfully uploaded files to dataset {dataset_handler.repo_id}")
        else:
            print("Failed to upload files to Hugging Face dataset")
        
        # Read the output image before cleaning up
        if os.path.exists(output_path):
            output_image = Image.open(output_path)
            output_array = np.array(output_image)
            # Clean up temp folder after reading the image
            shutil.rmtree(folder_path)
            return output_array
        else:
            print("Output image not found")
            if folder_path and os.path.exists(folder_path):
                shutil.rmtree(folder_path)
            return None
        
    except Exception as e:
        print(f"Error in face swap process: {str(e)}")
        if folder_path and os.path.exists(folder_path):
            shutil.rmtree(folder_path)
        raise gr.Error(f"Face swap failed: {str(e)}")

def create_interface():
    # Create custom style
    custom_css = """
    .container {
        max-width: 1200px;
        margin: auto;
        padding: 20px;
    }
    .output-image {
        min-height: 400px;
        border: 1px solid #ccc;
        border-radius: 8px;
        padding: 10px;
    }
    """

    # Gradio interface setup
    title = "Face - Integrator"
    description = r"""
    Please upload source and target images to begin the face swap process.
    """
    article = r"""
    <div style="text-align: center; max-width: 650px; margin: 40px auto;">
        <p>
            This tool performs face swapping with optional enhancement.
        </p>
    </div>
    """

    # Create Gradio interface with improved layout
    with gr.Blocks(title=title, css=custom_css) as app:
        gr.Markdown(f"<h1 style='text-align: center;'>{title}</h1>")
        gr.Markdown(description)
        
        with gr.Row():
            with gr.Column(scale=1):
                source_image = gr.Image(
                    label="Source Image",
                    type="numpy",
                    sources=["upload"]
                )
                
            with gr.Column(scale=1):
                target_image = gr.Image(
                    label="Target Image",
                    type="numpy",
                    sources=["upload"]
                )
                
            with gr.Column(scale=1):
                output_image = gr.Image(
                    label="Output Image",
                    type="numpy",
                    interactive=False,
                    elem_classes="output-image"
                )
        
        with gr.Row():
            enhance_checkbox = gr.Checkbox(
                label="Apply the algorithm?",
                info="Image Quality Improvement",
                value=False
            )
            
        with gr.Row():
            process_btn = gr.Button(
                "Process Face Swap",
                variant="primary",
                size="lg"
            )
        
        # Set up the processing event
        process_btn.click(
            fn=swap_face,
            inputs=[source_image, target_image, enhance_checkbox],
            outputs=output_image,
            api_name="swap_face"
        )
        
        gr.Markdown(article)
    
    return app

def main():
    # Create and launch the interface
    app = create_interface()
    app.launch(share=False)

if __name__ == "__main__":
    main()