Document final training decision and organize research history
Browse files- DEPLOYMENT.md +13 -29
- README.md +10 -0
- SHA256SUMS.json +15 -6
- app.py +58 -24
- decision_model.py +117 -0
- reports/FINAL_TRAINING_REVIEW.md +25 -0
- reports/README.md +11 -0
- reports/SPACE_QA_multi.json +15 -0
- reports/SPACE_QA_v3.json +11 -0
- RESEARCH_REPORT.md → reports/history/RESEARCH_REPORT.md +0 -0
- RESEARCH_ROUND2.md → reports/history/RESEARCH_ROUND2.md +0 -0
- colorization_project_log.md → reports/history/colorization_project_log.md +0 -0
- requirements-space.txt +1 -0
- training/final_config.json +21 -0
- training/train_final.py +368 -0
DEPLOYMENT.md
CHANGED
|
@@ -1,34 +1,18 @@
|
|
| 1 |
-
#
|
| 2 |
|
| 3 |
-
|
|
|
|
| 4 |
|
| 5 |
-
|
|
|
|
| 6 |
|
| 7 |
-
```
|
| 8 |
-
from inference import load_colorizer, colorize
|
| 9 |
-
model = load_colorizer('.', 'cpu')
|
| 10 |
-
```
|
| 11 |
|
| 12 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 13 |
|
| 14 |
-
|
| 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": "
|
| 3 |
"PIPELINE_EVAL.json": "595be88704d580c01a3c4dd642409b3dc5dd25e86ca1c7eb40a696605f181a19",
|
| 4 |
"QA.json": "b249b1d49045e31b0bf85f9e01515b91a72c4dd973963751c30d14272d262b90",
|
| 5 |
-
"README.md": "
|
| 6 |
"RELEASE_MANIFEST.json": "872eebfb322c8990b5e6ac493e4e4068b87f8daac0ddf219b8b55299514ca03c",
|
| 7 |
"RELEASE_REPORT.md": "69677ec5e13f5c1249f59b01c31ed1b0f1588bddc216bce9167b5d31423a5574",
|
| 8 |
-
"app.py": "
|
| 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": "
|
| 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
|
| 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
|
|
|
|
| 9 |
|
| 10 |
REPO=os.getenv('MODEL_ID','User-2468/mini-unet-colorizer')
|
| 11 |
-
REV=os.getenv('MODEL_REVISION','
|
|
|
|
| 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():
|
| 17 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 18 |
|
| 19 |
-
|
| 20 |
-
def
|
| 21 |
-
|
|
|
|
| 22 |
|
| 23 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 24 |
if image is None:raise gr.Error('Upload a photograph first.')
|
| 25 |
-
|
|
|
|
| 26 |
try:
|
| 27 |
-
|
| 28 |
-
|
|
|
|
|
|
|
| 29 |
result=render(light,alpha,coefficients,strength)
|
| 30 |
-
|
|
|
|
| 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('#
|
|
|
|
| 36 |
with gr.Row():
|
| 37 |
-
source=gr.Image(type='pil',image_mode='RGBA',label='
|
| 38 |
-
result=gr.Image(type='pil',label='Colourised photograph',format='png',height=
|
| 39 |
with gr.Row():
|
| 40 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 41 |
smoothing=gr.Radio(['Gentle','Balanced','Strong'],value='Balanced',label='Colour smoothing')
|
| 42 |
-
|
| 43 |
-
|
| 44 |
-
|
| 45 |
-
|
|
|
|
|
|
|
| 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)
|