Push model using huggingface_hub.

README.md

@@ -1,46 +1,16 @@
 ---
+license: apache-2.0
+pipeline_tag: image-to-image
 tags:
-- image-to-image
+- classification
 - colorization
+- image-to-image
+- model_hub_mixin
+- pytorch_model_hub_mixin
 - unet
-- pytorch
-license: apache-2.0
 ---
 
-# Mini U-Net Colorizer
-
-A 3,651,004-parameter U-Net that colorizes grayscale photos. Classification-style
-(Zhang et al., "Colorful Image Colorization"): it predicts a distribution over
-236 quantized CIE Lab `a`/`b` bins per pixel rather than regressing a
-single ab value directly, with color-bin loss weights derived from the real
-training-data distribution (rare/saturated colors weighted higher) so it
-doesn't just hedge toward desaturated averages. Decode with an annealed mean.
-
-- **Status:** Final (20 epochs complete)
-- **Input:** L channel normalized as `L/50 - 1` -> `[-1, 1]`, shape `(1, 256, 256)`
-- **Output:** logits over 236 ab bins, shape `(236, 256, 256)`
-- **Trained on:** `johnowhitaker/imagenette2-320` (None), warm-started from `User-2468/mini-unet-colorizer`
-- **Loss so far (weighted soft cross-entropy):** train 2.6135, val 2.7720
-
-## Usage
-
-```python
-import numpy as np, torch
-from skimage.color import rgb2lab, lab2rgb
-from PIL import Image
-# paste the SmallUNetColorizer class definition from the training script, then:
-model = SmallUNetColorizer.from_pretrained("User-2468/mini-unet-colorizer")
-model.eval()
-
-img = Image.open("photo.jpg").convert("RGB").resize((256, 256))
-lab = rgb2lab(np.asarray(img).astype("float32") / 255.0)
-L = torch.from_numpy(lab[:, :, 0:1] / 50.0 - 1.0).permute(2, 0, 1)[None]
-
-with torch.no_grad():
-    logits = model(L)
-    ab = model.decode(logits, temperature=0.38)[0].permute(1, 2, 0).numpy()
-
-L_out = (L[0, 0].numpy() + 1) * 50.0
-lab_out = np.concatenate([L_out[:, :, None], ab], axis=-1)
-rgb_out = np.clip(lab2rgb(lab_out), 0, 1)
-```
+This model has been pushed to the Hub using the [PytorchModelHubMixin](https://huggingface.co/docs/huggingface_hub/package_reference/mixins#huggingface_hub.PyTorchModelHubMixin) integration:
+- Code: [More Information Needed]
+- Paper: [More Information Needed]
+- Docs: [More Information Needed]

config.json

@@ -946,5 +946,11 @@
       -35.0
     ]
   ],
+  "context_dilations": [
+    2,
+    4,
+    8
+  ],
+  "context_mid_ch": 96,
   "in_ch": 1
 }

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:acffe0090532e4a6fa714b1e9488d6d2b7ffda4b25273c690e83deb9a55c2f43
-size 14629040
+oid sha256:c2b4ca6eb9b48a053dd64c7e38c7b84585fc81135f416b377409f701300301c9
+size 15909320