models.py

"""
Deep Analog — Model Definitions (deployment build)
Architecture matches training exactly. Backbone weights come from the
checkpoint, so no ImageNet download is needed at startup.
"""

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.models as tv_models


# =========================================================================
#  Trilinear Interpolation
# =========================================================================

class TrilinearInterpolation(nn.Module):
    def forward(self, lut, x):
        B, C, H, W = x.shape
        x = torch.clamp(x, 0, 1)
        output = []
        for c in range(C):
            lut_c = lut[:, c:c+1, :, :, :]
            grid = x[:, :, :, :].permute(0, 2, 3, 1)
            grid_norm = grid * 2 - 1
            grid_reshaped = grid_norm.unsqueeze(1)
            sampled = F.grid_sample(lut_c, grid_reshaped, mode='bilinear',
                                    padding_mode='border', align_corners=False)
            sampled = sampled.squeeze(1).squeeze(1)
            output.append(sampled)
        return torch.stack(output, dim=1)


# =========================================================================
#  StyleLUTNet — Conditional 3D LUT prediction (open-set color transfer)
# =========================================================================

class StyleLUTNet(nn.Module):
    def __init__(self, lut_dim=33, lut_dim_low=17):
        super().__init__()
        self.lut_dim = lut_dim
        self.lut_dim_low = lut_dim_low

        resnet18 = tv_models.resnet18(weights=None)  # checkpoint supplies weights
        self.encoder = nn.Sequential(*list(resnet18.children())[:-1])
        self.feat_dim = 512

        n_values = lut_dim_low ** 3 * 3
        self.decoder = nn.Sequential(
            nn.Linear(self.feat_dim, 1024), nn.LayerNorm(1024), nn.GELU(), nn.Dropout(0.1),
            nn.Linear(1024, 2048), nn.LayerNorm(2048), nn.GELU(), nn.Dropout(0.1),
            nn.Linear(2048, 4096), nn.LayerNorm(4096), nn.GELU(),
            nn.Linear(4096, n_values),
        )

        self.register_buffer('residual_scale', torch.tensor(1.0))
        self.register_buffer('identity_lut', self._make_identity_lut(lut_dim_low))
        self.trilinear = TrilinearInterpolation()

    @staticmethod
    def _make_identity_lut(dim):
        coords = torch.linspace(0, 1, dim)
        lut = torch.zeros(3, dim, dim, dim)
        lut[0] = coords.view(dim, 1, 1).expand(dim, dim, dim)
        lut[1] = coords.view(1, dim, 1).expand(dim, dim, dim)
        lut[2] = coords.view(1, 1, dim).expand(dim, dim, dim)
        return lut

    def predict_lut(self, reference_224):
        B = reference_224.shape[0]
        feat = self.encoder(reference_224).flatten(1)
        residual = self.decoder(feat)
        d = self.lut_dim_low
        residual = residual.view(B, 3, d, d, d)
        lut_low = self.identity_lut.unsqueeze(0) + self.residual_scale * residual
        if self.lut_dim != self.lut_dim_low:
            lut = F.interpolate(
                lut_low.view(B * 3, 1, d, d, d),
                size=(self.lut_dim, self.lut_dim, self.lut_dim),
                mode='trilinear', align_corners=True,
            ).view(B, 3, self.lut_dim, self.lut_dim, self.lut_dim)
        else:
            lut = lut_low
        return lut, feat

    def forward(self, x_full, reference_224):
        lut, feat = self.predict_lut(reference_224)
        output = self.trilinear(lut, x_full)
        return output, lut, feat


# =========================================================================
#  Film Physics Renderers (grain + halation) and tone matching
# =========================================================================

def _gaussian_blur_2d(x, sigma):
    if sigma < 0.3:
        return x
    ks = int(6 * sigma + 1) | 1
    ks = min(ks, 31)
    coords = torch.arange(ks, dtype=torch.float32, device=x.device) - ks // 2
    k = torch.exp(-0.5 * (coords / sigma) ** 2)
    k = k / k.sum()
    B, C, H, W = x.shape
    pad = ks // 2
    kh = k.view(1, 1, 1, ks).expand(C, -1, -1, -1)
    x = F.conv2d(F.pad(x, [pad, pad, 0, 0], 'reflect'), kh, groups=C)
    kv = k.view(1, 1, ks, 1).expand(C, -1, -1, -1)
    x = F.conv2d(F.pad(x, [0, 0, pad, pad], 'reflect'), kv, groups=C)
    return x


def _channel_blur(x1, radius):
    sigma = radius / 2.0
    if sigma < 0.5:
        return x1
    ks = int(6 * sigma + 1) | 1
    ks = min(ks, 181)
    coords = torch.arange(ks, dtype=torch.float32, device=x1.device) - ks // 2
    k = torch.exp(-0.5 * (coords / sigma) ** 2)
    k = k / k.sum()
    pad = ks // 2
    kh = k.view(1, 1, 1, ks)
    x1 = F.conv2d(F.pad(x1, [pad, pad, 0, 0], 'reflect'), kh)
    kv = k.view(1, 1, ks, 1)
    x1 = F.conv2d(F.pad(x1, [0, 0, pad, pad], 'reflect'), kv)
    return x1


def match_tone_curve(source, reference, strength=0.8, return_transfer=False):
    """Per-channel histogram matching with smooth transfer function."""
    B, C, H, W = source.shape
    dev = source.device
    result = torch.zeros_like(source)
    n_bins = 256
    transfer_curves = []

    for c in range(C):
        src_ch = source[0, c].flatten()
        ref_ch = reference[0, c].flatten()

        bins = torch.linspace(0, 1, n_bins + 1, device=dev)
        bin_centers = (bins[:-1] + bins[1:]) / 2

        src_cdf = torch.zeros(n_bins, device=dev)
        ref_cdf = torch.zeros(n_bins, device=dev)
        for i in range(n_bins):
            src_cdf[i] = (src_ch <= bins[i + 1]).float().mean()
            ref_cdf[i] = (ref_ch <= bins[i + 1]).float().mean()

        transfer = torch.zeros(n_bins, device=dev)
        for i in range(n_bins):
            target_cdf = src_cdf[i]
            idx = torch.searchsorted(ref_cdf, target_cdf.unsqueeze(0)).squeeze()
            idx = idx.clamp(0, n_bins - 1)
            transfer[i] = bin_centers[idx]

        kernel_size = 9
        sigma = 2.0
        coords = torch.arange(kernel_size, device=dev).float() - kernel_size // 2
        kernel = torch.exp(-0.5 * (coords / sigma) ** 2)
        kernel = kernel / kernel.sum()
        transfer_padded = F.pad(transfer.unsqueeze(0).unsqueeze(0),
                                [kernel_size // 2, kernel_size // 2], mode='reflect')
        transfer_smooth = F.conv1d(transfer_padded, kernel.view(1, 1, -1)).squeeze()
        transfer_curves.append(transfer_smooth)

        src_flat = src_ch.clamp(0, 1)
        idx_float = src_flat * (n_bins - 1)
        idx_low = idx_float.long().clamp(0, n_bins - 2)
        idx_high = (idx_low + 1).clamp(max=n_bins - 1)
        frac = idx_float - idx_low.float()
        mapped = transfer_smooth[idx_low] * (1 - frac) + transfer_smooth[idx_high] * frac

        result[0, c] = mapped.view(H, W)

    lum_w = torch.tensor([0.2126, 0.7152, 0.0722], device=dev).view(1, 3, 1, 1)
    src_lum = (source * lum_w).sum(dim=1, keepdim=True)

    highlight_protect = ((src_lum - 0.55) / (0.90 - 0.55)).clamp(0, 1)
    effective_strength = strength * (1.0 - highlight_protect * 0.7)

    out = source * (1 - effective_strength) + result * effective_strength

    dither = (torch.rand_like(out) + torch.rand_like(out) - 1.0) * (0.5 / 256.0)
    out = out + dither * strength

    out = torch.clamp(out, 0, 1)
    if return_transfer:
        return out, torch.stack(transfer_curves, dim=0)
    return out


def render_grain(img, sigma_val, grain_size_val, lum_params, grain_mult=1.0):
    B, C, H, W = img.shape
    dev = img.device
    lum_w = torch.tensor([0.2126, 0.7152, 0.0722], device=dev).view(1, 3, 1, 1)
    L = (img * lum_w).sum(dim=1, keepdim=True)
    a = lum_params[:, 0:1, None, None]
    b = lum_params[:, 1:2, None, None]
    c = lum_params[:, 2:3, None, None]
    lum_mask = torch.sigmoid(a * L**2 + b * L + c)

    res_factor = math.sqrt(H * W) / 480.0
    raw_gs = grain_size_val if isinstance(grain_size_val, float) else grain_size_val.mean().item()

    fine_sigma = 0.0
    mid_sigma = max(0.3, raw_gs * 0.12) * res_factor
    coarse_sigma = min(max(0.5, raw_gs * 0.25) * res_factor, 2.0 * res_factor)
    scales = [fine_sigma, mid_sigma, coarse_sigma]
    weights = [0.50, 0.33, 0.17]

    shared = torch.randn(B, 1, H, W, device=dev)
    noise = torch.cat([
        shared * 0.5 + torch.randn(B, 1, H, W, device=dev) * 0.5,
        shared * 0.5 + torch.randn(B, 1, H, W, device=dev) * 0.5,
        shared * 0.5 + torch.randn(B, 1, H, W, device=dev) * 0.5,
    ], dim=1)

    grain = torch.zeros(B, 3, H, W, device=dev)
    for s, w in zip(scales, weights):
        octave = _gaussian_blur_2d(noise, s) if s >= 0.3 else noise
        std = octave.std(dim=[2, 3], keepdim=True).clamp(min=1e-6)
        grain = grain + w * (octave / std)

    intensity = sigma_val * grain_mult
    if isinstance(intensity, torch.Tensor):
        intensity = intensity.view(-1, 1, 1, 1) if intensity.dim() >= 1 else intensity
    grain = intensity * lum_mask * grain
    return torch.clamp(img + grain, 0, 1), grain


def render_halation(img, threshold_val, radius_val, intensity_val, color_bias):
    B, C, H, W = img.shape
    lum_w = torch.tensor([0.2126, 0.7152, 0.0722], device=img.device).view(1, 3, 1, 1)
    L = (img * lum_w).sum(dim=1, keepdim=True)
    if isinstance(threshold_val, torch.Tensor):
        mask = torch.sigmoid(20.0 * (L - threshold_val.view(B, 1, 1, 1)))
    else:
        mask = torch.sigmoid(20.0 * (L - threshold_val))
    highlights = img * mask

    res_factor = math.sqrt(H * W) / 480.0
    base_r = (radius_val.mean().item() if isinstance(radius_val, torch.Tensor) else radius_val) * res_factor

    sr = _channel_blur(highlights[:, 0:1], base_r * 1.4)
    sg = _channel_blur(highlights[:, 1:2], base_r * 1.0)
    sb = _channel_blur(highlights[:, 2:3], base_r * 0.7)
    scattered = torch.cat([sr, sg, sb], dim=1)

    cw = torch.sigmoid(color_bias).view(B, 3, 1, 1)
    if isinstance(intensity_val, torch.Tensor):
        iv = intensity_val.view(B, 1, 1, 1)
    else:
        iv = intensity_val
    hmap = iv * cw * scattered
    return torch.clamp(img + hmap, 0, 1), hmap