import os
import gradio as gr
from scepter.modules.utils.file_system import FS

# Set necessary environment variables for ACE++
os.environ["FLUX_FILL_PATH"] = "hf://black-forest-labs/FLUX.1-Fill-dev"
os.environ["PORTRAIT_MODEL_PATH"] = "ms://iic/ACE_Plus@portrait/comfyui_portrait_lora64.safetensors"
os.environ["SUBJECT_MODEL_PATH"] = "ms://iic/ACE_Plus@subject/comfyui_subject_lora16.safetensors"
os.environ["LOCAL_MODEL_PATH"] = "ms://iic/ACE_Plus@local_editing/comfyui_local_lora16.safetensors"
os.environ["ACE_PLUS_FFT_MODEL"] = "hf://ali-vilab/ACE_Plus@ace_plus_fft.safetensors"

os.environ["FLUX_FILL_PATH"]=FS.get_dir_to_local_dir(os.environ["FLUX_FILL_PATH"])
os.environ["ACE_PLUS_FFT_MODEL"]=FS.get_from(os.environ["ACE_PLUS_FFT_MODEL"])

from inference.ace_plus_inference import ACEInference
from scepter.modules.utils.config import Config
from modules.flux import FluxMRModiACEPlus
from inference.registry import INFERENCES


config_path = os.path.join("config", "ace_plus_fft.yaml")
cfg = Config(load=True, cfg_file=config_path)

# Instantiate the ACEInference object.
ace_infer = ACEInference(cfg)

def face_swap_app(target_img, face_img):
    """
    Swaps the face in the target image using the provided face image via ACE++.

    Parameters:
      target_img: The image in which you want to swap a face.
      face_img: The reference face image to insert.

    Returns:
      The output image after applying ACE++ face swapping.
    """
    # For ACEInference, we pass:
    # - reference_image: the target image,
    # - edit_image: the new face image,
    # - edit_mask: set to None so the image processor will create it,
    # - prompt: "Face swap" instructs the model to perform face swapping.
    # Other parameters (output dimensions, sampler, etc.) are set here as desired.
    output_img, edit_image, change_image, mask, seed = ace_infer(
        reference_image=target_img,
        edit_image=face_img,
        edit_mask=None,          # No manual mask provided; let ACE++ handle it
        prompt="Face swap",
        output_height=1024,
        output_width=1024,
        sampler='flow_euler',
        sample_steps=28,
        guide_scale=50,
        seed=-1                  # Use a random seed if not specified
    )
    return output_img

# Create the Gradio interface.
iface = gr.Interface(
    fn=face_swap_app,
    inputs=[
        gr.Image(type="pil", label="Target Image"),
        gr.Image(type="pil", label="Face Image")
    ],
    outputs=gr.Image(type="pil", label="Swapped Face Output"),
    title="ACE++ Face Swap Demo",
    description="Upload a target image and a face image to swap the face using the ACE++ model."
)

if __name__ == "__main__":
    iface.launch()