Preserve original hue, saturation and brightness through restoration and finishing

Replace blurred RGB matching with per-pixel hue/saturation anchoring and broad value alignment. Reapply after finishing and default color preservation to 1.0. Convert embedded RGB ICC profiles to sRGB when available and tag PNG exports. All 22 automated checks pass, including vivid color recovery with restored texture retained.

README.md

@@ -38,7 +38,7 @@ is recoverable at commit `af15c36a9052b361c7a69f1bd210290125a9b07c`.
   Gaussian overlap blending, padded borders, and separately tiled VAE inference.
 - PNG downloads preserve alpha. The comparison uses RGB; the original alpha
   channel is resized separately and attached to the download.
-- Color preservation uses low-frequency alignment. Restoration amount blends
+- Color preservation defaults to 1.0: original hue and saturation are anchored per pixel, broad brightness is aligned, and restored fine texture is retained. The anchor is reapplied after finishing. Embedded RGB ICC profiles are converted to sRGB when available, and exported PNGs are tagged sRGB. Restoration amount blends
   with bicubic input resizing; it does not change the underlying model conditioning.
 - UI reports inference and processing time separately; ZeroGPU queue wait is additional.
 

app.py

@@ -18,7 +18,7 @@ if os.environ.get('TILE_V3_UI_ONLY') != '1':
     MODELS['Detail'] = RestorationModel('Detail', shared=MODELS['Restore'])
 
 
-def duration(image, mode='Restore', size='2×', seed=42, restoration=1., color_strength=.85,
+def duration(image, mode='Restore', size='2×', seed=42, restoration=1., color_strength=1.,
              batch_size=4, passes=3, pass2_strength=.5, pass3_strength=.3, texture_only=True,
              detail_boost=.25, local_contrast=.1, overlap=128, variations=1, steps=10,
              guidance=.5, weak_condition=.1, preset='Maximum detail', progress=None, **kwargs):
@@ -30,7 +30,7 @@ def duration(image, mode='Restore', size='2×', seed=42, restoration=1., color_s
 
 
 @spaces.GPU(duration=duration)
-def upscale(image, mode='Restore', size='2×', seed=42, restoration=1., color_strength=.85,
+def upscale(image, mode='Restore', size='2×', seed=42, restoration=1., color_strength=1.,
             batch_size=4, passes=3, pass2_strength=.5, pass3_strength=.3, texture_only=True,
             detail_boost=.25, local_contrast=.1, overlap=128, variations=1, steps=10,
             guidance=.5, weak_condition=.1, preset='Maximum detail', progress=gr.Progress()):
@@ -95,7 +95,8 @@ with gr.Blocks(title='Tile Upscaler v3', theme=gr.themes.Soft(primary_hue='indig
         with gr.Row():
             texture_only = gr.Checkbox(value=True, label='Texture-only refinements',
                                        info='Later passes add small detail while retaining broader shapes.')
-            color_strength = gr.Slider(0, 1, value=.85, step=.05, label='Original color preservation')
+            color_strength = gr.Slider(0, 1, value=1., step=.05, label='Original color preservation',
+                                       info='At 1, keeps original hue, saturation and broad brightness while retaining restored texture.')
         with gr.Row():
             detail_boost = gr.Slider(0, 1, value=.25, step=.05, label='Fine-detail boost',
                                     info='Luminance sharpening with local halo limits; increases crispness.')

pipeline.py

@@ -1,5 +1,6 @@
 """A single request: restore, refine, finish, and export distinct seed variants."""
 import json
+from io import BytesIO
 import math
 import shutil
 import tempfile
@@ -9,7 +10,7 @@ from pathlib import Path
 
 import cv2
 import numpy as np
-from PIL import Image, ImageOps
+from PIL import Image, ImageOps, ImageCms
 from geometry import output_size, padded_size, tile_coordinates
 
 cv2.setNumThreads(2)
@@ -17,11 +18,11 @@ PRESET_FIELDS = ('mode', 'passes', 'restoration', 'pass2_strength', 'pass3_stren
                  'texture_only', 'detail_boost', 'local_contrast', 'color_strength',
                  'overlap', 'variations', 'steps', 'guidance', 'weak_condition')
 PRESETS = {
-    'Maximum detail': ('Restore', 3, 1., .5, .3, True, .25, .1, .85, 128, 1, 10, .5, .1),
+    'Maximum detail': ('Restore', 3, 1., .5, .3, True, .25, .1, 1., 128, 1, 10, .5, .1),
     'Balanced': ('Restore', 2, 1., .4, .25, True, .15, .05, 1., 64, 1, 10, .5, .1),
     'Natural': ('Restore', 1, .8, .3, .2, True, .08, 0., 1., 64, 1, 10, 1., .2),
     'Fast': ('Fast', 1, 1., .4, .25, True, .12, .05, 1., 64, 1, 5, .5, .1),
-    'Creative detail': ('Detail', 2, 1., .35, .2, True, .15, .05, .8, 128, 1, 15, .5, .1),
+    'Creative detail': ('Detail', 2, 1., .35, .2, True, .15, .05, 1., 128, 1, 15, .5, .1),
 }
 DEFAULTS = dict(zip(PRESET_FIELDS, PRESETS['Maximum detail']))
 MODEL_REVISION = 'f24b3061b7f350b81e1907bdcfc27f71fb4ff3f3'
@@ -73,10 +74,43 @@ def blend_pass(base, prediction, strength, texture_only):
 
 
 def align_color(image, reference, strength):
+    """Anchor original hue/saturation and broad brightness; retain restored texture.
+
+    Matching blurred RGB alone cannot recover small saturated objects or color
+    edges. RGB ratios preserve original hue and saturation without histogram
+    matching, while a high-pass value residual retains the restored fine detail.
+    """
+    image = image.astype(np.float32)
     if not strength:
         return image
-    return np.clip(image + float(strength) * (cv2.GaussianBlur(reference, (0, 0), 5)
-                                             - cv2.GaussianBlur(image, (0, 0), 5)), 0, 255)
+    reference = reference.astype(np.float32)
+    value = image.max(axis=2)
+    source_value = reference.max(axis=2)
+    target_value = np.clip(value + cv2.GaussianBlur(source_value - value, (0, 0), 5), 0, 255)
+    ratios = np.divide(reference, source_value[:, :, None],
+                       out=np.ones_like(reference), where=source_value[:, :, None] > 1e-6)
+    anchored = ratios * target_value[:, :, None]
+    return np.clip(image + float(strength) * (anchored - image), 0, 255)
+
+
+def normalize_color(image):
+    """Convert embedded RGB profiles to sRGB before processing and PNG export."""
+    image = ImageOps.exif_transpose(image)
+    profile = image.info.get('icc_profile')
+    if profile:
+        try:
+            converted = ImageCms.profileToProfile(image.convert('RGB'),
+                ImageCms.ImageCmsProfile(BytesIO(profile)), ImageCms.createProfile('sRGB'), outputMode='RGB')
+            if 'A' in image.getbands() or 'transparency' in image.info:
+                converted.putalpha(image.convert('RGBA').getchannel('A'))
+            return converted
+        except (ImageCms.PyCMSError, ValueError, OSError):
+            # Malformed profiles must not make an otherwise readable image fail.
+            print('Ignoring unreadable input ICC profile', flush=True)
+    return image
+
+
+SRGB_PROFILE = ImageCms.ImageCmsProfile(ImageCms.createProfile('sRGB')).tobytes()
 
 
 def finish_detail(image, boost, contrast):
@@ -105,7 +139,7 @@ def output_directory():
 
 
 def run_pipeline(models, image, mode='Restore', size='2×', seed=42, restoration=1.,
-                 color_strength=.85, batch_size=4, passes=3, pass2_strength=.5,
+                 color_strength=1., batch_size=4, passes=3, pass2_strength=.5,
                  pass3_strength=.3, texture_only=True, detail_boost=.25, local_contrast=.1,
                  overlap=128, variations=1, steps=10, guidance=.5, weak_condition=.1,
                  preset='Maximum detail', progress=None):
@@ -116,7 +150,7 @@ def run_pipeline(models, image, mode='Restore', size='2×', seed=42, restoration
              guidance, weak_condition, batch_size, seed)
     if mode not in models:
         raise ValueError('Model inference is unavailable in the local UI preview.')
-    image = ImageOps.exif_transpose(image)
+    image = normalize_color(image)
     target = output_size(*image.size, size)
     _, estimate = work_estimate(target, mode, passes, overlap, variations, steps, batch_size)
     if estimate > 165:
@@ -129,7 +163,8 @@ def run_pipeline(models, image, mode='Restore', size='2×', seed=42, restoration
     pw, ph = padded_size(*target)
     folder = output_directory()
     reference_path = folder / 'before.png'
-    reference.save(reference_path, compress_level=3)
+    reference.info.pop("transparency", None)
+    reference.save(reference_path, compress_level=3, icc_profile=SRGB_PROFILE)
     settings = dict(mode=mode, size=size, seed=int(seed), restoration=restoration, color_strength=color_strength,
                     batch_size=batch_size, passes=passes, pass2_strength=pass2_strength, pass3_strength=pass3_strength,
                     texture_only=texture_only, detail_boost=detail_boost, local_contrast=local_contrast,
@@ -162,11 +197,13 @@ def run_pipeline(models, image, mode='Restore', size='2×', seed=42, restoration
             current = result.round().astype(np.uint8)
             timings.append(dict(pass_number=pass_index + 1, model=pass_mode, seed=(variant_seed + pass_index) % 2147483648, **timing))
         current = finish_detail(current, detail_boost, local_contrast)
+        # Finishing must not undo the color anchor through additive sharpening.
+        current = align_color(current, source, color_strength).round().astype(np.uint8)
         output = Image.fromarray(current)
         if alpha is not None:
             output.putalpha(alpha)
         path = folder / f'tile-v3-{mode.lower()}-{target[0]}x{target[1]}-{passes}pass-seed{variant_seed}.png'
-        output.save(path, compress_level=3)
+        output.save(path, compress_level=3, icc_profile=SRGB_PROFILE)
         gpu_seconds = sum(t['gpu_seconds'] for t in timings)
         record = dict(variant=variant + 1, seed=variant_seed, gpu_seconds=gpu_seconds, passes=timings, file=path.name)
         records.append(record)

test_pipeline.py

@@ -28,6 +28,10 @@ def test_variants_passes_alpha_and_reproducible_archive():
         metadata=json.loads(z.read('settings.json'))
         assert metadata['settings']['pass3_strength']==.3
         assert len(metadata['variants'][0]['passes'])==3
+    for transparent in [Image.new('LA',(40,24),(90,128)), Image.new('P',(40,24),0)]:
+        if transparent.mode=='P': transparent.info['transparency']=0
+        exported=run_pipeline({'Restore':Fake()},transparent,passes=1)
+        assert Image.open(exported[1]).mode=='RGBA'
     again=run_pipeline({'Restore':Fake()},image,variations=2)
     np.testing.assert_array_equal(first,Image.open(again[1]))
 
@@ -49,6 +53,36 @@ def test_settings_validation_and_capacity():
     for kwargs in ({'seed':float('nan')},{'guidance':3},{'overlap':100},{'detail_boost':float('nan')}):
         with pytest.raises(ValueError): run_pipeline({'Restore':Fake()},image,**kwargs)
     with pytest.raises(ValueError,match='exceeds'):
-        run_pipeline({'Detail':Fake()},image,mode='Detail',size='4096 px long edge',passes=3,variations=3,steps=25,overlap=192,batch_size=1)
+        run_pipeline({'Detail':Fake()},image,mode='Detail',size='4096 px long edge',passes=3,variations=1,steps=25,overlap=128,batch_size=4)
     assert DEFAULTS['mode']=='Restore' and DEFAULTS['passes']==3
     assert all(len(v)==len(PRESET_FIELDS) for v in PRESETS.values())
+
+
+def test_color_anchor_recovers_vivid_colors_and_keeps_restored_detail():
+    from pipeline import align_color
+    import cv2
+    reference=np.zeros((64,96,3),np.float32)
+    reference[:,:32]=[230,30,20]
+    reference[:,32:64]=[20,210,50]
+    reference[:,64:]=[30,50,240]
+    gray=reference.mean(2,keepdims=True)*.65+30
+    prediction=np.repeat(gray,3,axis=2)
+    # Add texture the source does not have: it should survive the color correction.
+    prediction[::2]+=8
+    corrected=align_color(prediction,reference,1.)
+    src_hsv=cv2.cvtColor(reference/255,cv2.COLOR_RGB2HSV)
+    dst_hsv=cv2.cvtColor(corrected/255,cv2.COLOR_RGB2HSV)
+    np.testing.assert_allclose(src_hsv[:,:,:2],dst_hsv[:,:,:2],atol=.001)
+    assert corrected[::2,16,0].mean()>corrected[1::2,16,0].mean()+5
+    np.testing.assert_array_equal(align_color(prediction,reference,0),prediction)
+    flat=np.broadcast_to([230,30,20],(32,32,3)).astype('float32')
+    np.testing.assert_allclose(align_color(np.full_like(flat,100),flat,1),flat,atol=.0001)
+
+
+def test_profile_normalization_and_export_tag():
+    from pipeline import normalize_color,SRGB_PROFILE
+    image=Image.new('RGBA',(16,12),(230,30,20,128));image.info['icc_profile']=SRGB_PROFILE
+    converted=normalize_color(image)
+    np.testing.assert_array_equal(image,converted)
+    r=run_pipeline({'Restore':Fake()},image,passes=1)
+    assert Image.open(r[1]).info.get('icc_profile')==SRGB_PROFILE