User-2468 commited on
Commit
a34407b
·
verified ·
1 Parent(s): 85b5823

Colorizer round6: source and write verification before training

Browse files
experiments/decision-20260928-v2/PROTOCOL.md ADDED
@@ -0,0 +1,47 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Penultimate architecture and curriculum decision — 28 September 2026
2
+
3
+ ## Objective and decision scope
4
+ Select one reproducible course for the user's next, longer training run. Prioritize coherent plausible colors, fewer within-surface patches, and retained color over recovering an unknowable original hue. Strictly fewer than 4,000,000 deployed learned parameters, including encoder and any selector. This screen cannot establish a universal architecture ceiling. No experimental model replaces root production weights during this pass.
5
+
6
+ ## Evidence audit
7
+ The uploaded project log describes an older classification U-Net. Current production is round6 palette step9000 at base revision `1a9eb8af2754ad2329a24cfe50d388cb559441d0`, 3,994,676 parameters. The authoritative current repo revision before this pass was `26d731862cc0ea531e7198fa823cb09db25bd243`.
8
+
9
+ The semantic encoder/palette change improved substantially over the old U-Net. Round7's easier-of-original-or-teacher objective selected original targets only 5.8–9.5% of the time; it did not learn explicit alternative outputs. Critic variants (including 27M, smaller, and dense critics) failed visual release review. A region auxiliary and repeating the same decoder did not produce a clear gain. These results motivate separating color ambiguity from spatial assignment, not another saturation/critic sweep.
10
+
11
+ ## Primary research and implications
12
+ - Zhang et al., Colorful Image Colorization: https://richzhang.github.io/colorization/ — classification/rebalancing addresses uncertainty; pixel-wise decoding still needs spatial coordination. Earlier log language that classification is universally required was an overstatement: later palette regression already outperformed that baseline.
13
+ - Deshpande et al., Learning Diverse Image Colorization: https://arxiv.org/abs/1612.01958 and https://arxiv.org/html/1612.01958v2 — low-dimensional color-field variables coordinate alternatives over long spatial distances. Motivates global modes instead of pixel noise.
14
+ - Wang et al., PalGAN: https://arxiv.org/html/2210.11204v1 — probabilistic palette prediction and assignment separate ambiguity from boundary handling. This project's small palette model is not a PalGAN reproduction.
15
+ - Lee et al., Stochastic Multiple Choice Learning: https://arxiv.org/abs/1606.07839 — winner-based learning supports distinct hypotheses, but oracle quality is not automatic-default quality. We explicitly measure the gap and train an image-only selector.
16
+ - Mehta and Rastegari, MobileViTv2: https://arxiv.org/abs/2206.02680 — efficient global attention in a mobile encoder. Test the ImageNet-pretrained 0.75 variant as a distinct backbone challenger, with all deployment parameters counted.
17
+ - Saharia et al., Palette: https://arxiv.org/abs/2111.05826 — diffusion supports colorization diversity. Not a justified final-run choice here without training/validating a compact denoiser; repeating the old decoder did not test diffusion. Excluded from this finite screen because establishing a from-scratch diffusion model and sampler is substantially larger work than these warm-started alternatives.
18
+
19
+ ## Controlled comparisons
20
+ Six arms: single palette vs four coherent palettes, crossed with teacher / mixed-target / staged curriculum. Same MobileNet v3 initialization, fixed8192 training images, seeds, image-index batches, augmentation draws, AdamW, batch24, 3000 updates. Encoder LR1e-5, decoder LR1e-4, warmup150 and cosine decay to15%, weight decay0.01, frozen encoder BatchNorm, gradient norm cap1, BF16 forwards/FP32 losses. New mode tensors and selector are included in the decoder group. The multi head adds only1,156 parameters (3,995,832 total).
21
+
22
+ Multi head: shared spatial masks and encoder; four latent palette variants; image-only score selects one complete output. Training loss is95% winner reconstruction +5% mean reconstruction +0.1 selector cross-entropy against the detached winning index. Reconstruction is L1/20 +0.25 sliced color distribution distance +0.1 boundary difference. Winner is selected per entire image, never per pixel. Report mode usage, between-mode diversity, default vs oracle losses and mode-switching under input resize. No original color is available to the runtime selector. No hue rotation, saturation bonus or unvalidated generated 'plausible' labels.
23
+
24
+ Curricula:
25
+ 1. Teacher: deterministic cached DDColor artistic targets.
26
+ 2. Mixed: independently sample original photograph or teacher target with probability0.5 per image each update. Never average two hues and never pick the easier target. Both represent complete color proposals; teacher remains fallible.
27
+ 3. Staged: first15% teacher, next35% ramp original probability to0.5; final third adds35% mild spectral/film/noise augmentation plus same-mode resize consistency (mean coefficient0.1, computed every fourth update). This tests a curriculum bundle, not isolated causal effects of each component. Anchored reconstruction remains active to guard against consistency-induced gray collapse.
28
+
29
+ Seventh arm: MobileViTv2-0.75 + single palette, copied compatible decoder weights and new lateral projections. 3,481,205 parameters.1000 extra teacher alignment updates, then3000 staged updates; encoder LR2e-5, decoder2e-4. Pretrained encoder pinned to the revision recorded in encoder_provenance.json. This is a practical challenger with unequal prior colorization training and an extra alignment stage, not a pure equal-initialization ablation. A negative short pilot does not prove the encoder incapable.
30
+
31
+ ## Data and evaluation
32
+ Training subset sampled from existing filtered16,230-image Imagenette/COCO training pool; no new teacher expense. Existing target cache hashes and row IDs checked. Pin dataset revisions. New128 development and160 test images from COCO val/0001 exclude earlier round7, region, recurrent evaluation IDs and exact-byte training duplicates. Teacher/encoder upstream exposure and near duplicates remain possible. Fix splits before training. Choose each arm's candidate from evaluations every1000 updates on development only; test opened after choices frozen. Keep all evaluated checkpoints.
33
+
34
+ Screen on ordinary and film-style grayscale. Record raw and guided patch metrics, chroma retention, missed color, spill, auxiliary original-color error, resize and flip sensitivity, latency, and image-bootstrap paired intervals. Gray output can win a smoothness metric; do not rank on patch score alone. Screening guardrails allow missed color/spill at most+1.5points, coverage at least95% baseline, and error at most110% baseline. These are diagnostic eligibility conditions, not proof of plausibility. All checkpoint scans are exploratory and intervals do not correct multiple comparisons.
35
+
36
+ Visual review: 32 fixed fresh examples per group, archival photos never used as neutral targets, seven familiar probes, and aspect-preserving production-pipeline outputs. Label comparisons anonymously with saved keys; assistant review is qualitative, not an independent human preference study. Archive source URLs and checksums retained. A modest historical panel is not a representative archival benchmark.
37
+
38
+ ## Targeted fine-tuning and RL
39
+ Late-stage robustness and consistency are tested rather than assumed helpful. No RL launched in this screen: no independently validated preference reward currently exists, and previous learned critics were exploitable. Prefer direct differentiable supervision for measurable consistency; if a later reward is proposed it must first rank coherent alternate hues above gray collapse and patchwork on held-out images. No claim that RL is intrinsically unsuitable.
40
+
41
+ ## Persistence, budget, and final selection
42
+ One L4 job with60min hard limit and50min internal training deadline; all outputs checkpointed to this main-branch experiment folder with SHA256 download verification. Source/access/preflight checks before training. The first attempt was canceled before training after about12min because of repeated NPZ decompression. Corrected version2 materializes each cache once and logs data preparation. A separate confirmation, if warranted, is capped18min; total maximum hardware charges at the verified$0.80/h rate approximately$1.20, excluding no separate service purchases. Current account balance is not exposed by the connector. No automatic paid reruns.
43
+
44
+ Choose architecture/curriculum from test plus visual review, with production baseline retained as an eligible fallback. If a candidate clearly wins, confirm against the matching control on a second seed or targeted ablation within the remaining cap. Save final_config.json and NEXT_PASS.md with exact checkpoints, source revisions, dataset/curriculum, learning rates, retention gates and validation instructions. The next pass starts the longer run and performs Space/model-card/repository work. This pass does not claim any longer run already completed.
45
+
46
+ ## Version2 corrections before training
47
+ NPZ cache members are decompressed once per shard, not once per selected image; a synthetic equivalence and timing check is saved. MobileViT now uses a separate post-warmup curriculum step. Six additional National Archives inputs supplement the original historical panel. Version1 was canceled before any training checkpoint; its source/preflight remains for audit.
experiments/decision-20260928-v2/REVIEW_HANDOFF.md ADDED
@@ -0,0 +1,30 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Experimental pass handoff — 2026-09-28
2
+
3
+ ## Status
4
+ Version1 job `6aba3fb36b030d633f69c572` was canceled before training to fix repeated NPZ decompression. The replacement runs the seven-arm suite. Source, architecture preflight, parameter counts and upload/download verification are already committed to `experiments/decision-20260928-v2` on main. At this handoff no trained candidate has been reviewed or approved. Do not describe the suite as complete or declare a winner without reading its outputs.
5
+
6
+ Root production remains v3 at `1a9eb8af2754ad2329a24cfe50d388cb559441d0`. The next pass should review this suite before launching the final longer run. Never overwrite root model files merely because an experimental score is lower.
7
+
8
+ ## Exact procedure for selecting the final course
9
+ 1. Read status.json, failure.json if present, summary.json and candidates.json. Verify every candidate's actual steps and parameter count; a partial/timed-out screen cannot establish the missing comparison. All model counts include encoder and selector.
10
+ 2. Read each development result and selected test result. The test is newly held out from project training and listed prior evaluations; upstream encoder/teacher exposure and near duplicates cannot be excluded. Use paired intervals as exploratory diagnostics, not multiple-comparison-corrected proof.
11
+ 3. View comparison JPEGs **before reading their key.json files**. Inspect both full frames and useful crops. Review all 32 fixed COCO examples for single/multi groups, the encoder group, and actual-aspect production_test images. Log visible findings per image: incoherent hue changes on one surface, missing colour, boundary leakage, implausible skin/sky/foliage, whole-image colour cast. An alternate plausible clothing colour is not a failure simply because it differs from the original.
12
+ 4. View the seven probes, historical comparison images, and actual-aspect production_archive images. Read archive_sources.json to identify failed downloads; do not report attempted images as evaluated. Supplement with the six National Archives inputs whose source records are in `supplement/nara_sources.json`. These are evaluation-only. Never supervise historical grayscale inputs with zero chroma.
13
+ 5. Inspect multi diagnostics: all mode frequencies, diversity, default-v-oracle losses and resize/flip mode switches. An oracle-only gain is insufficient for automatic deployment. An optional alternatives UI may be valuable but must not conceal a worse default. Inspect MobileViT visual outputs in light of its shorter colourisation history; do not declare a universally inferior backbone from a short warm-start test.
14
+ 6. Prefer a candidate only if images show a consistent benefit without unacceptable loss of colour, and quantitative diagnostics do not reveal collapse. Retain v3 as the fallback. If one method looks clearly better, a second-seed matched confirmation is authorized within the **remaining** pass budget; maximum additional L4 runtime18min, not an unbounded rerun. Previous attempt used about12min; replacement capped60min at$0.80/h. Preserve all results before exit.
15
+ 7. Write REVIEW.md and final_config.json with the actual checkpoint path, immutable commit, architecture, complete parameter count, recipe, seed, dataset revisions, teacher-cache revision and validation evidence. If nothing wins, state that plainly and use the established palette architecture. Do not promise that longer continuation necessarily improves it.
16
+
17
+ ## Specific implementation caveats to carry forward
18
+ - Single/multi arms: same initial encoder and decoder, index batches, learning rates and update counts. Staged augmentation changes random-number consumption late in training; its flip/target-choice draws are not identical to teacher/mix. This is a curriculum-bundle comparison.
19
+ - MobileViT receives1000 teacher alignment steps plus3000 curriculum steps. Version2 uses the post-warmup curriculum step for both augmentation/target ramps and consistency, correcting the canceled first attempt.
20
+ - Each arm's candidate can come from an earlier development checkpoint. optimizer.pt belongs to the **last** step, not necessarily that selected checkpoint. Start a new optimizer from selected candidate weights; never pair candidate weights with a mismatched optimizer.
21
+ - Existing r7 loss uses teacher/original targets for supervision; neither original colour error nor patch score alone measures plausibility.
22
+ - Do not add RL without an independently validated reward. Previous critic failures are project evidence for caution, not proof RL can never help.
23
+ - Keep original-image luminance and aspect ratio through the established inference pipeline. Single/multi DecisionColorizer needs its own loader; root load_semantic cannot silently load a new architecture.
24
+ - If deploying MobileViT, include its upstream Apple license/notice and appropriate pretrained-model attribution.
25
+
26
+ ## Inference compute experiment already inspected
27
+ Seven real historical photographs were evaluated with unchanged v3 weights at256/384/512 input size and the existing radius8 filter. On the six new National Archives images, larger size notably improves the children-in-a-field image (more coherent green foliage and blue sky), and slightly clarifies some people. It does not fix the near-monochrome casts on machinery, street and interior scenes. The familiar powerhouse pipe retains uneven colouring and changes hue with size. Therefore higher resolution is useful in some cases but is not a general cure. At fixed radius8, its physical smoothing footprint also changes: do not call this an isolated network-resolution ablation. Record measured CPU times as local indicative timings, not a deployment latency guarantee; they include a second diagnostic forward.
28
+
29
+ ## Final-pass scope
30
+ After evidence review, start one longer run on the selected course, with checkpoint retention, bounded runtime, early exit for non-finite losses, durable metrics and uploads. Final production promotion requires comparing its saved checkpoints visually to v3 and the selected pilot. Develop the Space/ZeroGPU app, requirements and model card, and clean the repository without removing experiments or the stable branch needed by the user. Never state that the model is mathematically optimal or that no improvement is possible under4M parameters.
experiments/decision-20260928-v2/SHA256SUMS.json ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "PROTOCOL.md": "7d73f476df74bf01b6ea93346fd08b98dde0456af62eda5b297b6a3044e33869",
3
+ "REVIEW_HANDOFF.md": "feae011a9a04136caf796224ae92e913eda5bc9e47c94a7dfea481e071bb58b2",
4
+ "source/data.py": "eb43b886fe45435841439c36646985f914386a48faa712fb6806accac2ad48e0",
5
+ "source/decision_model.py": "c7c4370d93bc52fdd363d6d4fcc684c317806d9b393bd969f9ee670a01487f27",
6
+ "source/inference.py": "0ac1f382205b42cef0afaa9173017fa660e6859194a65a61b9b9052c437cd476",
7
+ "source/legacy_eval.py": "dd7e8bd5f961dfb6f1c2ce98fac41824c5009188e306000b195e922e4ed03e58",
8
+ "source/metrics.py": "af0ac996f6be0a351a55e372fab637a0d06740afe4a744cb80d0eae56ec97f66",
9
+ "source/model.py": "4cc57f82ffd6378bdf23a088ce7a1ed56c09e5673de3ee5232b8ee1cacc9be0a",
10
+ "source/objectives.py": "133e85b6f2914c0e3803cc26a26fa29b779a49036219dd245c2ecfe0ba411e20",
11
+ "source/persistence.py": "8b26dd8acdc135ecbfcdd30b7f7723feb6008ac820c083403da1e820bdc15d33",
12
+ "source/r7.py": "019ea992052f0233954b86e06e31c85c74ca313e67a55ff7fb03248f8874aaf5",
13
+ "source/semantic_model.py": "b019358cbcaa214743cbd244d20eb6e59eeff0ab5264093ec0927be8529b25a2",
14
+ "source/spatial.py": "d13abe399ef049e21a6459a7003461afaae0c00e0c560262c6cf375de4c9884a",
15
+ "source/train_suite.py": "fe09629930b44a2ce4d630339b15f6cb3c09a4504de890df41016113163efb6a",
16
+ "status.json": "820e46a29d75ee397d8d2814ad0429cb3d300952f32030d83e06a4dc0b56f803",
17
+ "supplement/baseline_nara_resolution.json": "8920921e71ef95ed442495171a0ce229b5338170c5ccdf6d60e241fdd61cb570",
18
+ "supplement/baseline_resolution.json": "e770a0dff51460de711b05a2dcb107e128f8635f05284bc075816da8a44f51e2",
19
+ "supplement/cache_loader_check.json": "b7f282c99f9371b40f1a0b403bd31296ed1c93496d99ed60e144e088397f22f9",
20
+ "supplement/nara_sources.json": "d4551fb5857161ecd72be4f315dc9e643e33700827e7fcf51e90049a6ef0f6ba",
21
+ "supplement/resolution_review.py": "d942406942a5197b54db053ea696fc59ff8b5d0093d4e0b5817431db55e9d115"
22
+ }
experiments/decision-20260928-v2/source/data.py ADDED
@@ -0,0 +1,116 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Fixed color-bin semantics and reproducible dataset membership."""
2
+ import hashlib
3
+ import io
4
+ import json
5
+ from pathlib import Path
6
+ import numpy as np
7
+ from PIL import Image
8
+ import pyarrow.parquet as pq
9
+ from scipy.spatial import cKDTree
10
+ from skimage.color import rgb2lab
11
+ import torch
12
+ from torch.utils.data import Dataset
13
+
14
+ class ColorBins:
15
+ def __init__(self, centers, weights=None):
16
+ self.centers = np.asarray(centers, np.float32)
17
+ if self.centers.ndim != 2 or self.centers.shape[1] != 2:
18
+ raise ValueError("Expected Q by 2 color centers")
19
+ self.tree = cKDTree(self.centers)
20
+ self.weights = np.ones(len(self.centers), np.float32) if weights is None else np.asarray(weights,np.float32)
21
+
22
+ def encode(self, ab, k=5, sigma=5):
23
+ k = min(k, len(self.centers))
24
+ if sigma <= 0 or k < 1:
25
+ raise ValueError("Invalid encoding parameters")
26
+ dist, idx = self.tree.query(ab.reshape(-1,2), k=k)
27
+ dist, idx = dist.reshape(-1,k), idx.reshape(-1,k)
28
+ logw = -dist**2/(2*sigma**2)
29
+ logw -= logw.max(1, keepdims=True)
30
+ weight = np.exp(logw); weight /= weight.sum(1,keepdims=True)
31
+ shape = (*ab.shape[:-1], k)
32
+ return idx.reshape(shape).astype(np.int64), weight.reshape(shape).astype(np.float32)
33
+
34
+ def load_table(path):
35
+ table = pq.read_table(path)
36
+ return table
37
+
38
+ def validate_manifest(path, table, manifest):
39
+ with open(path,'rb') as f:
40
+ digest=hashlib.file_digest(f,'sha256').hexdigest()
41
+ if digest!=manifest['dataset_sha256'] or len(table)!=manifest['rows']:
42
+ raise ValueError('Dataset differs from pinned split manifest')
43
+ groups=[list(map(int,manifest[k])) for k in ['train','validation','test']]
44
+ if any(not g or len(g)!=len(set(g)) or min(g)<0 or max(g)>=len(table) for g in groups):
45
+ raise ValueError('Invalid or repeated indices in manifest')
46
+ if any(set(groups[i])&set(groups[j]) for i,j in [(0,1),(0,2),(1,2)]):
47
+ raise ValueError('Train/validation/test overlap')
48
+
49
+ def decode_image(table, index):
50
+ obj = table['image'][int(index)].as_py()
51
+ return Image.open(io.BytesIO(obj['bytes'])).convert('RGB')
52
+
53
+ def legacy_split(n, seed=0):
54
+ # Exactly reproduce datasets.Dataset.train_test_split(test_size=.05, seed=0).
55
+ order = np.random.default_rng(seed).permutation(n)
56
+ nval = int(np.ceil(.05*n))
57
+ return order[nval:], order[:nval]
58
+
59
+ def stratified_subset(indices, labels, per_class, seed):
60
+ rng = np.random.default_rng(seed)
61
+ selected = []
62
+ for label in sorted(set(labels)):
63
+ group = np.asarray([i for i in indices if labels[i] == label])
64
+ selected.extend(rng.choice(group, min(per_class,len(group)), replace=False).tolist())
65
+ return np.asarray(selected, dtype=np.int64)
66
+
67
+ class PhotoDataset(Dataset):
68
+ def __init__(self, table, indices, bins, size=256, augment=False):
69
+ self.table, self.indices, self.bins = table, list(map(int,indices)), bins
70
+ self.size, self.augment = size, augment
71
+
72
+ def __len__(self):
73
+ return len(self.indices)
74
+
75
+ def __getitem__(self, item):
76
+ # Keep square preprocessing for controlled comparison with historical runs.
77
+ # Production inference preserves aspect ratio separately.
78
+ image = decode_image(self.table,self.indices[item]).resize((self.size,self.size),Image.Resampling.BILINEAR)
79
+ arr = np.asarray(image,dtype=np.float32)/255
80
+ if self.augment and torch.rand(()) < .5:
81
+ arr = arr[:,::-1].copy()
82
+ lab = rgb2lab(arr).astype(np.float32)
83
+ idx, weight = self.bins.encode(lab[...,1:])
84
+ return (torch.from_numpy((lab[...,:1]/50-1).transpose(2,0,1).copy()),
85
+ torch.from_numpy(lab[...,1:].transpose(2,0,1).copy()),
86
+ torch.from_numpy(idx),torch.from_numpy(weight))
87
+
88
+ def estimate_weights(table, indices, bins, mix=.7, limit=1000, seed=0):
89
+ """Re-estimate the prior ON checkpoint bins; never replace the bin vocabulary."""
90
+ if not 0 <= mix <= 1:
91
+ raise ValueError("Rebalance mixture must be in [0,1]")
92
+ rng=np.random.default_rng(seed)
93
+ chosen=rng.choice(indices,min(limit,len(indices)),replace=False)
94
+ counts=np.ones(len(bins.centers),np.float64)*1e-3
95
+ for index in chosen:
96
+ image=decode_image(table,index).resize((32,32))
97
+ ab=rgb2lab(np.asarray(image,dtype=np.float32)/255)[...,1:]
98
+ # Low chroma is a heuristic, not proof that an image was originally B&W.
99
+ if np.linalg.norm(ab,axis=-1).mean()<3:
100
+ continue
101
+ idx,weight=bins.encode(ab)
102
+ np.add.at(counts,idx.ravel(),weight.ravel())
103
+ prior=counts/counts.sum()
104
+ weights=1/((1-mix)*prior+mix/len(prior))
105
+ weights/=np.sum(prior*weights)
106
+ return weights.astype(np.float32),prior.astype(np.float32)
107
+
108
+ def write_manifest(path, table, train, validation, test, dataset_sha256):
109
+ groups=[set(map(int,g)) for g in [train,validation,test]]
110
+ assert not groups[0]&groups[1] and not groups[0]&groups[2] and not groups[1]&groups[2]
111
+ content={'dataset':'johnowhitaker/imagenette2-320','dataset_sha256':dataset_sha256,
112
+ 'rows':len(table),'split_seed':0,'legacy_holdout_fraction':.05,
113
+ 'train':list(map(int,train)),'validation':list(map(int,validation)), 'test':list(map(int,test)),
114
+ 'caveat':'Matches the recent seed-0 holdout; older upstream training exposure is not established.'}
115
+ Path(path).write_text(json.dumps(content,indent=2))
116
+ return content
experiments/decision-20260928-v2/source/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()
experiments/decision-20260928-v2/source/inference.py ADDED
@@ -0,0 +1,111 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Aspect-preserving colourisation with a compact semantic model.
2
+
3
+ Input and output are PIL images. No teacher, critic or second learned model is
4
+ loaded. Original lightness and alpha are retained; out-of-gamut chroma is reduced.
5
+ """
6
+ import json
7
+ import math
8
+ from pathlib import Path
9
+ import numpy as np
10
+ import torch
11
+ import torch.nn.functional as F
12
+ from PIL import Image, ImageOps
13
+ from skimage.color import rgb2lab
14
+ from semantic_model import load_semantic
15
+
16
+ MAX_PIXELS = 12_000_000
17
+
18
+ def load_colorizer(path, device='cpu'):
19
+ model = load_semantic(path, device)
20
+ count = sum(p.numel() for p in model.parameters())
21
+ if count >= 4_000_000:
22
+ raise ValueError('Model exceeds the four-million-parameter limit')
23
+ return model
24
+
25
+ def box_mean(x, r):
26
+ return F.avg_pool2d(x, 2*r+1, 1, r, count_include_pad=False)
27
+
28
+ def prepare(image, size=256):
29
+ if not isinstance(image, Image.Image):
30
+ raise TypeError('Expected a PIL image')
31
+ if not 128 <= int(size) <= 512:
32
+ raise ValueError('Input size must be between 128 and 512')
33
+ image = ImageOps.exif_transpose(image)
34
+ if image.width * image.height > MAX_PIXELS:
35
+ raise ValueError('Please resize the image to at most 12 megapixels')
36
+ alpha = np.asarray(image.getchannel('A')).copy() if 'A' in image.getbands() else None
37
+ rgb = np.asarray(image.convert('RGB'), dtype=np.float32) / 255
38
+ light = rgb2lab(rgb)[..., 0].astype(np.float32)
39
+ scale = min(int(size)/max(image.size), 1.)
40
+ shape = (max(8, round(image.height*scale)), max(8, round(image.width*scale)))
41
+ x = torch.from_numpy(light)[None,None]/50-1
42
+ small = F.interpolate(x, size=shape, mode='bilinear', align_corners=False, antialias=True)
43
+ return light, alpha, small
44
+
45
+ @torch.inference_mode()
46
+ def chroma_coefficients(model, small, radius=8):
47
+ if radius not in (0,4,8,12,16):
48
+ raise ValueError('Unsupported smoothing radius')
49
+ device = next(model.parameters()).device
50
+ L = small.to(device)
51
+ ab = model(L).float()
52
+ if not torch.isfinite(ab).all():
53
+ raise RuntimeError('Model returned non-finite colours')
54
+ if radius == 0:
55
+ return torch.zeros_like(ab).cpu(), ab.cpu()
56
+ guide = (L.float()+1)/2
57
+ mi, mp = box_mean(guide,radius), box_mean(ab,radius)
58
+ var = (box_mean(guide*guide,radius)-mi*mi).clamp_min(0)
59
+ cov = box_mean(guide*ab,radius)-mi*mp
60
+ a = cov/(var+.001)
61
+ b = mp-a*mi
62
+ # Coefficients are upsampled, then evaluated against original-resolution L.
63
+ return box_mean(a,radius).cpu(), box_mean(b,radius).cpu()
64
+
65
+ def _linear_rgb(light, ab):
66
+ fy = (light+16)/116
67
+ f = np.stack([fy+ab[...,0]/500,fy,fy-ab[...,1]/200],axis=-1)
68
+ xyz = np.where(f>6/29,f**3,(f-4/29)*(3*(6/29)**2))
69
+ xyz *= np.array([.95047,1.,1.08883],np.float32)
70
+ matrix = np.array([[3.24048134,-1.53715152,-.49853633],[-.96925495,1.87599,.04155593],[.05564664,-.20404134,1.05731107]],np.float32)
71
+ return xyz @ matrix.T
72
+
73
+ def render(light, alpha, coefficients, saturation=1.):
74
+ saturation = float(saturation)
75
+ if not math.isfinite(saturation) or not 0 <= saturation <= 1.5:
76
+ raise ValueError('Colour strength must be between 0 and 1.5')
77
+ a,b = [F.interpolate(v.float(),size=light.shape,mode='bilinear',align_corners=False)[0].permute(1,2,0).numpy() for v in coefficients]
78
+ ab = (a*(light[...,None]/100)+b)*saturation
79
+ # Binary-search chroma compression retains Lab hue and lightness.
80
+ linear = _linear_rgb(light,ab)
81
+ invalid = ((linear < -1e-5)|(linear > 1+1e-5)).any(-1)
82
+ if invalid.any():
83
+ L = light[invalid]; colors=ab[invalid];lo=np.zeros(len(L),np.float32);hi=np.ones(len(L),np.float32)
84
+ for _ in range(9):
85
+ mid=(lo+hi)/2; candidate=_linear_rgb(L,colors*mid[:,None])
86
+ valid=((candidate>=-1e-5)&(candidate<=1+1e-5)).all(-1)
87
+ lo=np.where(valid,mid,lo);hi=np.where(valid,hi,mid)
88
+ ab[invalid]=colors*lo[:,None]
89
+ linear[invalid]=_linear_rgb(L,ab[invalid])
90
+ linear=np.clip(linear,0,1)
91
+ rgb=np.where(linear<=.0031308,12.92*linear,1.055*np.power(linear,1/2.4)-.055)
92
+ pixels=np.uint8(np.clip(np.rint(rgb*255),0,255))
93
+ if alpha is not None: pixels=np.concatenate([pixels,alpha[...,None]],axis=-1)
94
+ return Image.fromarray(pixels)
95
+
96
+ def colorize(model, image, size=256, radius=8, saturation=1.):
97
+ light,alpha,small=prepare(image,size)
98
+ return render(light,alpha,chroma_coefficients(model,small,radius),saturation)
99
+
100
+ def main():
101
+ import argparse
102
+ parser=argparse.ArgumentParser(description='Compact photo colouriser')
103
+ parser.add_argument('input');parser.add_argument('output')
104
+ parser.add_argument('--model',default='.');parser.add_argument('--device',default='cpu')
105
+ parser.add_argument('--size',type=int,default=256);parser.add_argument('--saturation',type=float,default=1.)
106
+ args=parser.parse_args()
107
+ model=load_colorizer(args.model,args.device)
108
+ with Image.open(args.input) as image:
109
+ colorize(model,image,args.size,saturation=args.saturation).save(args.output)
110
+
111
+ if __name__=='__main__':main()
experiments/decision-20260928-v2/source/legacy_eval.py ADDED
@@ -0,0 +1,150 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Round 4: matched multiscale structure and DDColor distillation experiments."""
2
+ import os,sys,json,time,random,hashlib,traceback,shutil
3
+ from pathlib import Path
4
+ import numpy as np
5
+ import cv2
6
+ from PIL import Image,ImageDraw
7
+ from skimage import data as sample_data
8
+ from skimage.color import rgb2lab,lab2rgb
9
+ import pyarrow as pa
10
+ import pyarrow.parquet as pq
11
+ import torch
12
+ import torch.nn.functional as F
13
+ from torch.utils.data import Dataset,DataLoader
14
+ from huggingface_hub import HfApi,hf_hub_download,PyTorchModelHubMixin
15
+ from model import load_model,save_model
16
+ from data import decode_image,ColorBins
17
+ from metrics import soft_ce,per_image_metrics
18
+ from spatial import guided_chroma
19
+
20
+ OUT=Path('round6');OUT.mkdir(exist_ok=True)
21
+ DURABLE=None
22
+ START=time.monotonic();DEADLINE=START+3.5*3600
23
+ DEVICE='cuda';SEED=409;STEPS=3000
24
+ BASE_REV='704fa80d792c3d759db91daa00b2dcfe6f0f6412'
25
+ BASE_HASH='0e4c417375684a044860f8af3ac3a2fb44e1a5729254ca33abda758ea71aea6e'
26
+ TEACHER_ID='piddnad/ddcolor_artistic'
27
+ torch.set_num_threads(6)
28
+
29
+ def deadline():
30
+ if time.monotonic()>DEADLINE:raise TimeoutError('Graceful export before remote hard timeout')
31
+
32
+ def gray_arrays(im,style='gray'):
33
+ im=im.resize((256,256),Image.Resampling.BILINEAR)
34
+ rgb=np.asarray(im,np.float32)/255
35
+ lab=rgb2lab(rgb).astype(np.float32)
36
+ if style=='film':
37
+ g=(rgb*np.array([.42,.45,.13],np.float32)).sum(-1)
38
+ g=np.clip((g**1.15-.5)*.8+.5,0,1)
39
+ else:g=np.asarray(im.convert('L'),np.float32)/255
40
+ # Feed the same neutral sRGB input to teacher and student.
41
+ gray=np.repeat(g[...,None],3,-1).astype(np.float32)
42
+ L=rgb2lab(gray)[...,0].astype(np.float32)
43
+ return (L[None]/50-1).copy(),lab[...,1:].transpose(2,0,1).copy(),gray
44
+
45
+ class Photos(Dataset):
46
+ def __init__(self,table,ids,bins=None,cache=None,augment=False,style='gray'):
47
+ self.table,self.ids,self.bins,self.cache,self.augment,self.style=table,ids,bins,cache,augment,style
48
+ def __len__(self):return len(self.ids)
49
+ def __getitem__(self,k):
50
+ L,ab,gray=gray_arrays(decode_image(self.table,self.ids[k]),self.style)
51
+ teach=np.array(self.cache[k],np.float32) if self.cache is not None else np.zeros((2,64,64),np.float32)
52
+ if self.augment:
53
+ if random.random()<.5:L=L[...,::-1].copy();ab=ab[...,::-1].copy();teach=teach[...,::-1].copy()
54
+ if random.random()<.3:L=np.clip(L*random.uniform(.85,1.15)+random.uniform(-.08,.08),-1,1)
55
+ if self.bins is not None:
56
+ idx,w=self.bins.encode(ab.transpose(1,2,0));return torch.from_numpy(L),torch.from_numpy(ab),torch.from_numpy(idx),torch.from_numpy(w),torch.from_numpy(teach)
57
+ return torch.from_numpy(L),torch.from_numpy(ab),torch.from_numpy(gray.transpose(2,0,1)),self.ids[k]
58
+
59
+ def patch_excess(p,t):
60
+ """Excess color differences across 4–64px regions, in target-flat areas."""
61
+ values=[]
62
+ for scale in [4,16]:
63
+ a=F.avg_pool2d(p,scale);b=F.avg_pool2d(t,scale)
64
+ for offset in [1,4]:
65
+ for dim in [-1,-2]:
66
+ x=a.narrow(dim,offset,a.shape[dim]-offset)-a.narrow(dim,0,a.shape[dim]-offset)
67
+ y=b.narrow(dim,offset,b.shape[dim]-offset)-b.narrow(dim,0,b.shape[dim]-offset)
68
+ dx=x.norm(dim=1);dy=y.norm(dim=1);mask=(dy<3).float()
69
+ values.append(((dx-dy-1).relu()*mask).sum((1,2))/mask.sum((1,2)).clamp_min(1))
70
+ return sum(values)/len(values)
71
+
72
+ def structural_loss(p,t):
73
+ terms=[]
74
+ for scale in [4,16]:
75
+ a=F.avg_pool2d(p,scale);b=F.avg_pool2d(t,scale)
76
+ for offset in [1,4]:
77
+ for dim in [-1,-2]:
78
+ x=a.narrow(dim,offset,a.shape[dim]-offset)-a.narrow(dim,0,a.shape[dim]-offset)
79
+ y=b.narrow(dim,offset,b.shape[dim]-offset)-b.narrow(dim,0,b.shape[dim]-offset)
80
+ terms.append(F.smooth_l1_loss(x/10,y/10,beta=.3))
81
+ return sum(terms)/len(terms)
82
+
83
+ def additional_losses(p,t,teacher):
84
+ small=F.avg_pool2d(p,4);target=F.avg_pool2d(t,4)
85
+ # Teacher is fallible; reduce guidance when it conflicts strongly with ground truth.
86
+ confidence=torch.exp(-(teacher-target).norm(dim=1,keepdim=True)/30).detach()
87
+ kd=(F.smooth_l1_loss(small/20,teacher/20,beta=.5,reduction='none')*confidence).mean()
88
+ neutral=(t.norm(dim=1)<3).float()
89
+ neutral_loss=(p.norm(dim=1)*neutral).sum()/neutral.sum().clamp_min(1)/20
90
+ return structural_loss(p,t),kd,neutral_loss
91
+
92
+ def measure(L,p,t,raw=None):
93
+ r=per_image_metrics(L,p,t);C=p.norm(dim=1);T=t.norm(dim=1)
94
+ for name,mask,event in [('missed_color',T>12,C<5),('neutral_spill',T<3,C>10)]:
95
+ r[name]=(mask&event).sum((1,2))/mask.sum((1,2)).clamp_min(1)
96
+ r['color_coverage']=(C>10).float().mean((1,2));r['patch_excess']=patch_excess(p,t)
97
+ r['raw_patch_excess']=patch_excess(p if raw is None else raw,t)
98
+ return r
99
+
100
+ @torch.inference_mode()
101
+ def teacher_predict(teacher,gray):
102
+ gray=F.interpolate(gray.to(DEVICE),size=(512,512),mode='bilinear',align_corners=False)
103
+ # Full precision is deliberate: spectral normalization / attention are not assumed BF16-safe.
104
+ out=teacher(gray).float()
105
+ assert out.shape[1]==2 and torch.isfinite(out).all()
106
+ return F.interpolate(out,size=(256,256),mode='bilinear',align_corners=False)
107
+
108
+ @torch.inference_mode()
109
+ def evaluate(model,table,ids,is_teacher=False):
110
+ result={};model.eval()
111
+ for style in ['gray','film']:
112
+ rows=[]
113
+ for L,t,gray,indices in DataLoader(Photos(table,ids,style=style),batch_size=2 if is_teacher else 8,num_workers=4):
114
+ deadline();L,t=L.to(DEVICE),t.to(DEVICE)
115
+ raw=teacher_predict(model,gray) if is_teacher else model.decode(model(L),.38)
116
+ pred=raw if is_teacher else guided_chroma(L,raw,8)
117
+ metrics=measure(L,pred,t,raw)
118
+ for j,index in enumerate(indices):rows.append({'index':int(index)}|{k:float(v[j]) for k,v in metrics.items()})
119
+ keep=[x for x in rows if x['target_chroma']>=5]
120
+ result[style]={'n_total':len(rows),'n_color':len(keep),'summary':{k:float(np.mean([x[k] for x in keep])) for k in keep[0] if k!='index'},'per_image':rows}
121
+ return result
122
+
123
+ def selection(result,baseline):
124
+ scores=[];eligible=True
125
+ for style in ['gray','film']:
126
+ s=result[style]['summary'];b=baseline[style]['summary']
127
+ eligible &= (s['ab_error']<=b['ab_error']*1.05 and s['neutral_spill']<=b['neutral_spill']+.01
128
+ and s['missed_color']<=b['missed_color']+.01 and s['color_coverage']>=.95*b['color_coverage']
129
+ and s['patch_excess']<.85*b['patch_excess'] and s['raw_patch_excess']<.9*b['raw_patch_excess'])
130
+ scores.append(s['patch_excess']/max(b['patch_excess'],1e-6)+.3*s['neutral_spill']+.2*s['ab_error']/b['ab_error'])
131
+ return float(np.mean(scores)),bool(eligible)
132
+
133
+ @torch.inference_mode()
134
+ def probes(model,tag,teacher=False):
135
+ names=['astronaut','coffee','chelsea','rocket','camera','coins','moon'];images=[]
136
+ for name in names:
137
+ original=Image.fromarray(getattr(sample_data,name)()).convert('RGB')
138
+ L,_,gray=gray_arrays(original);x=torch.from_numpy(L)[None].to(DEVICE)
139
+ ab=teacher_predict(model,torch.from_numpy(gray.transpose(2,0,1))[None]) if teacher else guided_chroma(x,model.decode(model(x),.38),8)
140
+ lab=np.concatenate([(L[0]*50+50)[...,None],ab[0].cpu().numpy().transpose(1,2,0)],-1)
141
+ rgb=np.uint8(np.clip(lab2rgb(lab),0,1)*255)
142
+ images.append((name,Image.fromarray(np.uint8(gray*255)),Image.fromarray(rgb)))
143
+ canvas=Image.new('RGB',(512,len(names)*280),'white');d=ImageDraw.Draw(canvas)
144
+ for i,(name,gray,col) in enumerate(images):
145
+ canvas.paste(gray,(0,i*280+24));canvas.paste(col,(256,i*280+24));d.text((4,i*280+4),name+' input | '+tag,fill='black')
146
+ canvas.save(OUT/(tag+'.jpg'),quality=90)
147
+
148
+ def write(name,obj):
149
+ (OUT/name).write_text(json.dumps(obj,indent=2));return obj
150
+
experiments/decision-20260928-v2/source/metrics.py ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Metrics are fidelity/artifact proxies; none certifies plausible color by itself."""
2
+ import numpy as np
3
+ import torch
4
+
5
+ def soft_ce(logits, idx, weights, class_weights=None):
6
+ # FP32 reduction under autocast; no full target distribution materialized.
7
+ logits=logits.float()
8
+ idx=idx.permute(0,3,1,2)
9
+ weights=weights.permute(0,3,1,2)
10
+ nll=-((logits.gather(1,idx)-logits.logsumexp(1,keepdim=True))*weights).sum(1)
11
+ if class_weights is not None:
12
+ nll=nll*class_weights[idx[:,0]]
13
+ return nll.mean()
14
+
15
+ def per_image_metrics(L, pred, target):
16
+ error=(pred-target).square().sum(1).sqrt().mean((1,2))
17
+ chroma=pred.square().sum(1).sqrt().mean((1,2))
18
+ true_chroma=target.square().sum(1).sqrt().mean((1,2))
19
+ seams=[]
20
+ for dim in [-1,-2]:
21
+ dp=pred.diff(dim=dim).square().sum(1).sqrt()
22
+ dt=target.diff(dim=dim).square().sum(1).sqrt()
23
+ dl=L.diff(dim=dim).abs()[:,0]*50
24
+ # Penalize excess chroma discontinuity where luminance is flat.
25
+ mask=(dl<2).float()
26
+ seams.append(((dp-dt).relu()*mask).sum((1,2))/mask.sum((1,2)).clamp_min(1))
27
+ return {'ab_error':error,'chroma':chroma,'target_chroma':true_chroma,
28
+ 'excess_chroma_edge':sum(seams)/2}
29
+
30
+ def summarize(rows):
31
+ keys=['ce','ab_error','chroma','target_chroma','excess_chroma_edge']
32
+ result={k:float(np.mean([r[k] for r in rows])) for k in keys}
33
+ result['n']=len(rows)
34
+ result['chroma_ratio']=result['chroma']/max(result['target_chroma'],1e-9)
35
+ return result
experiments/decision-20260928-v2/source/model.py ADDED
@@ -0,0 +1,122 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Shared, checkpoint-compatible Mini U-Net. Parameters remain under 4M."""
2
+ import json
3
+ import math
4
+ from pathlib import Path
5
+ import numpy as np
6
+ import torch
7
+ from torch import nn
8
+ import torch.nn.functional as F
9
+ from huggingface_hub import PyTorchModelHubMixin, snapshot_download
10
+ from safetensors.torch import load_file, save_file
11
+ def double_conv(in_ch, out_ch):
12
+ return nn.Sequential(
13
+ nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False),
14
+ nn.BatchNorm2d(out_ch),
15
+ nn.ReLU(inplace=True),
16
+ nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False),
17
+ nn.BatchNorm2d(out_ch),
18
+ nn.ReLU(inplace=True),
19
+ )
20
+
21
+
22
+ class DilatedContextBlock(nn.Module):
23
+ def __init__(self, channels, mid_ch=96, dilations=(2, 4, 8)):
24
+ super().__init__()
25
+ self.proj_in = nn.Sequential(
26
+ nn.Conv2d(channels, mid_ch, 1, bias=False),
27
+ nn.BatchNorm2d(mid_ch),
28
+ nn.ReLU(inplace=True),
29
+ )
30
+ layers = []
31
+ for d in dilations:
32
+ layers += [
33
+ nn.Conv2d(mid_ch, mid_ch, 3, padding=d, dilation=d, bias=False),
34
+ nn.BatchNorm2d(mid_ch),
35
+ nn.ReLU(inplace=True),
36
+ ]
37
+ self.dilated = nn.Sequential(*layers)
38
+ self.proj_out = nn.Sequential(
39
+ nn.Conv2d(mid_ch, channels, 1, bias=False),
40
+ nn.BatchNorm2d(channels),
41
+ )
42
+ nn.init.zeros_(self.proj_out[-1].weight)
43
+ self.relu = nn.ReLU(inplace=True)
44
+
45
+ def forward(self, x):
46
+ y = self.proj_in(x)
47
+ y = self.dilated(y)
48
+ y = self.proj_out(y)
49
+ return self.relu(x + y)
50
+
51
+
52
+ class SmallUNetColorizer(
53
+ nn.Module,
54
+ PyTorchModelHubMixin,
55
+ pipeline_tag="image-to-image",
56
+ license="apache-2.0",
57
+ tags=["colorization", "unet", "image-to-image", "classification"],
58
+ ):
59
+ def __init__(self, bin_centers, in_ch: int = 1, base: int = 44,
60
+ context_mid_ch: int = 96, context_dilations=(2, 4, 8)):
61
+ super().__init__()
62
+ self.in_ch, self.base = in_ch, base
63
+ num_bins = len(bin_centers)
64
+ self.num_bins = num_bins
65
+ self.register_buffer("bin_centers", torch.tensor(bin_centers, dtype=torch.float32))
66
+
67
+ self.enc1 = double_conv(in_ch, base)
68
+ self.enc2 = double_conv(base, base * 2)
69
+ self.enc3 = double_conv(base * 2, base * 4)
70
+ self.enc4 = double_conv(base * 4, base * 8)
71
+ self.pool = nn.MaxPool2d(2)
72
+ self.context = DilatedContextBlock(base * 8, mid_ch=context_mid_ch,
73
+ dilations=tuple(context_dilations))
74
+ self.up3 = nn.ConvTranspose2d(base * 8, base * 4, 2, stride=2)
75
+ self.dec3 = double_conv(base * 8, base * 4)
76
+ self.up2 = nn.ConvTranspose2d(base * 4, base * 2, 2, stride=2)
77
+ self.dec2 = double_conv(base * 4, base * 2)
78
+ self.up1 = nn.ConvTranspose2d(base * 2, base, 2, stride=2)
79
+ self.dec1 = double_conv(base * 2, base)
80
+ self.out_conv = nn.Conv2d(base, num_bins, 1)
81
+
82
+ def forward(self, x):
83
+ h, w = x.shape[-2:]
84
+ x = F.pad(x, (0, (-w) % 8, 0, (-h) % 8), mode="replicate")
85
+ e1 = self.enc1(x)
86
+ e2 = self.enc2(self.pool(e1))
87
+ e3 = self.enc3(self.pool(e2))
88
+ e4 = self.context(self.enc4(self.pool(e3)))
89
+ d3 = self.dec3(torch.cat([self.up3(e4), e3], dim=1))
90
+ d2 = self.dec2(torch.cat([self.up2(d3), e2], dim=1))
91
+ d1 = self.dec1(torch.cat([self.up1(d2), e1], dim=1))
92
+ return self.out_conv(d1)[..., :h, :w]
93
+
94
+ def decode(self, logits, temperature: float = 0.38):
95
+ if not math.isfinite(temperature) or temperature <= 0:
96
+ raise ValueError("temperature must be finite and positive")
97
+ probs_t = F.softmax(logits.float() / temperature, dim=1)
98
+ return torch.einsum("bqhw,qc->bchw", probs_t, self.bin_centers)
99
+
100
+
101
+ def load_model(source, revision=None, device="cpu"):
102
+ path = Path(source)
103
+ if not path.is_dir():
104
+ path = Path(snapshot_download(source, revision=revision,
105
+ allow_patterns=["config.json", "model.safetensors"]))
106
+ config = json.loads((path / "config.json").read_text())
107
+ state = load_file(str(path / "model.safetensors"))
108
+ centers = torch.tensor(config["bin_centers"], dtype=torch.float32)
109
+ if not torch.equal(centers, state["bin_centers"]):
110
+ raise ValueError("Checkpoint config and state color bins differ; refusing ambiguous decode")
111
+ model = SmallUNetColorizer(**config)
112
+ model.load_state_dict(state, strict=True)
113
+ model.to(device).eval()
114
+ return model
115
+
116
+ def save_model(model, path):
117
+ path = Path(path); path.mkdir(parents=True, exist_ok=True)
118
+ model.save_pretrained(path)
119
+ # Mixin config can retain constructor bins; use the actual authoritative buffer.
120
+ cfg = json.loads((path / "config.json").read_text())
121
+ cfg["bin_centers"] = model.bin_centers.detach().cpu().tolist()
122
+ (path / "config.json").write_text(json.dumps(cfg, indent=2))
experiments/decision-20260928-v2/source/objectives.py ADDED
@@ -0,0 +1,17 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+ import torch
3
+ import torch.nn.functional as F
4
+
5
+ def color_loss(pred,target):
6
+ """Fit one coherent target; distribution matching discourages gray averages."""
7
+ pred=F.adaptive_avg_pool2d(pred,(64,64)).float();target=target.float()
8
+ pixel=F.l1_loss(pred/20,target/20)
9
+ p=F.avg_pool2d(pred,2).flatten(2);t=F.avg_pool2d(target,2).flatten(2)
10
+ angle=torch.arange(8,device=p.device,dtype=p.dtype)*math.pi/8
11
+ directions=torch.stack([angle.cos(),angle.sin()],1)
12
+ a=torch.einsum('kc,bcn->bkn',directions,p).sort(-1).values
13
+ b=torch.einsum('kc,bcn->bkn',directions,t).sort(-1).values
14
+ distribution=F.l1_loss(a/20,b/20)
15
+ # Match boundaries in the teacher target, never minimize gradients toward zero.
16
+ edge=(F.l1_loss((pred[:,:,:,1:]-pred[:,:,:,:-1])/20,(target[:,:,:,1:]-target[:,:,:,:-1])/20)+F.l1_loss((pred[:,:,1:]-pred[:,:,:-1])/20,(target[:,:,1:]-target[:,:,:-1])/20))/2
17
+ return pixel+.25*distribution+.1*edge,{'pixel':pixel,'distribution':distribution,'edge':edge}
experiments/decision-20260928-v2/source/persistence.py ADDED
@@ -0,0 +1,42 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Commit run artifacts atomically and verify every changed file by SHA256."""
2
+ import hashlib, json, os, time
3
+ from pathlib import Path
4
+ from huggingface_hub import HfApi, CommitOperationAdd, hf_hub_download
5
+
6
+ REPO = 'User-2468/mini-unet-colorizer'
7
+ PREFIX = 'experiments/round6-20260928'
8
+
9
+ class DurableRun:
10
+ def __init__(self, root):
11
+ if not os.environ.get('HF_TOKEN'):
12
+ raise RuntimeError('A write token is required before starting')
13
+ self.api = HfApi(token=os.environ['HF_TOKEN'])
14
+ if self.api.whoami()['name'] != 'User-2468':
15
+ raise RuntimeError('Unexpected account')
16
+ self.root = Path(root)
17
+ self.saved = {}
18
+
19
+ def sync(self, reason):
20
+ files = [p for p in sorted(self.root.rglob('*')) if p.is_file()]
21
+ hashes = {p.relative_to(self.root).as_posix(): hashlib.sha256(p.read_bytes()).hexdigest() for p in files}
22
+ changed = [p for p in files if self.saved.get(p.relative_to(self.root).as_posix()) != hashes[p.relative_to(self.root).as_posix()]]
23
+ if not changed:
24
+ return
25
+ for attempt in range(3):
26
+ try:
27
+ head = self.api.model_info(REPO, revision='main').sha
28
+ operations = [CommitOperationAdd(path_in_repo=PREFIX+'/'+p.relative_to(self.root).as_posix(), path_or_fileobj=str(p)) for p in changed]
29
+ operations.append(CommitOperationAdd(path_in_repo=PREFIX+'/SHA256SUMS.json', path_or_fileobj=json.dumps(hashes,sort_keys=True,indent=2).encode()))
30
+ commit = self.api.create_commit(repo_id=REPO, revision='main', parent_commit=head, operations=operations, commit_message='Colorizer round6: '+reason)
31
+ for p in changed:
32
+ name = p.relative_to(self.root).as_posix()
33
+ downloaded = hf_hub_download(REPO, PREFIX+'/'+name, revision=commit.oid, force_download=True, token=os.environ['HF_TOKEN'])
34
+ if hashlib.sha256(Path(downloaded).read_bytes()).hexdigest() != hashes[name]:
35
+ raise RuntimeError('Remote checksum mismatch: '+name)
36
+ self.saved = hashes
37
+ print('PERSISTED',reason,commit.oid,'verified_files',len(changed),flush=True)
38
+ return
39
+ except Exception:
40
+ if attempt == 2:
41
+ raise
42
+ time.sleep(2**attempt)
experiments/decision-20260928-v2/source/r7.py ADDED
@@ -0,0 +1,234 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os,sys,json,random,time,hashlib,shutil,math,traceback,io
2
+ from pathlib import Path
3
+ import numpy as np
4
+ import torch
5
+ import torch.nn.functional as F
6
+ from torch.utils.data import Dataset,DataLoader
7
+ from huggingface_hub import hf_hub_download
8
+ import pyarrow as pa
9
+ import pyarrow.parquet as pq
10
+ from PIL import Image,ImageDraw
11
+ from skimage.color import rgb2lab,lab2rgb
12
+ from skimage.segmentation import slic
13
+ import cv2
14
+ import persistence
15
+ from semantic_model import load_semantic,save_semantic
16
+ from objectives import color_loss
17
+ from data import decode_image
18
+ import legacy_eval as ev
19
+ REPO='User-2468/mini-unet-colorizer'
20
+ REV='7b10485ee9d205ce2f7b660dfc6a90012d78b1cb'
21
+ OLD='experiments/round6-20260928'
22
+ OUT=Path('/work/results');OUT.mkdir(parents=True,exist_ok=True)
23
+ persistence.PREFIX='experiments/round7-20260928'
24
+ ev.OUT=OUT
25
+ START=time.monotonic()
26
+ SEED=709
27
+ STEPS=2000
28
+ DEVICE='cpu' if os.environ.get('SMOKE')=='1' else 'cuda'
29
+ ev.DEVICE=DEVICE
30
+ DURABLE=None
31
+ def write(name,obj):
32
+ p=OUT/name;p.parent.mkdir(parents=True,exist_ok=True);p.write_text(json.dumps(obj,indent=2));return obj
33
+ def deadline():
34
+ if time.monotonic()-START>9000:raise TimeoutError('Graceful deadline')
35
+ ev.deadline=deadline
36
+ def losses(p,t):
37
+ p=F.adaptive_avg_pool2d(p.float(),(64,64));t=t.float()
38
+ pixel=(p-t).abs().mean((1,2,3))/20
39
+ a=F.avg_pool2d(p,2).flatten(2);b=F.avg_pool2d(t,2).flatten(2)
40
+ angles=torch.arange(8,device=p.device)*math.pi/8
41
+ directions=torch.stack([angles.cos(),angles.sin()],1)
42
+ ap=torch.einsum('kc,bcn->bkn',directions,a).sort(-1).values
43
+ bp=torch.einsum('kc,bcn->bkn',directions,b).sort(-1).values
44
+ distribution=(ap-bp).abs().mean((1,2))/20
45
+ edge=((torch.diff(p,dim=-1)-torch.diff(t,dim=-1)).abs().mean((1,2,3))+(torch.diff(p,dim=-2)-torch.diff(t,dim=-2)).abs().mean((1,2,3)))/40
46
+ return pixel+.25*distribution+.1*edge
47
+ def select_loss(p,teacher,original,mode):
48
+ a=losses(p,teacher)
49
+ if mode=='teacher':return a.mean(),torch.zeros_like(a,dtype=torch.bool)
50
+ b=losses(p,original);choose=b<a
51
+ return torch.where(choose,b,a).mean(),choose
52
+ def smoke():
53
+ torch.set_num_threads(4);torch.manual_seed(SEED)
54
+ p=torch.randn(2,2,64,64,requires_grad=True);t=torch.randn_like(p)
55
+ assert torch.allclose(losses(p,t).mean(),color_loss(p,t)[0],atol=1e-6)
56
+ # Whole-image choice: exact red or blue is better than an average or a patchwork.
57
+ red=torch.zeros(1,2,64,64);red[:,0]=30
58
+ blue=-red;mix=red.clone();mix[:,:,:,32:]=blue[:,:,:,32:]
59
+ assert select_loss(red,red,blue,'set')[0]==0
60
+ assert select_loss(torch.zeros_like(red),red,blue,'set')[0]>0
61
+ assert select_loss(mix,red,blue,'set')[0]>0
62
+ m=load_semantic('/work/initial',DEVICE)
63
+ assert sum(p.numel() for p in m.parameters())==3994676
64
+ x=torch.rand(2,1,96,128,device=DEVICE)*2-1
65
+ pred=m(x);loss,_=select_loss(pred,t.to(DEVICE),-t.to(DEVICE),'set')
66
+ loss.backward()
67
+ assert torch.isfinite(loss) and all(torch.isfinite(p.grad).all() for p in m.parameters() if p.grad is not None)
68
+ save_semantic(m,OUT/'initial')
69
+ re=load_semantic(OUT/'initial',DEVICE)
70
+ with torch.no_grad():assert torch.allclose(m.eval()(x),re(x),atol=1e-5)
71
+ write('smoke.json',{'passed':True,'parameters':3994676,'loss_equivalence':True,'whole_image_choice':True,'finite_backward':True,'reload':True})
72
+ del m,re
73
+ class Photos(Dataset):
74
+ def __init__(self,table,ids,cache,archive):
75
+ self.table,self.ids,self.cache,self.archive=table,ids,cache,archive;self.epoch=0
76
+ def __len__(self):return len(self.ids)
77
+ def __getitem__(self,k):
78
+ im=decode_image(self.table,self.ids[k]);L,ab,gray=ev.gray_arrays(im)
79
+ # Independent per-example RNG; augmentation differences never alter image ordering/flips.
80
+ rng=np.random.default_rng(np.random.SeedSequence([SEED,self.epoch,k]))
81
+ target=np.array(self.cache[k],np.float32)
82
+ if self.archive and rng.random()<.5:
83
+ rgb=np.asarray(im.resize((256,256),Image.Resampling.BILINEAR),np.float32)/255
84
+ weights=rng.dirichlet(np.array([3.,5.,2.]))
85
+ g=(rgb*weights).sum(-1)
86
+ g=np.clip((g**rng.uniform(.7,1.5)-.5)*rng.uniform(.6,1.1)+.5+rng.uniform(-.08,.08),0,1)
87
+ if rng.random()<.5:g=cv2.GaussianBlur(g,(5,5),float(rng.uniform(.3,1.2)))
88
+ if rng.random()<.6:g=np.clip(g+rng.normal(0,rng.uniform(.003,.025),g.shape),0,1)
89
+ L=(rgb2lab(np.repeat(g[...,None],3,-1).astype(np.float32))[...,0][None]/50-1).astype(np.float32)
90
+ else:
91
+ if rng.random()<.5:L=np.clip(L*rng.uniform(.7,1.15)+rng.uniform(-.12,.12),-1,1)
92
+ flip=np.random.default_rng(np.random.SeedSequence([SEED+1,self.epoch,k])).random()<.5
93
+ if flip:L=L[...,::-1].copy();ab=ab[...,::-1].copy();target=target[...,::-1].copy()
94
+ original=cv2.resize(ab.transpose(1,2,0),(64,64),interpolation=cv2.INTER_AREA).transpose(2,0,1).copy()
95
+ return torch.from_numpy(L),torch.from_numpy(target),torch.from_numpy(original)
96
+ @torch.inference_mode()
97
+ def predict(m,L):
98
+ x=torch.from_numpy(L)[None].to(DEVICE);raw=m(x)
99
+ guided=ev.guided_chroma(x,raw,8)
100
+ return raw,guided
101
+ def toimage(L,ab):
102
+ rgb=lab2rgb(np.concatenate([(L[0]*50+50)[...,None],ab],-1))
103
+ return Image.fromarray(np.uint8(np.clip(rgb,0,1)*255))
104
+ @torch.inference_mode()
105
+ def review(m,table,ids,name):
106
+ m.eval();rows=[]
107
+ # Reference-derived surface regions for diagnostics only, never inference.
108
+ for i in ids:
109
+ L,target,_=ev.gray_arrays(decode_image(table,i))
110
+ raw,pred=predict(m,L);x=torch.from_numpy(L)[None].to(DEVICE)
111
+ resize=F.interpolate(x,(224,224),mode='bilinear',align_corners=False)
112
+ resized=F.interpolate(m(resize),(256,256),mode='bilinear',align_corners=False)
113
+ contrast=m((x*.95+.02).clamp(-1,1))
114
+ pab=pred[0].cpu().numpy().transpose(1,2,0)
115
+ rgb=np.asarray(decode_image(table,i).resize((256,256)),np.float32)/255
116
+ seg=slic(rgb,n_segments=100,compactness=10,start_label=0)
117
+ tab=target.transpose(1,2,0);variations=[]
118
+ for label in np.unique(seg):
119
+ mask=seg==label
120
+ if mask.sum()<64:continue
121
+ ground=tab[mask];center=np.median(ground,axis=0)
122
+ if np.percentile(np.linalg.norm(ground-center,axis=1),90)>5:continue
123
+ col=pab[mask];variations.append(float(np.mean(np.linalg.norm(col-np.median(col,axis=0),axis=1))))
124
+ rows.append({'id':int(i),'resize_ab_change':float((raw-resized).abs().mean()),'contrast_ab_change':float((raw-contrast).abs().mean()),'surface_chroma_variation':float(np.mean(variations)) if variations else None,'chroma_coverage':float((pred.norm(dim=1)>10).float().mean()),'mean_chroma':float(pred.norm(dim=1).mean())})
125
+ write(name+'_diagnostics.json',{'rows':rows,'summary':{k:float(np.mean([r[k] for r in rows if r[k] is not None])) for k in rows[0] if k!='id'},'warning':'Surface uniformity alone rewards grey collapse; never use alone to select model. Color differences from original are diagnostic only.'})
126
+ return rows
127
+ @torch.inference_mode()
128
+ def grid(models,table,ids,label):
129
+ order=np.random.default_rng(SEED+2).permutation(len(models))
130
+ names=list(models);key={chr(65+j):names[int(i)] for j,i in enumerate(order)}
131
+ write(label+'_key.json',key)
132
+ for page in range(0,min(len(ids),32),8):
133
+ chosen=ids[page:page+8];canvas=Image.new('RGB',(256*(len(models)+1),274*len(chosen)),'white');draw=ImageDraw.Draw(canvas)
134
+ for row,index in enumerate(chosen):
135
+ L,_,gray=ev.gray_arrays(decode_image(table,index));draw.text((2,row*274),f'{index}: gray | '+' | '.join(key),fill='black')
136
+ canvas.paste(Image.fromarray(np.uint8(gray*255)),(0,row*274+18))
137
+ for j,i in enumerate(order):
138
+ _,p=predict(models[names[int(i)]],L)
139
+ canvas.paste(toimage(L,p[0].cpu().numpy().transpose(1,2,0)),((j+1)*256,row*274+18))
140
+ canvas.save(OUT/f'{label}_{page//8}.jpg',quality=92)
141
+ def main():
142
+ global DURABLE
143
+ torch.set_num_threads(6);torch.manual_seed(SEED);random.seed(SEED);np.random.seed(SEED)
144
+ DURABLE=persistence.DurableRun(OUT)
145
+ shutil.copy2('/work/runner.py',OUT/'runner.py');shutil.copy2('/work/PROTOCOL.md',OUT/'PROTOCOL.md')
146
+ shutil.copytree('/work/source',OUT/'source',ignore=shutil.ignore_patterns('__pycache__'),dirs_exist_ok=True)
147
+ smoke();DURABLE.sync('round7 initial checkpoint and smoke verified')
148
+ if DEVICE=='cpu':return
149
+ torch.backends.cudnn.benchmark=True
150
+ tables=[]
151
+ for repo,rev,files in [('johnowhitaker/imagenette2-320','771c1310a2487e8076ede6b7d6307244aa8400af',['default/train/0000.parquet']),('detection-datasets/coco','26ddc382fe75dfc2a0655b5977e296ea10efebce',['default/train/0000.parquet','default/train/0001.parquet'])]:
152
+ for f in files:tables.append(pq.read_table(hf_hub_download(repo,f,repo_type='dataset',revision=rev),columns=['image']))
153
+ table=pa.concat_tables(tables);manifest=json.loads(Path('/work/source/previous_manifest.json').read_text());ids=manifest['train']
154
+ cache=np.empty((len(ids),2,64,64),np.float16)
155
+ sums=json.loads(Path(hf_hub_download(REPO,OLD+'/SHA256SUMS.json',revision=REV)).read_text())
156
+ for start in range(0,len(ids),2048):
157
+ f=f'teacher_cache/{start:05d}.npz';path=Path(hf_hub_download(REPO,OLD+'/'+f,revision=REV))
158
+ assert hashlib.sha256(path.read_bytes()).hexdigest()==sums[f]
159
+ with np.load(path) as z:
160
+ assert z['indices'].tolist()==ids[start:start+len(z['indices'])]
161
+ cache[start:start+len(z['target'])]=z['target']
162
+ fresh=pq.read_table(hf_hub_download('detection-datasets/coco','default/val/0001.parquet',repo_type='dataset',revision='26ddc382fe75dfc2a0655b5977e296ea10efebce'),columns=['image'])
163
+ hashes={hashlib.sha256(table['image'][int(i)].as_py()['bytes']).digest() for i in ids}
164
+ freshids=[]
165
+ for i in np.random.default_rng(SEED).permutation(len(fresh)):
166
+ if hashlib.sha256(fresh['image'][int(i)].as_py()['bytes']).digest() not in hashes:freshids.append(int(i))
167
+ if len(freshids)==160:break
168
+ write('manifest.json',{'train':ids,'development':freshids[:80],'heldout':freshids[80:],'fresh_file':'default/val/0001.parquet','revision':'26ddc382fe75dfc2a0655b5977e296ea10efebce','prior_model_revision':REV,'upstream_pretraining_overlap_unknown':True,'near_duplicates_not_excluded':True})
169
+ DURABLE.sync('data cache verified and split fixed')
170
+ m=load_semantic('/work/initial',DEVICE)
171
+ write('initial_validation.json',ev.evaluate(m,table,manifest['validation']))
172
+ review(m,fresh,freshids[:80],'initial_development')
173
+ review(m,fresh,freshids[80:],'initial_heldout')
174
+ write('initial_heldout_reference_metrics.json',ev.evaluate(m,fresh,freshids[80:]))
175
+ del m
176
+ arms=[('teacher_clean','teacher',False),('set_clean','set',False),('teacher_archive','teacher',True),('set_archive','set',True)]
177
+ for name,mode,archive in arms:
178
+ torch.manual_seed(SEED);np.random.seed(SEED);random.seed(SEED)
179
+ m=load_semantic('/work/initial',DEVICE)
180
+ opt=torch.optim.AdamW([{'params':m.encoder.parameters(),'lr':5e-6},{'params':[p for n,p in m.named_parameters() if not n.startswith('encoder.')],'lr':5e-5}],weight_decay=.01)
181
+ sched=torch.optim.lr_scheduler.LambdaLR(opt,lambda step:min((step+1)/100,1)*(.2+.8*(1+math.cos(math.pi*min(step,STEPS)/STEPS))/2))
182
+ ds=Photos(table,ids,cache,archive)
183
+ loader=DataLoader(ds,batch_size=32,shuffle=True,num_workers=4,pin_memory=True,generator=torch.Generator().manual_seed(SEED))
184
+ it=iter(loader);history=[];choices=0;seen=0
185
+ for step in range(1,STEPS+1):
186
+ deadline();m.train()
187
+ for layer in m.encoder.modules():
188
+ if isinstance(layer,torch.nn.BatchNorm2d):layer.eval()
189
+ try:L,t,o=next(it)
190
+ except StopIteration:ds.epoch+=1;it=iter(loader);L,t,o=next(it)
191
+ L,t,o=L.to(DEVICE),t.to(DEVICE),o.to(DEVICE);opt.zero_grad(set_to_none=True)
192
+ with torch.autocast('cuda',dtype=torch.bfloat16):p=m(L)
193
+ loss,choice=select_loss(p,t,o,mode)
194
+ assert torch.isfinite(loss)
195
+ loss.backward();torch.nn.utils.clip_grad_norm_(m.parameters(),1,error_if_nonfinite=True);opt.step();sched.step()
196
+ choices+=int(choice.sum());seen+=len(L)
197
+ if step%100==0:
198
+ history.append({'step':step,'loss':float(loss),'original_choice_fraction':choices/seen})
199
+ print('STEP',name,step,round(float(loss),4),flush=True)
200
+ if step%500==0:
201
+ path=OUT/name/f'step_{step:05d}';save_semantic(m,path)
202
+ torch.save({'optimizer':opt.state_dict(),'scheduler':sched.state_dict(),'step':step,'epoch':ds.epoch,'torch_rng':torch.get_rng_state(),'cuda_rng':torch.cuda.get_rng_state_all(),'numpy_rng':np.random.get_state(),'python_rng':random.getstate()},path/'training_state.pt')
203
+ write(name+'/history.json',history);write('status.json',{'status':'training','arm':name,'step':step,'elapsed_seconds':time.monotonic()-START})
204
+ DURABLE.sync(name+' checkpoint '+str(step))
205
+ write(name+'_validation.json',ev.evaluate(m,table,manifest['validation']))
206
+ review(m,fresh,freshids[:80],name+'_development')
207
+ review(m,fresh,freshids[80:],name+'_heldout')
208
+ write(name+'_heldout_reference_metrics.json',ev.evaluate(m,fresh,freshids[80:]))
209
+ ev.probes(m,name)
210
+ DURABLE.sync(name+' complete diagnostics')
211
+ del m,opt,sched,it,loader;torch.cuda.empty_cache()
212
+ models={'round6':load_semantic('/work/initial',DEVICE)}
213
+ for name,_,_ in arms:models[name]=load_semantic(OUT/name/f'step_{STEPS:05d}',DEVICE)
214
+ grid(models,fresh,freshids[:80],'blinded_development')
215
+ grid(models,fresh,freshids[80:],'blinded_heldout')
216
+ # Familiar historical development probes; never claim these are a held-out archival benchmark.
217
+ import requests
218
+ sources={'migrant_mother':'https://cdn.loc.gov/service/pnp/ppmsca/50200/50236v.jpg','power_house_mechanic':'https://www.archives.gov/exhibits/picturing_the_century/images/port_hine_022_v36.jpg'}
219
+ for name,url in sources.items():
220
+ try:
221
+ response=requests.get(url,timeout=45);response.raise_for_status()
222
+ im=Image.open(io.BytesIO(response.content)).convert('RGB')
223
+ tb=pa.table({'image':[{'bytes':response.content,'path':None}]})
224
+ grid(models,tb,[0],name)
225
+ write(name+'_source.json',{'url':url,'sha256':hashlib.sha256(response.content).hexdigest(),'use':'previously seen archival development probe; original colors unknown'})
226
+ except Exception as e:write(name+'_error.json',{'type':type(e).__name__,'message':str(e)})
227
+ write('status.json',{'status':'completed','elapsed_seconds':time.monotonic()-START,'production_approved':False,'next':'Review blinded grids for coherent plausible color and gray/sepia collapse. This pilot does not test a stochastic conditioned decoder.'})
228
+ DURABLE.sync('round7 four-arm experiment completed')
229
+ if __name__=='__main__':
230
+ try:main()
231
+ except BaseException as e:
232
+ write('failure.json',{'type':type(e).__name__,'traceback':traceback.format_exc().replace(os.environ.get('HF_TOKEN','__EMPTY__'),'[REDACTED]')})
233
+ if DURABLE:DURABLE.sync('round7 failure report and latest saved checkpoints')
234
+ raise
experiments/decision-20260928-v2/source/semantic_model.py ADDED
@@ -0,0 +1,94 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Compact pretrained semantic colorizer; dense and shared-palette variants."""
2
+ import json
3
+ from pathlib import Path
4
+ import torch
5
+ from torch import nn
6
+ import torch.nn.functional as F
7
+ from torchvision.models import mobilenet_v3_large, MobileNet_V3_Large_Weights
8
+ from safetensors.torch import save_file, load_file
9
+
10
+ class QueryBlock(nn.Module):
11
+ def __init__(self,d=96):
12
+ super().__init__()
13
+ self.self_attn=nn.MultiheadAttention(d,4,batch_first=True,dropout=0)
14
+ self.cross_attn=nn.MultiheadAttention(d,4,batch_first=True,dropout=0)
15
+ self.norms=nn.ModuleList([nn.LayerNorm(d) for _ in range(3)])
16
+ self.ff=nn.Sequential(nn.Linear(d,2*d),nn.GELU(),nn.Linear(2*d,d))
17
+ def forward(self,q,memory):
18
+ x=self.norms[0](q);q=q+self.self_attn(x,x,x,need_weights=False)[0]
19
+ x=self.norms[1](q);q=q+self.cross_attn(x,memory,memory,need_weights=False)[0]
20
+ return q+self.ff(self.norms[2](q))
21
+
22
+ def refine(d):
23
+ return nn.Sequential(nn.Conv2d(d,d,3,padding=1,bias=False),nn.GroupNorm(8,d),nn.SiLU())
24
+
25
+ class SemanticColorizer(nn.Module):
26
+ def __init__(self,head='palette',pretrained=False,width=128,queries=16):
27
+ super().__init__()
28
+ if head not in ['palette','dense']:raise ValueError(head)
29
+ self.config={'architecture':'SemanticColorizer','head':head,'width':width,'queries':queries,'format_version':1}
30
+ self.encoder=mobilenet_v3_large(weights=MobileNet_V3_Large_Weights.IMAGENET1K_V2 if pretrained else None,progress=False).features
31
+ self.lateral=nn.ModuleList([nn.Conv2d(c,width,1) for c in [24,40,112,960]])
32
+ self.refine=nn.ModuleList([refine(width) for _ in range(3)])
33
+ self.register_buffer('rgb_mean',torch.tensor([.485,.456,.406]).view(1,3,1,1))
34
+ self.register_buffer('rgb_std',torch.tensor([.229,.224,.225]).view(1,3,1,1))
35
+ if head=='palette':
36
+ self.queries=nn.Parameter(torch.randn(queries,width)*.2)
37
+ self.query_blocks=nn.ModuleList([QueryBlock(width) for _ in range(2)])
38
+ self.memory_norm=nn.LayerNorm(width)
39
+ self.query_norm=nn.LayerNorm(width)
40
+ self.pixel=nn.Conv2d(width,width,1)
41
+ self.palette=nn.Sequential(nn.Linear(width,width),nn.GELU(),nn.Linear(width,2))
42
+ self.residual=nn.Conv2d(width,2,1)
43
+ nn.init.normal_(self.palette[-1].weight,std=.01);nn.init.zeros_(self.palette[-1].bias)
44
+ nn.init.zeros_(self.residual.weight);nn.init.zeros_(self.residual.bias)
45
+ else:
46
+ self.dense=nn.Sequential(refine(width),nn.Conv2d(width,2,1))
47
+ nn.init.normal_(self.dense[-1].weight,std=.01);nn.init.zeros_(self.dense[-1].bias)
48
+ count=sum(p.numel() for p in self.parameters())
49
+ if count>=4_000_000:raise ValueError(f'Parameter budget exceeded: {count}')
50
+
51
+ @staticmethod
52
+ def neutral_rgb(L):
53
+ light=(L.float()*50+50).clamp(0,100)
54
+ y=torch.where(light>8,((light+16)/116)**3,light/903.296296)
55
+ g=torch.where(y<=.0031308,12.92*y,1.055*y.clamp_min(1e-8).pow(1/2.4)-.055)
56
+ return g.expand(-1,3,-1,-1)
57
+
58
+ def forward(self,L):
59
+ h,w=L.shape[-2:]
60
+ x=F.pad(self.neutral_rgb(L),(0,(-w)%32,0,(-h)%32),mode='replicate')
61
+ x=(x-self.rgb_mean)/self.rgb_std
62
+ features=[]
63
+ for i,layer in enumerate(self.encoder):
64
+ x=layer(x)
65
+ if i in [3,6,12,16]:features.append(x)
66
+ projected=[layer(f) for layer,f in zip(self.lateral,features)]
67
+ x=projected[-1]
68
+ for i in range(2,-1,-1):
69
+ x=self.refine[2-i](F.interpolate(x,size=projected[i].shape[-2:],mode='bilinear',align_corners=False)+projected[i])
70
+ if self.config['head']=='palette':
71
+ memory=torch.cat([F.adaptive_avg_pool2d(f,(8,8)).flatten(2).transpose(1,2) for f in projected[1:]],1)
72
+ memory=self.memory_norm(memory)
73
+ q=self.queries[None].expand(L.shape[0],-1,-1)
74
+ for block in self.query_blocks:q=block(q,memory)
75
+ q=self.query_norm(q)
76
+ palette=80*torch.tanh(self.palette(q))
77
+ masks=torch.einsum('bqd,bdhw->bqhw',q,self.pixel(x))/(self.config['width']**.5)
78
+ weights=F.softmax(masks.float(),dim=1)
79
+ ab=torch.einsum('bqhw,bqc->bchw',weights,palette.float())+2*torch.tanh(self.residual(x).float())
80
+ else:ab=80*torch.tanh(self.dense(x).float())
81
+ return F.interpolate(ab,size=(x.shape[-2]*4,x.shape[-1]*4),mode='bilinear',align_corners=False)[...,:h,:w]
82
+ def decode(self,z,temperature=.38):return z.float()
83
+
84
+ def save_semantic(model,path):
85
+ path=Path(path);path.mkdir(parents=True,exist_ok=True)
86
+ save_file({k:v.detach().cpu().contiguous() for k,v in model.state_dict().items()},str(path/'model.safetensors'))
87
+ (path/'config.json').write_text(json.dumps(model.config,indent=2))
88
+ (path/'README.md').write_text('# Experimental semantic colorizer\n\nHead: '+model.config['head']+'. Parameters: '+str(sum(p.numel() for p in model.parameters()))+'.\n\nNot approved for production. Requires semantic_model.py; incompatible with the old U-Net loader. Input is Lab lightness normalized to [-1,1]; output is Lab ab. See the run protocol, provenance, selection and visual comparisons. Predictions are plausible colors, not recovered historical truth.\n')
89
+
90
+ def load_semantic(path,device='cpu'):
91
+ path=Path(path);cfg=json.loads((path/'config.json').read_text())
92
+ model=SemanticColorizer(**{k:cfg[k] for k in ['head','width','queries']})
93
+ model.load_state_dict(load_file(str(path/'model.safetensors')),strict=True)
94
+ return model.to(device).eval()
experiments/decision-20260928-v2/source/spatial.py ADDED
@@ -0,0 +1,34 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Controlled spatial decoding alternatives; no learned parameters."""
2
+ import math
3
+ import torch
4
+ import torch.nn.functional as F
5
+
6
+ def box_mean(x,radius):
7
+ # Border-normalized local windows, also work on images smaller than kernel.
8
+ return F.avg_pool2d(x,2*radius+1,stride=1,padding=radius,count_include_pad=False)
9
+
10
+ def guided_chroma(L,ab,radius=4,epsilon=.001):
11
+ """Scalar luminance-guided local linear filter (He et al., ECCV 2010).
12
+
13
+ L is the network's [-1,1] luminance. Epsilon is in [0,1] luminance units.
14
+ Uses luminance only; reference colors never enter inference.
15
+ """
16
+ if not isinstance(radius,int) or radius<0 or not math.isfinite(epsilon) or epsilon<=0:
17
+ raise ValueError('Invalid guided-filter radius or epsilon')
18
+ if radius==0:return ab
19
+ I=(L.float()+1)/2;p=ab.float()
20
+ mi=box_mean(I,radius);mp=box_mean(p,radius)
21
+ var=(box_mean(I*I,radius)-mi*mi).clamp_min(0)
22
+ cov=box_mean(I*p,radius)-mi*mp
23
+ a=cov/(var+epsilon);b=mp-a*mi
24
+ return box_mean(a,radius)*I+box_mean(b,radius)
25
+
26
+ def spatial_decode(model,logits,L,temperature=.38,pool=1,radius=0,epsilon=.001):
27
+ if pool<1 or not isinstance(pool,int):raise ValueError('pool must be a positive integer')
28
+ if pool>1:
29
+ # Average evidence before annealing, then upsample chroma, not RGB.
30
+ small=F.avg_pool2d(logits,pool,ceil_mode=True,count_include_pad=False)
31
+ ab=model.decode(small,temperature)
32
+ ab=F.interpolate(ab,size=L.shape[-2:],mode='bilinear',align_corners=False)
33
+ else:ab=model.decode(logits,temperature)
34
+ return guided_chroma(L,ab,radius,epsilon)
experiments/decision-20260928-v2/source/train_suite.py ADDED
@@ -0,0 +1,359 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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='26d731862cc0ea531e7198fa823cb09db25bd243'
22
+ CACHE_REV='7b10485ee9d205ce2f7b660dfc6a90012d78b1cb'
23
+ COCO_REV='26ddc382fe75dfc2a0655b5977e296ea10efebce'
24
+ PREFIX=os.environ.get('EXPERIMENT_PREFIX','experiments/decision-20260928')
25
+ SEED=int(os.environ.get('SEED','92893'))
26
+ STEPS=int(os.environ.get('STEPS','3000'));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=api.model_info('timm/mobilevitv2_075.cvnets_in1k').sha
184
+ write('encoder_provenance.json',{'repo':'timm/mobilevitv2_075.cvnets_in1k','revision':encoder_rev,'license':'Apple sample code license; include upstream LICENSE if deployed','timm':'1.0.19'})
185
+ constructors=[('single',1,'mobilenet'),('multi',4,'mobilenet'),('mobilevit',1,'mobilevit')]
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);m.eval();count=sum(p.numel() for p in m.parameters());assert count<4_000_000
189
+ counts[name]=count
190
+ x=torch.rand(2,1,97,113,device='cuda')*2-1
191
+ with torch.no_grad():
192
+ v=m(x,True);assert v['all'].shape==(2,k,2,97,113) and torch.isfinite(v['all']).all()
193
+ parity=float((v['output']-baseline(x)).abs().max()) if name=='single' else None
194
+ if parity is not None:assert parity<1e-3,parity
195
+ loss,_,_=objective(m(x,True),torch.randn(2,2,64,64,device='cuda'),1)
196
+ loss.backward();assert torch.isfinite(loss) and all(torch.isfinite(p.grad).all() for p in m.parameters() if p.grad is not None)
197
+ m.zero_grad(set_to_none=True);save_decision(m,OUT/'initial'/name)
198
+ re=load_decision(OUT/'initial'/name,'cuda')
199
+ with torch.no_grad():assert torch.allclose(m(x),re(x),atol=1e-4)
200
+ smokes[name]={'parameters':count,'copied_tensors':len(copied),'odd_shape':True,'finite_backward':True,'reload':True,'baseline_max_abs':parity}
201
+ del m,re
202
+ write('smoke.json',smokes);DURABLE.sync('all three architectures pass preflight')
203
+ torch.backends.cuda.matmul.allow_tf32=True;torch.backends.cudnn.allow_tf32=True;torch.backends.cudnn.benchmark=True
204
+ tables=[]
205
+ 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'])]:
206
+ for f in files:
207
+ print('PREP downloading',repo,f,flush=True)
208
+ tables.append(pq.read_table(hf_hub_download(repo,f,repo_type='dataset',revision=rev),columns=['image']))
209
+ print('PREP loaded',repo,f,len(tables[-1]),flush=True)
210
+ table=pa.concat_tables(tables)
211
+ full_manifest=json.loads(dl('experiments/region-coherence-20260928/source/previous_manifest.json').read_text())
212
+ full_ids=full_manifest['train']
213
+ # Same fixed training subset for every arm, independent of new dev/test.
214
+ ids=np.random.default_rng(92893).choice(full_ids,8192,replace=False).tolist();pos={i:j for j,i in enumerate(ids)}
215
+ teacher=np.empty((len(ids),2,64,64),np.float16)
216
+ sums=json.loads(dl('experiments/round6-20260928/SHA256SUMS.json',CACHE_REV).read_text())
217
+ for start in range(0,len(full_ids),2048):
218
+ f=f'teacher_cache/{start:05d}.npz';p=dl('experiments/round6-20260928/'+f,CACHE_REV)
219
+ assert hashlib.sha256(p.read_bytes()).hexdigest()==sums[f]
220
+ with np.load(p) as z:
221
+ # NpzFile indexing decompresses a member on every access. Materialize once.
222
+ cache_ids=z['indices'];cache_targets=z['target']
223
+ assert cache_ids.tolist()==full_ids[start:start+len(cache_ids)]
224
+ src=[j for j,i in enumerate(cache_ids) if int(i) in pos]
225
+ dst=[pos[int(cache_ids[j])] for j in src]
226
+ teacher[dst]=cache_targets[src]
227
+ print('PREP cache',start,'selected',len(src),flush=True)
228
+ print('PREP decoding training images',len(ids),flush=True)
229
+ rgb=np.stack([np.asarray(decode_image(table,i).resize((256,256),Image.Resampling.BILINEAR)) for i in ids])
230
+ 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)
231
+ batcher=Batches(rgb,teacher);original=torch.from_numpy(original).cuda();del rgb,teacher
232
+ print('PREP training tensors ready',flush=True)
233
+ fresh=pq.read_table(hf_hub_download('detection-datasets/coco','default/val/0001.parquet',repo_type='dataset',revision=COCO_REV),columns=['image'])
234
+ exclude=set()
235
+ for path,keys in [('round7-20260928',['development','heldout']),('region-coherence-20260928',['development_regression','heldout_regression','new_holdout']),('recurrent-refinement-20260928',['fresh_holdout','development_regression'])]:
236
+ manifest=json.loads(dl('experiments/'+path+'/manifest.json').read_text())
237
+ for key in keys:exclude.update(manifest[key])
238
+ hashes={hashlib.sha256(table['image'][i].as_py()['bytes']).digest() for i in full_ids}
239
+ freshids=[int(i) for i in np.random.default_rng(92893).permutation(len(fresh)) if int(i) not in exclude and hashlib.sha256(fresh['image'][int(i)].as_py()['bytes']).digest() not in hashes][:288]
240
+ assert len(freshids)==288
241
+ dev,test=freshids[:128],freshids[128:]
242
+ 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})
243
+ base_dev=r.ev.evaluate(baseline,fresh,dev);write('baseline_development.json',base_dev)
244
+ DURABLE.sync('dataset checks and unseen development/test split frozen')
245
+ arms=[(h,c) for h in ['single','multi'] for c in ['teacher','mix','staged']]+[('mobilevit','staged')]
246
+ if os.environ.get('ARMS'):arms=[tuple(x.split(':')) for x in os.environ['ARMS'].split(',')]
247
+ selected={};histories={};dev_results={};begin_train=time.monotonic()
248
+ for head,recipe in arms:
249
+ name=head+'_'+recipe;seed(SEED)
250
+ m=load_decision(OUT/'initial'/head,'cuda')
251
+ # New encoder gets a separate alignment warmup, explicitly outside matched six arms.
252
+ warmup=1000 if head=='mobilevit' else 0
253
+ groups=[{'params':[p for n,p in m.named_parameters() if n.startswith('encoder.')],'lr':2e-5 if head=='mobilevit' else 1e-5},
254
+ {'params':[p for n,p in m.named_parameters() if not n.startswith('encoder.')],'lr':2e-4 if head=='mobilevit' else 1e-4}]
255
+ opt=torch.optim.AdamW(groups,weight_decay=.01)
256
+ for g in opt.param_groups:g['base_lr']=g['lr']
257
+ history=[];best=1e9;beststep=0;last=0;arm_start=time.monotonic();mode_counts=np.zeros(m.modes,np.int64)
258
+ for s in range(1,STEPS+warmup+1):
259
+ if time.monotonic()-START>TRAIN_LIMIT:break
260
+ step=max(1,s-warmup);active_recipe='teacher' if s<=warmup else recipe
261
+ L,target,_,fraction=batcher.batch(s,active_recipe,original,curriculum_step=step)
262
+ m.train()
263
+ for layer in m.encoder.modules():
264
+ if isinstance(layer,nn.BatchNorm2d):layer.eval()
265
+ opt.zero_grad(set_to_none=True)
266
+ factor=min(s/150,1)*(.15+.85*(1+math.cos(math.pi*max(0,s-warmup)/STEPS))/2)
267
+ for g in opt.param_groups:g['lr']=g['base_lr']*factor
268
+ with torch.autocast('cuda',dtype=torch.bfloat16):
269
+ details=m(L,True)
270
+ loss,winners,terms=objective(details,target,step)
271
+ consistency=loss.new_zeros(())
272
+ if active_recipe=='staged' and step>STEPS*2/3 and s%4==0:
273
+ small=m(F.interpolate(L,(192,192),mode='bilinear',align_corners=False),True)
274
+ a=F.adaptive_avg_pool2d(small['all'].flatten(0,1).float(),(64,64))
275
+ b=F.adaptive_avg_pool2d(details['all'].flatten(0,1).detach().float(),(64,64))
276
+ consistency=F.l1_loss(a/20,b/20)
277
+ if m.modes>1:consistency+=.1*F.kl_div(small['scores'].log_softmax(1),details['scores'].detach().softmax(1),reduction='batchmean')
278
+ loss=loss+.4*consistency # every fourth step -> mean weight .1
279
+ assert torch.isfinite(loss),name
280
+ loss.backward();torch.nn.utils.clip_grad_norm_(m.parameters(),1,error_if_nonfinite=True);opt.step();last=s
281
+ mode_counts+=np.bincount(winners.detach().cpu().numpy(),minlength=m.modes)
282
+ if s%100==0:
283
+ 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}
284
+ history.append(row);print('STEP',name,s,round(row['loss'],4),flush=True)
285
+ if s%1000==0 or s==STEPS+warmup:
286
+ m.eval();res=r.ev.evaluate(m,fresh,dev);val,okay=score(res,base_dev)
287
+ write(name+f'/development_{s:05d}.json',res)
288
+ write(name+'/history.json',history)
289
+ save_decision(m,OUT/name/f'step_{s:05d}')
290
+ # Ineligible models still retained for diagnosis, ranked after eligible.
291
+ ranked=val+(0 if okay else 100)
292
+ if ranked<best:
293
+ best=ranked;beststep=s;save_decision(m,OUT/name/'candidate');dev_results[name]=res
294
+ write('status.json',{'status':'training','arm':name,'step':s,'elapsed_seconds':time.monotonic()-START,'production_approved':False})
295
+ DURABLE.sync(name+' checkpoint '+str(s))
296
+ if not beststep:raise RuntimeError('Budget exhausted before first candidate for '+name)
297
+ write(name+'/history.json',history)
298
+ 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')
299
+ selected[name]={'step':beststep,'rank_score':best,'parameters':counts[head],'steps_completed':last,'warmup_steps':warmup,'seconds':time.monotonic()-arm_start}
300
+ write('candidates.json',selected);DURABLE.sync(name+' selected checkpoint on development only')
301
+ del m,opt;torch.cuda.empty_cache()
302
+ if last<STEPS+warmup:break
303
+ # Test is opened only after all checkpoint choices have been frozen.
304
+ write('candidates.json',selected);DURABLE.sync('freeze candidates before opening final test')
305
+ models={'v3':baseline};test_results={}
306
+ base_test=r.ev.evaluate(baseline,fresh,test);write('v3_test.json',base_test);test_results['v3']=base_test
307
+ for name in selected:
308
+ m=load_decision(OUT/name/'candidate','cuda');models[name]=m
309
+ res=r.ev.evaluate(m,fresh,test);test_results[name]=res
310
+ write(name+'_test.json',res);write(name+'_test_paired.json',paired_bootstrap(res,base_test))
311
+ write(name+'_diagnostics.json',diagnostics(m,fresh,test[:48]))
312
+ r.ev.probes(m,name+'_probes')
313
+ x=torch.zeros(1,1,256,256,device='cuda')
314
+ with torch.inference_mode():
315
+ for _ in range(5):m(x)
316
+ torch.cuda.synchronize();t0=time.monotonic()
317
+ for _ in range(30):m(x)
318
+ torch.cuda.synchronize()
319
+ write(name+'_latency.json',{'milliseconds':1000*(time.monotonic()-t0)/30,'hardware':torch.cuda.get_device_name(),'size':256,'precision':'FP32'})
320
+ DURABLE.sync(name+' final test and diagnostic outputs')
321
+ r.ev.probes(baseline,'v3_probes')
322
+ # Smaller comparison groups preserve enough image detail to inspect.
323
+ for head in ['single','multi']:
324
+ group={n:m for n,m in models.items() if n=='v3' or n.startswith(head+'_')}
325
+ r.grid(group,fresh,test,head+'_test')
326
+ if 'mobilevit_staged' in models:r.grid({'v3':baseline,'single_staged':models['single_staged'],'mobilevit_staged':models['mobilevit_staged']},fresh,test,'encoder_test')
327
+ # Historical originals deliberately never serve as neutral-color supervision.
328
+ archive=json.loads(Path('/work/archive_sources.json').read_text())
329
+ import requests
330
+ records=[];images=[]
331
+ for entry in archive:
332
+ try:
333
+ response=requests.get(entry['image_url'],timeout=20);response.raise_for_status()
334
+ im=Image.open(io.BytesIO(response.content)).convert('RGB');im.load()
335
+ images.append({'bytes':response.content,'path':None});records.append(entry|{'sha256':hashlib.sha256(response.content).hexdigest()})
336
+ except Exception as e:records.append(entry|{'failed':type(e).__name__})
337
+ if images:
338
+ ar=pa.table({'image':images})
339
+ resolution_grid(baseline,ar,[x for x in records if 'failed' not in x])
340
+ for head in ['single','multi']:
341
+ r.grid({n:m for n,m in models.items() if n=='v3' or n.startswith(head+'_')},ar,list(range(len(images))),head+'_archive')
342
+ if 'mobilevit_staged' in models:r.grid({'v3':baseline,'mobilevit_staged':models['mobilevit_staged']},ar,list(range(len(images))),'encoder_archive')
343
+ pipeline_grid({'v3':baseline,**{n:m for n,m in models.items() if n.endswith('staged')}},ar,list(range(len(images))),'production_archive')
344
+ pipeline_grid({'v3':baseline,**{n:m for n,m in models.items() if n.endswith('staged')}},fresh,test[:16],'production_test')
345
+ write('archive_sources.json',records)
346
+ 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),'elapsed_seconds':time.monotonic()-START,'estimated_compute_usd':(time.monotonic()-START)/3600*.8,'warning':'No metric proves realism. All outputs need visual review. MobileViT received extra alignment warmup and different encoder LR; not a pure head ablation.'})
347
+ write('status.json',{'status':'completed','all_arms_completed':len(selected)==len(arms),'elapsed_seconds':time.monotonic()-START,'production_approved':False})
348
+ DURABLE.sync('architecture and curriculum suite completed')
349
+
350
+ if __name__=='__main__':
351
+ try:main()
352
+ except BaseException:
353
+ msg=traceback.format_exc();token=os.environ.get('HF_TOKEN')
354
+ if token:msg=msg.replace(token,'[REDACTED]')
355
+ write('failure.json',{'traceback':msg});write('status.json',{'status':'failed','elapsed_seconds':time.monotonic()-START,'production_approved':False})
356
+ if DURABLE:
357
+ try:DURABLE.sync('failure report and saved checkpoints')
358
+ except Exception:pass
359
+ print(msg,flush=True);raise SystemExit(1)
experiments/decision-20260928-v2/status.json ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ {
2
+ "status": "preflight",
3
+ "seed": 92893,
4
+ "production_approved": false
5
+ }
experiments/decision-20260928-v2/supplement/baseline_nara_resolution.json ADDED
@@ -0,0 +1,128 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [
2
+ {
3
+ "name": "nara_42.jpg",
4
+ "size": 256,
5
+ "seconds_cpu": 0.14755021000019042,
6
+ "mask_entropy": 2.0555808544158936,
7
+ "output_chroma": 12.437393188476562
8
+ },
9
+ {
10
+ "name": "nara_42.jpg",
11
+ "size": 384,
12
+ "seconds_cpu": 0.23218034499996065,
13
+ "mask_entropy": 2.090886116027832,
14
+ "output_chroma": 14.703699111938477
15
+ },
16
+ {
17
+ "name": "nara_42.jpg",
18
+ "size": 512,
19
+ "seconds_cpu": 0.3470577670000239,
20
+ "mask_entropy": 2.164964199066162,
21
+ "output_chroma": 13.910662651062012
22
+ },
23
+ {
24
+ "name": "nara_43.jpg",
25
+ "size": 256,
26
+ "seconds_cpu": 0.11568649200034997,
27
+ "mask_entropy": 2.221789598464966,
28
+ "output_chroma": 22.151323318481445
29
+ },
30
+ {
31
+ "name": "nara_43.jpg",
32
+ "size": 384,
33
+ "seconds_cpu": 0.17595398799994655,
34
+ "mask_entropy": 2.151303291320801,
35
+ "output_chroma": 23.234233856201172
36
+ },
37
+ {
38
+ "name": "nara_43.jpg",
39
+ "size": 512,
40
+ "seconds_cpu": 0.2790908069996476,
41
+ "mask_entropy": 2.1258065700531006,
42
+ "output_chroma": 20.545372009277344
43
+ },
44
+ {
45
+ "name": "nara_44.jpg",
46
+ "size": 256,
47
+ "seconds_cpu": 0.09606875300005413,
48
+ "mask_entropy": 1.8140740394592285,
49
+ "output_chroma": 7.230224609375
50
+ },
51
+ {
52
+ "name": "nara_44.jpg",
53
+ "size": 384,
54
+ "seconds_cpu": 0.1647310340003969,
55
+ "mask_entropy": 1.9336862564086914,
56
+ "output_chroma": 8.660061836242676
57
+ },
58
+ {
59
+ "name": "nara_44.jpg",
60
+ "size": 512,
61
+ "seconds_cpu": 0.2678544229997897,
62
+ "mask_entropy": 2.003618001937866,
63
+ "output_chroma": 10.750614166259766
64
+ },
65
+ {
66
+ "name": "nara_45.jpg",
67
+ "size": 256,
68
+ "seconds_cpu": 0.12318589300002714,
69
+ "mask_entropy": 1.9286333322525024,
70
+ "output_chroma": 10.617255210876465
71
+ },
72
+ {
73
+ "name": "nara_45.jpg",
74
+ "size": 384,
75
+ "seconds_cpu": 0.2194160090002697,
76
+ "mask_entropy": 1.786165475845337,
77
+ "output_chroma": 13.136454582214355
78
+ },
79
+ {
80
+ "name": "nara_45.jpg",
81
+ "size": 512,
82
+ "seconds_cpu": 0.3068678010004078,
83
+ "mask_entropy": 1.711348056793213,
84
+ "output_chroma": 15.462327003479004
85
+ },
86
+ {
87
+ "name": "nara_46.jpg",
88
+ "size": 256,
89
+ "seconds_cpu": 0.14160436099973595,
90
+ "mask_entropy": 1.984742522239685,
91
+ "output_chroma": 6.706954002380371
92
+ },
93
+ {
94
+ "name": "nara_46.jpg",
95
+ "size": 384,
96
+ "seconds_cpu": 0.20283023299998604,
97
+ "mask_entropy": 2.0790905952453613,
98
+ "output_chroma": 8.468958854675293
99
+ },
100
+ {
101
+ "name": "nara_46.jpg",
102
+ "size": 512,
103
+ "seconds_cpu": 0.3308731010001793,
104
+ "mask_entropy": 2.1869609355926514,
105
+ "output_chroma": 9.939047813415527
106
+ },
107
+ {
108
+ "name": "nara_49.jpg",
109
+ "size": 256,
110
+ "seconds_cpu": 0.14086192399963693,
111
+ "mask_entropy": 2.110356330871582,
112
+ "output_chroma": 15.332934379577637
113
+ },
114
+ {
115
+ "name": "nara_49.jpg",
116
+ "size": 384,
117
+ "seconds_cpu": 0.2311230479999722,
118
+ "mask_entropy": 2.158229112625122,
119
+ "output_chroma": 19.6121768951416
120
+ },
121
+ {
122
+ "name": "nara_49.jpg",
123
+ "size": 512,
124
+ "seconds_cpu": 0.2948662769999828,
125
+ "mask_entropy": 2.1546900272369385,
126
+ "output_chroma": 16.430278778076172
127
+ }
128
+ ]
experiments/decision-20260928-v2/supplement/baseline_resolution.json ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [
2
+ {
3
+ "image": "07.jpg",
4
+ "size": 256,
5
+ "mask_entropy": 1.855709433555603,
6
+ "palette_chroma": 30.90544319152832,
7
+ "output_chroma": 19.500049591064453
8
+ },
9
+ {
10
+ "image": "07.jpg",
11
+ "size": 384,
12
+ "mask_entropy": 1.8696938753128052,
13
+ "palette_chroma": 20.328649520874023,
14
+ "output_chroma": 11.265676498413086
15
+ },
16
+ {
17
+ "image": "07.jpg",
18
+ "size": 512,
19
+ "mask_entropy": 1.9025498628616333,
20
+ "palette_chroma": 20.42593765258789,
21
+ "output_chroma": 12.631791114807129
22
+ }
23
+ ]
experiments/decision-20260928-v2/supplement/cache_loader_check.json ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "old_seconds": 0.6009589720001713,
3
+ "new_seconds": 0.010625808999975561,
4
+ "speedup": 56.556538142324364,
5
+ "exact_same_targets": true,
6
+ "synthetic_rows": 128,
7
+ "selected_rows": 64
8
+ }
experiments/decision-20260928-v2/supplement/nara_sources.json ADDED
@@ -0,0 +1,50 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [
2
+ {
3
+ "name": "nara_42.jpg",
4
+ "title": "Power station",
5
+ "source": "https://www.archives.gov/exhibits/picturing_the_century/greatdep/greatdep_img42.html",
6
+ "image_url": "https://www.archives.gov/exhibits/picturing_the_century/images/greatdep_042_v71.jpg",
7
+ "sha256": "a73186b25853d95a8e1785d1497923c1052747e8b045c5eac4c05dade60b2e87",
8
+ "bytes": 47847
9
+ },
10
+ {
11
+ "name": "nara_43.jpg",
12
+ "title": "Boxing team",
13
+ "source": "https://www.archives.gov/exhibits/picturing_the_century/greatdep/greatdep_img43.html",
14
+ "image_url": "https://www.archives.gov/exhibits/picturing_the_century/images/greatdep_043_v69.jpg",
15
+ "sha256": "f22dc07c4bdd0f3dac4b65ff0548395dc9aa88d8169a714fc40524249233b977",
16
+ "bytes": 40864
17
+ },
18
+ {
19
+ "name": "nara_44.jpg",
20
+ "title": "Street corner",
21
+ "source": "https://www.archives.gov/exhibits/picturing_the_century/greatdep/greatdep_img44.html",
22
+ "image_url": "https://www.archives.gov/exhibits/picturing_the_century/images/greatdep_044_v73.jpg",
23
+ "sha256": "c5e9a3fa57a6a0789e4b1529910e0294d98f245b24d205a1caccd08a5a3069b2",
24
+ "bytes": 68857
25
+ },
26
+ {
27
+ "name": "nara_45.jpg",
28
+ "title": "Children and sugar beets",
29
+ "source": "https://www.archives.gov/exhibits/picturing_the_century/greatdep/greatdep_img45.html",
30
+ "image_url": "https://www.archives.gov/exhibits/picturing_the_century/images/greatdep_045_v75.jpg",
31
+ "sha256": "3a3e3900cdb1384ca028a0a8843dcbcf499ffdd34832bc6b2e3f477d520403ae",
32
+ "bytes": 66706
33
+ },
34
+ {
35
+ "name": "nara_46.jpg",
36
+ "title": "Abandoned house",
37
+ "source": "https://www.archives.gov/exhibits/picturing_the_century/greatdep/greatdep_img46.html",
38
+ "image_url": "https://www.archives.gov/exhibits/picturing_the_century/images/greatdep_046_v76.jpg",
39
+ "sha256": "e7f2881c420c51c1855eb64219dc8c31e9a112976a3c81bfe0bcacb690d6b6e8",
40
+ "bytes": 74972
41
+ },
42
+ {
43
+ "name": "nara_49.jpg",
44
+ "title": "Man in his home",
45
+ "source": "https://www.archives.gov/exhibits/picturing_the_century/greatdep/greatdep_img49.html",
46
+ "image_url": "https://www.archives.gov/exhibits/picturing_the_century/images/greatdep_049_v72.jpg",
47
+ "sha256": "6d46bcc386aabf2dc090b2a05baf2e3ead04bc1274c38d87a60e1806ecfb4bab",
48
+ "bytes": 54941
49
+ }
50
+ ]
experiments/decision-20260928-v2/supplement/resolution_review.py ADDED
@@ -0,0 +1,27 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Resolution-only review on fixed historical inputs; no training or pixel editing."""
2
+ import json,time
3
+ from pathlib import Path
4
+ import torch
5
+ from PIL import Image,ImageDraw,ImageOps
6
+ from decision_model import load_decision
7
+ from inference import colorize,prepare
8
+ root=Path(__file__).resolve().parent
9
+ torch.set_num_threads(4)
10
+ m=load_decision(root/'smoke/mobilenet1')
11
+ records=json.loads((root/'archive_inputs/nara_sources.json').read_text())
12
+ rows=[]
13
+ for start in range(0,len(records),3):
14
+ canvas=Image.new('RGB',(1280,344*len(records[start:start+3])),'white');draw=ImageDraw.Draw(canvas)
15
+ for j,item in enumerate(records[start:start+3]):
16
+ im=Image.open(root/'archive_inputs'/item['name']).convert('L').convert('RGB')
17
+ draw.text((3,j*344+3),item['title']+' | Input | 256 | 384 | 512',fill='black')
18
+ canvas.paste(ImageOps.pad(im,(320,320),color='white'),(0,j*344+24))
19
+ for col,size in enumerate([256,384,512],1):
20
+ t=time.monotonic();pred=colorize(m,im,size=size,radius=8)
21
+ canvas.paste(ImageOps.pad(pred,(320,320),color='white'),(col*320,j*344+24))
22
+ with torch.no_grad():
23
+ details=m(prepare(im,size)[2],True);w=details['weights'].clamp_min(1e-8)
24
+ rows.append({'name':item['name'],'size':size,'seconds_cpu':time.monotonic()-t,'mask_entropy':float(-(w*w.log()).sum(1).mean()),'output_chroma':float(details['output'].norm(dim=1).mean())})
25
+ canvas.save(root/f'baseline_nara_resolution_{start//3}.jpg',quality=94)
26
+ (root/'baseline_nara_resolution.json').write_text(json.dumps(rows,indent=2))
27
+ print('Rendered',len(records),'historical photographs at three resolutions.')