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