BokehFlow Code Audit — Issues Found and Fixed
CRITICAL BUGS (Will cause training failure or incorrect results)
1. ❌ Compositing double-multiplication bug (render_bokeh, line ~825)
Problem: output = output + blurred * visible / (blurred_mask + 1e-6) * visible
This multiplies by visible TWICE — blurred * visible / blurred_mask * visible = wrong alpha compositing.
Fix: output = output + blurred / (blurred_mask + 1e-6) * visible
2. ❌ CoC map computation doesn't handle focus distance == depth correctly
Problem: When D == S₁ (pixel at focus distance), CoC should be exactly 0. The formula computes abs(D - S₁) which is correct, but the S1.clamp(min=f+1.0) can produce NaN gradients when f is a learnable parameter.
Fix: Detach f from the clamp or use a fixed minimum.
3. ❌ BatchNorm in ConvStem will break at batch_size=1 during inference
Problem: nn.BatchNorm2d computes running stats during training but fails with batch_size=1 if model is in training mode.
Fix: Use nn.GroupNorm(num_groups=8, num_channels=...) or nn.InstanceNorm2d instead.
STABILITY ISSUES (May cause NaN/Inf during training)
4. ⚠️ No gradient clipping mentioned in training config
Problem: The GatedDeltaNet recurrence compounds matrix operations. Without gradient clipping, gradients can explode.
Fix: Add max_grad_norm=1.0 to training config.
5. ⚠️ Key L2-normalization — correct but needs epsilon
Problem: F.normalize(k, p=2, dim=-1) can produce NaN if k is all zeros.
Fix: Add eps: k = F.normalize(k, p=2, dim=-1, eps=1e-8)
6. ⚠️ State explosion risk
Problem: The state update state = a_t * (state - b_t * (state @ kk_t)) + b_t * vk_t has matrix products that can grow unbounded if α≈1 and β≈0 for many steps.
Fix: Add periodic state normalization: state = state / (state.norm() + 1e-6).clamp(min=0.1) every 256 steps.
7. ⚠️ Softplus depth output has no upper bound
Problem: nn.Softplus() can output arbitrarily large values, causing CoC explosion.
Fix: depth = F.softplus(raw_depth).clamp(max=100.0) (100 meters max).
LOGICAL ISSUES
8. ⚠️ embed_dim mismatch for base variant
Problem: num_heads=6, head_dim=32 means inner_dim=192 but embed_dim=192, so the linear projections to_qkv project 192→3*192=576. This is correct but the output gate also projects 192→192. No bug but very heavy for base variant.
9. ⚠️ Direction fusion uses outputs before normalization
Problem: The adaptive direction fusion softmax(W_γ · [o_→;...]) operates on raw scan outputs, then the result is LayerNorm'd. The softmax inputs can have different scales per direction, potentially making one direction always dominate.
Fix: Apply LayerNorm to each scan output BEFORE fusion, or use a temperature in the softmax.
10. ⚠️ TSP state shape mismatch
Problem: self.S_init has shape (1, num_heads, head_dim, head_dim) but BiGDR returns a list of states (one per scan direction), not a single state. The propagate function iterates over block_states which are lists, not tensors.
Fix: S_init should match the per-direction state shape, and propagation should handle the list structure properly.
DATASET CONFIRMED COMPATIBLE ✅
RealBokeh has paired data:
- Input:
{split}/in/{id}_f22.JPG(sharp, f/22) - GT:
{split}/gt/{id}/{id}_f{fstop}.JPG(variable bokeh) - Metadata:
{split}/metadata/{id}.jsonwith focal_length, focus_plane_distance, target_avs
This maps perfectly to BokehFlow's inputs: image, f_number, focal_length_mm, focus_distance_m.