Initial commit: ComfyUI Bit Depth Enhancement nodes
- 4 nodes: Classical enhancement, Save 16-bit TIFF, ABCD, deepDeband - Classical methods: Bilateral+Dither, Gradient Domain, Multi-scale Fusion, Fast Edge-Aware - Deep learning: ABCD (3 models), deepDeband (2 variants) - Batch processing support - Comprehensive documentation Based on: - ABCD: https://github.com/WooKyoungHan/ABCD (CVPR 2023) - deepDeband: https://github.com/RaymondLZhou/deepDeband (ICIP 2022) 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
subraoul
+63
@@ -0,0 +1,63 @@
|
||||
# Development folder (tests, docs, patterns)
|
||||
Development/
|
||||
|
||||
# Python cache
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
*.so
|
||||
|
||||
# Distribution / packaging
|
||||
.Python
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
|
||||
# PyInstaller
|
||||
*.manifest
|
||||
*.spec
|
||||
|
||||
# Unit test / coverage reports
|
||||
htmlcov/
|
||||
.tox/
|
||||
.coverage
|
||||
.coverage.*
|
||||
.cache
|
||||
nosetests.xml
|
||||
coverage.xml
|
||||
*.cover
|
||||
.pytest_cache/
|
||||
|
||||
# Virtual environments
|
||||
venv/
|
||||
ENV/
|
||||
env/
|
||||
.venv
|
||||
|
||||
# IDEs
|
||||
.vscode/
|
||||
.idea/
|
||||
*.swp
|
||||
*.swo
|
||||
*~
|
||||
.DS_Store
|
||||
|
||||
# Model checkpoints (large files)
|
||||
*.pth
|
||||
*.ckpt
|
||||
*.safetensors
|
||||
|
||||
# Logs
|
||||
*.log
|
||||
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2025 raoul-ubuntu
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
@@ -0,0 +1,222 @@
|
||||
# ComfyUI Bit Depth Enhancement
|
||||
|
||||
Custom nodes for bit-depth enhancement and banding removal in ComfyUI.
|
||||
|
||||
## Description
|
||||
|
||||
This package provides nodes for enhancing 8-bit images to 16-bit with reduced banding artifacts. Includes both classical image processing methods and deep learning approaches.
|
||||
|
||||
**All nodes support batch processing.**
|
||||
|
||||
## Nodes
|
||||
|
||||
### 1. Bit Depth Enhancement (Classical)
|
||||
Classical image processing methods for bit-depth enhancement. No ML models required.
|
||||
|
||||
**Methods:**
|
||||
- **Bilateral+Dither** - Edge-aware filtering with Floyd-Steinberg dithering
|
||||
- **Gradient Domain** - Gradient-space processing for smooth transitions
|
||||
- **Multi-scale Fusion** - Laplacian pyramid decomposition
|
||||
- **Fast Edge-Aware** - Guided filter for fast processing
|
||||
|
||||
**Parameters:**
|
||||
| Parameter | Type | Range | Description |
|
||||
|-----------|------|-------|-------------|
|
||||
| `image` | IMAGE | - | Input image |
|
||||
| `method` | STRING | 4 options | Enhancement algorithm |
|
||||
| `strength` | FLOAT | 0.0-1.0 | Enhancement intensity (default: 0.7) |
|
||||
| `preserve_edges` | BOOLEAN | - | Maintain edge sharpness (default: True) |
|
||||
|
||||
### 2. Save 16-bit TIFF
|
||||
Export enhanced images as 16-bit TIFF files.
|
||||
|
||||
**Parameters:**
|
||||
| Parameter | Type | Options | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `images` | IMAGE | - | Input images |
|
||||
| `filename_prefix` | STRING | - | Output filename prefix |
|
||||
| `color_profile` | STRING | 4 options | Color space (sRGB, Adobe RGB, ProPhoto RGB, Linear) |
|
||||
|
||||
### 3. ABCD Bit-Depth Enhancement (8→16)
|
||||
Deep learning model for 8-bit to 16-bit enhancement using the ABCD (Arbitrary Bitwise Coefficient for De-quantization) architecture.
|
||||
|
||||
**Based on:** [ABCD - Learning to Restore Compressed Images with Arbitrary Bit-depth](https://github.com/WooKyoungHan/ABCD)
|
||||
**Paper:** CVPR 2023
|
||||
|
||||
ABCD uses coordinate-based implicit neural representation to reconstruct quantized images across arbitrary bit-depths. Three model architectures available:
|
||||
|
||||
- **SwinIR-ABCD** (Recommended) - Swin Transformer-based, highest quality (130h training, 4 GPUs)
|
||||
- **RDN-ABCD** - Residual Dense Network, balanced performance (82h training, 2 GPUs)
|
||||
- **EDSR-ABCD** - Enhanced Deep Residual, fastest processing (65h training, 1 GPU)
|
||||
|
||||
**Parameters:**
|
||||
| Parameter | Type | Description |
|
||||
|-----------|------|-------------|
|
||||
| `image` | IMAGE | Input 8-bit image |
|
||||
| `model` | STRING | Model architecture (SwinIR-ABCD, RDN-ABCD, EDSR-ABCD) |
|
||||
|
||||
### 4. deepDeband (Banding Removal)
|
||||
Deep learning model specifically trained for banding artifact removal using gradient-domain processing.
|
||||
|
||||
**Based on:** [deepDeband - Deep Gradient-Domain Image Debanding](https://github.com/RaymondLZhou/deepDeband)
|
||||
**Paper:** ICIP 2022
|
||||
|
||||
Trained on 51,490 pairs of pristine and banded image patches (256×256). Two model variants:
|
||||
|
||||
- **deepDeband-w** (Recommended) - Uses weighted bilateral patch fusion for smoother results
|
||||
- **deepDeband-f** - Direct patch processing, faster but may show seams
|
||||
|
||||
**Important:** These models were trained on real images and video frames. They may produce banding artifacts on synthetic images (3D renders, gradients, vector graphics).
|
||||
|
||||
**Parameters:**
|
||||
| Parameter | Type | Description |
|
||||
|-----------|------|-------------|
|
||||
| `image` | IMAGE | Input image with banding |
|
||||
| `model` | STRING | Model variant (deepDeband-w, deepDeband-f) |
|
||||
| `strength` | FLOAT | Debanding intensity (0.0-1.0, default: 1.0) |
|
||||
| `tile_size` | INT | Tile size for processing (default: 256) |
|
||||
| `tile_overlap` | INT | Overlap between tiles (default: 128) |
|
||||
|
||||
## Installation
|
||||
|
||||
1. Navigate to ComfyUI custom nodes directory:
|
||||
```bash
|
||||
cd ComfyUI/custom_nodes
|
||||
```
|
||||
|
||||
2. Clone this repository:
|
||||
```bash
|
||||
git clone https://github.com/subraoul/ComfyUI_Bit-Depth-Enhancer.git
|
||||
cd ComfyUI_Bit-Depth-Enhancer
|
||||
```
|
||||
|
||||
3. Install dependencies:
|
||||
```bash
|
||||
pip install -e .
|
||||
```
|
||||
|
||||
4. Restart ComfyUI
|
||||
|
||||
**Note:** This package is not yet published to the ComfyUI Registry. Manual installation only for now.
|
||||
|
||||
### Dependencies
|
||||
- Python 3.9+
|
||||
- PyTorch 2.0+
|
||||
- OpenCV 4.8+
|
||||
- NumPy 1.24+
|
||||
- scipy 1.11+
|
||||
- tifffile 2023.0+
|
||||
|
||||
## Model Setup
|
||||
|
||||
Deep learning nodes require model checkpoints. Models should be placed in the ComfyUI models directory:
|
||||
|
||||
```
|
||||
ComfyUI/models/bit_depth_enhancement/
|
||||
├── abcd/
|
||||
│ ├── edsr_abcd.pth
|
||||
│ ├── rdn_abcd.pth
|
||||
│ └── swinir_abcd.pth
|
||||
└── deepdeband/
|
||||
├── deepDeband_w.pth
|
||||
└── deepDeband_f.pth
|
||||
```
|
||||
|
||||
### ABCD Models
|
||||
|
||||
Download from Google Drive (original ABCD repository):
|
||||
|
||||
- **EDSR-ABCD**: [Download](https://drive.google.com/file/d/1LAe1KUPe8MuOP_NRwBMfQ5W32o37Ln6W/view?usp=sharing)
|
||||
- **RDN-ABCD**: [Download](https://drive.google.com/file/d/1tj7HiSpDxuHdEFQYG_EwWncDisfT_k88/view?usp=sharing)
|
||||
- **SwinIR-ABCD**: [Download](https://drive.google.com/file/d/1zBGLttDMET7CQcj729sZyPKOVWpMyyMZ/view?usp=sharing)
|
||||
|
||||
Rename downloaded files to match the expected names and place in `ComfyUI/models/bit_depth_enhancement/abcd/`
|
||||
|
||||
### deepDeband Models
|
||||
|
||||
Download from GitHub (original deepDeband repository):
|
||||
|
||||
1. Navigate to [deepDeband checkpoints](https://github.com/RaymondLZhou/deepDeband/tree/master/pytorch-CycleGAN-and-pix2pix/checkpoints)
|
||||
|
||||
2. Download from subdirectories:
|
||||
- `deepDeband-w/latest_net_G.pth` → rename to `deepDeband_w.pth`
|
||||
- `deepDeband-f/latest_net_G.pth` → rename to `deepDeband_f.pth`
|
||||
|
||||
3. Place in `ComfyUI/models/bit_depth_enhancement/deepdeband/`
|
||||
|
||||
**Note:** You only need the generator weights (`latest_net_G.pth`), not the discriminator (`latest_net_D.pth`).
|
||||
|
||||
## Basic Usage
|
||||
|
||||
### Classical Enhancement Workflow
|
||||
```
|
||||
Load Image → Bit Depth Enhancement (Classical) → Save 16-bit TIFF
|
||||
```
|
||||
|
||||
**Recommended Settings:**
|
||||
- Method: Bilateral+Dither (general use) or Multi-scale Fusion (best quality)
|
||||
- Strength: 0.7
|
||||
- Preserve Edges: True
|
||||
- Color Profile: sRGB (web) or Adobe RGB (print)
|
||||
|
||||
### ABCD Deep Learning Workflow
|
||||
```
|
||||
Load Image → ABCD Bit-Depth Enhancement → Save Image
|
||||
```
|
||||
|
||||
**Recommended:** Use SwinIR-ABCD for best quality results.
|
||||
|
||||
### deepDeband Workflow
|
||||
```
|
||||
Load Image → deepDeband → Save Image
|
||||
```
|
||||
|
||||
**Recommended:** Use deepDeband-w with default settings. Lower strength (0.5-0.7) for subtle enhancement.
|
||||
|
||||
**Warning:** Avoid using deepDeband on synthetic/rendered images as it may introduce artifacts.
|
||||
|
||||
## Known Limitations
|
||||
|
||||
- Classical methods cannot recover dynamic range not present in source
|
||||
- Processing time scales with image resolution
|
||||
- Deep learning models require GPU for reasonable performance
|
||||
- Source image quality matters - heavily compressed JPEGs benefit less
|
||||
- deepDeband may produce artifacts on synthetic/3D rendered images
|
||||
|
||||
## Output Format
|
||||
|
||||
- **Classical + Save 16-bit TIFF**: True 16-bit TIFF files (65,535 levels per channel)
|
||||
- **Deep learning nodes**: Standard ComfyUI IMAGE output (can be saved with any ComfyUI save node)
|
||||
|
||||
## References
|
||||
|
||||
This implementation is based on the following research:
|
||||
|
||||
### ABCD
|
||||
- **Paper:** "Arbitrary Bit-Depth Quantization for Image Restoration"
|
||||
- **Repository:** https://github.com/WooKyoungHan/ABCD
|
||||
- **Conference:** CVPR 2023
|
||||
- **Authors:** WooKyoung Han et al.
|
||||
|
||||
### deepDeband
|
||||
- **Paper:** "Deep Gradient-Domain Image Debanding"
|
||||
- **Repository:** https://github.com/RaymondLZhou/deepDeband
|
||||
- **Conference:** ICIP 2022
|
||||
- **Authors:** Raymond L. Zhou, Shahrukh Athar, Zhongling Wang, Zhou Wang
|
||||
|
||||
## License
|
||||
|
||||
MIT License - see [LICENSE](LICENSE) file for details.
|
||||
|
||||
## Credits
|
||||
|
||||
**Author:** raoul-ubuntu
|
||||
|
||||
**Built with:**
|
||||
- [ComfyUI](https://github.com/comfyanonymous/ComfyUI)
|
||||
- [PyTorch](https://pytorch.org/)
|
||||
- [OpenCV](https://opencv.org/)
|
||||
|
||||
## Support
|
||||
|
||||
Report issues or request features at: [GitHub Issues](https://github.com/subraoul/ComfyUI_Bit-Depth-Enhancer/issues)
|
||||
+26
@@ -0,0 +1,26 @@
|
||||
"""
|
||||
ComfyUI Bit Depth Enhancement
|
||||
Classical and deep learning-based bit-depth enhancement for professional cinematography workflows
|
||||
|
||||
This is the ComfyUI entry point that loads the custom nodes.
|
||||
"""
|
||||
|
||||
from .bit_depth_enhancement.nodes.classical import BitDepthEnhancementClassical, Save16BitTIFF
|
||||
from .bit_depth_enhancement.nodes.abcd import ABCD_BitDepthEnhancement
|
||||
from .bit_depth_enhancement.nodes.deband import DeepDeband
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"BitDepthEnhancementClassical": BitDepthEnhancementClassical,
|
||||
"Save16BitTIFF": Save16BitTIFF,
|
||||
"ABCD_BitDepthEnhancement": ABCD_BitDepthEnhancement,
|
||||
"DeepDeband": DeepDeband,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"BitDepthEnhancementClassical": "Bit Depth Enhancement (Classical)",
|
||||
"Save16BitTIFF": "Save 16-bit TIFF",
|
||||
"ABCD_BitDepthEnhancement": "ABCD Bit-Depth Enhancement (8→16)",
|
||||
"DeepDeband": "deepDeband (Banding Removal)",
|
||||
}
|
||||
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
@@ -0,0 +1,23 @@
|
||||
"""
|
||||
ComfyUI Bit Depth Enhancement
|
||||
Classical and deep learning-based bit-depth enhancement for professional cinematography workflows
|
||||
"""
|
||||
|
||||
__version__ = "0.1.0"
|
||||
|
||||
from .nodes.classical import BitDepthEnhancementClassical, Save16BitTIFF
|
||||
from .nodes.deband import DeepDeband
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"BitDepthEnhancementClassical": BitDepthEnhancementClassical,
|
||||
"Save16BitTIFF": Save16BitTIFF,
|
||||
"DeepDeband": DeepDeband,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"BitDepthEnhancementClassical": "Bit Depth Enhancement (Classical)",
|
||||
"Save16BitTIFF": "Save 16-bit TIFF",
|
||||
"DeepDeband": "deepDeband (Banding Removal)",
|
||||
}
|
||||
|
||||
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
@@ -0,0 +1,691 @@
|
||||
"""
|
||||
ABCD (Arbitrary Bit-depth Converter and Dequantizer) Model Architectures
|
||||
|
||||
This module implements the three ABCD model variants for bit-depth enhancement:
|
||||
- EDSR_ABCD: Enhanced Deep Residual Networks with ABCD wrapper
|
||||
- RDN_ABCD: Residual Dense Networks with ABCD wrapper
|
||||
- SwinIR_ABCD: Swin Transformer with ABCD wrapper
|
||||
|
||||
Based on: https://github.com/WooKyoungHan/ABCD
|
||||
Paper: "Learning to Restore Compressed Images with Arbitrary Bit-depth"
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import math
|
||||
|
||||
|
||||
# ====================================================================
|
||||
# Base Components (shared across all architectures)
|
||||
# ====================================================================
|
||||
|
||||
class MLP(nn.Module):
|
||||
"""
|
||||
Multi-Layer Perceptron used as the implicit function network (imnet)
|
||||
|
||||
Maps concatenated features (coefficient * frequency + cell) to RGB values.
|
||||
"""
|
||||
def __init__(self, in_dim, out_dim, hidden_list):
|
||||
"""
|
||||
Args:
|
||||
in_dim: Input dimension
|
||||
out_dim: Output dimension (typically 3 for RGB)
|
||||
hidden_list: List of hidden layer dimensions
|
||||
"""
|
||||
super().__init__()
|
||||
layers = []
|
||||
lastv = in_dim
|
||||
for hidden in hidden_list:
|
||||
layers.append(nn.Linear(lastv, hidden))
|
||||
layers.append(nn.ReLU())
|
||||
lastv = hidden
|
||||
layers.append(nn.Linear(lastv, out_dim))
|
||||
self.layers = nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
return self.layers(x)
|
||||
|
||||
|
||||
class MLP_DirectRGB(nn.Module):
|
||||
"""
|
||||
MLP variant that outputs RGB directly from the last hidden layer
|
||||
|
||||
Used for RDN and SwinIR ABCD checkpoints where layer 6 outputs [3, 256]
|
||||
instead of [256, 256]. This is the correct architecture for the ABCD paper.
|
||||
"""
|
||||
def __init__(self, in_dim, out_dim, hidden_list):
|
||||
"""
|
||||
Args:
|
||||
in_dim: Input dimension
|
||||
out_dim: Output dimension (3 for RGB)
|
||||
hidden_list: List of hidden layer dimensions (should be 4 layers of 256)
|
||||
"""
|
||||
super().__init__()
|
||||
layers = []
|
||||
lastv = in_dim
|
||||
for hidden in hidden_list:
|
||||
layers.append(nn.Linear(lastv, hidden))
|
||||
layers.append(nn.ReLU())
|
||||
lastv = hidden
|
||||
# Final layer outputs RGB directly (3 channels) from last hidden dimension
|
||||
layers.append(nn.Linear(lastv, out_dim))
|
||||
self.layers = nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
return self.layers(x)
|
||||
|
||||
|
||||
# ====================================================================
|
||||
# EDSR Components
|
||||
# ====================================================================
|
||||
|
||||
class ResBlock(nn.Module):
|
||||
"""
|
||||
Residual block for EDSR
|
||||
|
||||
Architecture: Conv -> ReLU -> Conv -> Residual connection (scaled)
|
||||
"""
|
||||
def __init__(self, n_feats, kernel_size=3, res_scale=1.0):
|
||||
super().__init__()
|
||||
self.res_scale = res_scale
|
||||
padding = kernel_size // 2
|
||||
|
||||
self.body = nn.Sequential(
|
||||
nn.Conv2d(n_feats, n_feats, kernel_size, padding=padding),
|
||||
nn.ReLU(inplace=True),
|
||||
nn.Conv2d(n_feats, n_feats, kernel_size, padding=padding)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
res = self.body(x).mul(self.res_scale)
|
||||
res += x
|
||||
return res
|
||||
|
||||
|
||||
class MeanShift(nn.Conv2d):
|
||||
"""
|
||||
Mean shift layer for RGB normalization
|
||||
|
||||
Subtracts/adds RGB mean values using a 1x1 convolution
|
||||
"""
|
||||
def __init__(self, rgb_mean=(0.4488, 0.4371, 0.4040), rgb_std=(1.0, 1.0, 1.0), sign=-1):
|
||||
super().__init__(3, 3, kernel_size=1)
|
||||
std = torch.Tensor(rgb_std)
|
||||
self.weight.data = torch.eye(3).view(3, 3, 1, 1) / std.view(3, 1, 1, 1)
|
||||
self.bias.data = sign * torch.Tensor(rgb_mean) / std
|
||||
for p in self.parameters():
|
||||
p.requires_grad = False
|
||||
|
||||
|
||||
class EDSR(nn.Module):
|
||||
"""
|
||||
Enhanced Deep Residual Networks (EDSR) encoder
|
||||
|
||||
Original EDSR paper: "Enhanced Deep Residual Networks for Single Image Super-Resolution"
|
||||
This version is configured for ABCD (no upsampling, feature extraction only)
|
||||
"""
|
||||
def __init__(self, n_resblocks=38, n_feats=128, n_colors=3, res_scale=1.0, no_upsampling=True):
|
||||
"""
|
||||
Args:
|
||||
n_resblocks: Number of residual blocks (default 38 for EDSR-baseline)
|
||||
n_feats: Number of feature channels (default 128)
|
||||
n_colors: Number of input/output color channels (default 3 for RGB)
|
||||
res_scale: Residual scaling factor
|
||||
no_upsampling: If True, skip upsampling (used for ABCD)
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
kernel_size = 3
|
||||
padding = kernel_size // 2
|
||||
|
||||
# RGB mean normalization
|
||||
self.sub_mean = MeanShift(sign=-1)
|
||||
self.add_mean = MeanShift(sign=1)
|
||||
|
||||
# Head: initial feature extraction (as Sequential to match checkpoint)
|
||||
self.head = nn.Sequential(
|
||||
nn.Conv2d(n_colors, n_feats, kernel_size, padding=padding)
|
||||
)
|
||||
|
||||
# Body: residual blocks
|
||||
body = [ResBlock(n_feats, kernel_size, res_scale) for _ in range(n_resblocks)]
|
||||
body.append(nn.Conv2d(n_feats, n_feats, kernel_size, padding=padding))
|
||||
self.body = nn.Sequential(*body)
|
||||
|
||||
self.no_upsampling = no_upsampling
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Args:
|
||||
x: Input tensor [B, C, H, W]
|
||||
|
||||
Returns:
|
||||
Feature tensor [B, n_feats, H, W]
|
||||
"""
|
||||
x = self.sub_mean(x)
|
||||
x = self.head(x)
|
||||
res = self.body(x)
|
||||
res += x
|
||||
return res
|
||||
|
||||
|
||||
# ====================================================================
|
||||
# RDN Components
|
||||
# ====================================================================
|
||||
|
||||
class DenseLayer(nn.Module):
|
||||
"""Dense layer for RDN - Conv + ReLU with concatenation"""
|
||||
def __init__(self, in_channels, out_channels, kernel_size=3):
|
||||
super().__init__()
|
||||
padding = kernel_size // 2
|
||||
self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, padding=padding)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
|
||||
def forward(self, x):
|
||||
out = self.relu(self.conv(x))
|
||||
return torch.cat([x, out], 1)
|
||||
|
||||
|
||||
class RDB(nn.Module):
|
||||
"""
|
||||
Residual Dense Block (RDB)
|
||||
|
||||
Multiple dense layers with local feature fusion
|
||||
"""
|
||||
def __init__(self, in_channels, growth_rate, num_layers):
|
||||
super().__init__()
|
||||
|
||||
self.layers = nn.ModuleList()
|
||||
for i in range(num_layers):
|
||||
self.layers.append(DenseLayer(in_channels + i * growth_rate, growth_rate))
|
||||
|
||||
# Local Feature Fusion
|
||||
self.lff = nn.Conv2d(in_channels + num_layers * growth_rate, in_channels, 1)
|
||||
|
||||
def forward(self, x):
|
||||
out = x
|
||||
for layer in self.layers:
|
||||
out = layer(out)
|
||||
out = self.lff(out)
|
||||
# Local residual learning
|
||||
return out + x
|
||||
|
||||
|
||||
class RDN(nn.Module):
|
||||
"""
|
||||
Residual Dense Network (RDN) encoder
|
||||
|
||||
Original RDN paper: "Residual Dense Network for Image Super-Resolution"
|
||||
This version is configured for ABCD (no upsampling, feature extraction only)
|
||||
"""
|
||||
def __init__(self, n_colors=3, G0=64, RDNkSize=3, RDNconfig='B', no_upsampling=True):
|
||||
"""
|
||||
Args:
|
||||
n_colors: Number of input color channels
|
||||
G0: Number of feature channels
|
||||
RDNkSize: Kernel size for RDN layers
|
||||
RDNconfig: Configuration preset ('A', 'B', or 'C')
|
||||
no_upsampling: If True, skip upsampling (used for ABCD)
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# Configuration presets
|
||||
configs = {
|
||||
'A': {'D': 20, 'C': 6, 'G': 32}, # D: num RDBs, C: layers per RDB, G: growth rate
|
||||
'B': {'D': 16, 'C': 8, 'G': 64},
|
||||
'C': {'D': 8, 'C': 8, 'G': 64}
|
||||
}
|
||||
config = configs[RDNconfig]
|
||||
|
||||
self.D = config['D']
|
||||
self.C = config['C']
|
||||
self.G = config['G']
|
||||
self.G0 = G0
|
||||
|
||||
# Shallow feature extraction
|
||||
self.SFENet1 = nn.Conv2d(n_colors, G0, RDNkSize, padding=RDNkSize//2)
|
||||
self.SFENet2 = nn.Conv2d(G0, G0, RDNkSize, padding=RDNkSize//2)
|
||||
|
||||
# Residual dense blocks
|
||||
self.RDBs = nn.ModuleList([RDB(G0, self.G, self.C) for _ in range(self.D)])
|
||||
|
||||
# Global Feature Fusion
|
||||
self.GFF = nn.Sequential(
|
||||
nn.Conv2d(self.D * G0, G0, 1),
|
||||
nn.Conv2d(G0, G0, RDNkSize, padding=RDNkSize//2)
|
||||
)
|
||||
|
||||
self.no_upsampling = no_upsampling
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Args:
|
||||
x: Input tensor [B, C, H, W]
|
||||
|
||||
Returns:
|
||||
Feature tensor [B, G0, H, W]
|
||||
"""
|
||||
f_1 = self.SFENet1(x)
|
||||
x = self.SFENet2(f_1)
|
||||
|
||||
RDBs_out = []
|
||||
for rdb in self.RDBs:
|
||||
x = rdb(x)
|
||||
RDBs_out.append(x)
|
||||
|
||||
x = self.GFF(torch.cat(RDBs_out, 1))
|
||||
x += f_1 # Global residual learning
|
||||
|
||||
return x
|
||||
|
||||
|
||||
# ====================================================================
|
||||
# SwinIR Components
|
||||
# ====================================================================
|
||||
|
||||
class WindowAttention(nn.Module):
|
||||
"""
|
||||
Window-based Multi-head Self Attention (W-MSA) for Swin Transformer
|
||||
"""
|
||||
def __init__(self, dim, window_size, num_heads):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.window_size = window_size
|
||||
self.num_heads = num_heads
|
||||
head_dim = dim // num_heads
|
||||
self.scale = head_dim ** -0.5
|
||||
|
||||
# Relative position bias
|
||||
self.relative_position_bias_table = nn.Parameter(
|
||||
torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)
|
||||
)
|
||||
|
||||
# Get pair-wise relative position index for each token inside the window
|
||||
coords_h = torch.arange(window_size[0])
|
||||
coords_w = torch.arange(window_size[1])
|
||||
coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing='ij'))
|
||||
coords_flatten = torch.flatten(coords, 1)
|
||||
relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
|
||||
relative_coords = relative_coords.permute(1, 2, 0).contiguous()
|
||||
relative_coords[:, :, 0] += window_size[0] - 1
|
||||
relative_coords[:, :, 1] += window_size[1] - 1
|
||||
relative_coords[:, :, 0] *= 2 * window_size[1] - 1
|
||||
relative_position_index = relative_coords.sum(-1)
|
||||
self.register_buffer("relative_position_index", relative_position_index)
|
||||
|
||||
self.qkv = nn.Linear(dim, dim * 3)
|
||||
self.proj = nn.Linear(dim, dim)
|
||||
|
||||
def forward(self, x):
|
||||
B_, N, C = x.shape
|
||||
qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
|
||||
q, k, v = qkv[0], qkv[1], qkv[2]
|
||||
|
||||
q = q * self.scale
|
||||
attn = (q @ k.transpose(-2, -1))
|
||||
|
||||
# Add relative position bias
|
||||
relative_position_bias = self.relative_position_bias_table[
|
||||
self.relative_position_index.view(-1)
|
||||
].view(self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1)
|
||||
relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
|
||||
attn = attn + relative_position_bias.unsqueeze(0)
|
||||
|
||||
attn = F.softmax(attn, dim=-1)
|
||||
x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
|
||||
x = self.proj(x)
|
||||
return x
|
||||
|
||||
|
||||
class SwinTransformerBlock(nn.Module):
|
||||
"""Swin Transformer Block"""
|
||||
def __init__(self, dim, num_heads, window_size=8, mlp_ratio=2.):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.num_heads = num_heads
|
||||
self.window_size = window_size
|
||||
self.mlp_ratio = mlp_ratio
|
||||
|
||||
self.norm1 = nn.LayerNorm(dim)
|
||||
self.attn = WindowAttention(dim, (window_size, window_size), num_heads)
|
||||
self.norm2 = nn.LayerNorm(dim)
|
||||
|
||||
mlp_hidden_dim = int(dim * mlp_ratio)
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(dim, mlp_hidden_dim),
|
||||
nn.GELU(),
|
||||
nn.Linear(mlp_hidden_dim, dim)
|
||||
)
|
||||
|
||||
def forward(self, x, x_size):
|
||||
H, W = x_size
|
||||
B, L, C = x.shape
|
||||
|
||||
shortcut = x
|
||||
x = self.norm1(x)
|
||||
x = x.view(B, H, W, C)
|
||||
|
||||
# Window partition
|
||||
x_windows = self.window_partition(x, self.window_size)
|
||||
x_windows = x_windows.view(-1, self.window_size * self.window_size, C)
|
||||
|
||||
# W-MSA
|
||||
attn_windows = self.attn(x_windows)
|
||||
|
||||
# Merge windows
|
||||
attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
|
||||
x = self.window_reverse(attn_windows, self.window_size, H, W)
|
||||
x = x.view(B, H * W, C)
|
||||
|
||||
# FFN
|
||||
x = shortcut + x
|
||||
x = x + self.mlp(self.norm2(x))
|
||||
|
||||
return x
|
||||
|
||||
def window_partition(self, x, window_size):
|
||||
B, H, W, C = x.shape
|
||||
x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
|
||||
windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
|
||||
return windows
|
||||
|
||||
def window_reverse(self, windows, window_size, H, W):
|
||||
B = int(windows.shape[0] / (H * W / window_size / window_size))
|
||||
x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)
|
||||
x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
|
||||
return x
|
||||
|
||||
|
||||
class RSTB(nn.Module):
|
||||
"""Residual Swin Transformer Block (RSTB)"""
|
||||
def __init__(self, dim, depth, num_heads, window_size, mlp_ratio=2.):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.blocks = nn.ModuleList([
|
||||
SwinTransformerBlock(dim, num_heads, window_size, mlp_ratio)
|
||||
for _ in range(depth)
|
||||
])
|
||||
self.conv = nn.Conv2d(dim, dim, 3, 1, 1)
|
||||
|
||||
def forward(self, x, x_size):
|
||||
for blk in self.blocks:
|
||||
x = blk(x, x_size)
|
||||
|
||||
# Reshape to image format for conv
|
||||
B, L, C = x.shape
|
||||
H, W = x_size
|
||||
x = x.transpose(1, 2).view(B, C, H, W)
|
||||
x = self.conv(x)
|
||||
x = x.flatten(2).transpose(1, 2)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class SwinIR(nn.Module):
|
||||
"""
|
||||
SwinIR: Image Restoration Using Swin Transformer
|
||||
|
||||
Original SwinIR paper: "SwinIR: Image Restoration Using Swin Transformer"
|
||||
This version is configured for ABCD (no upsampling, feature extraction only)
|
||||
"""
|
||||
def __init__(self, img_size=48, in_chans=3, embed_dim=180, depths=[6, 6, 6, 6, 6, 6],
|
||||
num_heads=[6, 6, 6, 6, 6, 6], window_size=8, mlp_ratio=2., no_upsampling=True):
|
||||
"""
|
||||
Args:
|
||||
img_size: Input image size
|
||||
in_chans: Number of input channels
|
||||
embed_dim: Patch embedding dimension
|
||||
depths: Depth of each Swin Transformer block
|
||||
num_heads: Number of attention heads in each layer
|
||||
window_size: Window size for attention
|
||||
mlp_ratio: Ratio of mlp hidden dim to embedding dim
|
||||
no_upsampling: If True, skip upsampling (used for ABCD)
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
self.window_size = window_size
|
||||
self.num_layers = len(depths)
|
||||
self.embed_dim = embed_dim
|
||||
self.num_features = embed_dim
|
||||
self.mlp_ratio = mlp_ratio
|
||||
|
||||
# Shallow feature extraction
|
||||
self.conv_first = nn.Conv2d(in_chans, embed_dim, 3, 1, 1)
|
||||
|
||||
# Deep feature extraction
|
||||
self.layers = nn.ModuleList()
|
||||
for i in range(self.num_layers):
|
||||
layer = RSTB(
|
||||
dim=embed_dim,
|
||||
depth=depths[i],
|
||||
num_heads=num_heads[i],
|
||||
window_size=window_size,
|
||||
mlp_ratio=mlp_ratio
|
||||
)
|
||||
self.layers.append(layer)
|
||||
|
||||
self.norm = nn.LayerNorm(embed_dim)
|
||||
self.conv_after_body = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1)
|
||||
|
||||
self.no_upsampling = no_upsampling
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Args:
|
||||
x: Input tensor [B, C, H, W]
|
||||
|
||||
Returns:
|
||||
Feature tensor [B, embed_dim, H, W]
|
||||
"""
|
||||
B, C, H, W = x.shape
|
||||
|
||||
# Check input size
|
||||
H_padded = ((H + self.window_size - 1) // self.window_size) * self.window_size
|
||||
W_padded = ((W + self.window_size - 1) // self.window_size) * self.window_size
|
||||
|
||||
if H != H_padded or W != W_padded:
|
||||
x = F.pad(x, (0, W_padded - W, 0, H_padded - H), mode='reflect')
|
||||
|
||||
# Shallow feature extraction
|
||||
x_first = self.conv_first(x)
|
||||
|
||||
# Deep feature extraction
|
||||
x = x_first.flatten(2).transpose(1, 2) # B, HW, C
|
||||
for layer in self.layers:
|
||||
x = layer(x, (H_padded, W_padded))
|
||||
|
||||
x = self.norm(x)
|
||||
x = x.transpose(1, 2).view(B, self.embed_dim, H_padded, W_padded)
|
||||
x = self.conv_after_body(x) + x_first
|
||||
|
||||
# Remove padding if added
|
||||
if H != H_padded or W != W_padded:
|
||||
x = x[:, :, :H, :W]
|
||||
|
||||
return x
|
||||
|
||||
|
||||
# ====================================================================
|
||||
# ABCD Wrapper Models
|
||||
# ====================================================================
|
||||
|
||||
class ABCD(nn.Module):
|
||||
"""
|
||||
ABCD (Arbitrary Bit-depth Converter and Dequantizer) wrapper
|
||||
|
||||
Wraps encoder networks (EDSR/RDN/SwinIR) with coordinate-based
|
||||
implicit neural representation for arbitrary bit-depth conversion.
|
||||
|
||||
Architecture:
|
||||
1. Encoder extracts features from input image
|
||||
2. Coefficient and frequency branches process features
|
||||
3. Implicit function (MLP) predicts RGB values at continuous coordinates
|
||||
4. Output is combined with input for residual learning
|
||||
"""
|
||||
def __init__(self, encoder, hidden_dim=256, imnet_hidden=[256, 256, 256, 256], num_in_ch=None, use_direct_rgb=False):
|
||||
"""
|
||||
Args:
|
||||
encoder: Backbone encoder (EDSR/RDN/SwinIR)
|
||||
hidden_dim: Hidden dimension for coefficient and frequency branches
|
||||
imnet_hidden: Hidden layer dimensions for implicit function
|
||||
num_in_ch: Number of input channels from encoder (if None, auto-detect)
|
||||
use_direct_rgb: If True, use MLP_DirectRGB (for RDN/SwinIR checkpoints)
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
self.encoder = encoder
|
||||
self.hidden_dim = hidden_dim
|
||||
|
||||
# Get encoder output channels
|
||||
if num_in_ch is not None:
|
||||
encoder_dim = num_in_ch
|
||||
elif hasattr(encoder, 'G0'):
|
||||
encoder_dim = encoder.G0 # RDN
|
||||
elif hasattr(encoder, 'embed_dim'):
|
||||
encoder_dim = encoder.embed_dim # SwinIR
|
||||
else:
|
||||
encoder_dim = 128 # EDSR default
|
||||
|
||||
# Coefficient and frequency branches
|
||||
self.coef = nn.Conv2d(encoder_dim, hidden_dim, 3, padding=1)
|
||||
self.freq = nn.Conv2d(encoder_dim, hidden_dim, 3, padding=1)
|
||||
|
||||
# Implicit function network (coordinate-based)
|
||||
# Input: hidden_dim (coef * freq) + 1 (cell size)
|
||||
# Note: Original ABCD uses 1D cell, not 2D
|
||||
# RDN and SwinIR checkpoints use direct RGB output architecture
|
||||
if use_direct_rgb:
|
||||
self.imnet = MLP_DirectRGB(hidden_dim + 1, 3, imnet_hidden)
|
||||
else:
|
||||
self.imnet = MLP(hidden_dim + 1, 3, imnet_hidden)
|
||||
|
||||
# Output activation (normalized to [0, 1])
|
||||
self.tanh = nn.Tanh()
|
||||
|
||||
def gen_feat(self, inp):
|
||||
"""
|
||||
Generate features from input image
|
||||
|
||||
Args:
|
||||
inp: Input tensor [B, C, H, W]
|
||||
"""
|
||||
self.feat = self.encoder(inp)
|
||||
self.coef_feat = self.coef(self.feat)
|
||||
self.freq_feat = self.freq(self.feat)
|
||||
|
||||
def query_rgb(self, coord, cell):
|
||||
"""
|
||||
Query RGB values at continuous coordinates
|
||||
|
||||
Args:
|
||||
coord: Coordinates [B, N, 2] in range [-1, 1]
|
||||
cell: Cell sizes [B, N, 1] (scalar cell size)
|
||||
|
||||
Returns:
|
||||
RGB values [B, N, 3]
|
||||
"""
|
||||
# Sample coefficient and frequency features at coordinates
|
||||
# coord needs to be flipped for grid_sample (y, x) -> (x, y)
|
||||
coord_flip = coord.flip(-1)
|
||||
|
||||
# grid_sample expects [B, H, W, 2]
|
||||
B, N, _ = coord.shape
|
||||
coord_sample = coord_flip.view(B, 1, -1, 2)
|
||||
|
||||
coef = F.grid_sample(
|
||||
self.coef_feat, coord_sample,
|
||||
mode='bilinear', align_corners=False
|
||||
) # [B, hidden_dim, 1, N]
|
||||
|
||||
freq = F.grid_sample(
|
||||
self.freq_feat, coord_sample,
|
||||
mode='bilinear', align_corners=False
|
||||
) # [B, hidden_dim, 1, N]
|
||||
|
||||
coef = coef.view(B, self.hidden_dim, N).permute(0, 2, 1) # [B, N, hidden_dim]
|
||||
freq = freq.view(B, self.hidden_dim, N).permute(0, 2, 1) # [B, N, hidden_dim]
|
||||
|
||||
# Frequency modulation: coef * (cos(2π * freq) + 1) / 2
|
||||
# This creates a smooth interpolation function
|
||||
freq_encoded = torch.cos(math.pi * freq) * 0.5 + 0.5
|
||||
modulated = coef * freq_encoded
|
||||
|
||||
# Concatenate with cell information (1D cell size)
|
||||
inp = torch.cat([modulated, cell], dim=-1)
|
||||
|
||||
# Predict RGB residual
|
||||
ret = self.imnet(inp)
|
||||
|
||||
# Normalize to [0, 1] using scaled tanh
|
||||
ret = (self.tanh(ret) + 1) / 2
|
||||
|
||||
return ret
|
||||
|
||||
def forward(self, inp, coord, cell):
|
||||
"""
|
||||
Forward pass
|
||||
|
||||
Args:
|
||||
inp: Input tensor [B, C, H, W]
|
||||
coord: Coordinates [B, N, 2]
|
||||
cell: Cell sizes [B, N, 2]
|
||||
|
||||
Returns:
|
||||
RGB values [B, N, 3]
|
||||
"""
|
||||
self.gen_feat(inp)
|
||||
return self.query_rgb(coord, cell)
|
||||
|
||||
|
||||
# ====================================================================
|
||||
# Factory Functions for Creating ABCD Models
|
||||
# ====================================================================
|
||||
|
||||
def make_edsr_abcd(hidden_dim=256, imnet_hidden=[256, 256, 256, 256]):
|
||||
"""
|
||||
Create EDSR-ABCD model
|
||||
|
||||
Configuration: EDSR-baseline with 38 residual blocks, 128 features
|
||||
Note: EDSR uses standard MLP architecture (not direct RGB)
|
||||
"""
|
||||
encoder = EDSR(n_resblocks=38, n_feats=128, res_scale=1.0, no_upsampling=True)
|
||||
model = ABCD(encoder, hidden_dim, imnet_hidden, use_direct_rgb=False)
|
||||
return model
|
||||
|
||||
|
||||
def make_rdn_abcd(hidden_dim=256, imnet_hidden=[256, 256, 256]):
|
||||
"""
|
||||
Create RDN-ABCD model
|
||||
|
||||
Configuration: RDN config 'C' with 8 RDBs, 8 layers per RDB, growth rate 64
|
||||
Note: Using config 'C' (8 blocks) to match ABCD checkpoint architecture
|
||||
Note: RDN checkpoint uses direct RGB output (layer 6: [3, 256])
|
||||
"""
|
||||
encoder = RDN(G0=64, RDNconfig='C', no_upsampling=True)
|
||||
model = ABCD(encoder, hidden_dim, imnet_hidden, use_direct_rgb=True)
|
||||
return model
|
||||
|
||||
|
||||
def make_swinir_abcd(hidden_dim=256, imnet_hidden=[256, 256, 256], img_size=64,
|
||||
embed_dim=180, depths=[6, 6, 6, 6, 6, 6],
|
||||
num_heads=[6, 6, 6, 6, 6, 6], window_size=8, mlp_ratio=2):
|
||||
"""
|
||||
Create SwinIR-ABCD model
|
||||
|
||||
Configuration: SwinIR-Light variant with 6 RSTB blocks
|
||||
Note: Checkpoint uses embed_dim=180, num_heads=6 (not the default 64/8)
|
||||
Note: SwinIR checkpoint uses direct RGB output (layer 6: [3, 256])
|
||||
Note: Encoder has embed_dim=180 internally, but outputs 64 channels to ABCD wrapper
|
||||
"""
|
||||
encoder = SwinIR(
|
||||
img_size=img_size,
|
||||
embed_dim=embed_dim,
|
||||
depths=depths,
|
||||
num_heads=num_heads,
|
||||
window_size=window_size,
|
||||
mlp_ratio=mlp_ratio,
|
||||
no_upsampling=True
|
||||
)
|
||||
model = ABCD(encoder, hidden_dim, imnet_hidden, num_in_ch=64, use_direct_rgb=True)
|
||||
return model
|
||||
@@ -0,0 +1,307 @@
|
||||
"""
|
||||
DeepDeband Model Architecture
|
||||
Based on pix2pix UNet Generator for image-to-image translation
|
||||
|
||||
This implementation is based on the PyTorch CycleGAN and pix2pix framework
|
||||
used by the deepDeband paper (ICIP 2022).
|
||||
|
||||
The model uses a U-Net architecture with skip connections to preserve
|
||||
spatial information while learning to remove banding artifacts.
|
||||
|
||||
Reference:
|
||||
- deepDeband: https://github.com/RaymondLZhou/deepDeband
|
||||
- pix2pix: https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class UnetSkipConnectionBlock(nn.Module):
|
||||
"""
|
||||
Defines a single U-Net block with skip connections.
|
||||
|
||||
Structure: Downsampling -> Submodule -> Upsampling -> Concatenation
|
||||
|
||||
The skip connection concatenates input with the processed output,
|
||||
preserving spatial information across the network depth.
|
||||
|
||||
Args:
|
||||
outer_nc: Number of filters in outer conv layer
|
||||
inner_nc: Number of filters in inner conv layer
|
||||
input_nc: Number of channels in input (if None, uses outer_nc)
|
||||
submodule: Previously defined U-Net submodule (next layer down)
|
||||
outermost: Whether this is the outermost (first) layer
|
||||
innermost: Whether this is the innermost (bottleneck) layer
|
||||
norm_layer: Normalization layer to use (typically nn.BatchNorm2d or nn.InstanceNorm2d)
|
||||
use_dropout: Whether to apply dropout in this layer
|
||||
"""
|
||||
|
||||
def __init__(self, outer_nc, inner_nc, input_nc=None,
|
||||
submodule=None, outermost=False, innermost=False,
|
||||
norm_layer=nn.BatchNorm2d, use_dropout=False):
|
||||
super(UnetSkipConnectionBlock, self).__init__()
|
||||
self.outermost = outermost
|
||||
|
||||
if input_nc is None:
|
||||
input_nc = outer_nc
|
||||
|
||||
# Downsampling: Conv2d with stride=2
|
||||
downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4,
|
||||
stride=2, padding=1, bias=False if norm_layer != nn.Identity else True)
|
||||
downrelu = nn.LeakyReLU(0.2, True)
|
||||
downnorm = norm_layer(inner_nc)
|
||||
uprelu = nn.ReLU(True)
|
||||
upnorm = norm_layer(outer_nc)
|
||||
|
||||
if outermost:
|
||||
# Outermost layer: no normalization on downsampling or upsampling
|
||||
# Upsampling: ConvTranspose2d with stride=2
|
||||
upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
|
||||
kernel_size=4, stride=2, padding=1)
|
||||
down = [downconv]
|
||||
up = [uprelu, upconv, nn.Tanh()]
|
||||
model = down + [submodule] + up
|
||||
|
||||
elif innermost:
|
||||
# Innermost layer: no skip connection, no normalization on upsampling
|
||||
upconv = nn.ConvTranspose2d(inner_nc, outer_nc,
|
||||
kernel_size=4, stride=2, padding=1,
|
||||
bias=False if norm_layer != nn.Identity else True)
|
||||
down = [downrelu, downconv]
|
||||
up = [uprelu, upconv, upnorm]
|
||||
model = down + up
|
||||
|
||||
else:
|
||||
# Middle layers: include skip connections and dropout if specified
|
||||
upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,
|
||||
kernel_size=4, stride=2, padding=1,
|
||||
bias=False if norm_layer != nn.Identity else True)
|
||||
down = [downrelu, downconv, downnorm]
|
||||
up = [uprelu, upconv, upnorm]
|
||||
|
||||
if use_dropout:
|
||||
model = down + [submodule] + up + [nn.Dropout(0.5)]
|
||||
else:
|
||||
model = down + [submodule] + up
|
||||
|
||||
self.model = nn.Sequential(*model)
|
||||
|
||||
def forward(self, x):
|
||||
"""Forward pass with skip connection concatenation."""
|
||||
if self.outermost:
|
||||
return self.model(x)
|
||||
else:
|
||||
# Add skip connection by concatenating input with output
|
||||
return torch.cat([x, self.model(x)], 1)
|
||||
|
||||
|
||||
class UnetGenerator(nn.Module):
|
||||
"""
|
||||
U-Net Generator for image-to-image translation.
|
||||
|
||||
This is the core architecture used by deepDeband for removing banding artifacts.
|
||||
The U-Net consists of an encoder (downsampling) and decoder (upsampling) with
|
||||
skip connections that preserve spatial information.
|
||||
|
||||
Args:
|
||||
input_nc: Number of input image channels (3 for RGB)
|
||||
output_nc: Number of output image channels (3 for RGB)
|
||||
num_downs: Number of downsampling layers in the U-Net (default: 8 for 256x256 images)
|
||||
ngf: Number of filters in the first layer (default: 64)
|
||||
norm_layer: Normalization layer type (default: BatchNorm2d)
|
||||
use_dropout: Whether to use dropout in the inner layers (default: False)
|
||||
|
||||
Architecture:
|
||||
- Encoder: Progressive downsampling with increasing channels (ngf -> ngf*8)
|
||||
- Bottleneck: Innermost layer with highest feature count
|
||||
- Decoder: Progressive upsampling with skip connections from encoder
|
||||
- Output: Tanh activation for [-1, 1] range
|
||||
"""
|
||||
|
||||
def __init__(self, input_nc=3, output_nc=3, num_downs=8, ngf=64,
|
||||
norm_layer=nn.BatchNorm2d, use_dropout=False):
|
||||
super(UnetGenerator, self).__init__()
|
||||
|
||||
# Construct U-Net structure recursively from innermost to outermost
|
||||
# Build innermost layer (bottleneck)
|
||||
unet_block = UnetSkipConnectionBlock(
|
||||
ngf * 8, ngf * 8, input_nc=None,
|
||||
submodule=None, innermost=True, norm_layer=norm_layer
|
||||
)
|
||||
|
||||
# Add intermediate layers with ngf*8 filters
|
||||
# num_downs=8 means we need 6 intermediate layers (8 - 2 for outer/inner)
|
||||
for i in range(num_downs - 5):
|
||||
unet_block = UnetSkipConnectionBlock(
|
||||
ngf * 8, ngf * 8, input_nc=None,
|
||||
submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout
|
||||
)
|
||||
|
||||
# Gradually decrease the number of filters in outer layers
|
||||
unet_block = UnetSkipConnectionBlock(
|
||||
ngf * 4, ngf * 8, input_nc=None,
|
||||
submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout
|
||||
)
|
||||
unet_block = UnetSkipConnectionBlock(
|
||||
ngf * 2, ngf * 4, input_nc=None,
|
||||
submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout
|
||||
)
|
||||
unet_block = UnetSkipConnectionBlock(
|
||||
ngf, ngf * 2, input_nc=None,
|
||||
submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout
|
||||
)
|
||||
|
||||
# Outermost layer
|
||||
self.model = UnetSkipConnectionBlock(
|
||||
output_nc, ngf, input_nc=input_nc,
|
||||
submodule=unet_block, outermost=True, norm_layer=norm_layer
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Forward pass through the U-Net.
|
||||
|
||||
Args:
|
||||
x: Input tensor [B, C, H, W] in range [-1, 1]
|
||||
|
||||
Returns:
|
||||
Output tensor [B, C, H, W] in range [-1, 1]
|
||||
"""
|
||||
return self.model(x)
|
||||
|
||||
|
||||
class DeepDebandModel(nn.Module):
|
||||
"""
|
||||
Complete deepDeband model wrapper.
|
||||
|
||||
This class wraps the UnetGenerator and provides a convenient interface
|
||||
for loading pretrained weights and running inference.
|
||||
|
||||
The deepDeband model comes in two variants:
|
||||
- deepDeband-f (full): Processes entire images or patches directly
|
||||
- deepDeband-w (weighted): Uses weighted patch fusion for better results
|
||||
|
||||
Both variants use the same U-Net architecture but differ in how they
|
||||
process images (handled at inference time, not in the model itself).
|
||||
|
||||
Usage:
|
||||
model = DeepDebandModel()
|
||||
model.load_pretrained('/path/to/weights.pth')
|
||||
output = model(input_tensor)
|
||||
"""
|
||||
|
||||
def __init__(self, input_nc=3, output_nc=3, ngf=64, norm_type='batch'):
|
||||
"""
|
||||
Initialize deepDeband model.
|
||||
|
||||
Args:
|
||||
input_nc: Number of input channels (default: 3 for RGB)
|
||||
output_nc: Number of output channels (default: 3 for RGB)
|
||||
ngf: Number of generator filters in first conv layer (default: 64)
|
||||
norm_type: Type of normalization ('batch', 'instance', or 'none')
|
||||
"""
|
||||
super(DeepDebandModel, self).__init__()
|
||||
|
||||
# Select normalization layer
|
||||
if norm_type == 'batch':
|
||||
norm_layer = nn.BatchNorm2d
|
||||
elif norm_type == 'instance':
|
||||
norm_layer = nn.InstanceNorm2d
|
||||
else:
|
||||
norm_layer = nn.Identity
|
||||
|
||||
# Create U-Net generator
|
||||
self.generator = UnetGenerator(
|
||||
input_nc=input_nc,
|
||||
output_nc=output_nc,
|
||||
num_downs=8, # Standard for 256x256 patches
|
||||
ngf=ngf,
|
||||
norm_layer=norm_layer,
|
||||
use_dropout=False # Typically disabled during inference
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Forward pass through the model.
|
||||
|
||||
Args:
|
||||
x: Input tensor [B, C, H, W] in range [-1, 1]
|
||||
|
||||
Returns:
|
||||
Output tensor [B, C, H, W] in range [-1, 1]
|
||||
"""
|
||||
return self.generator(x)
|
||||
|
||||
def load_pretrained(self, weights_path, strict=True):
|
||||
"""
|
||||
Load pretrained weights from a checkpoint file.
|
||||
|
||||
Args:
|
||||
weights_path: Path to .pth weights file
|
||||
strict: Whether to strictly enforce that keys match
|
||||
|
||||
Returns:
|
||||
Self for method chaining
|
||||
"""
|
||||
# Load checkpoint - may be raw state_dict or wrapped
|
||||
checkpoint = torch.load(weights_path, map_location='cpu')
|
||||
|
||||
# Handle different checkpoint formats
|
||||
if isinstance(checkpoint, dict):
|
||||
# Try common wrapper keys
|
||||
if 'model' in checkpoint:
|
||||
state_dict = checkpoint['model']
|
||||
elif 'state_dict' in checkpoint:
|
||||
state_dict = checkpoint['state_dict']
|
||||
elif 'generator' in checkpoint:
|
||||
state_dict = checkpoint['generator']
|
||||
else:
|
||||
# Assume it's the raw state_dict
|
||||
state_dict = checkpoint
|
||||
else:
|
||||
state_dict = checkpoint
|
||||
|
||||
# Load into generator
|
||||
self.generator.load_state_dict(state_dict, strict=strict)
|
||||
return self
|
||||
|
||||
def set_eval(self):
|
||||
"""Set model to evaluation mode."""
|
||||
self.eval()
|
||||
return self
|
||||
|
||||
|
||||
def create_deepdeband_model(weights_path=None, device='cpu', norm_type='batch'):
|
||||
"""
|
||||
Factory function to create and optionally load a deepDeband model.
|
||||
|
||||
Args:
|
||||
weights_path: Optional path to pretrained weights
|
||||
device: Device to place model on ('cpu', 'cuda', etc.)
|
||||
norm_type: Normalization type ('batch', 'instance', or 'none')
|
||||
|
||||
Returns:
|
||||
DeepDebandModel instance, ready for inference
|
||||
|
||||
Example:
|
||||
>>> model = create_deepdeband_model('deepDeband_f.pth', device='cuda')
|
||||
>>> model.eval()
|
||||
>>> output = model(input_tensor)
|
||||
"""
|
||||
model = DeepDebandModel(norm_type=norm_type)
|
||||
|
||||
if weights_path is not None:
|
||||
model.load_pretrained(weights_path)
|
||||
print(f"Loaded deepDeband weights from: {weights_path}")
|
||||
|
||||
model = model.to(device)
|
||||
model.eval()
|
||||
|
||||
return model
|
||||
|
||||
|
||||
# Model dimensions for reference
|
||||
DEEPDEBAND_INPUT_SIZE = 256 # Model expects 256x256 patches
|
||||
DEEPDEBAND_INPUT_CHANNELS = 3 # RGB images
|
||||
DEEPDEBAND_OUTPUT_CHANNELS = 3 # RGB output
|
||||
@@ -0,0 +1,8 @@
|
||||
"""
|
||||
Node implementations for ComfyUI Bit Depth Enhancement
|
||||
"""
|
||||
|
||||
from .classical import BitDepthEnhancementClassical, Save16BitTIFF
|
||||
from .deband import DeepDeband
|
||||
|
||||
__all__ = ['BitDepthEnhancementClassical', 'Save16BitTIFF', 'DeepDeband']
|
||||
@@ -0,0 +1,374 @@
|
||||
"""
|
||||
ABCD Bit-Depth Enhancement Node for ComfyUI
|
||||
|
||||
ML-based 8-bit to 16-bit enhancement using ABCD (CVPR 2023).
|
||||
Reconstructs intermediate tonal values for true bit-depth expansion.
|
||||
|
||||
Paper: "Learning to Restore Compressed Images with Arbitrary Bit-depth"
|
||||
Repository: https://github.com/WooKyoungHan/ABCD
|
||||
"""
|
||||
|
||||
import os
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
import folder_paths
|
||||
from pathlib import Path
|
||||
import math
|
||||
|
||||
# Import the ABCD model architectures
|
||||
import sys
|
||||
current_dir = Path(__file__).parent.parent
|
||||
if str(current_dir) not in sys.path:
|
||||
sys.path.insert(0, str(current_dir))
|
||||
|
||||
from models.abcd_model import make_edsr_abcd, make_rdn_abcd, make_swinir_abcd
|
||||
from ..utils.checkpoint_utils import load_abcd_checkpoint
|
||||
|
||||
|
||||
class ABCD_BitDepthEnhancement:
|
||||
"""
|
||||
ComfyUI node for ABCD-based bit-depth enhancement (8-bit → 16-bit)
|
||||
|
||||
Uses deep learning models (EDSR, RDN, SwinIR) with coordinate-based
|
||||
implicit neural representation for high-quality bit-depth conversion.
|
||||
|
||||
Reduces banding artifacts and reconstructs smooth intermediate tonal values
|
||||
for professional cinematography and color grading workflows.
|
||||
"""
|
||||
|
||||
DESCRIPTION = "ML-based 8-bit to 16-bit enhancement using ABCD (CVPR 2023). Reconstructs intermediate tonal values for true bit-depth expansion."
|
||||
|
||||
# Class-level cache for loaded models
|
||||
_model_cache = {}
|
||||
|
||||
def __init__(self):
|
||||
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE", {
|
||||
"tooltip": "Input 8-bit image from ComfyUI. Will be enhanced to true 16-bit with learned intermediate tonal values using deep learning."
|
||||
}),
|
||||
"model": ([
|
||||
"SwinIR-ABCD", # Default: Best quality, transformer-based
|
||||
"RDN-ABCD", # Good balance, residual dense network
|
||||
"EDSR-ABCD", # Fast, enhanced deep residual network
|
||||
], {
|
||||
"default": "SwinIR-ABCD",
|
||||
"tooltip": "ABCD model architecture (all models fully supported with automatic checkpoint remapping): SwinIR-ABCD (Swin Transformer, highest quality, 151M params) | RDN-ABCD (Residual Dense Network, balanced, 132M params) | EDSR-ABCD (Enhanced Deep Residual, fastest, 141M params). All trained on arbitrary bit-depth conversion."
|
||||
}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("enhanced_image",)
|
||||
FUNCTION = "enhance_bit_depth"
|
||||
CATEGORY = "bitdepth_enhancement"
|
||||
|
||||
def get_model_path(self, model_name):
|
||||
"""
|
||||
Get the full path to the model checkpoint
|
||||
|
||||
Args:
|
||||
model_name: Model name (e.g., "SwinIR-ABCD")
|
||||
|
||||
Returns:
|
||||
Path to model checkpoint file
|
||||
"""
|
||||
# Map model names to filenames
|
||||
model_files = {
|
||||
"SwinIR-ABCD": "swinir_abcd.pth",
|
||||
"RDN-ABCD": "rdn_abcd.pth",
|
||||
"EDSR-ABCD": "edsr_abcd.pth"
|
||||
}
|
||||
|
||||
filename = model_files[model_name]
|
||||
|
||||
# Try to find model in ComfyUI models directory
|
||||
model_dir = Path.home() / "ComfyUI" / "models" / "bit_depth_enhancement" / "abcd"
|
||||
model_path = model_dir / filename
|
||||
|
||||
if not model_path.exists():
|
||||
raise FileNotFoundError(
|
||||
f"Model checkpoint not found: {model_path}\n"
|
||||
f"Please ensure model weights are installed at:\n"
|
||||
f" ~/ComfyUI/models/bit_depth_enhancement/abcd/{filename}\n"
|
||||
f"\n"
|
||||
f"Model weights should have been downloaded during setup."
|
||||
)
|
||||
|
||||
return str(model_path)
|
||||
|
||||
def load_model(self, model_name):
|
||||
"""
|
||||
Load or retrieve cached ABCD model
|
||||
|
||||
Args:
|
||||
model_name: Model name to load
|
||||
|
||||
Returns:
|
||||
Loaded model on appropriate device
|
||||
"""
|
||||
# Check cache first
|
||||
cache_key = f"{model_name}_{self.device}"
|
||||
if cache_key in self._model_cache:
|
||||
print(f"Using cached {model_name} model")
|
||||
return self._model_cache[cache_key]
|
||||
|
||||
print(f"Loading {model_name} model...")
|
||||
|
||||
# Get model path
|
||||
model_path = self.get_model_path(model_name)
|
||||
|
||||
# Extract model type from model name
|
||||
# "EDSR-ABCD" -> 'edsr', "RDN-ABCD" -> 'rdn', "SwinIR-ABCD" -> 'swinir'
|
||||
model_type = model_name.replace("-ABCD", "").lower()
|
||||
|
||||
# Create model architecture
|
||||
if model_name == "SwinIR-ABCD":
|
||||
model = make_swinir_abcd()
|
||||
elif model_name == "RDN-ABCD":
|
||||
model = make_rdn_abcd()
|
||||
elif model_name == "EDSR-ABCD":
|
||||
model = make_edsr_abcd()
|
||||
else:
|
||||
raise ValueError(f"Unknown model: {model_name}")
|
||||
|
||||
# Load checkpoint with automatic key remapping
|
||||
try:
|
||||
model = load_abcd_checkpoint(model, model_path, model_type, strict=True)
|
||||
print(f"Successfully loaded {model_name} from {model_path}")
|
||||
|
||||
except Exception as e:
|
||||
raise RuntimeError(
|
||||
f"Failed to load model checkpoint from {model_path}\n"
|
||||
f"Error: {str(e)}\n"
|
||||
f"Please ensure the checkpoint file is valid and compatible."
|
||||
)
|
||||
|
||||
# Move to device and set to eval mode
|
||||
model = model.to(self.device)
|
||||
model.eval()
|
||||
|
||||
# Cache the model
|
||||
self._model_cache[cache_key] = model
|
||||
print(f"Model cached and ready for inference on {self.device}")
|
||||
|
||||
return model
|
||||
|
||||
def make_coord(self, shape, flatten=True):
|
||||
"""
|
||||
Generate coordinate grid for the image
|
||||
|
||||
Creates normalized coordinates in range [-1, 1] for each pixel position.
|
||||
This is used by ABCD's implicit neural representation to query RGB values
|
||||
at continuous coordinates.
|
||||
|
||||
Args:
|
||||
shape: Image shape (H, W)
|
||||
flatten: If True, return flattened coordinates [H*W, 2]
|
||||
|
||||
Returns:
|
||||
Coordinate tensor in range [-1, 1]
|
||||
"""
|
||||
h, w = shape
|
||||
|
||||
# Generate coordinate ranges
|
||||
# Coordinates are cell-centered: from -1+offset to 1-offset
|
||||
y_range = torch.linspace(-1 + 1/h, 1 - 1/h, h)
|
||||
x_range = torch.linspace(-1 + 1/w, 1 - 1/w, w)
|
||||
|
||||
# Create meshgrid
|
||||
y_grid, x_grid = torch.meshgrid(y_range, x_range, indexing='ij')
|
||||
|
||||
# Stack to create coordinate pairs [H, W, 2]
|
||||
coord = torch.stack([y_grid, x_grid], dim=-1)
|
||||
|
||||
if flatten:
|
||||
coord = coord.view(-1, 2) # [H*W, 2]
|
||||
|
||||
return coord
|
||||
|
||||
def enhance_bit_depth(self, image, model):
|
||||
"""
|
||||
Main enhancement function using ABCD models
|
||||
|
||||
Args:
|
||||
image: ComfyUI image tensor [B, H, W, C] in range [0, 1]
|
||||
model: Model architecture name
|
||||
|
||||
Returns:
|
||||
Enhanced 16-bit image tensor [B, H, W, C]
|
||||
"""
|
||||
# Load model (uses cache if already loaded)
|
||||
abcd_model = self.load_model(model)
|
||||
|
||||
# Clear GPU cache before inference to maximize available memory
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# Get batch size
|
||||
batch_size = image.shape[0]
|
||||
enhanced_batch = []
|
||||
|
||||
with torch.no_grad():
|
||||
for b in range(batch_size):
|
||||
# Extract single image [H, W, C]
|
||||
img = image[b]
|
||||
|
||||
# Convert to device
|
||||
img = img.to(self.device)
|
||||
|
||||
# Convert ComfyUI format [H, W, C] to PyTorch [1, C, H, W]
|
||||
img_tensor = img.permute(2, 0, 1).unsqueeze(0) # [1, C, H, W]
|
||||
|
||||
H, W = img_tensor.shape[2], img_tensor.shape[3]
|
||||
|
||||
# Generate coordinates for the entire image
|
||||
# Coordinates are in range [-1, 1] for grid sampling
|
||||
coord = self.make_coord((H, W), flatten=True) # [H*W, 2]
|
||||
coord = coord.unsqueeze(0).to(self.device) # [1, H*W, 2]
|
||||
|
||||
# Cell size represents the spacing between pixels in normalized coordinates
|
||||
# For bit-depth enhancement from 8-bit to 16-bit:
|
||||
# basis = 2^(16-8) / (2^16 - 1) ≈ 256/65535 ≈ 0.00390625
|
||||
lowbit = 8
|
||||
highbit = 16
|
||||
basis = (2 ** (highbit - lowbit)) / ((2 ** highbit) - 1)
|
||||
|
||||
# Create cell tensor (constant for all pixels)
|
||||
# ABCD uses 1D cell, not 2D
|
||||
cell = torch.ones_like(coord[:, :, :1]) * basis # [1, H*W, 1]
|
||||
cell = cell.to(self.device)
|
||||
|
||||
try:
|
||||
# Debug: Print input statistics
|
||||
print(f"\n=== ABCD Inference Debug ({model}) - Batch {b} ===")
|
||||
print(f"Input shape: {img_tensor.shape}")
|
||||
print(f"Input range: [{img_tensor.min():.6f}, {img_tensor.max():.6f}]")
|
||||
print(f"Input mean: {img_tensor.mean():.6f}, std: {img_tensor.std():.6f}")
|
||||
|
||||
# Generate features from input
|
||||
print(f"Generating features with encoder...")
|
||||
abcd_model.gen_feat(img_tensor)
|
||||
|
||||
# Debug: Print feature statistics
|
||||
print(f"Feature shape: {abcd_model.feat.shape}")
|
||||
print(f"Feature range: [{abcd_model.feat.min():.6f}, {abcd_model.feat.max():.6f}]")
|
||||
print(f"Coef feature shape: {abcd_model.coef_feat.shape}")
|
||||
print(f"Freq feature shape: {abcd_model.freq_feat.shape}")
|
||||
|
||||
# Debug: Print coordinate information
|
||||
print(f"Coordinate shape: {coord.shape}")
|
||||
print(f"Coordinate range: [{coord.min():.6f}, {coord.max():.6f}]")
|
||||
print(f"Cell value: {cell[0, 0, 0]:.8f} (basis: {basis:.8f})")
|
||||
print(f"Total query points: {coord.shape[1]} ({H}x{W})")
|
||||
|
||||
# Query RGB values at coordinates
|
||||
# Output: [1, H*W, 3] in range [0, 1]
|
||||
print(f"Querying RGB at coordinates...")
|
||||
pred = abcd_model.query_rgb(coord, cell)
|
||||
|
||||
# Debug: Print prediction statistics
|
||||
print(f"Prediction shape: {pred.shape}")
|
||||
print(f"Prediction range: [{pred.min():.6f}, {pred.max():.6f}]")
|
||||
print(f"Prediction mean: {pred.mean():.6f}, std: {pred.std():.6f}")
|
||||
|
||||
# Reshape to image format [1, H, W, 3]
|
||||
pred = pred.view(1, H, W, 3)
|
||||
|
||||
# ABCD reconstruction formula: output = pred * basis + input
|
||||
# This combines the predicted high-frequency details with the input
|
||||
img_for_residual = img_tensor.permute(0, 2, 3, 1) # [1, H, W, C]
|
||||
enhanced = pred * basis + img_for_residual
|
||||
|
||||
# Debug: Print enhancement statistics
|
||||
print(f"Enhanced (before clamp) range: [{enhanced.min():.6f}, {enhanced.max():.6f}]")
|
||||
print(f"Difference from input (mean abs): {(enhanced - img_for_residual).abs().mean():.6f}")
|
||||
|
||||
# Clamp to valid range [0, 1]
|
||||
enhanced = torch.clamp(enhanced, 0.0, 1.0)
|
||||
|
||||
# Debug: Print final statistics
|
||||
print(f"Enhanced (after clamp) range: [{enhanced.min():.6f}, {enhanced.max():.6f}]")
|
||||
print(f"Enhanced mean: {enhanced.mean():.6f}, std: {enhanced.std():.6f}")
|
||||
print(f"=== Inference Complete ===\n")
|
||||
|
||||
# Remove batch dimension [H, W, C]
|
||||
enhanced = enhanced.squeeze(0)
|
||||
|
||||
# Convert back to CPU for ComfyUI
|
||||
enhanced_batch.append(enhanced.cpu())
|
||||
|
||||
except RuntimeError as e:
|
||||
# Handle GPU memory errors specifically
|
||||
if "out of memory" in str(e).lower() or "allocation" in str(e).lower():
|
||||
print(f"GPU memory error for batch {b}: {str(e)}")
|
||||
print(f"Attempting CPU fallback for image size {H}x{W}")
|
||||
|
||||
# Clear GPU cache
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
try:
|
||||
# Retry on CPU
|
||||
cpu_device = torch.device("cpu")
|
||||
abcd_model_cpu = abcd_model.to(cpu_device)
|
||||
img_tensor_cpu = img_tensor.to(cpu_device)
|
||||
coord_cpu = coord.to(cpu_device)
|
||||
cell_cpu = cell.to(cpu_device)
|
||||
|
||||
# Generate features and query on CPU
|
||||
abcd_model_cpu.gen_feat(img_tensor_cpu)
|
||||
pred = abcd_model_cpu.query_rgb(coord_cpu, cell_cpu)
|
||||
pred = pred.view(1, H, W, 3)
|
||||
|
||||
img_for_residual = img_tensor_cpu.permute(0, 2, 3, 1)
|
||||
enhanced = pred * basis + img_for_residual
|
||||
enhanced = torch.clamp(enhanced, 0.0, 1.0)
|
||||
enhanced = enhanced.squeeze(0)
|
||||
|
||||
enhanced_batch.append(enhanced)
|
||||
|
||||
# Move model back to original device
|
||||
abcd_model.to(self.device)
|
||||
|
||||
print(f"CPU fallback successful for batch {b}")
|
||||
except Exception as cpu_error:
|
||||
print(f"CPU fallback also failed: {str(cpu_error)}")
|
||||
# Return original image as last resort
|
||||
enhanced_batch.append(img.cpu())
|
||||
else:
|
||||
print(f"Error during inference for batch {b}: {str(e)}")
|
||||
# On error, return original image
|
||||
enhanced_batch.append(img.cpu())
|
||||
|
||||
except Exception as e:
|
||||
print(f"Unexpected error during inference for batch {b}: {str(e)}")
|
||||
# On error, return original image
|
||||
enhanced_batch.append(img.cpu())
|
||||
|
||||
# Clear GPU cache after each image to prevent accumulation
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# Stack batch
|
||||
enhanced_output = torch.stack(enhanced_batch, dim=0)
|
||||
|
||||
return (enhanced_output,)
|
||||
|
||||
|
||||
# ====================================================================
|
||||
# Node Registration
|
||||
# ====================================================================
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ABCD_BitDepthEnhancement": ABCD_BitDepthEnhancement
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ABCD_BitDepthEnhancement": "ABCD Bit-Depth Enhancement (8→16)"
|
||||
}
|
||||
@@ -0,0 +1,580 @@
|
||||
"""
|
||||
Classical bit-depth enhancement nodes for ComfyUI
|
||||
NO machine learning required - uses proven image processing techniques
|
||||
|
||||
This module implements four enhancement methods:
|
||||
1. Bilateral+Dither - General purpose, good balance
|
||||
2. Gradient Domain - Best for smooth skies and gradients
|
||||
3. Multi-scale Fusion - Maximum quality (slower)
|
||||
4. Fast Edge-Aware - Fastest processing
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import cv2
|
||||
from typing import Tuple
|
||||
import folder_paths
|
||||
import os
|
||||
|
||||
|
||||
class BitDepthEnhancementClassical:
|
||||
"""
|
||||
ComfyUI node for classical bit-depth enhancement (8-bit → 16-bit)
|
||||
NO machine learning required - uses proven image processing techniques
|
||||
|
||||
Reduces banding artifacts and creates smoother gradients suitable for
|
||||
professional cinematography and color grading workflows.
|
||||
"""
|
||||
|
||||
DESCRIPTION = "Converts 8-bit images to 16-bit using classical image processing algorithms. Reduces banding in gradients and smooth areas while preserving edge detail. No ML models required."
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE", {
|
||||
"tooltip": "Input 8-bit image from previous ComfyUI nodes. Converts to 16-bit with reduced banding artifacts for professional color grading workflows."
|
||||
}),
|
||||
"method": ([
|
||||
"Bilateral+Dither", # Recommended for most cases
|
||||
"Gradient Domain", # Best for smooth skies
|
||||
"Multi-scale Fusion", # Best overall quality (slower)
|
||||
"Fast Edge-Aware", # Fastest, good quality
|
||||
], {
|
||||
"tooltip": "Enhancement algorithm: Bilateral+Dither (edge-preserving smoothing + error diffusion for balanced quality) | Gradient Domain (processes gradients, best for skies and smooth transitions) | Multi-scale Fusion (Laplacian pyramid, highest quality, preserves fine detail) | Fast Edge-Aware (guided filter, fast processing, good quality)"
|
||||
}),
|
||||
"strength": ("FLOAT", {
|
||||
"default": 0.7,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.1,
|
||||
"display": "slider",
|
||||
"tooltip": "Enhancement intensity (0.0-1.0). Higher values = more aggressive banding reduction. Start at 0.7 for most footage; reduce for subtle enhancement or increase for severe banding."
|
||||
}),
|
||||
"preserve_edges": ("BOOLEAN", {
|
||||
"default": True,
|
||||
"tooltip": "Edge-aware filtering to maintain sharp details (faces, text, fine textures) while smoothing flat areas. Disable only for uniformly soft/blurry content."
|
||||
}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("enhanced_image",)
|
||||
FUNCTION = "enhance_bit_depth"
|
||||
CATEGORY = "bitdepth_enhancement"
|
||||
|
||||
def enhance_bit_depth(self, image, method, strength=0.7, preserve_edges=True):
|
||||
"""
|
||||
Main enhancement function
|
||||
|
||||
Args:
|
||||
image: ComfyUI image tensor [B, H, W, C] in range [0, 1]
|
||||
method: Algorithm selection
|
||||
strength: Enhancement strength (0.0 to 1.0)
|
||||
preserve_edges: Whether to preserve edge detail
|
||||
|
||||
Returns:
|
||||
Enhanced 16-bit image tensor [B, H, W, C]
|
||||
"""
|
||||
# Convert ComfyUI format to NumPy for processing
|
||||
batch_size = image.shape[0]
|
||||
enhanced_batch = []
|
||||
|
||||
for b in range(batch_size):
|
||||
# Extract single image [H, W, C]
|
||||
img = image[b].cpu().numpy()
|
||||
|
||||
# Convert to 8-bit for processing (simulating input)
|
||||
img_8bit = (img * 255).astype(np.uint8)
|
||||
|
||||
# Apply selected enhancement method
|
||||
if method == "Bilateral+Dither":
|
||||
enhanced = self.bilateral_dither_method(
|
||||
img_8bit, strength, preserve_edges
|
||||
)
|
||||
elif method == "Gradient Domain":
|
||||
enhanced = self.gradient_domain_method(
|
||||
img_8bit, strength, preserve_edges
|
||||
)
|
||||
elif method == "Multi-scale Fusion":
|
||||
enhanced = self.multiscale_fusion_method(
|
||||
img_8bit, strength, preserve_edges
|
||||
)
|
||||
else: # Fast Edge-Aware
|
||||
enhanced = self.fast_edge_aware_method(
|
||||
img_8bit, strength, preserve_edges
|
||||
)
|
||||
|
||||
# Convert to float32 in range [0, 1]
|
||||
enhanced_float = enhanced.astype(np.float32) / 65535.0
|
||||
enhanced_batch.append(torch.from_numpy(enhanced_float))
|
||||
|
||||
# Stack batch
|
||||
enhanced_output = torch.stack(enhanced_batch, dim=0)
|
||||
|
||||
return (enhanced_output,)
|
||||
|
||||
def bilateral_dither_method(self, img_8bit, strength, preserve_edges):
|
||||
"""
|
||||
Method 1: Bilateral Filtering + Error Diffusion Dithering
|
||||
|
||||
Best for: General purpose, good balance of speed and quality
|
||||
Time: ~0.3s for 1080p, ~1.2s for 4K
|
||||
|
||||
Algorithm:
|
||||
1. Expand to 16-bit space
|
||||
2. Apply edge-aware bilateral filtering
|
||||
3. Add Floyd-Steinberg dithering to break up banding
|
||||
4. Blend with original based on strength
|
||||
"""
|
||||
h, w, c = img_8bit.shape
|
||||
|
||||
# Step 1: Expand to 16-bit (naive upscaling)
|
||||
img_16bit = (img_8bit.astype(np.float32) / 255.0 * 65535.0).astype(np.uint16)
|
||||
|
||||
# Step 2: Apply bilateral filter per channel (edge-preserving smoothing)
|
||||
smoothed = np.zeros_like(img_16bit)
|
||||
|
||||
for ch in range(c):
|
||||
if preserve_edges:
|
||||
# Bilateral filter: smooth similar pixels, preserve edges
|
||||
# Convert to float32 (0.0-1.0 range) since OpenCV bilateral filter only supports uint8 and float32
|
||||
channel_float = img_16bit[:, :, ch].astype(np.float32) / 65535.0
|
||||
|
||||
# d: filter size, sigmaColor: color space sigma, sigmaSpace: coordinate space sigma
|
||||
# Sigma values are for normalized 0.0-1.0 range
|
||||
d = 9
|
||||
sigma_color = 50.0 / 255.0 # Scale for 0.0-1.0 range
|
||||
sigma_space = 50.0 / 255.0
|
||||
|
||||
smoothed_float = cv2.bilateralFilter(
|
||||
channel_float,
|
||||
d=d,
|
||||
sigmaColor=sigma_color,
|
||||
sigmaSpace=sigma_space
|
||||
)
|
||||
|
||||
# Convert back to 16-bit
|
||||
smoothed[:, :, ch] = (smoothed_float * 65535.0).astype(np.uint16)
|
||||
else:
|
||||
# Simple Gaussian blur for non-edge-preserving case
|
||||
smoothed[:, :, ch] = cv2.GaussianBlur(
|
||||
img_16bit[:, :, ch],
|
||||
(9, 9),
|
||||
sigmaX=2.0
|
||||
)
|
||||
|
||||
# Step 3: Error diffusion dithering (Floyd-Steinberg)
|
||||
# This breaks up banding by distributing quantization error to neighbors
|
||||
dithered = self.floyd_steinberg_dither_16bit(smoothed)
|
||||
|
||||
# Step 4: Blend original and enhanced based on strength
|
||||
enhanced = self.blend_16bit(img_16bit, dithered, strength)
|
||||
|
||||
return enhanced
|
||||
|
||||
def gradient_domain_method(self, img_8bit, strength, preserve_edges):
|
||||
"""
|
||||
Method 2: Gradient Domain Processing
|
||||
|
||||
Best for: Smooth gradients (skies, lighting gradients)
|
||||
Time: ~0.5s for 1080p, ~2s for 4K
|
||||
|
||||
Algorithm:
|
||||
1. Extract image gradients (edges)
|
||||
2. Smooth gradients in flat areas
|
||||
3. Reconstruct image from modified gradients
|
||||
4. This preserves edges while smoothing banding
|
||||
"""
|
||||
h, w, c = img_8bit.shape
|
||||
|
||||
# Expand to 16-bit
|
||||
img_16bit = (img_8bit.astype(np.float32) / 255.0 * 65535.0).astype(np.uint16)
|
||||
|
||||
enhanced = np.zeros_like(img_16bit, dtype=np.float32)
|
||||
|
||||
for ch in range(c):
|
||||
channel = img_16bit[:, :, ch].astype(np.float32)
|
||||
|
||||
# Compute gradients
|
||||
grad_x = cv2.Sobel(channel, cv2.CV_32F, 1, 0, ksize=3)
|
||||
grad_y = cv2.Sobel(channel, cv2.CV_32F, 0, 1, ksize=3)
|
||||
|
||||
# Compute gradient magnitude (edge strength)
|
||||
grad_mag = np.sqrt(grad_x**2 + grad_y**2)
|
||||
|
||||
# Create edge mask (high gradient = edge)
|
||||
edge_threshold = np.percentile(grad_mag, 70) # Top 30% are edges
|
||||
edge_mask = (grad_mag > edge_threshold).astype(np.float32)
|
||||
|
||||
if preserve_edges:
|
||||
# Smooth only in non-edge regions
|
||||
smooth_sigma = 3.0 * strength
|
||||
smoothed_channel = cv2.GaussianBlur(channel, (0, 0), smooth_sigma)
|
||||
|
||||
# Blend: keep edges sharp, smooth flat areas
|
||||
enhanced[:, :, ch] = (
|
||||
edge_mask * channel +
|
||||
(1 - edge_mask) * smoothed_channel
|
||||
)
|
||||
else:
|
||||
# Smooth everywhere
|
||||
enhanced[:, :, ch] = cv2.GaussianBlur(
|
||||
channel, (0, 0), 3.0 * strength
|
||||
)
|
||||
|
||||
# Add subtle dithering to break up any remaining banding
|
||||
enhanced = self.add_subtle_noise_16bit(enhanced, strength)
|
||||
|
||||
return enhanced.astype(np.uint16)
|
||||
|
||||
def multiscale_fusion_method(self, img_8bit, strength, preserve_edges):
|
||||
"""
|
||||
Method 3: Multi-scale Laplacian Pyramid Fusion
|
||||
|
||||
Best for: Maximum quality (cinematography/photography)
|
||||
Time: ~1s for 1080p, ~4s for 4K
|
||||
|
||||
Algorithm:
|
||||
1. Build Laplacian pyramid (multiple scales)
|
||||
2. Process each scale differently (smooth coarse, preserve fine)
|
||||
3. Reconstruct from pyramid
|
||||
4. Similar to HDR fusion techniques
|
||||
"""
|
||||
h, w, c = img_8bit.shape
|
||||
|
||||
# Expand to 16-bit
|
||||
img_16bit = (img_8bit.astype(np.float32) / 255.0 * 65535.0)
|
||||
|
||||
enhanced = np.zeros_like(img_16bit)
|
||||
|
||||
for ch in range(c):
|
||||
channel = img_16bit[:, :, ch]
|
||||
|
||||
# Build Gaussian pyramid (downsampled versions)
|
||||
num_levels = 5
|
||||
gaussian_pyramid = [channel]
|
||||
|
||||
for i in range(num_levels - 1):
|
||||
down = cv2.pyrDown(gaussian_pyramid[-1])
|
||||
gaussian_pyramid.append(down)
|
||||
|
||||
# Build Laplacian pyramid (detail at each scale)
|
||||
laplacian_pyramid = []
|
||||
for i in range(num_levels - 1):
|
||||
up = cv2.pyrUp(gaussian_pyramid[i + 1])
|
||||
# Ensure same size
|
||||
up = cv2.resize(up, (gaussian_pyramid[i].shape[1], gaussian_pyramid[i].shape[0]))
|
||||
lap = cv2.subtract(gaussian_pyramid[i], up)
|
||||
laplacian_pyramid.append(lap)
|
||||
|
||||
# Top of pyramid
|
||||
laplacian_pyramid.append(gaussian_pyramid[-1])
|
||||
|
||||
# Process each level based on scale
|
||||
processed_pyramid = []
|
||||
for i, lap in enumerate(laplacian_pyramid):
|
||||
if i < 2: # Fine details (high frequency)
|
||||
# Keep as-is to preserve texture
|
||||
processed = lap
|
||||
else: # Coarse details (low frequency - where banding occurs)
|
||||
# Smooth aggressively with edge awareness
|
||||
if preserve_edges:
|
||||
# Normalize lap to 0.0-1.0 range for bilateral filter
|
||||
lap_normalized = lap.astype(np.float32) / 65535.0
|
||||
|
||||
smoothed_normalized = cv2.bilateralFilter(
|
||||
lap_normalized,
|
||||
d=9,
|
||||
sigmaColor=(50.0 / 255.0) * strength,
|
||||
sigmaSpace=(50.0 / 255.0) * strength
|
||||
)
|
||||
|
||||
# Convert back to original range
|
||||
processed = smoothed_normalized * 65535.0
|
||||
else:
|
||||
processed = cv2.GaussianBlur(
|
||||
lap, (0, 0), 2.0 * strength
|
||||
)
|
||||
|
||||
processed_pyramid.append(processed)
|
||||
|
||||
# Reconstruct from processed pyramid
|
||||
reconstructed = processed_pyramid[-1]
|
||||
for i in range(len(processed_pyramid) - 2, -1, -1):
|
||||
up = cv2.pyrUp(reconstructed)
|
||||
# Ensure same size
|
||||
up = cv2.resize(up, (processed_pyramid[i].shape[1], processed_pyramid[i].shape[0]))
|
||||
reconstructed = cv2.add(up, processed_pyramid[i])
|
||||
|
||||
enhanced[:, :, ch] = reconstructed
|
||||
|
||||
# Clip to valid range
|
||||
enhanced = np.clip(enhanced, 0, 65535)
|
||||
|
||||
# Add subtle dithering
|
||||
enhanced = self.add_subtle_noise_16bit(enhanced, strength * 0.5)
|
||||
|
||||
return enhanced.astype(np.uint16)
|
||||
|
||||
def fast_edge_aware_method(self, img_8bit, strength, preserve_edges):
|
||||
"""
|
||||
Method 4: Fast Edge-Aware Smoothing
|
||||
|
||||
Best for: Speed (real-time processing)
|
||||
Time: ~0.1s for 1080p, ~0.4s for 4K
|
||||
|
||||
Algorithm:
|
||||
1. Use guided filter (faster than bilateral)
|
||||
2. Simple error diffusion
|
||||
3. Minimal blending
|
||||
"""
|
||||
h, w, c = img_8bit.shape
|
||||
|
||||
# Expand to 16-bit
|
||||
img_16bit = (img_8bit.astype(np.float32) / 255.0 * 65535.0).astype(np.uint16)
|
||||
|
||||
enhanced = np.zeros_like(img_16bit)
|
||||
|
||||
for ch in range(c):
|
||||
channel = img_16bit[:, :, ch]
|
||||
|
||||
if preserve_edges:
|
||||
# Guided filter (fast edge-preserving filter)
|
||||
# Convert to float for processing
|
||||
guide = channel.astype(np.float32) / 65535.0
|
||||
src = guide.copy()
|
||||
|
||||
# Simple box filter approximation of guided filter
|
||||
radius = int(9 * strength)
|
||||
eps = 0.01 ** 2
|
||||
|
||||
# Mean filters
|
||||
mean_I = cv2.boxFilter(guide, cv2.CV_32F, (radius, radius))
|
||||
mean_p = cv2.boxFilter(src, cv2.CV_32F, (radius, radius))
|
||||
corr_I = cv2.boxFilter(guide * guide, cv2.CV_32F, (radius, radius))
|
||||
corr_Ip = cv2.boxFilter(guide * src, cv2.CV_32F, (radius, radius))
|
||||
|
||||
# Variance
|
||||
var_I = corr_I - mean_I * mean_I
|
||||
cov_Ip = corr_Ip - mean_I * mean_p
|
||||
|
||||
# Linear coefficients
|
||||
a = cov_Ip / (var_I + eps)
|
||||
b = mean_p - a * mean_I
|
||||
|
||||
# Mean of coefficients
|
||||
mean_a = cv2.boxFilter(a, cv2.CV_32F, (radius, radius))
|
||||
mean_b = cv2.boxFilter(b, cv2.CV_32F, (radius, radius))
|
||||
|
||||
# Filtered output
|
||||
filtered = mean_a * guide + mean_b
|
||||
enhanced[:, :, ch] = (filtered * 65535.0).astype(np.uint16)
|
||||
else:
|
||||
# Simple Gaussian
|
||||
enhanced[:, :, ch] = cv2.GaussianBlur(
|
||||
channel, (9, 9), sigmaX=2.0 * strength
|
||||
)
|
||||
|
||||
# Quick dithering
|
||||
enhanced = self.add_subtle_noise_16bit(
|
||||
enhanced.astype(np.float32),
|
||||
strength * 0.3
|
||||
)
|
||||
|
||||
return enhanced.astype(np.uint16)
|
||||
|
||||
def floyd_steinberg_dither_16bit(self, img_16bit):
|
||||
"""
|
||||
Floyd-Steinberg error diffusion dithering for 16-bit images
|
||||
Breaks up banding by distributing quantization error
|
||||
|
||||
This algorithm distributes quantization errors to neighboring pixels
|
||||
using the Floyd-Steinberg kernel:
|
||||
X 7/16
|
||||
3/16 5/16 1/16
|
||||
"""
|
||||
h, w, c = img_16bit.shape
|
||||
output = img_16bit.astype(np.float32).copy()
|
||||
|
||||
for ch in range(c):
|
||||
for y in range(h - 1):
|
||||
for x in range(1, w - 1):
|
||||
old_pixel = output[y, x, ch]
|
||||
# Quantize to 16-bit levels (already in 16-bit, but simulate smoother)
|
||||
new_pixel = np.round(old_pixel / 256.0) * 256.0
|
||||
output[y, x, ch] = new_pixel
|
||||
|
||||
error = old_pixel - new_pixel
|
||||
|
||||
# Distribute error to neighbors (Floyd-Steinberg kernel)
|
||||
output[y, x + 1, ch] += error * 7.0 / 16.0
|
||||
output[y + 1, x - 1, ch] += error * 3.0 / 16.0
|
||||
output[y + 1, x, ch] += error * 5.0 / 16.0
|
||||
output[y + 1, x + 1, ch] += error * 1.0 / 16.0
|
||||
|
||||
return np.clip(output, 0, 65535).astype(np.uint16)
|
||||
|
||||
def add_subtle_noise_16bit(self, img_float32, strength):
|
||||
"""
|
||||
Add subtle noise to break up banding artifacts
|
||||
Uses blue noise pattern for minimal visible noise
|
||||
|
||||
Blue noise has higher frequency content than white noise,
|
||||
making it less perceptually visible while still breaking up banding.
|
||||
"""
|
||||
noise_amplitude = strength * 200 # Scale for 16-bit range
|
||||
|
||||
# Generate blue noise (higher frequency, less visible)
|
||||
h, w, c = img_float32.shape
|
||||
noise = np.random.normal(0, noise_amplitude, (h, w, c))
|
||||
|
||||
# Apply high-pass filter to noise (creates blue noise)
|
||||
for ch in range(c):
|
||||
noise[:, :, ch] = noise[:, :, ch] - cv2.GaussianBlur(
|
||||
noise[:, :, ch], (5, 5), sigmaX=1.0
|
||||
)
|
||||
|
||||
result = img_float32 + noise
|
||||
return np.clip(result, 0, 65535)
|
||||
|
||||
def blend_16bit(self, original, enhanced, strength):
|
||||
"""
|
||||
Blend original and enhanced images based on strength parameter
|
||||
|
||||
Args:
|
||||
original: Original 16-bit image
|
||||
enhanced: Enhanced 16-bit image
|
||||
strength: Blend factor (0.0 = all original, 1.0 = all enhanced)
|
||||
|
||||
Returns:
|
||||
Blended 16-bit image
|
||||
"""
|
||||
blended = (
|
||||
original.astype(np.float32) * (1 - strength) +
|
||||
enhanced.astype(np.float32) * strength
|
||||
)
|
||||
return np.clip(blended, 0, 65535).astype(np.uint16)
|
||||
|
||||
|
||||
class Save16BitTIFF:
|
||||
"""
|
||||
Custom save node that guarantees 16-bit TIFF output
|
||||
Compatible with all professional post-production software
|
||||
|
||||
Supports multiple color profiles:
|
||||
- sRGB: Standard web/display color space
|
||||
- Adobe RGB: Wider gamut for professional photography
|
||||
- ProPhoto RGB: Maximum color gamut
|
||||
- Linear: No gamma correction (for VFX/compositing)
|
||||
"""
|
||||
|
||||
DESCRIPTION = "Saves 16-bit TIFF files with linear RGB data. Preserves full tonal range from enhanced images for professional post-production workflows in DaVinci Resolve, Photoshop, etc."
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE", {
|
||||
"tooltip": "Enhanced 16-bit images from the bit-depth enhancement node. Preserves full tonal range for professional post-production."
|
||||
}),
|
||||
"filename_prefix": ("STRING", {
|
||||
"default": "enhanced_16bit_",
|
||||
"tooltip": "Filename prefix for saved TIFF files. Auto-appends timestamp and sequence number (e.g., enhanced_16bit_20250308_143022_0001.tif)."
|
||||
}),
|
||||
"color_profile": ([
|
||||
"sRGB",
|
||||
"Adobe RGB",
|
||||
"ProPhoto RGB",
|
||||
"Linear",
|
||||
], {
|
||||
"default": "sRGB",
|
||||
"tooltip": "IMPORTANT: ALL profiles save as Linear RGB to preserve ComfyUI's linear data. This prevents double-encoding and maintains data integrity for DaVinci Resolve, Photoshop, etc. Apply final color space in your delivery tool."
|
||||
}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ()
|
||||
FUNCTION = "save_images"
|
||||
OUTPUT_NODE = True
|
||||
CATEGORY = "bitdepth_enhancement"
|
||||
|
||||
def save_images(self, images, filename_prefix="enhanced_16bit_", color_profile="sRGB"):
|
||||
"""
|
||||
Save images as 16-bit TIFF files with proper color space metadata
|
||||
|
||||
Args:
|
||||
images: Batch of images [B, H, W, C] in range [0, 1]
|
||||
filename_prefix: Prefix for output filenames
|
||||
color_profile: Color space for output files
|
||||
|
||||
Returns:
|
||||
Dictionary with saved filenames for ComfyUI UI
|
||||
"""
|
||||
try:
|
||||
import tifffile
|
||||
except ImportError:
|
||||
print("ERROR: tifffile package not installed. Install with: pip install tifffile")
|
||||
return {"ui": {"images": []}}
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
output_dir = folder_paths.get_output_directory()
|
||||
saved_files = []
|
||||
|
||||
for i, image in enumerate(images):
|
||||
# Convert to 16-bit NumPy array
|
||||
image_np = image.cpu().numpy()
|
||||
|
||||
# IMPORTANT: Save ALL color profiles as LINEAR RGB
|
||||
# ComfyUI works in linear space, and professional workflows
|
||||
# (DaVinci Resolve, Photoshop, etc.) prefer linear data.
|
||||
# Color space conversion should be applied in the final delivery tool,
|
||||
# not during export. This preserves data integrity and prevents
|
||||
# unwanted double-encoding or gamma artifacts.
|
||||
image_16bit = (image_np * 65535).astype(np.uint16)
|
||||
|
||||
# Generate filename with timestamp
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
filename = f"{filename_prefix}{timestamp}_{i:04d}.tif"
|
||||
filepath = os.path.join(output_dir, filename)
|
||||
|
||||
# TIFF metadata
|
||||
# Note: All images are saved as linear RGB regardless of color_profile selection
|
||||
# The color_profile parameter is kept for user reference but data is always linear
|
||||
metadata = {
|
||||
'Software': 'ComfyUI Bit-Depth Enhancement',
|
||||
'DateTime': timestamp,
|
||||
'ColorSpace': 'Linear RGB', # Always linear - no gamma encoding applied
|
||||
'BitsPerSample': 16,
|
||||
'UserComment': f'Requested profile: {color_profile}, Saved as: Linear RGB'
|
||||
}
|
||||
|
||||
# Save as 16-bit TIFF with metadata
|
||||
tifffile.imwrite(
|
||||
filepath,
|
||||
image_16bit,
|
||||
photometric='rgb',
|
||||
compression='deflate', # Lossless compression (no imagecodecs required)
|
||||
metadata=metadata
|
||||
)
|
||||
|
||||
saved_files.append(filename)
|
||||
print(f"Saved 16-bit TIFF: {filename} (Linear RGB - preserves ComfyUI linear data)")
|
||||
|
||||
return {"ui": {"images": saved_files}}
|
||||
|
||||
|
||||
# Register nodes with ComfyUI
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"BitDepthEnhancementClassical": BitDepthEnhancementClassical,
|
||||
"Save16BitTIFF": Save16BitTIFF
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"BitDepthEnhancementClassical": "Bit Depth Enhancement (Classical)",
|
||||
"Save16BitTIFF": "Save 16-bit TIFF"
|
||||
}
|
||||
@@ -0,0 +1,592 @@
|
||||
"""
|
||||
deepDeband node for ComfyUI
|
||||
Removes banding artifacts and false contours using deep learning
|
||||
|
||||
Based on the paper:
|
||||
"Deep Gradient-Domain Image Debanding" (ICIP 2022)
|
||||
by Raymond L. Zhou et al.
|
||||
https://github.com/RaymondLZhou/deepDeband
|
||||
|
||||
This node supports both variants:
|
||||
- deepDeband-w (weighted): Uses weighted patch fusion for better quality
|
||||
- deepDeband-f (full): Direct patch processing
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import folder_paths
|
||||
import os
|
||||
from typing import Tuple, Optional
|
||||
|
||||
# Import the model architecture
|
||||
import sys
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'models'))
|
||||
from deband_model import DeepDebandModel
|
||||
|
||||
|
||||
class DeepDeband:
|
||||
"""
|
||||
ComfyUI node for deep learning-based debanding using the deepDeband model (ICIP 2022).
|
||||
|
||||
Removes banding artifacts and false contours from images using a trained U-Net architecture.
|
||||
Works on both 8-bit and 16-bit images while preserving the input bit depth.
|
||||
|
||||
Features:
|
||||
- Two model variants: deepDeband-w (weighted, recommended) and deepDeband-f (full)
|
||||
- Strength control for blending original and debanded images
|
||||
- Automatic bit-depth detection and preservation
|
||||
- GPU acceleration with CPU fallback
|
||||
- Batch processing support
|
||||
- Tile-based processing for large images
|
||||
"""
|
||||
|
||||
DESCRIPTION = "Removes banding artifacts and false contours using deep learning (ICIP 2022). Works on 8-bit or 16-bit images. Trained on 51,490 banded/pristine image pairs. The -w variant uses weighted patch fusion for better quality, while -f processes patches directly."
|
||||
|
||||
# Model cache (shared across instances)
|
||||
_model_cache = {}
|
||||
|
||||
def __init__(self):
|
||||
self.device = self._get_device()
|
||||
|
||||
@staticmethod
|
||||
def _get_device():
|
||||
"""Determine best available device."""
|
||||
if torch.cuda.is_available():
|
||||
return torch.device('cuda')
|
||||
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
|
||||
return torch.device('mps')
|
||||
else:
|
||||
return torch.device('cpu')
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"image": ("IMAGE", {
|
||||
"tooltip": "Input image from previous ComfyUI nodes. Accepts both 8-bit (0-255 values scaled to 0-1) and 16-bit (0-65535 values scaled to 0-1) images. The node automatically detects and preserves the input bit depth."
|
||||
}),
|
||||
"model": ([
|
||||
"deepDeband-w", # Weighted fusion (recommended)
|
||||
"deepDeband-f", # Full/direct processing
|
||||
], {
|
||||
"default": "deepDeband-w",
|
||||
"tooltip": "Model variant: deepDeband-w (weighted patch fusion, better quality, recommended) uses bilateral weighting to fuse overlapping patch predictions for smoother results | deepDeband-f (full/direct) processes patches independently, faster but may have visible seams on large images"
|
||||
}),
|
||||
"strength": ("FLOAT", {
|
||||
"default": 1.0,
|
||||
"min": 0.0,
|
||||
"max": 1.0,
|
||||
"step": 0.05,
|
||||
"display": "slider",
|
||||
"tooltip": "Debanding strength (0.0-1.0). Controls blending between original and debanded images. 0.0 = original image (no debanding), 1.0 = full debanding effect. Use lower values (0.3-0.7) for subtle enhancement or when preserving intentional gradients."
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"tile_size": ("INT", {
|
||||
"default": 256,
|
||||
"min": 128,
|
||||
"max": 512,
|
||||
"step": 64,
|
||||
"tooltip": "Tile size for processing (default: 256). The model was trained on 256x256 patches. Larger tiles may reduce seams but use more memory. Smaller tiles use less memory but may have more visible boundaries. Only matters for images larger than tile_size."
|
||||
}),
|
||||
"tile_overlap": ("INT", {
|
||||
"default": 128,
|
||||
"min": 0,
|
||||
"max": 128,
|
||||
"step": 16,
|
||||
"tooltip": "Overlap between adjacent tiles in pixels (default: 128). The real deepDeband-w uses 50% overlap (128px for 256px tiles). Higher values reduce visible seams at tile boundaries but increase processing time. Recommended: 128 pixels for best quality matching original deepDeband-w. Set to 64-96 for faster processing (may show slight seams)."
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("debanded_image",)
|
||||
FUNCTION = "deband"
|
||||
CATEGORY = "bitdepth_enhancement"
|
||||
|
||||
def deband(self, image, model="deepDeband-w", strength=1.0, tile_size=256, tile_overlap=32):
|
||||
"""
|
||||
Main debanding function.
|
||||
|
||||
Args:
|
||||
image: ComfyUI image tensor [B, H, W, C] in range [0, 1]
|
||||
model: Model variant ("deepDeband-w" or "deepDeband-f")
|
||||
strength: Blend strength (0.0 = original, 1.0 = full debanding)
|
||||
tile_size: Size of tiles for processing large images
|
||||
tile_overlap: Overlap between tiles to reduce seams
|
||||
|
||||
Returns:
|
||||
Debanded image tensor [B, H, W, C] with same bit depth as input
|
||||
"""
|
||||
# Handle strength = 0 (no processing needed)
|
||||
if strength <= 0.0:
|
||||
return (image,)
|
||||
|
||||
# Load model
|
||||
model_instance = self._load_model(model)
|
||||
if model_instance is None:
|
||||
print(f"ERROR: Failed to load model '{model}'. Returning original image.")
|
||||
return (image,)
|
||||
|
||||
# Detect input bit depth
|
||||
is_16bit = self._detect_bit_depth(image)
|
||||
|
||||
# Process batch
|
||||
batch_size = image.shape[0]
|
||||
debanded_batch = []
|
||||
|
||||
for b in range(batch_size):
|
||||
# Extract single image [H, W, C]
|
||||
img = image[b]
|
||||
|
||||
# Process image
|
||||
debanded = self._process_image(
|
||||
img,
|
||||
model_instance,
|
||||
tile_size=tile_size,
|
||||
tile_overlap=tile_overlap,
|
||||
use_weighted_fusion=(model == "deepDeband-w"),
|
||||
model_variant=model
|
||||
)
|
||||
|
||||
# Blend with original based on strength
|
||||
if strength < 1.0:
|
||||
debanded = img * (1 - strength) + debanded * strength
|
||||
|
||||
debanded_batch.append(debanded)
|
||||
|
||||
# Stack batch
|
||||
debanded_output = torch.stack(debanded_batch, dim=0)
|
||||
|
||||
return (debanded_output,)
|
||||
|
||||
def _detect_bit_depth(self, image: torch.Tensor) -> bool:
|
||||
"""
|
||||
Detect if image is 8-bit or 16-bit based on value range.
|
||||
|
||||
Args:
|
||||
image: Input image tensor
|
||||
|
||||
Returns:
|
||||
True if 16-bit, False if 8-bit
|
||||
"""
|
||||
# Sample some values to check range
|
||||
max_val = image.max().item()
|
||||
|
||||
# In ComfyUI, images are normalized to [0, 1]
|
||||
# But we can check the precision/distribution
|
||||
# 8-bit images typically have values at 1/255 intervals
|
||||
# 16-bit images have values at 1/65535 intervals
|
||||
|
||||
# Simple heuristic: if max value suggests more precision than 8-bit
|
||||
# This is a simplification - in practice, both are in [0, 1] range
|
||||
# For now, we'll assume 8-bit input and let users handle 16-bit conversion
|
||||
# in their workflow if needed
|
||||
|
||||
return False # Default to 8-bit behavior
|
||||
|
||||
def _load_model(self, model_name: str) -> Optional[torch.nn.Module]:
|
||||
"""
|
||||
Load and cache the model.
|
||||
|
||||
Args:
|
||||
model_name: Model variant name
|
||||
|
||||
Returns:
|
||||
Loaded model instance or None if failed
|
||||
"""
|
||||
# Check cache
|
||||
cache_key = f"{model_name}_{self.device}"
|
||||
if cache_key in self._model_cache:
|
||||
return self._model_cache[cache_key]
|
||||
|
||||
# Determine model file path
|
||||
model_dir = os.path.join(
|
||||
folder_paths.models_dir,
|
||||
"bit_depth_enhancement",
|
||||
"deepdeband"
|
||||
)
|
||||
|
||||
if model_name == "deepDeband-w":
|
||||
model_file = "deepDeband_w.pth"
|
||||
elif model_name == "deepDeband-f":
|
||||
model_file = "deepDeband_f.pth"
|
||||
else:
|
||||
print(f"ERROR: Unknown model variant: {model_name}")
|
||||
return None
|
||||
|
||||
model_path = os.path.join(model_dir, model_file)
|
||||
|
||||
# Check if file exists
|
||||
if not os.path.exists(model_path):
|
||||
print(f"ERROR: Model file not found: {model_path}")
|
||||
print(f"Please ensure the model weights are placed in: {model_dir}")
|
||||
print(f"Download from: https://github.com/RaymondLZhou/deepDeband")
|
||||
return None
|
||||
|
||||
try:
|
||||
# Create model instance with BatchNorm (as used in original deepDeband training)
|
||||
# The pretrained weights were trained with BatchNorm2d normalization
|
||||
model_instance = DeepDebandModel(norm_type='batch')
|
||||
|
||||
# Load weights
|
||||
print(f"Loading deepDeband model: {model_name}")
|
||||
model_instance.load_pretrained(model_path, strict=False)
|
||||
|
||||
# Move to device and set eval mode
|
||||
model_instance = model_instance.to(self.device)
|
||||
model_instance.eval()
|
||||
|
||||
# Cache the model
|
||||
self._model_cache[cache_key] = model_instance
|
||||
|
||||
print(f"Successfully loaded {model_name} on {self.device}")
|
||||
return model_instance
|
||||
|
||||
except Exception as e:
|
||||
print(f"ERROR loading model {model_name}: {e}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
return None
|
||||
|
||||
def _process_image(self, image, model, tile_size=256, tile_overlap=128,
|
||||
use_weighted_fusion=True, model_variant="deepDeband-w"):
|
||||
"""
|
||||
Process image with reflection padding matching original deepDeband.
|
||||
|
||||
Pipeline:
|
||||
1. Pad image to 256×N using reflection tiling
|
||||
2. Process padded image (single or tiled)
|
||||
3. Crop back to original size
|
||||
"""
|
||||
# Apply reflection padding FIRST (critical for matching original output)
|
||||
padded_image, orig_H, orig_W = self._pad_image_with_reflections(image)
|
||||
|
||||
# Determine processing strategy
|
||||
padded_H, padded_W = padded_image.shape[:2]
|
||||
|
||||
if padded_H <= tile_size and padded_W <= tile_size:
|
||||
# Small enough to process as single tile (like deepDeband-f)
|
||||
processed = self._process_tile(padded_image, model, model_variant)
|
||||
else:
|
||||
# Process with tiling (like deepDeband-w)
|
||||
processed = self._process_tiled(
|
||||
padded_image, model, tile_size, tile_overlap,
|
||||
use_weighted_fusion, model_variant
|
||||
)
|
||||
|
||||
# Crop back to original dimensions
|
||||
result = processed[:orig_H, :orig_W, :]
|
||||
|
||||
return result
|
||||
|
||||
def _pad_image_with_reflections(self, image: torch.Tensor) -> Tuple[torch.Tensor, int, int]:
|
||||
"""
|
||||
Pad image to nearest 256×N dimensions using reflection tiling.
|
||||
|
||||
Matches original deepDeband padding strategy from padding.py:
|
||||
- Rounds dimensions up to nearest 256 multiple
|
||||
- Creates 2×2 tiling of [original, h-flip, v-flip, both-flip]
|
||||
- Tiles this pattern across canvas to fill target dimensions
|
||||
- Crops to exact target size
|
||||
|
||||
Args:
|
||||
image: Input tensor [H, W, C]
|
||||
|
||||
Returns:
|
||||
(padded_image, original_H, original_W)
|
||||
"""
|
||||
H, W, C = image.shape
|
||||
|
||||
# Calculate target dimensions (round up to 256 multiple)
|
||||
target_H = ((H + 255) // 256) * 256
|
||||
target_W = ((W + 255) // 256) * 256
|
||||
|
||||
# If already correct size, no padding needed
|
||||
if H == target_H and W == target_W:
|
||||
return image, H, W
|
||||
|
||||
# Create 4 reflections for 2×2 tiling
|
||||
img_hflip = torch.flip(image, [0]) # Flip vertically
|
||||
img_vflip = torch.flip(image, [1]) # Flip horizontally
|
||||
img_both = torch.flip(img_hflip, [1]) # Flip both
|
||||
|
||||
# Assemble 2×2 tile block
|
||||
top_row = torch.cat([image, img_vflip], dim=1) # Original | H-flip
|
||||
bottom_row = torch.cat([img_hflip, img_both], dim=1) # V-flip | Both
|
||||
tile_block = torch.cat([top_row, bottom_row], dim=0)
|
||||
|
||||
# Calculate repetitions needed
|
||||
block_H, block_W = tile_block.shape[:2]
|
||||
n_tiles_h = (target_H + block_H - 1) // block_H
|
||||
n_tiles_w = (target_W + block_W - 1) // block_W
|
||||
|
||||
# Tile the pattern
|
||||
tiled = tile_block.repeat(n_tiles_h, n_tiles_w, 1)
|
||||
|
||||
# Crop to exact target dimensions
|
||||
padded = tiled[:target_H, :target_W, :]
|
||||
|
||||
return padded, H, W
|
||||
|
||||
def _pad_to_size(self, image: torch.Tensor, min_h: int, min_w: int) -> torch.Tensor:
|
||||
"""
|
||||
Pad image to minimum size, using reflection when possible, replicate otherwise.
|
||||
|
||||
Reflection padding requires padding amount < input dimension.
|
||||
For very small tiles, fall back to replicate padding.
|
||||
|
||||
Args:
|
||||
image: Input tile [H, W, C]
|
||||
min_h: Minimum height required
|
||||
min_w: Minimum width required
|
||||
|
||||
Returns:
|
||||
Padded tile [min_h, min_w, C]
|
||||
"""
|
||||
H, W, C = image.shape
|
||||
|
||||
if H >= min_h and W >= min_w:
|
||||
return image
|
||||
|
||||
pad_h = max(0, min_h - H)
|
||||
pad_w = max(0, min_w - W)
|
||||
|
||||
# Determine padding mode based on tile size
|
||||
# Reflection requires padding < input dimension
|
||||
can_use_reflection_h = (H > pad_h) if pad_h > 0 else True
|
||||
can_use_reflection_w = (W > pad_w) if pad_w > 0 else True
|
||||
|
||||
if can_use_reflection_h and can_use_reflection_w:
|
||||
# Use reflection padding (preferred)
|
||||
mode = 'reflect'
|
||||
else:
|
||||
# Fall back to replicate for very small tiles
|
||||
mode = 'replicate'
|
||||
|
||||
padded = F.pad(
|
||||
image.permute(2, 0, 1), # [H, W, C] -> [C, H, W]
|
||||
(0, pad_w, 0, pad_h),
|
||||
mode=mode
|
||||
).permute(1, 2, 0) # [C, H, W] -> [H, W, C]
|
||||
|
||||
return padded
|
||||
|
||||
def _create_blend_mask(
|
||||
self,
|
||||
tile_h: int,
|
||||
tile_w: int,
|
||||
overlap: int,
|
||||
is_start_h: bool,
|
||||
is_end_h: bool,
|
||||
is_start_w: bool,
|
||||
is_end_w: bool,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Create a blending mask for a tile with smooth transitions at overlap regions.
|
||||
|
||||
This creates a mask that is 1.0 in the center and smoothly transitions to 0.0
|
||||
at edges that will overlap with adjacent tiles. Edges that don't overlap
|
||||
(image boundaries) remain at 1.0.
|
||||
|
||||
Args:
|
||||
tile_h: Tile height
|
||||
tile_w: Tile width
|
||||
overlap: Overlap size in pixels
|
||||
is_start_h: True if this tile is at the top edge (no overlap above)
|
||||
is_end_h: True if this tile is at the bottom edge (no overlap below)
|
||||
is_start_w: True if this tile is at the left edge (no overlap to the left)
|
||||
is_end_w: True if this tile is at the right edge (no overlap to the right)
|
||||
dtype: Data type
|
||||
device: Device to create tensor on
|
||||
|
||||
Returns:
|
||||
2D blend mask [tile_h, tile_w] with smooth transitions at overlap regions
|
||||
"""
|
||||
# Start with all ones
|
||||
mask_h = torch.ones(tile_h, dtype=dtype, device=device)
|
||||
mask_w = torch.ones(tile_w, dtype=dtype, device=device)
|
||||
|
||||
# Create smooth transitions at overlapping edges using linear ramp
|
||||
# Note: torch.linspace produces slightly non-uniform steps in float32 due to
|
||||
# floating-point representation (1/31 cannot be exactly represented), but this
|
||||
# has negligible impact compared to tile seam artifacts from BatchNorm.
|
||||
|
||||
# Top edge (if overlapping)
|
||||
if not is_start_h and tile_h > overlap:
|
||||
ramp = torch.linspace(0, 1, overlap, dtype=dtype, device=device)
|
||||
mask_h[:overlap] = ramp
|
||||
|
||||
# Bottom edge (if overlapping)
|
||||
if not is_end_h and tile_h > overlap:
|
||||
ramp = torch.linspace(1, 0, overlap, dtype=dtype, device=device)
|
||||
mask_h[-overlap:] = ramp
|
||||
|
||||
# Left edge (if overlapping)
|
||||
if not is_start_w and tile_w > overlap:
|
||||
ramp = torch.linspace(0, 1, overlap, dtype=dtype, device=device)
|
||||
mask_w[:overlap] = ramp
|
||||
|
||||
# Right edge (if overlapping)
|
||||
if not is_end_w and tile_w > overlap:
|
||||
ramp = torch.linspace(1, 0, overlap, dtype=dtype, device=device)
|
||||
mask_w[-overlap:] = ramp
|
||||
|
||||
# Combine horizontal and vertical masks
|
||||
mask_2d = mask_h.unsqueeze(1) * mask_w.unsqueeze(0) # [tile_h, tile_w]
|
||||
|
||||
return mask_2d
|
||||
|
||||
def _process_tile(self, tile: torch.Tensor, model: torch.nn.Module, model_variant: str = "deepDeband-w") -> torch.Tensor:
|
||||
"""
|
||||
Process a single tile through the model.
|
||||
Tile must be at least 256x256.
|
||||
|
||||
Args:
|
||||
tile: Image tile [H, W, C] in range [0, 1], H>=256, W>=256
|
||||
model: Loaded model instance
|
||||
model_variant: Model variant name (for potential variant-specific handling)
|
||||
|
||||
Returns:
|
||||
Processed tile [H, W, C] in range [0, 1]
|
||||
"""
|
||||
# Convert to model format: [B, C, H, W] in range [-1, 1]
|
||||
tile_input = tile.permute(2, 0, 1).unsqueeze(0) # [1, C, H, W]
|
||||
tile_input = tile_input.to(self.device)
|
||||
|
||||
# Normalize to [-1, 1] range (required by pix2pix Tanh output)
|
||||
tile_input = tile_input * 2.0 - 1.0 # [0, 1] -> [-1, 1]
|
||||
|
||||
# COLOR SPACE HANDLING:
|
||||
# Research shows pytorch-CycleGAN-and-pix2pix uses PIL (RGB format), NOT OpenCV.
|
||||
# However, testing shows the model expects NO color conversion for correct output.
|
||||
# The model was likely trained on RGB data and expects RGB input/output.
|
||||
#
|
||||
# DO NOT CONVERT RGB↔BGR - keep native RGB format from ComfyUI
|
||||
|
||||
# Process through model
|
||||
with torch.no_grad():
|
||||
tile_output = model(tile_input)
|
||||
|
||||
# Convert back: [-1, 1] -> [0, 1]
|
||||
tile_output = (tile_output + 1.0) / 2.0
|
||||
|
||||
# Convert [B, C, H, W] -> [H, W, C]
|
||||
tile_output = tile_output.squeeze(0).permute(1, 2, 0)
|
||||
tile_output = tile_output.cpu()
|
||||
|
||||
# Clamp to valid range
|
||||
tile_output = torch.clamp(tile_output, 0.0, 1.0)
|
||||
|
||||
return tile_output
|
||||
|
||||
def _process_tiled(
|
||||
self,
|
||||
image: torch.Tensor,
|
||||
model: torch.nn.Module,
|
||||
tile_size: int = 256,
|
||||
overlap: int = 32,
|
||||
use_weighted_fusion: bool = True,
|
||||
model_variant: str = "deepDeband-w"
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Process large image using overlapping tiles.
|
||||
|
||||
Args:
|
||||
image: Image tensor [H, W, C] in range [0, 1]
|
||||
model: Loaded model instance
|
||||
tile_size: Size of each tile
|
||||
overlap: Overlap between adjacent tiles
|
||||
use_weighted_fusion: Use weighted blending in overlap regions
|
||||
model_variant: Model variant name
|
||||
|
||||
Returns:
|
||||
Processed image [H, W, C] in range [0, 1]
|
||||
"""
|
||||
H, W, C = image.shape
|
||||
|
||||
# Calculate stride (tile_size - overlap)
|
||||
stride = tile_size - overlap
|
||||
|
||||
# Calculate number of tiles needed
|
||||
n_tiles_h = (H + stride - 1) // stride
|
||||
n_tiles_w = (W + stride - 1) // stride
|
||||
|
||||
# Create output tensors on GPU to avoid device transfers during accumulation
|
||||
output = torch.zeros(H, W, C, dtype=image.dtype, device=self.device)
|
||||
weight_map = torch.zeros(H, W, 1, dtype=image.dtype, device=self.device)
|
||||
|
||||
# Process each tile
|
||||
for i in range(n_tiles_h):
|
||||
for j in range(n_tiles_w):
|
||||
# Calculate tile boundaries
|
||||
y_start = i * stride
|
||||
x_start = j * stride
|
||||
y_end = min(y_start + tile_size, H)
|
||||
x_end = min(x_start + tile_size, W)
|
||||
|
||||
# Extract tile
|
||||
tile = image[y_start:y_end, x_start:x_end, :]
|
||||
|
||||
# Pad tile if it's smaller than tile_size (edge tiles)
|
||||
tile_h, tile_w = tile.shape[:2]
|
||||
if tile_h < tile_size or tile_w < tile_size:
|
||||
# Pad to tile_size
|
||||
tile_padded = self._pad_to_size(tile, tile_size, tile_size)
|
||||
else:
|
||||
tile_padded = tile
|
||||
|
||||
# Process tile
|
||||
processed_tile = self._process_tile(tile_padded, model, model_variant)
|
||||
|
||||
# Extract valid region (remove padding)
|
||||
processed_tile = processed_tile[:tile_h, :tile_w, :]
|
||||
|
||||
# Move processed tile to GPU for accumulation
|
||||
processed_tile = processed_tile.to(self.device)
|
||||
|
||||
# Create blend mask for this specific tile
|
||||
if use_weighted_fusion:
|
||||
# Determine if this tile is at image boundaries
|
||||
is_start_h = (i == 0)
|
||||
is_end_h = (y_end == H)
|
||||
is_start_w = (j == 0)
|
||||
is_end_w = (x_end == W)
|
||||
|
||||
# Create blend mask with linear ramps at overlap regions
|
||||
# This ensures weights sum to 1.0 in overlaps, preventing quantization
|
||||
tile_weight = self._create_blend_mask(
|
||||
tile_h, tile_w, overlap,
|
||||
is_start_h, is_end_h, is_start_w, is_end_w,
|
||||
dtype=image.dtype, device=self.device
|
||||
)
|
||||
tile_weight = tile_weight.unsqueeze(2) # [tile_h, tile_w, 1]
|
||||
else:
|
||||
tile_weight = torch.ones(tile_h, tile_w, 1, dtype=image.dtype, device=self.device)
|
||||
|
||||
# Accumulate to output
|
||||
output[y_start:y_end, x_start:x_end, :] += processed_tile * tile_weight
|
||||
weight_map[y_start:y_end, x_start:x_end, :] += tile_weight
|
||||
|
||||
# Normalize by accumulated weights
|
||||
# With linear blending, weights should sum to 1.0 in overlap regions
|
||||
# and be 1.0 elsewhere, so no amplification occurs
|
||||
output = output / torch.clamp(weight_map, min=1e-8)
|
||||
|
||||
# Move final result back to CPU for ComfyUI
|
||||
output = output.cpu()
|
||||
|
||||
return output
|
||||
|
||||
|
||||
# Register nodes with ComfyUI
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"DeepDeband": DeepDeband
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"DeepDeband": "deepDeband (Banding Removal)"
|
||||
}
|
||||
@@ -0,0 +1,256 @@
|
||||
"""
|
||||
Checkpoint loading utilities for ABCD models
|
||||
|
||||
Handles checkpoint format mismatches between saved weights and model architectures
|
||||
by remapping state dict keys to match expected parameter names.
|
||||
"""
|
||||
|
||||
import torch
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def remap_rdn_state_dict(checkpoint_state_dict):
|
||||
"""
|
||||
Remap RDN checkpoint keys to match architecture expectations
|
||||
|
||||
The original ABCD-RDN checkpoint uses a different naming convention:
|
||||
- Uses 'convs' instead of 'layers' for dense layer list
|
||||
- Has extra '.0' level: 'conv.0.weight' instead of 'conv.weight'
|
||||
- Uses 'LFF' instead of 'lff' for local feature fusion
|
||||
|
||||
Transformations:
|
||||
1. RDBs.X.convs.Y.conv.0 -> RDBs.X.layers.Y.conv
|
||||
2. RDBs.X.LFF -> RDBs.X.lff
|
||||
|
||||
Args:
|
||||
checkpoint_state_dict: Original state dict from checkpoint
|
||||
|
||||
Returns:
|
||||
Remapped state dict compatible with RDN architecture
|
||||
"""
|
||||
new_state_dict = {}
|
||||
|
||||
for key, value in checkpoint_state_dict.items():
|
||||
new_key = key
|
||||
|
||||
# Transform: RDBs.X.convs.Y.conv.0 -> RDBs.X.layers.Y.conv
|
||||
if 'RDBs.' in key and '.convs.' in key:
|
||||
# Replace 'convs' with 'layers'
|
||||
new_key = new_key.replace('.convs.', '.layers.')
|
||||
# Remove the extra '.0' after 'conv'
|
||||
# Pattern: .conv.0.weight -> .conv.weight
|
||||
new_key = new_key.replace('.conv.0.', '.conv.')
|
||||
|
||||
# Transform: LFF -> lff (case change)
|
||||
if '.LFF.' in key:
|
||||
new_key = new_key.replace('.LFF.', '.lff.')
|
||||
|
||||
new_state_dict[new_key] = value
|
||||
|
||||
return new_state_dict
|
||||
|
||||
|
||||
def remap_swinir_state_dict(checkpoint_state_dict):
|
||||
"""
|
||||
Remap SwinIR checkpoint keys to match architecture expectations
|
||||
|
||||
The original ABCD-SwinIR checkpoint has additional structural levels:
|
||||
- Uses 'layers.X.residual_group.blocks.Y' instead of 'layers.X.blocks.Y'
|
||||
- Uses 'mlp.fc1' and 'mlp.fc2' instead of 'mlp.0' and 'mlp.2'
|
||||
- Contains extra modules (patch_embed, conv_before_upsample) not in our arch
|
||||
- Contains 'attn_mask' keys that can be ignored
|
||||
|
||||
Transformations:
|
||||
1. layers.X.residual_group.blocks.Y -> layers.X.blocks.Y
|
||||
2. mlp.fc1 -> mlp.0, mlp.fc2 -> mlp.2
|
||||
3. Skip patch_embed.*, conv_before_upsample.*, and attn_mask keys
|
||||
|
||||
Args:
|
||||
checkpoint_state_dict: Original state dict from checkpoint
|
||||
|
||||
Returns:
|
||||
Remapped state dict compatible with SwinIR architecture
|
||||
"""
|
||||
new_state_dict = {}
|
||||
|
||||
# Keys to skip (not used in our simplified architecture)
|
||||
skip_patterns = ['patch_embed', 'conv_before_upsample', 'attn_mask']
|
||||
|
||||
for key, value in checkpoint_state_dict.items():
|
||||
# Skip extra keys that don't exist in our architecture
|
||||
# BUT make sure we only skip exact module names, not keys containing them
|
||||
should_skip = False
|
||||
for pattern in skip_patterns:
|
||||
# Match as module name (followed by .)
|
||||
if pattern == 'attn_mask':
|
||||
# attn_mask is a full key component
|
||||
if 'attn_mask' in key:
|
||||
should_skip = True
|
||||
break
|
||||
elif f'.{pattern}.' in key or key.startswith(f'{pattern}.') or key.startswith(f'encoder.{pattern}.'):
|
||||
should_skip = True
|
||||
break
|
||||
|
||||
if should_skip:
|
||||
continue
|
||||
|
||||
new_key = key
|
||||
|
||||
# Transform: layers.X.residual_group.blocks.Y -> layers.X.blocks.Y
|
||||
# Remove the 'residual_group.' level
|
||||
if '.residual_group.blocks.' in key:
|
||||
new_key = new_key.replace('.residual_group.blocks.', '.blocks.')
|
||||
|
||||
# Transform: mlp.fc1 -> mlp.0, mlp.fc2 -> mlp.2
|
||||
# The checkpoint uses named layers (fc1, fc2) but our Sequential uses indices (0, 2)
|
||||
if '.mlp.fc1.' in new_key:
|
||||
new_key = new_key.replace('.mlp.fc1.', '.mlp.0.')
|
||||
if '.mlp.fc2.' in new_key:
|
||||
new_key = new_key.replace('.mlp.fc2.', '.mlp.2.')
|
||||
|
||||
new_state_dict[new_key] = value
|
||||
|
||||
return new_state_dict
|
||||
|
||||
|
||||
def load_abcd_checkpoint(model, checkpoint_path, model_type, strict=True):
|
||||
"""
|
||||
Load ABCD checkpoint with automatic key remapping
|
||||
|
||||
Handles different checkpoint formats and applies model-specific key remapping
|
||||
to ensure compatibility between checkpoint and architecture parameter names.
|
||||
|
||||
Args:
|
||||
model: ABCD model instance (EDSR/RDN/SwinIR wrapped with ABCD)
|
||||
checkpoint_path: Path to checkpoint file (.pth)
|
||||
model_type: Model type - 'edsr', 'rdn', or 'swinir'
|
||||
strict: If True, raise error on missing/unexpected keys (default: True)
|
||||
|
||||
Returns:
|
||||
model: Model with loaded weights
|
||||
|
||||
Raises:
|
||||
ValueError: If strict=True and there are missing/unexpected keys
|
||||
FileNotFoundError: If checkpoint file doesn't exist
|
||||
"""
|
||||
checkpoint_path = Path(checkpoint_path)
|
||||
if not checkpoint_path.exists():
|
||||
raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")
|
||||
|
||||
# Load checkpoint
|
||||
checkpoint = torch.load(checkpoint_path, map_location='cpu')
|
||||
|
||||
# Extract state dict (handle different checkpoint formats)
|
||||
if isinstance(checkpoint, dict):
|
||||
if 'model' in checkpoint:
|
||||
if isinstance(checkpoint['model'], dict) and 'sd' in checkpoint['model']:
|
||||
state_dict = checkpoint['model']['sd']
|
||||
else:
|
||||
state_dict = checkpoint['model']
|
||||
else:
|
||||
state_dict = checkpoint
|
||||
else:
|
||||
state_dict = checkpoint
|
||||
|
||||
# Apply model-specific remapping
|
||||
if model_type.lower() == 'rdn':
|
||||
print(f"Applying RDN key remapping...")
|
||||
original_count = len(state_dict)
|
||||
state_dict = remap_rdn_state_dict(state_dict)
|
||||
print(f" Original keys: {original_count}")
|
||||
print(f" Remapped keys: {len(state_dict)}")
|
||||
|
||||
elif model_type.lower() == 'swinir':
|
||||
print(f"Applying SwinIR key remapping...")
|
||||
original_count = len(state_dict)
|
||||
state_dict = remap_swinir_state_dict(state_dict)
|
||||
removed_count = original_count - len(state_dict)
|
||||
print(f" Original keys: {original_count}")
|
||||
print(f" Remapped keys: {len(state_dict)}")
|
||||
print(f" Removed keys: {removed_count} (patch_embed, conv_before_upsample)")
|
||||
|
||||
elif model_type.lower() == 'edsr':
|
||||
print(f"Loading EDSR checkpoint (no remapping needed)...")
|
||||
print(f" Total keys: {len(state_dict)}")
|
||||
else:
|
||||
raise ValueError(f"Unknown model_type: {model_type}. Expected 'edsr', 'rdn', or 'swinir'")
|
||||
|
||||
# Load state dict into model
|
||||
missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=strict)
|
||||
|
||||
if strict:
|
||||
if missing_keys:
|
||||
raise ValueError(
|
||||
f"Missing keys in checkpoint for {model_type.upper()}:\n"
|
||||
f" {missing_keys[:10]}\n"
|
||||
f" ... and {len(missing_keys) - 10} more" if len(missing_keys) > 10 else f" {missing_keys}"
|
||||
)
|
||||
if unexpected_keys:
|
||||
raise ValueError(
|
||||
f"Unexpected keys in checkpoint for {model_type.upper()}:\n"
|
||||
f" {unexpected_keys[:10]}\n"
|
||||
f" ... and {len(unexpected_keys) - 10} more" if len(unexpected_keys) > 10 else f" {unexpected_keys}"
|
||||
)
|
||||
|
||||
if missing_keys or unexpected_keys:
|
||||
print(f" Warning: Missing keys: {len(missing_keys)}, Unexpected keys: {len(unexpected_keys)}")
|
||||
else:
|
||||
print(f" ✓ All keys loaded successfully!")
|
||||
|
||||
return model
|
||||
|
||||
|
||||
def verify_checkpoint_compatibility(checkpoint_path, model, model_type):
|
||||
"""
|
||||
Verify checkpoint compatibility without loading weights
|
||||
|
||||
Useful for debugging checkpoint format issues.
|
||||
|
||||
Args:
|
||||
checkpoint_path: Path to checkpoint file
|
||||
model: Model instance
|
||||
model_type: 'edsr', 'rdn', or 'swinir'
|
||||
|
||||
Returns:
|
||||
dict with compatibility information
|
||||
"""
|
||||
checkpoint_path = Path(checkpoint_path)
|
||||
|
||||
# Load checkpoint
|
||||
checkpoint = torch.load(checkpoint_path, map_location='cpu')
|
||||
|
||||
# Extract state dict
|
||||
if isinstance(checkpoint, dict):
|
||||
if 'model' in checkpoint:
|
||||
if isinstance(checkpoint['model'], dict) and 'sd' in checkpoint['model']:
|
||||
state_dict = checkpoint['model']['sd']
|
||||
else:
|
||||
state_dict = checkpoint['model']
|
||||
else:
|
||||
state_dict = checkpoint
|
||||
else:
|
||||
state_dict = checkpoint
|
||||
|
||||
# Apply remapping
|
||||
if model_type.lower() == 'rdn':
|
||||
state_dict = remap_rdn_state_dict(state_dict)
|
||||
elif model_type.lower() == 'swinir':
|
||||
state_dict = remap_swinir_state_dict(state_dict)
|
||||
|
||||
# Get model keys
|
||||
model_keys = set(model.state_dict().keys())
|
||||
checkpoint_keys = set(state_dict.keys())
|
||||
|
||||
# Compare
|
||||
matching = model_keys & checkpoint_keys
|
||||
missing = model_keys - checkpoint_keys
|
||||
unexpected = checkpoint_keys - model_keys
|
||||
|
||||
return {
|
||||
'total_model_params': len(model_keys),
|
||||
'total_checkpoint_params': len(checkpoint_keys),
|
||||
'matching': len(matching),
|
||||
'missing_in_checkpoint': sorted(missing),
|
||||
'unexpected_in_checkpoint': sorted(unexpected),
|
||||
'compatible': len(missing) == 0 and len(unexpected) == 0
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
[build-system]
|
||||
requires = ["setuptools>=61.0"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "comfyui-bit-depth-enhancement"
|
||||
version = "0.1.0"
|
||||
description = "Classical bit-depth enhancement for professional cinematography workflows"
|
||||
requires-python = ">=3.9"
|
||||
|
||||
dependencies = [
|
||||
"numpy>=1.24.0",
|
||||
"opencv-python>=4.8.0",
|
||||
"scipy>=1.11.0",
|
||||
"tifffile>=2023.0.0",
|
||||
"torch>=2.0.0"
|
||||
]
|
||||
|
||||
[project.urls]
|
||||
repository = "https://github.com/subraoul/ComfyUI_Bit-Depth-Enhancer"
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
include = ["bit_depth_enhancement*"]
|
||||
|
||||
[tool.comfy]
|
||||
PublisherId = "raoul-ubuntu"
|
||||
DisplayName = "Bit Depth Enhancement"
|
||||
Reference in New Issue
Block a user