app.py

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST come before torch / any CUDA-touching import
import torch
import gradio as gr
import glob
from PIL import Image
from huggingface_hub import snapshot_download

from mirrorppr.data.image_ops import round_to_multiple
from diffsynth import load_state_dict
from diffsynth.pipelines.qwen_image import ModelConfig, QwenImagePipeline

MODEL_ID = "SJTU-DENG-Lab/MirrorPPR-Face"


def _glob_required(pattern):
    files = sorted(glob.glob(pattern))
    if not files:
        raise FileNotFoundError(f"No files matched: {pattern}")
    return files


def _build_paths(weights_root, qwen_root):
    qwen = qwen_root or os.path.join(weights_root, "qwen_image_edit")
    face = os.path.join(weights_root, "mirrorppr_face")
    return {
        "dit": _glob_required(os.path.join(qwen, "transformer", "diffusion_pytorch_model*.safetensors")),
        "text_encoder": _glob_required(os.path.join(qwen, "text_encoder", "model*.safetensors")),
        "vae": os.path.join(qwen, "vae", "diffusion_pytorch_model.safetensors"),
        "processor": os.path.join(qwen, "processor"),
        "mae": os.path.join(face, "mae", "mae_pretrained.safetensors"),
        "rformer": os.path.join(face, "rformer", "rformer.safetensors"),
        "connector": os.path.join(face, "connector", "connector.safetensors"),
        "lora": os.path.join(face, "lora", "lora.safetensors"),
    }


print("Downloading model weights from Hugging Face Hub...")
_local_root = snapshot_download(repo_id=MODEL_ID)
_paths = _build_paths(_local_root, None)
print(f"Model downloaded to: {_local_root}")

pipe = QwenImagePipeline.from_pretrained(
    torch_dtype=torch.bfloat16,
    device="cuda",
    model_configs=[
        ModelConfig(path=_paths["dit"]),
        ModelConfig(path=_paths["text_encoder"]),
        ModelConfig(path=_paths["vae"]),
        ModelConfig(path=_paths["mae"]),
        ModelConfig(path=_paths["rformer"]),
        ModelConfig(path=_paths["connector"]),
    ],
    tokenizer_config=None,
    processor_config=ModelConfig(path=_paths["processor"]),
)
if pipe.rformer is None:
    raise RuntimeError("R-Former module failed to load.")
if not hasattr(pipe, "connector") or pipe.connector is None:
    raise RuntimeError("Connector module failed to load.")
pipe.rformer.load_state_dict(load_state_dict(_paths["rformer"], device="cpu"))
pipe.connector.load_state_dict(load_state_dict(_paths["connector"], device="cpu"))
pipe.load_lora(pipe.dit, _paths["lora"])
print("MirrorPPR-Face pipeline loaded successfully.")


# Pre-packaged exemplar pairs for quick selection
EXEMPLAR_PAIRS = [
    {
        "name": "Style 1: Eye enlargement + mouth adjustments",
        "origin": "assets/exemplar_origin_0.png",
        "retouched": "assets/exemplar_retouched_0.png",
    },
    {
        "name": "Style 2: Eye enlargement + nose lengthening",
        "origin": "assets/exemplar_origin_1.png",
        "retouched": "assets/exemplar_retouched_1.png",
    },
    {
        "name": "Style 3: Eye enlargement + lip plump",
        "origin": "assets/exemplar_origin_2.png",
        "retouched": "assets/exemplar_retouched_2.png",
    },
]


def _on_exemplar_select(evt: gr.SelectData):
    """Load a pre-packaged exemplar pair when the user clicks a gallery item."""
    idx = evt.index
    if isinstance(idx, list):
        idx = idx[0] if idx else 0
    idx = int(idx)
    if 0 <= idx < len(EXEMPLAR_PAIRS):
        pair = EXEMPLAR_PAIRS[idx]
        return pair["origin"], pair["retouched"]
    return None, None


@spaces.GPU(duration=180)
def retouch(
    query_image,
    exemplar_origin,
    exemplar_retouched,
    steps=40,
    seed=123,
    cfg_scale=4.0,
):
    """Apply exemplar-based portrait photo retouching to a query image.

    Given an exemplar pair (an original face and its retouched version),
    this function infers the retouching operations and applies them to
    a new query face image.

    Args:
        query_image: The face image to be retouched.
        exemplar_origin: The original (pre-retouch) exemplar image.
        exemplar_retouched: The retouched exemplar image.
        steps: Number of diffusion inference steps (default 40).
        seed: Random seed for reproducibility (default 123).
        cfg_scale: Classifier-free guidance scale (default 4.0).

    Returns:
        The retouched query image.
    """
    if query_image is None:
        raise gr.Error("Please provide a query image.")
    if exemplar_origin is None or exemplar_retouched is None:
        raise gr.Error("Please provide both exemplar images (origin and retouched).")

    query = Image.fromarray(query_image).convert("RGB")
    ex_origin = Image.fromarray(exemplar_origin).convert("RGB")
    ex_target = Image.fromarray(exemplar_retouched).convert("RGB")

    width, height = query.size
    width = round_to_multiple(width, 16)
    height = round_to_multiple(height, 16)

    result = pipe(
        "",
        example_origin=ex_origin,
        example_target=ex_target,
        edit_image=query,
        seed=int(seed),
        num_inference_steps=int(steps),
        height=height,
        width=width,
        edit_image_auto_resize=False,
        cfg_scale=cfg_scale,
    )
    return result


CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
    gr.Markdown(
        """
        # MirrorPPR: Exemplar-Based Portrait Photo Retouching

        Upload a face image (query) and provide an exemplar pair (original → retouched).
        The model infers the retouching operations from the exemplar pair and applies
        them to your query image. Try a pre-packaged exemplar from the gallery below.

        [Paper](https://arxiv.org/abs/2606.29308) · [GitHub](https://github.com/SJTU-DENG-Lab/MirrorPPR) · [Model](https://huggingface.co/SJTU-DENG-Lab/MirrorPPR-Face)
        """
    )

    with gr.Row():
        # Left column: inputs
        with gr.Column(scale=1):
            gr.Markdown("### Query Image (to retouch)")
            query_img = gr.Image(
                label="Query Image",
                type="numpy",
                height=300,
            )
            gr.Markdown("### Exemplar Pair (reference retouching style)")
            ex_origin_img = gr.Image(
                label="Exemplar Original",
                type="numpy",
                height=200,
            )
            ex_retouched_img = gr.Image(
                label="Exemplar Retouched",
                type="numpy",
                height=200,
            )

            gr.Markdown("### Quick Exemplar Templates")
            exemplar_gallery = gr.Gallery(
                label="Click a template to load an exemplar pair",
                value=[
                    (pair["origin"], pair["name"])
                    for pair in EXEMPLAR_PAIRS
                ],
                columns=3,
                height=150,
                show_label=False,
                allow_preview=False,
            )

            with gr.Accordion("Advanced settings", open=False):
                steps_slider = gr.Slider(
                    label="Inference steps",
                    minimum=10,
                    maximum=80,
                    value=40,
                    step=1,
                )
                seed_input = gr.Number(
                    label="Seed",
                    value=123,
                    precision=0,
                )
                cfg_slider = gr.Slider(
                    label="CFG scale",
                    minimum=1.0,
                    maximum=10.0,
                    value=4.0,
                    step=0.5,
                )

            run_btn = gr.Button("Retouch", variant="primary", size="lg")

        # Right column: output
        with gr.Column(scale=1):
            gr.Markdown("### Retouched Result")
            output_img = gr.Image(
                label="Retouched Query",
                type="pil",
                height=400,
            )

    # Wire up exemplar gallery selection
    exemplar_gallery.select(
        fn=_on_exemplar_select,
        outputs=[ex_origin_img, ex_retouched_img],
    )

    # Wire up the run button
    run_btn.click(
        fn=retouch,
        inputs=[query_img, ex_origin_img, ex_retouched_img, steps_slider, seed_input, cfg_slider],
        outputs=output_img,
        api_name="retouch",
    )

    gr.Examples(
        examples=[
            [
                "assets/query_0.png",
                "assets/exemplar_origin_0.png",
                "assets/exemplar_retouched_0.png",
                40,
                123,
                4.0,
            ],
            [
                "assets/query_1.png",
                "assets/exemplar_origin_1.png",
                "assets/exemplar_retouched_1.png",
                40,
                123,
                4.0,
            ],
            [
                "assets/query_2.png",
                "assets/exemplar_origin_2.png",
                "assets/exemplar_retouched_2.png",
                40,
                123,
                4.0,
            ],
        ],
        inputs=[query_img, ex_origin_img, ex_retouched_img, steps_slider, seed_input, cfg_slider],
        outputs=output_img,
        fn=retouch,
        cache_examples=True,
        cache_mode="lazy",
    )

demo.launch(mcp_server=True)