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)