app.py

from __future__ import annotations

import gc
import os
import sys
import threading
import time
import warnings
from pathlib import Path
from typing import Any

warnings.filterwarnings(
    "ignore",
    message=r".*HTTP_422_UNPROCESSABLE_ENTITY.*",
    module=r"gradio\.routes",
)

os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"

import gradio as gr
import numpy as np
from huggingface_hub import snapshot_download
from PIL import Image


MODEL_REPO_ID = os.getenv("MODEL_REPO_ID", "heyh97791/NCF")
HF_TOKEN = os.environ.get("HF_TOKEN")
SPACE_DIR = Path(__file__).resolve().parent
SEGMENTATION_MODEL_DIR = SPACE_DIR / "segformer-b5-finetuned-ade-640-640"

RESOLUTION_MAP = {
    "Original": 0,
    "4K": 3840 * 2160,
    "2K": 2560 * 1440,
    "1080p": 1920 * 1080,
    "720p": 1280 * 720,
    "512": 512 * 512,
}

_RUN_LOCK = threading.Lock()
_FLOW_MODULES: dict[str, Any] = {}


I18N = {
    "lang_btn": {
        "en": "中文",
        "zh": "English",
    },
    "title": {
        "en": "# 🎨 ColorFM-O [CPU]",
        "zh": "# 🎨 ColorFM-O [CPU]",
    },
    "paper_code": {
        "en": """
        📄 **Paper:** [ColorFM: An Optimization-to-Learning Framework for Color Transfer via Flow Matching (ECCV 2026)](https://github.com/cszn/ColorFM)  |  💻 **Code:** [GitHub](https://github.com/cszn/ColorFM)  |  ⚙️ **Model:** Optimization method for color transfer
        """,
        "zh": """
        📄 **论文:** [ColorFM: An Optimization-to-Learning Framework for Color Transfer via Flow Matching (ECCV 2026)](https://github.com/cszn/ColorFM)  |  💻 **代码:** [GitHub](https://github.com/cszn/ColorFM)  |  ⚙️ **模型:** 基于优化的色彩转移方法
        """,
    },
    "desc_header": {
        "en": "💡 What is Color Transfer?",
        "zh": "💡 什么是追色?",
    },
    "desc_content": {
        "en": """
        **1. What is "Color Transfer"?**
        It is a technique that extracts the **color palette and atmosphere** from a reference image (Style) and applies it to your target image (Content), while preserving the original structure.
        
        **2. Use Cases:**
        * 📷 **Photography:** Instantly mimic color grading styles from master photographers.
        * 🎨 **Art & Design:** Unify the color theme of different assets quickly.
        
        **3. Tips for Best Results:**
        * ✅ **Match Content:** Results are best when the content and style images share similar scenes (e.g., Landscape to Landscape).
        * 🎨 **Hue Similarity:** Images with closer hues between content and style may produce better color transfer results.
        * ⚡ **High Resolution:** Supports processing high-resolution images (use the Settings menu to optimize speed).
        
        **🛡️ Privacy Disclaimer:**
        This demo runs entirely on the cloud instance. Your uploaded images are processed in memory and are **NOT saved** or stored permanently on our servers.
        """,
        "zh": """
        **1. 什么是“追色”?**
        追色（Color Transfer）是指从一张参考图（风格图）中提取**色调与氛围**，并将其迁移到你的目标图片（内容图）上，同时保留原图的细节纹理。
        
        **2. 主要用途:**
        * 📷 **摄影后期:** 一键复刻摄影大师的调色风格，无需手动调参。
        * 🎨 **设计创作:** 快速统一多张素材的色调风格，提高创作效率。
       
        **3. 如何获得最佳效果?**
        * ✅ **内容匹配:** 当内容图与风格图的场景相似时（例如都是风景），效果通常最好。
        * 🎨 **色相相似:** 内容图像和风格图像之间色相相似，可能会呈现更好的结果。
        * ⚡ **高清支持:** 支持高分辨率图片处理（可在“设置”中调整分辨率以获得更快速度）。
        
        **🛡️ 免责声明:**
        本 Demo 仅供演示体验。您的图片仅在内存中进行临时处理，**不会被保存**、存储或用于任何其他用途，处理结束后即刻销毁。
        """,
    },
    "label_content": {
        "en": "Content Image",
        "zh": "内容图",
    },
    "label_style": {
        "en": "Style Reference",
        "zh": "色彩参考图",
    },
    "settings_header": {
        "en": "⚙️ Settings",
        "zh": "⚙️ 设置",
    },
    "steps_label": {
        "en": "Fit Steps",
        "zh": "迭代步数",
    },
    "seg_label": {
        "en": "Use Semantic Segmentation",
        "zh": "使用语义分割",
    },
    "res_label": {
        "en": "Max Resolution Limit",
        "zh": "最大分辨率限制",
    },
    "res_info": {
        "en": "Limits the input image size before optimization. Original keeps the uploaded size.",
        "zh": "在优化前限制输入图像尺寸；Original 保持上传尺寸。",
    },
    "btn_run": {
        "en": "🚀 Start Color Transfer",
        "zh": "🚀 开始追色",
    },
    "label_result": {
        "en": "Result Image",
        "zh": "追色结果",
    },
}


def default_device() -> str:
    try:
        import torch

        return "cuda" if torch.cuda.is_available() else "cpu"
    except Exception:
        return "cpu"




def normalize_image(image: np.ndarray) -> np.ndarray:
    if image.ndim == 2:
        image = np.stack([image, image, image], axis=-1)
    if image.shape[-1] == 4:
        image = image[..., :3]
    if image.dtype != np.uint8:
        max_value = float(np.nanmax(image)) if image.size else 0.0
        if np.issubdtype(image.dtype, np.floating) and max_value <= 1.0:
            image = image * 255.0
        image = np.clip(image, 0, 255).astype(np.uint8)
    return image


def resize_to_limit(image: np.ndarray, resolution_choice: str) -> np.ndarray:
    image = normalize_image(image)
    max_pixels = RESOLUTION_MAP.get(str(resolution_choice), 0)
    if max_pixels <= 0:
        return image

    height, width = image.shape[:2]
    pixels = height * width
    if pixels <= max_pixels:
        return image

    scale = (max_pixels / pixels) ** 0.5
    new_width = max(1, int(width * scale))
    new_height = max(1, int(height * scale))
    return np.asarray(Image.fromarray(image).resize((new_width, new_height), Image.Resampling.LANCZOS))


def missing_segmentation_files() -> list[str]:
    required_files = [
        "config.json",
        "preprocessor_config.json",
        "pytorch_model.bin",
    ]
    return [name for name in required_files if not (SEGMENTATION_MODEL_DIR / name).exists()]

def load_flow_modules():
    if MODEL_REPO_ID in _FLOW_MODULES:
        return _FLOW_MODULES[MODEL_REPO_ID]

    try:
        repo_dir = Path(
            snapshot_download(
                repo_id=MODEL_REPO_ID,
                token=HF_TOKEN,
                allow_patterns=[
                    "flow_interface.py",
                    "config/flow.yaml",
                    "dataset/**",
                    "models/**",
                    "solvers/**",
                ],
            )
        )
    except Exception as exc:
        raise gr.Error(
            "Failed to download the model repository. Make sure HF_TOKEN has read access. "
            f"Original error: {exc}"
        ) from exc

    repo_path = str(repo_dir)
    if repo_path not in sys.path:
        sys.path.insert(0, repo_path)

    try:
        from flow_interface import ColorFlowOptimizer, FlowRunOptions
    except Exception as exc:
        raise gr.Error(f"Failed to import flow_interface from the model repository: {exc}") from exc

    optimizer = ColorFlowOptimizer(repo_dir / "config" / "flow.yaml")
    _FLOW_MODULES[MODEL_REPO_ID] = {
        "optimizer": optimizer,
        "FlowRunOptions": FlowRunOptions,
    }
    return _FLOW_MODULES[MODEL_REPO_ID]


def run_optimization(
    content_image: np.ndarray | None,
    style_image: np.ndarray | None,
    fit_steps: int,
    use_segmentation: bool,
    resolution_choice: str,
    progress=gr.Progress(track_tqdm=True),
) -> np.ndarray | None:
    if content_image is None or style_image is None:
        raise gr.Error("Please upload both content and style/reference images.")
    fit_steps = int(fit_steps)
    if fit_steps <= 0:
        raise gr.Error("Fit steps must be greater than 0.")
    if use_segmentation:
        missing_files = missing_segmentation_files()
        if missing_files:
            raise gr.Error(
                "Segmentation model files are missing. Put these files under "
                f"{SEGMENTATION_MODEL_DIR}: {', '.join(missing_files)}"
            )

    def update_progress(step: int, total_steps: int, stage: str) -> None:
        total_steps = max(int(total_steps or 1), 1)
        step = max(0, min(int(step), total_steps))
        progress((step, total_steps), desc=stage)

    progress(0, desc="Waiting for worker")
    with _RUN_LOCK:
        try:
            start_time = time.time()
            modules = load_flow_modules()
            FlowRunOptions = modules["FlowRunOptions"]
            optimizer = modules["optimizer"]

            content = resize_to_limit(content_image, resolution_choice)
            style = resize_to_limit(style_image, resolution_choice)
            device_choice = default_device()

            options = FlowRunOptions(
                total_steps=fit_steps,
                max_epochs=1,
                num_workers=0,
                accelerator=device_choice,
                devices="auto" if device_choice == "cuda" else None,
                full=True,
                seg_mode=use_segmentation,
                allow_downloads=False,
                segmentation_model_name=str(SEGMENTATION_MODEL_DIR) if use_segmentation else None,
                show_progress=False,
                progress_callback=update_progress,
            )

            result_pil = optimizer.optimize(content, style, options=options)
            output = np.asarray(result_pil.convert("RGB"))
            print(
                f"Finished in {time.time() - start_time:.2f}s | "
                f"steps={fit_steps} | segmentation={'on' if use_segmentation else 'off'} | "
                f"output={output.shape[1]}x{output.shape[0]}"
            )
            progress(1.0, desc="Finished")
            return output
        except gr.Error:
            raise
        except Exception as exc:
            if "out of memory" in str(exc).lower():
                raise gr.Error("Out of memory. Try fewer fit steps or a lower resolution.") from exc
            raise gr.Error(str(exc)) from exc
        finally:
            gc.collect()
            try:
                import torch

                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
            except Exception:
                pass


def toggle_language(current_lang):
    target_lang = "zh" if current_lang == "en" else "en"
    return (
        target_lang,
        gr.update(value=I18N["lang_btn"][target_lang]),
        gr.update(value=I18N["title"][target_lang]),
        gr.update(label=I18N["desc_header"][target_lang]),
        gr.update(value=I18N["desc_content"][target_lang]),
        gr.update(label=I18N["label_content"][target_lang]),
        gr.update(label=I18N["label_style"][target_lang]),
        gr.update(label=I18N["settings_header"][target_lang]),
        gr.update(label=I18N["steps_label"][target_lang]),
        gr.update(label=I18N["seg_label"][target_lang]),
        gr.update(label=I18N["res_label"][target_lang], info=I18N["res_info"][target_lang]),
        gr.update(value=I18N["btn_run"][target_lang]),
        gr.update(label=I18N["label_result"][target_lang]),
    )


custom_css = """
#col-container {
    margin: 0 auto;
    max-width: 1100px;
}
.gallery-container img{
    object-fit: contain;
}
"""

with gr.Blocks(css=custom_css, title="Color-Transfer") as demo:
    lang_state = gr.State("en")
    with gr.Column(elem_id="col-container"):
        with gr.Row():
            with gr.Column(scale=5):
                md_title = gr.Markdown(I18N["title"]["en"])
            with gr.Column(scale=0, min_width=80):
                btn_lang = gr.Button(
                    value=I18N["lang_btn"]["en"],
                    variant="secondary",
                    size="sm",
                )

        gr.Markdown(I18N["paper_code"]["en"])
        with gr.Accordion(label=I18N["desc_header"]["en"], open=False) as acc_desc:
            md_desc = gr.Markdown(I18N["desc_content"]["en"])

        with gr.Row():
            with gr.Column():
                input_content = gr.Image(label=I18N["label_content"]["en"], type="numpy", height=300, image_mode="RGB")
            with gr.Column():
                input_style = gr.Image(label=I18N["label_style"]["en"], type="numpy", height=300, image_mode="RGB")

        with gr.Accordion(label=I18N["settings_header"]["en"], open=True) as acc_settings:
            with gr.Row():
                fit_steps = gr.Slider(1, 1000, value=300, step=1, label=I18N["steps_label"]["en"])
                chk_seg = gr.Checkbox(value=True, label=I18N["seg_label"]["en"])
                radio_res = gr.Radio(
                    choices=list(RESOLUTION_MAP.keys()),
                    value="Original",
                    label=I18N["res_label"]["en"],
                    info=I18N["res_info"]["en"],
                    interactive=True,
                )

        with gr.Row():
            btn_run = gr.Button(I18N["btn_run"]["en"], variant="primary")

        output_result = gr.Image(label=I18N["label_result"]["en"], type="numpy", interactive=False, height=450, format="png")

        btn_run.click(
            fn=lambda: None,
            inputs=None,
            outputs=output_result,
        ).then(
            fn=run_optimization,
            inputs=[input_content, input_style, fit_steps, chk_seg, radio_res],
            outputs=output_result,
        )

        btn_lang.click(
            fn=toggle_language,
            inputs=[lang_state],
            outputs=[
                lang_state,
                btn_lang,
                md_title,
                acc_desc,
                md_desc,
                input_content,
                input_style,
                acc_settings,
                fit_steps,
                chk_seg,
                radio_res,
                btn_run,
                output_result,
            ],
        )

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