tiny model, tiny runtime

requirements.txt

@@ -1,7 +1,5 @@
 gradio
-torch
-transformers
 matplotlib
 numpy
+tflite-runtime
 huggingface_hub
-

utils/color_model.py

@@ -1,37 +1,43 @@
-import torch
+import numpy as np
 from huggingface_hub import hf_hub_download
+import tflite_runtime.interpreter as tflite
 
-# Download model file from HF
+
+# Download TFLite model file
 model_path = hf_hub_download(
     repo_id="danielritchie/vibe-color-model",
-    filename="pytorch_model.bin"
+    filename="vibe_model.tflite"
 )
 
-# Load full torch model
-model = torch.load(model_path, map_location=torch.device("cpu"))
-model.eval()
+# Load interpreter
+interpreter = tflite.Interpreter(model_path=model_path)
+interpreter.allocate_tensors()
+
+input_details = interpreter.get_input_details()
+output_details = interpreter.get_output_details()
 
 
 def infer_color(vad):
-    input_tensor = torch.tensor([[
+    input_data = np.array([[
         vad["V"],
         vad["A"],
         vad["D"],
         vad["Cx"],
         vad["Co"]
-    ]], dtype=torch.float32)
+    ]], dtype=np.float32)
 
-    with torch.no_grad():
-        output = model(input_tensor)
+    interpreter.set_tensor(input_details[0]["index"], input_data)
+    interpreter.invoke()
+    output_data = interpreter.get_tensor(output_details[0]["index"])
 
-    r, g, b, e, i = output[0].tolist()
+    r, g, b, e, i = output_data[0]
 
     return {
-        "R": r,
-        "G": g,
-        "B": b,
-        "E": e,
-        "I": i
+        "R": float(r),
+        "G": float(g),
+        "B": float(b),
+        "E": float(e),
+        "I": float(i)
     }