Add Color Match feature (Reinhard RGB transfer)

app.py

@@ -193,6 +193,19 @@ def run_rmbg(pil_image: Image.Image) -> Image.Image:
     return result
 
 
+def color_match_reinhard(source: Image.Image, result: Image.Image) -> Image.Image:
+    """Reinhard RGB 均值/标准差色彩迁移：将 result 的色调对齐 source。"""
+    src = np.array(source.convert("RGB")).astype(np.float32)
+    res = np.array(result.convert("RGB")).astype(np.float32)
+    out = np.zeros_like(res)
+    for c in range(3):
+        s_mean, s_std = src[:, :, c].mean(), src[:, :, c].std()
+        r_mean, r_std = res[:, :, c].mean(), res[:, :, c].std()
+        ratio = s_std / (r_std + 1e-6) if r_std > 0.5 else 1.0
+        out[:, :, c] = (res[:, :, c] - r_mean) * ratio + s_mean
+    return Image.fromarray(np.clip(out, 0, 255).astype(np.uint8))
+
+
 def add_image_watermark(result: Image.Image, ref: Image.Image, size: int = 200, padding: int = 16) -> Image.Image:
     result = result.copy().convert("RGBA")
     thumb = ref.convert("RGBA")
@@ -376,6 +389,7 @@ def infer(
     randomize_seed,
     guidance_scale,
     steps,
+    color_match,
     progress=gr.Progress(track_tqdm=True),
 ):
     gc.collect()
@@ -432,6 +446,20 @@ def infer(
 
         if ref_image is not None and len(pil_images) > 1:
             result = add_image_watermark(result, pil_images[1])
+
+        if color_match:
+            # 用原图背景（无笔迹）作为色彩参考
+            if isinstance(source_image, dict):
+                bg = source_image.get("background")
+                if bg is not None:
+                    ref_pil = Image.fromarray(bg).convert("RGB") if isinstance(bg, np.ndarray) else bg.convert("RGB")
+                else:
+                    ref_pil = src_pil
+            else:
+                ref_pil = src_pil
+            ref_pil_resized = ref_pil.resize(result.size, Image.LANCZOS)
+            result = color_match_reinhard(ref_pil_resized, result)
+
         return result, seed
 
     finally:
@@ -563,6 +591,8 @@ with gr.Blocks(
                         lines=3,
                     )
 
+                    color_match_input = gr.Checkbox(label="Color Match — 色彩对齐原图", value=True)
+
                     with gr.Accordion("Advanced Settings", open=False):
                         seed_input = gr.Slider(
                             label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0,
@@ -593,6 +623,7 @@ with gr.Blocks(
                 inputs=[
                     source_input, ref_input, prompt_input, negative_prompt_input,
                     seed_input, randomize_seed_input, guidance_scale_input, steps_input,
+                    color_match_input,
                 ],
                 outputs=[output_image, seed_input],
             ).then(