utils/color_model.py

import json
import numpy as np
from huggingface_hub import hf_hub_download

# Download weights
weights_path = hf_hub_download(
    repo_id="danielritchie/vibe-color-model",
    filename="vibe_weights.json"
)

with open(weights_path, "r") as f:
    weights = json.load(f)

# Extract layers
W0 = np.array(weights["layer_0"]["weights"], dtype=np.float32)
b0 = np.array(weights["layer_0"]["biases"], dtype=np.float32)

W1 = np.array(weights["layer_1"]["weights"], dtype=np.float32)
b1 = np.array(weights["layer_1"]["biases"], dtype=np.float32)

W2 = np.array(weights["layer_2"]["weights"], dtype=np.float32)
b2 = np.array(weights["layer_2"]["biases"], dtype=np.float32)


def relu(x):
    return np.maximum(0, x)


def sigmoid(x):
    return 1 / (1 + np.exp(-x))


def infer_color(vad):
    x = np.array([
        vad["V"],
        vad["A"],
        vad["D"],
        vad["Cx"],
        vad["Co"]
    ], dtype=np.float32)

    # Layer 0
    x = relu(np.dot(x, W0) + b0)

    # Layer 1
    x = relu(np.dot(x, W1) + b1)

    # Layer 2
    x = sigmoid(np.dot(x, W2) + b2)

    r, g, b, e, i = x.tolist()

    return {
        "R": r,
        "G": g,
        "B": b,
        "E": e,
        "I": i
    }