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

def resolve_hf_path(path):
    if isinstance(path, str) and path.startswith("hf://"):
        # The expected format is: hf://{repo_id}@{filename}
        parts = path[len("hf://"):].split("@")
        if len(parts) != 2:
            raise ValueError(f"Invalid HF URI format: {path}")
        repo_id = parts[0]
        filename = parts[1]
        print(f"Downloading {filename} from {repo_id} ...")
        local_path = hf_hub_download(repo_id=repo_id, filename=filename)
        return local_path
    return path

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"

flux_fill_path = resolve_hf_path(os.environ["FLUX_FILL_PATH"])
ace_plus_fft_model_path = resolve_hf_path(os.environ["ACE_PLUS_FFT_MODEL"])

# Update the environment variables with the resolved local file paths.
os.environ["FLUX_FILL_PATH"] = flux_fill_path
os.environ["ACE_PLUS_FFT_MODEL"] = ace_plus_fft_model_path

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()