User-2468 commited on
Commit
1cca2dd
·
verified ·
1 Parent(s): 7ec6b12

Document final training decision and organize research history

Browse files
DEPLOYMENT.md CHANGED
@@ -1,34 +1,18 @@
1
- # v3 deployment
2
 
3
- ## Local Python
 
4
 
5
- Use Python 3.10 or 3.11 with the supplied requirements. CPU inference works without a GPU. Install a CUDA-compatible PyTorch build for GPU execution. Load once, put the model in evaluation mode, and reuse it across requests. The supplied loader does this automatically.
 
6
 
7
- ```python
8
- from inference import load_colorizer, colorize
9
- model = load_colorizer('.', 'cpu')
10
- ```
11
 
12
- The runtime loads a single 3,994,676-parameter network. The training-only teacher and discriminator are not required or downloaded.
 
 
 
 
 
13
 
14
- ## Hugging Face Space
15
-
16
- Copy `app.py`, `inference.py`, `semantic_model.py`, and `requirements-space.txt` into a Gradio Space. Rename the latter to `requirements.txt`. Use Gradio 6.28.0 in the Space README and choose ZeroGPU hardware.
17
-
18
- The released app pins `MODEL_REVISION` to the verified model commit. Setting this environment variable deliberately overrides that pin. `MODEL_ID` can be a local directory for offline use or a Hub repository ID.
19
-
20
- Import `spaces` before PyTorch. The model is placed on CUDA once at startup, and only the network stage is inside `@spaces.GPU`. Image preparation and full-resolution rendering happen on CPU. `ZEROGPU_CPU_TEST=1` enables the local CPU path for verification.
21
-
22
- ## Input handling
23
-
24
- The image pipeline applies EXIF orientation, preserves alpha, rejects images over 12 megapixels, and restores colour at the original output dimensions. Network processing preserves aspect ratio with a maximum 256-pixel side. Errors are returned explicitly rather than replaced by a blank image.
25
-
26
- The default API `/colorize` takes an image, colour strength (0–1.5), and smoothing (`Gentle`, `Balanced`, or `Strong`). It returns a PNG. Requests are serialized to bound memory usage.
27
-
28
- ## Versioning and rollback
29
-
30
- Model and source are committed atomically. The app pins the model commit rather than following moving main. `RELEASE_MANIFEST.json` records selection and previous commits. To roll back, use the recorded prior Space commit and model revision. The user-managed stable branch is not modified.
31
-
32
- ## Verification scope
33
-
34
- `QA.json` records tests for strict reload, parameter count, finite outputs, dimensions, transparency, black/white inputs, grayscale strength zero, input validation, a Gradio HTTP request, and ONNX parity. `RELEASE_REPORT.md` records visual selection and the separate live Space check. These checks do not imply factual colour reconstruction or eliminate every visual failure case.
 
1
+ # Deployment and final-run promotion
2
 
3
+ The live Space is https://huggingface.co/spaces/User-2468/mini-unet-colorizer.
4
+ Default model revision remains `1a9eb8af2754ad2329a24cfe50d388cb559441d0` (v3.0). The final longer job is https://huggingface.co/jobs/User-2468/6abaef6552d0dbd7f1da7286.
5
 
6
+ ## ZeroGPU
7
+ Use `app.py`, `inference.py`, `semantic_model.py`, `decision_model.py`, and `requirements-space.txt` (renamed `requirements.txt` in the Space). Space README SDK is Gradio6.28.0; requirements use the same version. Select ZeroGPU hardware. Import `spaces` before `torch`; load/place the model on CUDA at startup. Only chroma inference runs inside `@spaces.GPU(duration=15)`. CPU handles image preparation, guided rendering, and cached strength changes. Do not enable `torch.compile` in ZeroGPU.
8
 
9
+ For local CPU validation set `ZEROGPU_CPU_TEST=1`; `MODEL_ID` can be a local checkpoint directory. Image dimensions are preserved, maximum12MP, with PNG alpha. Cache is session-scoped with a10-minute lifetime. Upload changes invalidate cached predictions. Standard/Detailed/Maximum correspond to maximum network dimensions256/384/512. Strength changes after inference do not use another GPU request.
 
 
 
10
 
11
+ ## Final checkpoint review
12
+ 1. Inspect `experiments/final-20260929/status.json`, `candidates.json`, `summary.json`, final test/paired diagnostics and image grids. A candidate with step0 is the retained pilot, not a failed upload.
13
+ 2. Compare actual images to v3 and the pilot, especially historical originals and colour retention. Do not promote solely on patch score or original-colour error.
14
+ 3. Verify strict parameter count, output dimensions/alpha/lightness, CPU/CUDA loading and all palette modes. Use `decision_model.load_decision` for a DecisionColorizer checkpoint.
15
+ 4. Pin Space `MODEL_REVISION` to the evaluated immutable commit and `MODEL_SUBFOLDER` to the chosen directory. Alternative controls become visible for four-palette models.
16
+ 5. For a new root release, replace weights/config/loader together, update card and hashes, regenerate and verify ONNX or archive the old graph with an explicit version. The current v3 ONNX must never be described as the new four-palette model. Rebuild the release ZIP, tag the release, and run a real ZeroGPU request.
17
 
18
+ Training outputs are saved on main, but production promotion requires this review. Preserve the user's stable branch.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
README.md CHANGED
@@ -15,6 +15,16 @@ library_name: pytorch
15
 
16
  A compact automatic photo colouriser with **3,994,676 learned parameters in total**, including the complete deployed MobileNetV3 encoder. It predicts plausible colours from grayscale photographs. No teacher, critic, external semantic model, retrieval service or ensemble is needed at inference.
17
 
 
 
 
 
 
 
 
 
 
 
18
  ## Release contents
19
 
20
  - `model.safetensors` and `config.json`: the selected checkpoint.
 
15
 
16
  A compact automatic photo colouriser with **3,994,676 learned parameters in total**, including the complete deployed MobileNetV3 encoder. It predicts plausible colours from grayscale photographs. No teacher, critic, external semantic model, retrieval service or ensemble is needed at inference.
17
 
18
+ ## Current status — 29 September 2026
19
+
20
+ [Try the ZeroGPU Space](https://huggingface.co/spaces/User-2468/mini-unet-colorizer).
21
+
22
+ Root weights and ONNX remain the reviewed **v3.0** release. A final continuation of the3,995,832-parameter four-palette challenger is running under [`experiments/final-20260929`](experiments/final-20260929); it is **not yet approved as a production replacement**. See the [visual review and training decision](reports/FINAL_TRAINING_REVIEW.md) and [exact configuration](training/final_config.json). The pilot showed modest patch reductions, not elimination of blotches.
23
+
24
+ The updated Space offers Standard, Detailed and Maximum input detail, full-resolution PNG output, and cached strength adjustments. Higher input detail can help some photos but can also change their colours. The default stays at256 pixels. The Space can load future multi-palette checkpoints after review without downloading arbitrary Python code.
25
+
26
+ Historical research documents have moved to `reports/history/`. See [the report index](reports/README.md). The `stable` branch and existing release weights are preserved.
27
+
28
  ## Release contents
29
 
30
  - `model.safetensors` and `config.json`: the selected checkpoint.
SHA256SUMS.json CHANGED
@@ -1,23 +1,32 @@
1
  {
2
- "DEPLOYMENT.md": "201879644b48093363f025589641bc2a264ec8097512b4b4c30b621f3b9d9ee6",
3
  "PIPELINE_EVAL.json": "595be88704d580c01a3c4dd642409b3dc5dd25e86ca1c7eb40a696605f181a19",
4
  "QA.json": "b249b1d49045e31b0bf85f9e01515b91a72c4dd973963751c30d14272d262b90",
5
- "README.md": "5ade5fd9022a03937b91fbd3b0fa653ae93c068034b54b35677c3db47a9ccc59",
6
  "RELEASE_MANIFEST.json": "872eebfb322c8990b5e6ac493e4e4068b87f8daac0ddf219b8b55299514ca03c",
7
  "RELEASE_REPORT.md": "69677ec5e13f5c1249f59b01c31ed1b0f1588bddc216bce9167b5d31423a5574",
8
- "app.py": "58be79b62642c2f3393860791a1a7b39d8f2e83ac79caf6545839124ec787404",
9
- "colorization_project_log.md": "699238b4078fca93b15f4f2d951e439651502338dac4bf3a5b5a6d0b0449de21",
10
  "colorize_image.py": "9a85b7633bdf8a0d63f740ae69ed76644bb2087da4b21f8c1e01ca804cd4e6b2",
11
  "colorize_onnx.py": "cc106f066878a7372f4260279620238d46a3841baf09fbbe954f61a684199d05",
12
  "colorizer.json": "7d444d5f560b2ea6a070f73750f75f66aa2c61d3dc21a8e4300ed5f5196da179",
13
  "colorizer.onnx": "911fb777328fcfbeeb89cc25dca70092cf2de4686ede1edd44a8ccebc89ea341",
14
  "config.json": "eaeb49b39f9116bfed0d954e81852dc93476298752787dba88a63c0eec97ccec",
 
15
  "export_onnx.py": "f46a413c2b1b2a6e130cb7d2ad8d6d69c16a16955bf830efdc4c04c30d47ca60",
16
  "inference.py": "0ac1f382205b42cef0afaa9173017fa660e6859194a65a61b9b9052c437cd476",
17
  "model.py": "31628744fb5d362d91342063c90eb38b2b97fda8c91e303e4a58b6f388f53c55",
18
  "model.safetensors": "ec1f27d74533adc83f7ab3639a091fc4d8738a434dafc7d172c7873c28a9e715",
 
 
 
 
 
 
 
19
  "requirements-onnx.txt": "9bcca098a9da63b31fb9d803dfd7f99b753ad125ee9bce3a5bbe52e368739e11",
20
- "requirements-space.txt": "4c7f918aadc4afb31ec6808131570f9978cf0bcf9e2af562e93adc749e3245cf",
21
  "requirements.txt": "c855eb70020dc89f71fd38e44a0be2cfb80f2fe3d442d5ba45c91398c336f83f",
22
- "semantic_model.py": "b019358cbcaa214743cbd244d20eb6e59eeff0ab5264093ec0927be8529b25a2"
 
 
23
  }
 
1
  {
2
+ "DEPLOYMENT.md": "dac782d3f07838731ced3c0b7a21cfd24ca9943ea9e9660ba6e0dfa5f86de5c0",
3
  "PIPELINE_EVAL.json": "595be88704d580c01a3c4dd642409b3dc5dd25e86ca1c7eb40a696605f181a19",
4
  "QA.json": "b249b1d49045e31b0bf85f9e01515b91a72c4dd973963751c30d14272d262b90",
5
+ "README.md": "cdca79dbc42aeb31be15b61adc67f50746096b8fab9d95963232622cc8f691b5",
6
  "RELEASE_MANIFEST.json": "872eebfb322c8990b5e6ac493e4e4068b87f8daac0ddf219b8b55299514ca03c",
7
  "RELEASE_REPORT.md": "69677ec5e13f5c1249f59b01c31ed1b0f1588bddc216bce9167b5d31423a5574",
8
+ "app.py": "27832b99d910c1ff12d728d00e460060d061ab4c30aae2d8b2d256d4c894b8fa",
 
9
  "colorize_image.py": "9a85b7633bdf8a0d63f740ae69ed76644bb2087da4b21f8c1e01ca804cd4e6b2",
10
  "colorize_onnx.py": "cc106f066878a7372f4260279620238d46a3841baf09fbbe954f61a684199d05",
11
  "colorizer.json": "7d444d5f560b2ea6a070f73750f75f66aa2c61d3dc21a8e4300ed5f5196da179",
12
  "colorizer.onnx": "911fb777328fcfbeeb89cc25dca70092cf2de4686ede1edd44a8ccebc89ea341",
13
  "config.json": "eaeb49b39f9116bfed0d954e81852dc93476298752787dba88a63c0eec97ccec",
14
+ "decision_model.py": "c7c4370d93bc52fdd363d6d4fcc684c317806d9b393bd969f9ee670a01487f27",
15
  "export_onnx.py": "f46a413c2b1b2a6e130cb7d2ad8d6d69c16a16955bf830efdc4c04c30d47ca60",
16
  "inference.py": "0ac1f382205b42cef0afaa9173017fa660e6859194a65a61b9b9052c437cd476",
17
  "model.py": "31628744fb5d362d91342063c90eb38b2b97fda8c91e303e4a58b6f388f53c55",
18
  "model.safetensors": "ec1f27d74533adc83f7ab3639a091fc4d8738a434dafc7d172c7873c28a9e715",
19
+ "reports/FINAL_TRAINING_REVIEW.md": "2c354152ed07c22852401cff66865782840ffb143941a8f92a7ed4fd8542988f",
20
+ "reports/README.md": "b50c4af1398d0a8c56bd6bf1a1561f95bd8187c68f057f803ef5fddf4c718f90",
21
+ "reports/SPACE_QA_multi.json": "990353b85936624c24a71b4c2183d158251cd122ff23b1ee026eae94bca5715b",
22
+ "reports/SPACE_QA_v3.json": "78e8d283a7c7c0913705ee4a26fceae61c444fd836d29f784ec1bc88e46b53a4",
23
+ "reports/history/RESEARCH_REPORT.md": "6730cf20f7a584707fb0c084b882ac87db1ef858d72e7d5e716e5e0d382ddca1",
24
+ "reports/history/RESEARCH_ROUND2.md": "19f42c72003cb685f46b315105ad2d4e62a10a7e0cd1455ab11db51bd5083728",
25
+ "reports/history/colorization_project_log.md": "699238b4078fca93b15f4f2d951e439651502338dac4bf3a5b5a6d0b0449de21",
26
  "requirements-onnx.txt": "9bcca098a9da63b31fb9d803dfd7f99b753ad125ee9bce3a5bbe52e368739e11",
27
+ "requirements-space.txt": "bc5b3edd619278a37e0cdc499ab9a87057cdbd2023f7c22c08b6a437c5701c65",
28
  "requirements.txt": "c855eb70020dc89f71fd38e44a0be2cfb80f2fe3d442d5ba45c91398c336f83f",
29
+ "semantic_model.py": "b019358cbcaa214743cbd244d20eb6e59eeff0ab5264093ec0927be8529b25a2",
30
+ "training/final_config.json": "39502272a7bbc05c1ad15e70de4e93550b53cdd9fbea399192d1b400311205dd",
31
+ "training/train_final.py": "ba42f07649c05cfcf8d7a6cee78088a4ed0838f4e09e24227cf9fb3509b34b5b"
32
  }
app.py CHANGED
@@ -1,47 +1,81 @@
1
- """Gradio and ZeroGPU entry point. The model revision is pinned at release."""
2
  import os
3
  from pathlib import Path
4
  import spaces
5
  import torch
6
  import gradio as gr
7
  from huggingface_hub import snapshot_download
8
- from inference import load_colorizer, prepare, chroma_coefficients, render
 
9
 
10
  REPO=os.getenv('MODEL_ID','User-2468/mini-unet-colorizer')
11
- REV=os.getenv('MODEL_REVISION','v3.0.0')
 
12
  CPU=os.getenv('ZEROGPU_CPU_TEST')=='1'
13
  DEVICE='cpu' if CPU else 'cuda'
14
  if CPU:torch.set_num_threads(min(4,os.cpu_count() or 1))
15
  path=Path(REPO)
16
- if not path.is_dir():path=Path(snapshot_download(REPO,revision=REV,allow_patterns=['model.safetensors','config.json']))
17
- MODEL=load_colorizer(path,DEVICE)
 
 
 
 
 
 
 
 
 
 
 
18
 
19
- @spaces.GPU(duration=20)
20
- def infer(small,radius):
21
- return chroma_coefficients(MODEL,small,int(radius))
 
22
 
23
- def run(image,strength,smoothing):
 
 
 
 
24
  if image is None:raise gr.Error('Upload a photograph first.')
25
- radius={'Balanced':8,'Gentle':4,'Strong':16}[smoothing]
 
26
  try:
27
- light,alpha,small=prepare(image,256)
28
- coefficients=infer(small,radius)
 
 
29
  result=render(light,alpha,coefficients,strength)
30
- except (ValueError,RuntimeError,TypeError) as exc:
 
31
  raise gr.Error(str(exc)) from exc
32
- return result
 
 
 
 
 
33
 
34
- with gr.Blocks(title='Mini Photo Colorizer') as demo:
35
- gr.Markdown('# Mini Photo Colorizer\nAdd plausible colour to a black-and-white photograph. Original colours may be unknowable.')
 
36
  with gr.Row():
37
- source=gr.Image(type='pil',image_mode='RGBA',label='Original photograph',sources=['upload','clipboard'],height=440)
38
- result=gr.Image(type='pil',label='Colourised photograph',format='png',height=440)
39
  with gr.Row():
40
- strength=gr.Slider(0,1.5,value=1.,step=.05,label='Colour strength')
 
 
 
 
41
  smoothing=gr.Radio(['Gentle','Balanced','Strong'],value='Balanced',label='Colour smoothing')
42
- button=gr.Button('Colourise',variant='primary')
43
- button.click(run,[source,strength,smoothing],result,api_name='colorize',concurrency_limit=1)
44
- gr.ClearButton([source,result])
45
- gr.Markdown('3.995 million parameters · Full-resolution output, up to 12 megapixels · PNG download\n\nBest suited to ordinary photographs. Colours can be muted or inaccurate on unusual scenes, tiny objects and heavily damaged scans.')
 
 
46
  demo.queue(default_concurrency_limit=1,max_size=12)
47
- if __name__=='__main__':demo.launch()
 
1
+ """Pinned-release Gradio app. Import spaces before torch for ZeroGPU emulation."""
2
  import os
3
  from pathlib import Path
4
  import spaces
5
  import torch
6
  import gradio as gr
7
  from huggingface_hub import snapshot_download
8
+ from inference import prepare, chroma_coefficients, render, load_colorizer
9
+ import json
10
 
11
  REPO=os.getenv('MODEL_ID','User-2468/mini-unet-colorizer')
12
+ REV=os.getenv('MODEL_REVISION','1a9eb8af2754ad2329a24cfe50d388cb559441d0')
13
+ SUBFOLDER=os.getenv('MODEL_SUBFOLDER','').strip('/')
14
  CPU=os.getenv('ZEROGPU_CPU_TEST')=='1'
15
  DEVICE='cpu' if CPU else 'cuda'
16
  if CPU:torch.set_num_threads(min(4,os.cpu_count() or 1))
17
  path=Path(REPO)
18
+ if not path.is_dir():
19
+ prefix=SUBFOLDER+'/' if SUBFOLDER else ''
20
+ path=Path(snapshot_download(REPO,revision=REV,allow_patterns=[prefix+'model.safetensors',prefix+'config.json']))
21
+ if SUBFOLDER:path=path/SUBFOLDER
22
+ cfg=json.loads((path/'config.json').read_text())
23
+ if cfg.get('architecture')=='DecisionColorizer':
24
+ from decision_model import load_decision
25
+ MODEL=load_decision(path,DEVICE)
26
+ else:MODEL=load_colorizer(path,DEVICE)
27
+ MODEL.eval().requires_grad_(False)
28
+ PARAMETERS=sum(p.numel() for p in MODEL.parameters())
29
+ if PARAMETERS>=4_000_000:raise RuntimeError('Model exceeds the four-million-parameter limit.')
30
+ MODES=getattr(MODEL,'modes',1)
31
 
32
+ class PaletteView(torch.nn.Module):
33
+ def __init__(self,model,mode):super().__init__();self.model=model;self.mode=mode
34
+ def forward(self,L):
35
+ return self.model(L) if self.mode<0 else self.model(L,True)['all'][:,self.mode]
36
 
37
+ @spaces.GPU(duration=15)
38
+ def infer(small,radius,mode):
39
+ return chroma_coefficients(PaletteView(MODEL,mode),small,radius)
40
+
41
+ def run(image,strength,smoothing,detail,mode):
42
  if image is None:raise gr.Error('Upload a photograph first.')
43
+ mode=int(mode)
44
+ if mode < -1 or mode>=MODES or (MODES==1 and mode!=-1):raise gr.Error('That colour interpretation is unavailable.')
45
  try:
46
+ radius={'Gentle':4,'Balanced':8,'Strong':16}[smoothing]
47
+ size={'Standard':256,'Detailed':384,'Maximum':512}[detail]
48
+ light,alpha,small=prepare(image,size)
49
+ coefficients=infer(small,radius,mode)
50
  result=render(light,alpha,coefficients,strength)
51
+ cache=(light,alpha,coefficients)
52
+ except (ValueError,RuntimeError,TypeError,KeyError) as exc:
53
  raise gr.Error(str(exc)) from exc
54
+ return result,cache
55
+
56
+ def recolour(strength,cache):
57
+ if cache is None:return gr.skip()
58
+ light,alpha,coefficients=cache
59
+ return render(light,alpha,coefficients,strength)
60
 
61
+ with gr.Blocks(title='Mini Photo Colorizer',delete_cache=(3600,3600)) as demo:
62
+ gr.Markdown('# Give an old photograph new colour\nPlausible colour, with the original detail and proportions preserved.')
63
+ cache=gr.State(value=None,time_to_live=600)
64
  with gr.Row():
65
+ source=gr.Image(type='pil',image_mode='RGBA',label='Your photograph',sources=['upload','clipboard'],height=460)
66
+ result=gr.Image(type='pil',label='Colourised photograph',format='png',height=460,interactive=False)
67
  with gr.Row():
68
+ button=gr.Button('Colourise',variant='primary',scale=2)
69
+ clear=gr.ClearButton([source,result,cache],value='Clear',scale=1)
70
+ strength=gr.Slider(0,1.5,value=1.,step=.05,label='Colour strength',info='Adjust after colourising without another GPU request.')
71
+ with gr.Accordion('Fine-tune the result',open=False):
72
+ detail=gr.Radio(['Standard','Detailed','Maximum'],value='Standard',label='Detail',info='Higher detail can change the colours. Try Detailed if small objects are missed.')
73
  smoothing=gr.Radio(['Gentle','Balanced','Strong'],value='Balanced',label='Colour smoothing')
74
+ palette=gr.Dropdown([('Automatic',-1)]+[(f'Alternative {i+1}',i) for i in range(MODES)] if MODES>1 else [('Automatic',-1)],value=-1,label='Colour interpretation',visible=MODES>1)
75
+ gr.Markdown('After changing these settings, select **Colourise** again.')
76
+ button.click(run,[source,strength,smoothing,detail,palette],[result,cache],api_name='colorize',concurrency_limit=1)
77
+ strength.release(recolour,[strength,cache],result,api_name=False,concurrency_limit=1)
78
+ source.change(lambda:None,outputs=cache,api_name=False,queue=False)
79
+ gr.Markdown('Download the result as PNG using its download button. Up to 12 megapixels; transparency is preserved.\n\nColours are interpretations, not recovered historical facts. Unusual scenes, damaged scans and tiny objects can still produce muted colours or colour bleeding. [About the model](https://huggingface.co/User-2468/mini-unet-colorizer).')
80
  demo.queue(default_concurrency_limit=1,max_size=12)
81
+ if __name__=='__main__':demo.launch(max_file_size='30mb')
decision_model.py ADDED
@@ -0,0 +1,117 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Single / multiple coherent palettes, optional compact global-attention encoder.
2
+
3
+ Original experiment inspired by diverse colorization and multiple-choice learning;
4
+ not a reproduction. All deployed parameters, including the selector, are counted.
5
+ """
6
+ import json
7
+ from pathlib import Path
8
+ import torch
9
+ from torch import nn
10
+ import torch.nn.functional as F
11
+ from safetensors.torch import save_file, load_file
12
+ from semantic_model import SemanticColorizer
13
+
14
+
15
+ class DecisionColorizer(SemanticColorizer):
16
+ def __init__(self, modes=1, backbone='mobilenet', pretrained=False,
17
+ encoder_revision=None, **kwargs):
18
+ super().__init__(pretrained=False, **kwargs)
19
+ self.modes = modes
20
+ self.backbone = backbone
21
+ d = self.config['width']; q = self.config['queries']
22
+ if backbone == 'mobilevit':
23
+ import timm
24
+ from huggingface_hub import hf_hub_download
25
+ name = 'mobilevitv2_075.cvnets_in1k'
26
+ overlay = None
27
+ if pretrained:
28
+ assert encoder_revision, 'Pin pretrained encoder revision'
29
+ filename = hf_hub_download('timm/'+name, 'model.safetensors', revision=encoder_revision)
30
+ overlay = {'file': filename, 'hf_hub_id': None}
31
+ self.encoder = timm.create_model(name, pretrained=pretrained,
32
+ pretrained_cfg_overlay=overlay, features_only=True, out_indices=(1,2,3,4))
33
+ self.lateral = nn.ModuleList([nn.Conv2d(c,d,1) for c in self.encoder.feature_info.channels()])
34
+ cfg = self.encoder.pretrained_cfg
35
+ self.rgb_mean.copy_(torch.tensor(cfg['mean']).view(1,3,1,1))
36
+ self.rgb_std.copy_(torch.tensor(cfg['std']).view(1,3,1,1))
37
+ if modes > 1:
38
+ self.mode_embeddings = nn.Parameter(torch.randn(modes,d)*.03)
39
+ self.mode_offsets = nn.Parameter(torch.zeros(modes,q,2))
40
+ # Whole-image palette alternatives; no independently sampled pixels.
41
+ with torch.no_grad():
42
+ a=torch.arange(modes)*2*torch.pi/modes
43
+ self.mode_offsets[:,:,0]=.025*a.cos()[:,None]
44
+ self.mode_offsets[:,:,1]=.025*a.sin()[:,None]
45
+ self.mode_score = nn.Linear(d,modes)
46
+ nn.init.zeros_(self.mode_score.weight);nn.init.zeros_(self.mode_score.bias)
47
+ self.config.update(architecture='DecisionColorizer', format_version=3,
48
+ modes=modes, backbone=backbone, encoder_revision=encoder_revision)
49
+ count = sum(p.numel() for p in self.parameters())
50
+ assert count < 4_000_000, count
51
+
52
+ def features(self,L):
53
+ h,w=L.shape[-2:]
54
+ x=F.pad(self.neutral_rgb(L),(0,(-w)%32,0,(-h)%32),mode='replicate')
55
+ x=(x-self.rgb_mean)/self.rgb_std
56
+ if self.backbone=='mobilevit':
57
+ fs=self.encoder(x)
58
+ else:
59
+ fs=[]
60
+ for i,layer in enumerate(self.encoder):
61
+ x=layer(x)
62
+ if i in [3,6,12,16]:fs.append(x)
63
+ projected=[layer(f) for layer,f in zip(self.lateral,fs)]
64
+ x=projected[-1]
65
+ for i in range(2,-1,-1):
66
+ x=self.refine[2-i](F.interpolate(x,size=projected[i].shape[-2:],mode='bilinear',align_corners=False)+projected[i])
67
+ memory=self.memory_norm(torch.cat([F.adaptive_avg_pool2d(f,(8,8)).flatten(2).transpose(1,2) for f in projected[1:]],1))
68
+ q=self.queries[None].expand(L.shape[0],-1,-1)
69
+ for block in self.query_blocks:q=block(q,memory)
70
+ q=self.query_norm(q)
71
+ masks=torch.einsum('bqd,bdhw->bqhw',q,self.pixel(x))/(self.config['width']**.5)
72
+ weights=F.softmax(masks.float(),dim=1)
73
+ residual=2*torch.tanh(self.residual(x).float())
74
+ return q,weights,residual,memory.mean(1)
75
+
76
+ def forward(self,L,return_details=False):
77
+ q,weights,residual,context=self.features(L)
78
+ if self.modes==1:
79
+ palette=80*torch.tanh(self.palette(q))[:,None]
80
+ scores=q.new_zeros((len(L),1))
81
+ else:
82
+ mq=q[:,None]+self.mode_embeddings[None,:,None]
83
+ palette=80*torch.tanh(self.palette(mq)+self.mode_offsets[None])
84
+ scores=self.mode_score(context)
85
+ outputs=torch.einsum('bqhw,bkqc->bkchw',weights,palette.float())+residual[:,None]
86
+ b,k,c,h,w=outputs.shape
87
+ full=F.interpolate(outputs.reshape(b*k,c,h,w),size=(h*4,w*4),mode='bilinear',align_corners=False)
88
+ full=full[...,:L.shape[-2],:L.shape[-1]].reshape(b,k,c,*L.shape[-2:])
89
+ default=full[torch.arange(b,device=L.device),scores.argmax(1)]
90
+ if return_details:
91
+ return {'output':default,'all':full,'scores':scores.float(),
92
+ 'weights':weights,'palette':palette.float()}
93
+ return default
94
+
95
+
96
+ def initialize_from_v3(model,baseline):
97
+ old=baseline.state_dict();state=model.state_dict()
98
+ copied=[]
99
+ for k,v in old.items():
100
+ if model.backbone=='mobilevit' and (k.startswith(('encoder.','lateral.')) or k in ['rgb_mean','rgb_std']):continue
101
+ if k in state and v.shape==state[k].shape:
102
+ state[k]=v;copied.append(k)
103
+ model.load_state_dict(state,strict=True)
104
+ return copied
105
+
106
+
107
+ def save_decision(m,path):
108
+ p=Path(path);p.mkdir(parents=True,exist_ok=True)
109
+ save_file({k:v.detach().cpu().contiguous() for k,v in m.state_dict().items()},str(p/'model.safetensors'))
110
+ (p/'config.json').write_text(json.dumps(m.config,indent=2))
111
+
112
+
113
+ def load_decision(path,device='cpu'):
114
+ p=Path(path);cfg=json.loads((p/'config.json').read_text())
115
+ m=DecisionColorizer(**{k:cfg[k] for k in ['modes','backbone','width','queries','head','encoder_revision']})
116
+ m.load_state_dict(load_file(str(p/'model.safetensors')),strict=True)
117
+ return m.to(device).eval()
reports/FINAL_TRAINING_REVIEW.md ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Final training decision — 29 September 2026
2
+
3
+ All seven pilot arms completed. Reviewed all 32 fixed COCO examples in each of the single, multi and encoder comparison groups, all14 historical photos in production-pipeline grids, and the14-photo multi-palette archive comparison. This is qualitative assistant review, not an independent human preference study.
4
+
5
+ ## Decision
6
+ Continue the MobileNetV3 four-palette model with mixed complete-image teacher/original targets, initialized from `experiments/decision-20260928-v2/multi_mix/candidate` (pilot step2000). Complete deployed model:3,995,832 parameters. This is a modest best-compromise candidate, not a proven major advance over v3. Retain released v3 until final visual review.
7
+
8
+ ## Evidence
9
+ Relative to v3, multi_mix guided patch excess decreased4.69% on ordinary grayscale (0.81598→0.77772) and9.37% on film grayscale (0.77398→0.70148). Exploratory paired95% intervals for the absolute changes were[-0.07208,-0.00611] and[-0.10446,-0.04150]. These intervals do not correct for selecting among seven experiments.
10
+ Mean chroma was almost retained:14.0299→13.9184, and13.7452→13.6813. Coverage nevertheless fell from54.44%→52.63% and52.99%→52.30%; missed-colour rates slightly improved, with uncertain paired differences. Original-colour error changed little. No metric is a direct perceptual-quality guarantee.
11
+
12
+ MobileViT had the lowest patch metric but its images often had a yellow/brown cast: skateboarder background, tennis player and court, lake, hand holding food, kitchen hand, and indoor scenes. It also lost useful colour distinctions in the dance scene and food images. It improved some landscape/elephant details and the historical children/sky, but the overall trade-off is unsuitable for release. Its selected step1000 precedes its staged curriculum, so those scores do not validate that curriculum.
13
+
14
+ Single mixed/staged outputs were generally close to v3 and sometimes smoother, but reduced colour coverage and saturation. Single_staged selected step1000, before the late robustness phase; it is not evidence that the late film/consistency phase worked. Multi_staged retained more colour but its ordinary-grayscale patch score worsened. Thus the final course uses mixed targets without the unvalidated late-stage bundle.
15
+
16
+ Multi_mix preserves recognisable blue sky, green foliage and skin distinctions in the modern examples, with less evidence of the broad sepia substitution seen in MobileViT. Historical images remain substantially similar to v3. The old man's clothes and the children's field remain muted; the powerhouse pipe still changes hue along its length. The power station still receives a pink cast. No claim that blotches are eliminated.
17
+
18
+ The multi selector chose modes[39,8,1,0] on48 diagnostic images; alternatives are unevenly used. Default target loss0.5174 versus oracle0.4234 shows room for selection improvement, but oracle performance is not the deployable result. Resize changed the selected mode in1/48 examples. Keep all parameters and the learned default; do not choose an output using unknown ground truth at inference.
19
+
20
+ ## Final run
21
+ Up to24,000 additional updates, batch24, all16,230 existing filtered training images, cached teacher/original targets sampled50/50 per image. Fresh optimizer; encoder LR2e-6 and decoder LR2e-5, warmup150, cosine floor15%, frozen encoder batch statistics, clipping1. No RL, critic, new teacher generation, or architecture expansion. Save every2000 updates and retain the initial pilot as a valid candidate. Stop after six evaluations without a development-score improvement of0.002. The development score retains colour/error guardrails.
22
+
23
+ New development/test samples exclude the pilot's development/test IDs and the earlier evaluation IDs. Test opens after development checkpoint selection. Final outputs and image comparisons are pushed under `experiments/final-20260929` on main; root production weights are not automatically overwritten. A lower metric is insufficient for promotion.
24
+
25
+ Budget: one L4,60min hard timeout ($0.80 maximum at the currently documented rate),50min internal training limit. Expected about25–40min from measured pilot throughput, with setup and evaluation overhead. No automatic retry jobs.
reports/README.md ADDED
@@ -0,0 +1,11 @@
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Reports and experiment history
2
+
3
+ - [Final training decision, visual review and limitations](FINAL_TRAINING_REVIEW.md)
4
+ - [Initial research report](history/RESEARCH_REPORT.md)
5
+ - [Second research report](history/RESEARCH_ROUND2.md)
6
+ - [Historical project log](history/colorization_project_log.md)
7
+ - [Current v3 release report](../RELEASE_REPORT.md)
8
+ - [Seven-arm architecture/curriculum comparison](../experiments/decision-20260928-v2/PROTOCOL.md)
9
+ - [Final continuation outputs](../experiments/final-20260929/)
10
+
11
+ The historical log includes earlier U-Net versions and superseded conclusions. For deployed architecture, parameter counts and inference, use the root model card and its pinned release. Experiment checkpoints are not automatically approved releases.
reports/SPACE_QA_multi.json ADDED
@@ -0,0 +1,15 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "parameters": 3995832,
3
+ "gradio": "6.28.0",
4
+ "rgba_and_dimensions": true,
5
+ "cached_strength": true,
6
+ "max_lightness_quantization_error": 0.23923175575065514,
7
+ "blocks_constructed": true,
8
+ "tested_modes": [
9
+ -1,
10
+ 0,
11
+ 1,
12
+ 2,
13
+ 3
14
+ ]
15
+ }
reports/SPACE_QA_v3.json ADDED
@@ -0,0 +1,11 @@
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "parameters": 3994676,
3
+ "gradio": "6.28.0",
4
+ "rgba_and_dimensions": true,
5
+ "cached_strength": true,
6
+ "max_lightness_quantization_error": 0.24549842656874432,
7
+ "blocks_constructed": true,
8
+ "tested_modes": [
9
+ -1
10
+ ]
11
+ }
RESEARCH_REPORT.md → reports/history/RESEARCH_REPORT.md RENAMED
File without changes
RESEARCH_ROUND2.md → reports/history/RESEARCH_ROUND2.md RENAMED
File without changes
colorization_project_log.md → reports/history/colorization_project_log.md RENAMED
File without changes
requirements-space.txt CHANGED
@@ -7,3 +7,4 @@ safetensors==0.5.3
7
  huggingface-hub>=1.0,<2
8
  gradio==6.28.0
9
  spaces==0.51.3
 
 
7
  huggingface-hub>=1.0,<2
8
  gradio==6.28.0
9
  spaces==0.51.3
10
+ timm==1.0.19
training/final_config.json ADDED
@@ -0,0 +1,21 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architecture": "DecisionColorizer",
3
+ "backbone": "mobilenet",
4
+ "modes": 4,
5
+ "parameters": 3995832,
6
+ "initial_revision": "12dfd8111622f692a7f69d0a961ce9b65f044c01",
7
+ "initial_subfolder": "experiments/decision-20260928-v2/multi_mix/candidate",
8
+ "output_subfolder": "experiments/final-20260929",
9
+ "seed": 9292026,
10
+ "max_steps": 24000,
11
+ "batch_size": 24,
12
+ "encoder_lr": 2e-06,
13
+ "decoder_lr": 2e-05,
14
+ "original_target_probability": 0.5,
15
+ "checkpoint_interval": 2000,
16
+ "early_stopping_patience_evaluations": 6,
17
+ "early_stopping_min_delta": 0.002,
18
+ "automatic_production_promotion": false,
19
+ "job_id": "6abaef6552d0dbd7f1da7286",
20
+ "job_url": "https://huggingface.co/jobs/User-2468/6abaef6552d0dbd7f1da7286"
21
+ }
training/train_final.py ADDED
@@ -0,0 +1,368 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Budget-bounded architecture/curriculum screening. Sources pinned by bootstrap."""
2
+ import os,sys,json,time,math,random,hashlib,shutil,traceback,copy,io
3
+ from pathlib import Path
4
+ import numpy as np
5
+ import torch
6
+ from torch import nn
7
+ import torch.nn.functional as F
8
+ from PIL import Image,ImageDraw
9
+ from skimage.color import rgb2lab
10
+ import cv2,pyarrow as pa,pyarrow.parquet as pq
11
+ from huggingface_hub import HfApi,hf_hub_download
12
+ from safetensors.torch import load_file
13
+ from semantic_model import load_semantic
14
+ from decision_model import DecisionColorizer,initialize_from_v3,save_decision,load_decision
15
+ import persistence
16
+ import r7 as r
17
+ from data import decode_image
18
+
19
+ REPO='User-2468/mini-unet-colorizer'
20
+ BASE='1a9eb8af2754ad2329a24cfe50d388cb559441d0'
21
+ ARTIFACTS='12dfd8111622f692a7f69d0a961ce9b65f044c01'
22
+ CACHE_REV='7b10485ee9d205ce2f7b660dfc6a90012d78b1cb'
23
+ COCO_REV='26ddc382fe75dfc2a0655b5977e296ea10efebce'
24
+ PREFIX=os.environ.get('EXPERIMENT_PREFIX','experiments/final-20260929')
25
+ SEED=int(os.environ.get('SEED','9292026'))
26
+ STEPS=int(os.environ.get('STEPS','24000'));BATCH=24
27
+ START=time.monotonic();TRAIN_LIMIT=float(os.environ.get('TRAIN_LIMIT_SECONDS','3000'))
28
+ OUT=Path('/work/results');OUT.mkdir(exist_ok=True,parents=True)
29
+ persistence.PREFIX=PREFIX;r.OUT=OUT;r.ev.OUT=OUT;r.DEVICE='cuda';r.ev.DEVICE='cuda';r.ev.deadline=lambda:None
30
+ api=HfApi();DURABLE=None
31
+ def write(n,obj):
32
+ p=OUT/n;p.parent.mkdir(parents=True,exist_ok=True);p.write_text(json.dumps(obj,indent=2));return obj
33
+ def dl(f,rev=ARTIFACTS):return Path(hf_hub_download(REPO,f,revision=rev))
34
+ def seed(s):torch.manual_seed(s);np.random.seed(s);random.seed(s)
35
+
36
+ def objective(details,target,step):
37
+ allout=details['all'];b,k=allout.shape[:2]
38
+ losses=r.losses(allout.flatten(0,1),target[:,None].expand(-1,k,-1,-1,-1).flatten(0,1)).reshape(b,k)
39
+ if k==1:return losses.mean(),losses.argmin(1),losses
40
+ # One winning complete image, never a per-pixel best-of-K patchwork.
41
+ winner=losses.detach().argmin(1)
42
+ selected=losses.gather(1,winner[:,None]).mean()
43
+ # Small nonwinner term prevents completely untrained branches; no saturation reward.
44
+ loss=.95*selected+.05*losses.mean()+.1*F.cross_entropy(details['scores'],winner)
45
+ return loss,winner,losses
46
+
47
+ class Batches:
48
+ def __init__(self,rgb,teacher):
49
+ self.rgb=torch.from_numpy(rgb).cuda().permute(0,3,1,2)
50
+ self.teacher=torch.from_numpy(teacher).cuda().float()
51
+ self.n=len(rgb)
52
+ @staticmethod
53
+ def light(g):
54
+ y=torch.where(g<=.04045,g/12.92,((g+.055)/1.055).pow(2.4))
55
+ l=torch.where(y>.008856,116*y.clamp_min(1e-9).pow(1/3)-16,903.296296*y)
56
+ return l/50-1
57
+ def batch(self,step,recipe,original,curriculum_step=None):
58
+ phase=step if curriculum_step is None else curriculum_step
59
+ gen=torch.Generator(device='cuda').manual_seed(SEED*100000+step)
60
+ ids=torch.randint(self.n,(BATCH,),generator=gen,device='cuda')
61
+ rgb=self.rgb[ids].float()/255;t=self.teacher[ids];o=original[ids]
62
+ g=(rgb*rgb.new_tensor([.299,.587,.114])[None,:,None,None]).sum(1,keepdim=True)
63
+ # Emulate standard 8-bit PIL grayscale conversion before Lab conversion.
64
+ g=(g*255).round()/255
65
+ L=self.light(g)
66
+ if recipe=='staged' and phase>STEPS*2/3:
67
+ take=torch.rand(BATCH,1,1,1,generator=gen,device='cuda')<.35
68
+ weights=torch.rand(BATCH,3,1,1,generator=gen,device='cuda')+rgb.new_tensor([.8,1.2,.2])[None,:,None,None]
69
+ weights/=weights.sum(1,keepdim=True)
70
+ film=(rgb*weights).sum(1,keepdim=True)
71
+ gamma=.85+.4*torch.rand(BATCH,1,1,1,generator=gen,device='cuda')
72
+ film=(film.pow(gamma)-.5)*(.8+.3*torch.rand(BATCH,1,1,1,generator=gen,device='cuda'))+.5
73
+ film=film+.005*torch.randn(film.shape,generator=gen,device='cuda')
74
+ L=torch.where(take,self.light(film.clamp(0,1)),L)
75
+ else:
76
+ gain=.9+.2*torch.rand(BATCH,1,1,1,generator=gen,device='cuda')
77
+ L=(L*gain).clamp(-1,1)
78
+ flips=torch.rand(BATCH,1,1,1,generator=gen,device='cuda')<.5
79
+ L=torch.where(flips,L.flip(-1),L);t=torch.where(flips,t.flip(-1),t);o=torch.where(flips,o.flip(-1),o)
80
+ prob=0 if recipe=='teacher' else .5
81
+ if recipe=='staged':prob*=min(max((phase/STEPS-.15)/.35,0),1)
82
+ pick=torch.rand(BATCH,1,1,1,generator=gen,device='cuda')<prob
83
+ # Stochastic full targets: never choose the easier one, never average hues.
84
+ target=torch.where(pick,o,t)
85
+ return L,target,ids,float(pick.float().mean())
86
+
87
+ def summary(result):return {s:result[s]['summary'] for s in ['gray','film']}
88
+ def score(result,baseline):
89
+ val=0;okay=True
90
+ for style in ['gray','film']:
91
+ s=result[style]['summary'];b=baseline[style]['summary']
92
+ okay &= s['missed_color']<=b['missed_color']+.015 and s['neutral_spill']<=b['neutral_spill']+.015 and s['color_coverage']>=.95*b['color_coverage'] and s['ab_error']<=b['ab_error']*1.10
93
+ val+=s['patch_excess']/max(b['patch_excess'],1e-6)+.2*s['ab_error']/b['ab_error']+2*s['missed_color']+s['neutral_spill']
94
+ return val/2,bool(okay)
95
+
96
+ @torch.inference_mode()
97
+ def diagnostics(model,table,ids):
98
+ model.eval();rows=[]
99
+ for idx in ids:
100
+ L,t,_=r.ev.gray_arrays(decode_image(table,idx));x=torch.from_numpy(L)[None].cuda()
101
+ v=model(x,True);a=v['all'];sel=int(v['scores'].argmax(1))
102
+ resized=model(F.interpolate(x,(224,224),mode='bilinear',align_corners=False),True)
103
+ p=F.interpolate(resized['output'],(256,256),mode='bilinear',align_corners=False)
104
+ flip=model(x.flip(-1),True)
105
+ colors=a[0].mean((-1,-2));k=model.modes
106
+ diff=(a[:,:,None]-a[:,None]).abs().mean((0,3,4,5))
107
+ target=torch.from_numpy(t)[None].cuda();target64=F.avg_pool2d(target,4)
108
+ losses=r.losses(a.flatten(0,1),target64.expand(k,-1,-1,-1))
109
+ mask=v['weights'].clamp_min(1e-8)
110
+ rows.append({'index':int(idx),'selected_mode':sel,'resized_mode':int(resized['scores'].argmax(1)),
111
+ 'flip_mode':int(flip['scores'].argmax(1)),'resize_ab_change':float((p-v['output']).abs().mean()),
112
+ 'flip_ab_change':float((flip['output'].flip(-1)-v['output']).abs().mean()),
113
+ 'mode_diversity':float(diff.sum()/max(k*(k-1),1)),
114
+ 'default_target_loss':float(losses[sel]),'oracle_target_loss':float(losses.min()),
115
+ 'mask_entropy':float(-(mask*mask.log()).sum(1).mean()),
116
+ 'palette_chroma':float(v['palette'].norm(dim=-1).mean()),'output_chroma':float(v['output'].norm(dim=1).mean())})
117
+ keys=['resize_ab_change','flip_ab_change','mode_diversity','default_target_loss','oracle_target_loss','mask_entropy','palette_chroma','output_chroma']
118
+ return {'rows':rows,'summary':{k:float(np.mean([x[k] for x in rows])) for k in keys},
119
+ 'mode_frequency':np.bincount([x['selected_mode'] for x in rows],minlength=model.modes).tolist(),
120
+ 'resize_switch_fraction':float(np.mean([x['selected_mode']!=x['resized_mode'] for x in rows])),
121
+ 'note':'Oracle uses unknown original only for diagnosis. Default mode is selected by the learned image-only score, never by ground truth.'}
122
+
123
+ def paired_bootstrap(result,base):
124
+ rng=np.random.default_rng(937);out={}
125
+ for style in ['gray','film']:
126
+ br={z['index']:z for z in base[style]['per_image']};rows=[z for z in result[style]['per_image'] if z['target_chroma']>=5]
127
+ out[style]={}
128
+ for key in ['patch_excess','raw_patch_excess','missed_color','neutral_spill','color_coverage','ab_error']:
129
+ d=np.array([z[key]-br[z['index']][key] for z in rows]);boot=d[rng.integers(len(d),size=(3000,len(d)))].mean(1)
130
+ out[style][key]={'delta':float(d.mean()),'ci95':np.quantile(boot,[.025,.975]).tolist()}
131
+ return out
132
+
133
+ @torch.inference_mode()
134
+ def pipeline_grid(models,table,ids,label):
135
+ from inference import colorize
136
+ from PIL import ImageOps
137
+ names=list(models);order=np.random.default_rng(9184).permutation(len(names))
138
+ key={chr(65+j):names[int(i)] for j,i in enumerate(order)};write(label+'_key.json',key)
139
+ for start in range(0,min(32,len(ids)),4):
140
+ canvas=Image.new('RGB',(320*(len(names)+1),344*len(ids[start:start+4])),'white');d=ImageDraw.Draw(canvas)
141
+ for row,idx in enumerate(ids[start:start+4]):
142
+ original=decode_image(table,idx).convert('L').convert('RGB')
143
+ original.thumbnail((800,800))
144
+ d.text((2,row*344),f'{idx}: grayscale | '+' | '.join(key),fill='black')
145
+ canvas.paste(ImageOps.pad(original,(320,320),color='white'),(0,row*344+24))
146
+ for col,j in enumerate(order):
147
+ output=colorize(models[names[int(j)]].eval(),original,size=256,radius=8)
148
+ canvas.paste(ImageOps.pad(output,(320,320),color='white'),((col+1)*320,row*344+24))
149
+ canvas.save(OUT/f'{label}_{start//4}.jpg',quality=93)
150
+
151
+ @torch.inference_mode()
152
+ def resolution_grid(model,table,records):
153
+ from inference import colorize
154
+ from PIL import ImageOps
155
+ for start in range(0,len(table),3):
156
+ count=min(3,len(table)-start)
157
+ canvas=Image.new('RGB',(1280,344*count),'white');d=ImageDraw.Draw(canvas)
158
+ for row,idx in enumerate(range(start,start+count)):
159
+ im=decode_image(table,idx).convert('L').convert('RGB');im.thumbnail((1000,1000))
160
+ d.text((2,row*344),records[idx]['name']+' | input | 256 | 384 | 512',fill='black')
161
+ canvas.paste(ImageOps.pad(im,(320,320),color='white'),(0,row*344+24))
162
+ for col,size in enumerate([256,384,512],1):
163
+ pred=colorize(model,im,size=size,radius=8)
164
+ canvas.paste(ImageOps.pad(pred,(320,320),color='white'),(col*320,row*344+24))
165
+ canvas.save(OUT/f'v3_resolution_{start//3}.jpg',quality=94)
166
+
167
+ def main():
168
+ global DURABLE
169
+ assert api.whoami()['name']=='User-2468'
170
+ DURABLE=persistence.DurableRun(OUT)
171
+ shutil.copytree('/work/source',OUT/'source',dirs_exist_ok=True,ignore=shutil.ignore_patterns('__pycache__'))
172
+ shutil.copy2('/work/PROTOCOL.md',OUT/'PROTOCOL.md')
173
+ if Path('/work/REVIEW_HANDOFF.md').exists():shutil.copy2('/work/REVIEW_HANDOFF.md',OUT/'REVIEW_HANDOFF.md')
174
+ if Path('/work/supplement').exists():shutil.copytree('/work/supplement',OUT/'supplement',dirs_exist_ok=True)
175
+ write('status.json',{'status':'preflight','seed':SEED,'production_approved':False})
176
+ DURABLE.sync('source and write verification before training')
177
+ torch.set_num_threads(6);seed(SEED)
178
+ torch.backends.cuda.matmul.allow_tf32=False;torch.backends.cudnn.allow_tf32=False
179
+ initial=Path('/work/initial');initial.mkdir(exist_ok=True)
180
+ for f in ['model.safetensors','config.json']:shutil.copy2(dl(f,BASE),initial/f)
181
+ baseline=load_semantic(initial,'cuda').requires_grad_(False)
182
+ counts={};smokes={}
183
+ encoder_rev=None
184
+ write('encoder_provenance.json',{'backbone':'MobileNetV3 Large','initial_revision':ARTIFACTS,'initial_subfolder':'experiments/decision-20260928-v2/multi_mix/candidate'})
185
+ constructors=[('multi',4,'mobilenet')]
186
+ for name,k,backbone in constructors:
187
+ seed(SEED);m=DecisionColorizer(modes=k,backbone=backbone,pretrained=backbone=='mobilevit',encoder_revision=encoder_rev if backbone=='mobilevit' else None).cuda()
188
+ copied=initialize_from_v3(m,baseline)
189
+ seedpath=Path('/work/pilot');seedpath.mkdir(exist_ok=True)
190
+ for f in ['config.json','model.safetensors']:shutil.copy2(dl('experiments/decision-20260928-v2/multi_mix/candidate/'+f),seedpath/f)
191
+ m=load_decision(seedpath,'cuda');m.eval();count=sum(p.numel() for p in m.parameters());assert count<4_000_000
192
+ counts[name]=count
193
+ x=torch.rand(2,1,97,113,device='cuda')*2-1
194
+ with torch.no_grad():
195
+ v=m(x,True);assert v['all'].shape==(2,k,2,97,113) and torch.isfinite(v['all']).all()
196
+ parity=float((v['output']-baseline(x)).abs().max()) if name=='single' else None
197
+ if parity is not None:assert parity<1e-3,parity
198
+ loss,_,_=objective(m(x,True),torch.randn(2,2,64,64,device='cuda'),1)
199
+ loss.backward();assert torch.isfinite(loss) and all(torch.isfinite(p.grad).all() for p in m.parameters() if p.grad is not None)
200
+ m.zero_grad(set_to_none=True);save_decision(m,OUT/'initial'/name)
201
+ re=load_decision(OUT/'initial'/name,'cuda')
202
+ with torch.no_grad():assert torch.allclose(m(x),re(x),atol=1e-4)
203
+ smokes[name]={'parameters':count,'copied_tensors':len(copied),'odd_shape':True,'finite_backward':True,'reload':True,'baseline_max_abs':parity}
204
+ del m,re
205
+ write('smoke.json',smokes);DURABLE.sync('all three architectures pass preflight')
206
+ torch.backends.cuda.matmul.allow_tf32=True;torch.backends.cudnn.allow_tf32=True;torch.backends.cudnn.benchmark=True
207
+ tables=[]
208
+ for repo,rev,files in [('johnowhitaker/imagenette2-320','771c1310a2487e8076ede6b7d6307244aa8400af',['default/train/0000.parquet']),('detection-datasets/coco',COCO_REV,['default/train/0000.parquet','default/train/0001.parquet'])]:
209
+ for f in files:
210
+ print('PREP downloading',repo,f,flush=True)
211
+ tables.append(pq.read_table(hf_hub_download(repo,f,repo_type='dataset',revision=rev),columns=['image']))
212
+ print('PREP loaded',repo,f,len(tables[-1]),flush=True)
213
+ table=pa.concat_tables(tables)
214
+ full_manifest=json.loads(dl('experiments/region-coherence-20260928/source/previous_manifest.json').read_text())
215
+ full_ids=full_manifest['train']
216
+ # Same fixed training subset for every arm, independent of new dev/test.
217
+ ids=list(full_ids);pos={i:j for j,i in enumerate(ids)}
218
+ teacher=np.empty((len(ids),2,64,64),np.float16)
219
+ sums=json.loads(dl('experiments/round6-20260928/SHA256SUMS.json',CACHE_REV).read_text())
220
+ for start in range(0,len(full_ids),2048):
221
+ f=f'teacher_cache/{start:05d}.npz';p=dl('experiments/round6-20260928/'+f,CACHE_REV)
222
+ assert hashlib.sha256(p.read_bytes()).hexdigest()==sums[f]
223
+ with np.load(p) as z:
224
+ # NpzFile indexing decompresses a member on every access. Materialize once.
225
+ cache_ids=z['indices'];cache_targets=z['target']
226
+ assert cache_ids.tolist()==full_ids[start:start+len(cache_ids)]
227
+ src=[j for j,i in enumerate(cache_ids) if int(i) in pos]
228
+ dst=[pos[int(cache_ids[j])] for j in src]
229
+ teacher[dst]=cache_targets[src]
230
+ print('PREP cache',start,'selected',len(src),flush=True)
231
+ print('PREP decoding training images',len(ids),flush=True)
232
+ rgb=np.stack([np.asarray(decode_image(table,i).resize((256,256),Image.Resampling.BILINEAR)) for i in ids])
233
+ original=np.stack([cv2.resize(rgb2lab(im.astype(np.float32)/255)[...,1:],(64,64),interpolation=cv2.INTER_AREA).transpose(2,0,1) for im in rgb]).astype(np.float32)
234
+ batcher=Batches(rgb,teacher);original=torch.from_numpy(original).cuda();del rgb,teacher
235
+ print('PREP training tensors ready',flush=True)
236
+ fresh=pq.read_table(hf_hub_download('detection-datasets/coco','default/val/0001.parquet',repo_type='dataset',revision=COCO_REV),columns=['image'])
237
+ exclude=set()
238
+ for path,keys in [('decision-20260928-v2',['development','test']),('round7-20260928',['development','heldout']),('region-coherence-20260928',['development_regression','heldout_regression','new_holdout']),('recurrent-refinement-20260928',['fresh_holdout','development_regression'])]:
239
+ manifest=json.loads(dl('experiments/'+path+'/manifest.json').read_text())
240
+ for key in keys:exclude.update(manifest[key])
241
+ hashes={hashlib.sha256(table['image'][i].as_py()['bytes']).digest() for i in full_ids}
242
+ freshids=[int(i) for i in np.random.default_rng(SEED).permutation(len(fresh)) if int(i) not in exclude and hashlib.sha256(fresh['image'][int(i)].as_py()['bytes']).digest() not in hashes][:288]
243
+ assert len(freshids)==288
244
+ dev,test=freshids[:128],freshids[128:]
245
+ write('manifest.json',{'train':ids,'development':dev,'test':test,'prior_eval_excluded':sorted(exclude),'coco_revision':COCO_REV,'fresh_file':'default/val/0001.parquet','teacher_revision':CACHE_REV,'base_revision':BASE,'seed':SEED,'near_duplicates_and_upstream_overlap_not_excluded':True})
246
+ base_dev=r.ev.evaluate(baseline,fresh,dev);write('baseline_development.json',base_dev)
247
+ DURABLE.sync('dataset checks and unseen development/test split frozen')
248
+ arms=[('multi','mix')]
249
+ if os.environ.get('ARMS'):arms=[tuple(x.split(':')) for x in os.environ['ARMS'].split(',')]
250
+ selected={};histories={};dev_results={};begin_train=time.monotonic()
251
+ for head,recipe in arms:
252
+ name=head+'_'+recipe;seed(SEED)
253
+ m=load_decision(OUT/'initial'/head,'cuda')
254
+ # New encoder gets a separate alignment warmup, explicitly outside matched six arms.
255
+ warmup=1000 if head=='mobilevit' else 0
256
+ groups=[{'params':[p for n,p in m.named_parameters() if n.startswith('encoder.')],'lr':2e-5 if head=='mobilevit' else 2e-6},
257
+ {'params':[p for n,p in m.named_parameters() if not n.startswith('encoder.')],'lr':2e-4 if head=='mobilevit' else 2e-5}]
258
+ opt=torch.optim.AdamW(groups,weight_decay=.01)
259
+ for g in opt.param_groups:g['base_lr']=g['lr']
260
+ history=[];beststep=0;last=0;arm_start=time.monotonic();mode_counts=np.zeros(m.modes,np.int64);stale=0
261
+ initial_res=r.ev.evaluate(m,fresh,dev);val,okay=score(initial_res,base_dev);best=val+(0 if okay else 100)
262
+ save_decision(m,OUT/name/'candidate');dev_results[name]=initial_res
263
+ write(name+'/development_00000.json',initial_res);DURABLE.sync('retain initial pilot candidate before continuation')
264
+ for s in range(1,STEPS+warmup+1):
265
+ if time.monotonic()-START>TRAIN_LIMIT:break
266
+ step=max(1,s-warmup);active_recipe='teacher' if s<=warmup else recipe
267
+ L,target,_,fraction=batcher.batch(s,active_recipe,original,curriculum_step=step)
268
+ m.train()
269
+ for layer in m.encoder.modules():
270
+ if isinstance(layer,nn.BatchNorm2d):layer.eval()
271
+ opt.zero_grad(set_to_none=True)
272
+ factor=min(s/150,1)*(.15+.85*(1+math.cos(math.pi*max(0,s-warmup)/STEPS))/2)
273
+ for g in opt.param_groups:g['lr']=g['base_lr']*factor
274
+ with torch.autocast('cuda',dtype=torch.bfloat16):
275
+ details=m(L,True)
276
+ loss,winners,terms=objective(details,target,step)
277
+ consistency=loss.new_zeros(())
278
+ if active_recipe=='staged' and step>STEPS*2/3 and s%4==0:
279
+ small=m(F.interpolate(L,(192,192),mode='bilinear',align_corners=False),True)
280
+ a=F.adaptive_avg_pool2d(small['all'].flatten(0,1).float(),(64,64))
281
+ b=F.adaptive_avg_pool2d(details['all'].flatten(0,1).detach().float(),(64,64))
282
+ consistency=F.l1_loss(a/20,b/20)
283
+ if m.modes>1:consistency+=.1*F.kl_div(small['scores'].log_softmax(1),details['scores'].detach().softmax(1),reduction='batchmean')
284
+ loss=loss+.4*consistency # every fourth step -> mean weight .1
285
+ assert torch.isfinite(loss),name
286
+ loss.backward();torch.nn.utils.clip_grad_norm_(m.parameters(),1,error_if_nonfinite=True);opt.step();last=s
287
+ mode_counts+=np.bincount(winners.detach().cpu().numpy(),minlength=m.modes)
288
+ if s%100==0:
289
+ row={'step':s,'loss':float(loss.detach()),'reconstruction':float(terms.mean().detach()),'consistency':float(consistency.detach()),'photo_fraction_this_batch':fraction,'winner_counts':mode_counts.tolist(),'seconds':time.monotonic()-arm_start}
290
+ history.append(row);print('STEP',name,s,round(row['loss'],4),flush=True)
291
+ if s%2000==0 or s==STEPS+warmup:
292
+ m.eval();res=r.ev.evaluate(m,fresh,dev);val,okay=score(res,base_dev)
293
+ write(name+f'/development_{s:05d}.json',res)
294
+ write(name+'/history.json',history)
295
+ save_decision(m,OUT/name/f'step_{s:05d}')
296
+ # Ineligible models still retained for diagnosis, ranked after eligible.
297
+ ranked=val+(0 if okay else 100)
298
+ if ranked<best-.002:
299
+ stale=0;best=ranked;beststep=s;save_decision(m,OUT/name/'candidate');dev_results[name]=res
300
+ if ranked>=best-.002 and beststep!=s:stale+=1
301
+ write('status.json',{'status':'training','arm':name,'step':s,'elapsed_seconds':time.monotonic()-START,'production_approved':False})
302
+ DURABLE.sync(name+' checkpoint '+str(s))
303
+ if stale>=6:
304
+ print('EARLY_STOP development score unchanged for six evaluations',flush=True);break
305
+ if not last:raise RuntimeError('Budget exhausted before any final training update')
306
+ write(name+'/history.json',history)
307
+ torch.save({'optimizer':opt.state_dict(),'step':last,'seed':SEED,'torch_rng':torch.get_rng_state(),'cuda_rng':torch.cuda.get_rng_state_all(),'exact_resume':False},OUT/name/'optimizer.pt')
308
+ selected[name]={'step':beststep,'rank_score':best,'parameters':counts[head],'steps_completed':last,'warmup_steps':warmup,'seconds':time.monotonic()-arm_start}
309
+ write('candidates.json',selected);DURABLE.sync(name+' selected checkpoint on development only')
310
+ del m,opt;torch.cuda.empty_cache()
311
+ if last<STEPS+warmup:break
312
+ # Test is opened only after all checkpoint choices have been frozen.
313
+ write('candidates.json',selected);DURABLE.sync('freeze candidates before opening final test')
314
+ models={'v3':baseline};test_results={}
315
+ base_test=r.ev.evaluate(baseline,fresh,test);write('v3_test.json',base_test);test_results['v3']=base_test
316
+ for name in selected:
317
+ m=load_decision(OUT/name/'candidate','cuda');models[name]=m
318
+ res=r.ev.evaluate(m,fresh,test);test_results[name]=res
319
+ write(name+'_test.json',res);write(name+'_test_paired.json',paired_bootstrap(res,base_test))
320
+ write(name+'_diagnostics.json',diagnostics(m,fresh,test[:48]))
321
+ r.ev.probes(m,name+'_probes')
322
+ x=torch.zeros(1,1,256,256,device='cuda')
323
+ with torch.inference_mode():
324
+ for _ in range(5):m(x)
325
+ torch.cuda.synchronize();t0=time.monotonic()
326
+ for _ in range(30):m(x)
327
+ torch.cuda.synchronize()
328
+ write(name+'_latency.json',{'milliseconds':1000*(time.monotonic()-t0)/30,'hardware':torch.cuda.get_device_name(),'size':256,'precision':'FP32'})
329
+ DURABLE.sync(name+' final test and diagnostic outputs')
330
+ r.ev.probes(baseline,'v3_probes')
331
+ # Smaller comparison groups preserve enough image detail to inspect.
332
+ for head in ['single','multi']:
333
+ group={n:m for n,m in models.items() if n=='v3' or n.startswith(head+'_')}
334
+ r.grid(group,fresh,test,head+'_test')
335
+ if 'mobilevit_staged' in models:r.grid({'v3':baseline,'single_staged':models['single_staged'],'mobilevit_staged':models['mobilevit_staged']},fresh,test,'encoder_test')
336
+ # Historical originals deliberately never serve as neutral-color supervision.
337
+ archive=json.loads(Path('/work/archive_sources.json').read_text())
338
+ import requests
339
+ records=[];images=[]
340
+ for entry in archive:
341
+ try:
342
+ response=requests.get(entry['image_url'],timeout=20);response.raise_for_status()
343
+ im=Image.open(io.BytesIO(response.content)).convert('RGB');im.load()
344
+ images.append({'bytes':response.content,'path':None});records.append(entry|{'sha256':hashlib.sha256(response.content).hexdigest()})
345
+ except Exception as e:records.append(entry|{'failed':type(e).__name__})
346
+ if images:
347
+ ar=pa.table({'image':images})
348
+ resolution_grid(baseline,ar,[x for x in records if 'failed' not in x])
349
+ for head in ['single','multi']:
350
+ r.grid({n:m for n,m in models.items() if n=='v3' or n.startswith(head+'_')},ar,list(range(len(images))),head+'_archive')
351
+ if 'mobilevit_staged' in models:r.grid({'v3':baseline,'mobilevit_staged':models['mobilevit_staged']},ar,list(range(len(images))),'encoder_archive')
352
+ pipeline_grid({'v3':baseline,**{n:m for n,m in models.items() if n!='v3'}},ar,list(range(len(images))),'production_archive')
353
+ pipeline_grid({'v3':baseline,**{n:m for n,m in models.items() if n!='v3'}},fresh,test[:16],'production_test')
354
+ write('archive_sources.json',records)
355
+ write('summary.json',{'parameters':counts,'candidates':selected,'test':{n:summary(v) for n,v in test_results.items()},'all_arms_completed':len(selected)==len(arms),'planned_max_steps':STEPS,'elapsed_seconds':time.monotonic()-START,'estimated_compute_usd':(time.monotonic()-START)/3600*.8,'warning':'No metric proves realism. All outputs need visual review. Candidate can be the retained step0 pilot; selection does not automatically approve production deployment.'})
356
+ write('status.json',{'status':'completed','all_arms_completed':len(selected)==len(arms),'planned_max_steps':STEPS,'elapsed_seconds':time.monotonic()-START,'production_approved':False})
357
+ DURABLE.sync('architecture and curriculum suite completed')
358
+
359
+ if __name__=='__main__':
360
+ try:main()
361
+ except BaseException:
362
+ msg=traceback.format_exc();token=os.environ.get('HF_TOKEN')
363
+ if token:msg=msg.replace(token,'[REDACTED]')
364
+ write('failure.json',{'traceback':msg});write('status.json',{'status':'failed','elapsed_seconds':time.monotonic()-START,'production_approved':False})
365
+ if DURABLE:
366
+ try:DURABLE.sync('failure report and saved checkpoints')
367
+ except Exception:pass
368
+ print(msg,flush=True);raise SystemExit(1)