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