jonstreeterandClaude
36860b09dd
This fixes the ModuleNotFoundError issues that occurred during clean installations. Changes: - Embedded colormnet_source dependencies directly into repository as: - colormnet_model/ (model architecture) - colormnet_inference/ (inference pipeline) - colormnet_util/ (utility functions) - colormnet_dataset/ (data transforms and utilities) - Updated all internal imports to use new embedded paths - Removed external colormnet_source path dependency from __init__.py - Added graceful fallback for spatial_correlation_sampler CUDA extension - Fixed shape calculation bugs in PyTorch fallback implementation - Added missing dependencies (hickle, scikit-image) to auto-install The repository is now self-contained and can be installed without requiring the separate colormnet_source directory. Tested and verified: - Node imports work correctly - Colorization pipeline runs successfully - Model loads and processes video frames Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
17 lines
412 B
Python
17 lines
412 B
Python
import torch
|
|
import torch.nn.functional as F
|
|
|
|
|
|
# Soft aggregation from STM
|
|
def aggregate(prob, dim, return_logits=False):
|
|
new_prob = torch.cat([
|
|
torch.prod(1-prob, dim=dim, keepdim=True),
|
|
prob
|
|
], dim).clamp(1e-7, 1-1e-7)
|
|
logits = torch.log((new_prob /(1-new_prob)))
|
|
prob = F.softmax(logits, dim=dim)
|
|
|
|
if return_logits:
|
|
return logits, prob
|
|
else:
|
|
return prob |