Update app.py

app.py

@@ -20,6 +20,24 @@ moondream = AutoModelForCausalLM.from_pretrained(
 )
 moondream.eval()
 
+control_vectors = torch.load("control_vectors.pt")
+control_vectors = [t.to('cuda', dtype=torch.bfloat16) for t in control_vectors]
+
+class LayerWrapper(torch.nn.Module):
+    def __init__(self, og_layer, control_vectors, scale=4.2):
+        super().__init__()
+        self.og_layer = og_layer
+        self.control_vectors = control_vectors
+        self.scale = scale
+
+    def forward(self, *args, **kwargs):
+        layer_outputs = self.og_layer(*args, **kwargs)
+        layer_outputs = (layer_outputs[0] + self.scale * self.control_vectors, *layer_outputs[1:])
+        return layer_outputs
+
+moondream.text_model.transformer.h = torch.nn.ModuleList([
+    LayerWrapper(layer, vector, 4.2) for layer, vector in zip(og_h, control_vectors)
+])
 
 @spaces.GPU(duration=10)
 def answer_question(img, prompt):
@@ -74,8 +92,8 @@ def process_answer(img, answer):
 with gr.Blocks() as demo:
     gr.Markdown(
         """
-        # 🌔 moondream2
-        A tiny vision language model. [GitHub](https://github.com/vikhyat/moondream)
+        # 🌔 contemplative moondream
+        a demo of [moondream](http://moondream.ai) steered to discuss the meaning of life using [activation vectors](https://github.com/vikhyat/moondream/blob/main/notebooks/RepEng.ipynb)
         """
     )
     with gr.Row():