Upload folder using huggingface_hub

Dockerfile

@@ -0,0 +1,36 @@
+FROM nvidia/cuda:12.3.2-cudnn9-devel-ubuntu22.04
+
+# Set environment variables
+ENV PYTHONUNBUFFERED=1 \
+    DEBIAN_FRONTEND=noninteractive \
+    CUDA_HOME=/usr/local/cuda \
+    PATH=/usr/local/cuda/bin:$PATH \
+    LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH \
+    NVIDIA_VISIBLE_DEVICES=all \
+    NVIDIA_DRIVER_CAPABILITIES=compute,utility
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    python3 \
+    python3-pip \
+    python3-dev \
+    build-essential \
+    ffmpeg \
+    libsndfile1 \
+    curl \
+    git \
+    && rm -rf /var/lib/apt/lists/*
+
+# Upgrade pip and install build tools
+RUN python3 -m pip install --upgrade pip setuptools wheel
+
+WORKDIR /app
+
+COPY . .
+
+# Install requirements
+RUN pip3 install --no-cache-dir -r requirements.txt
+
+EXPOSE 8000
+
+CMD ["python3", "server.py"]

README.md

@@ -0,0 +1,13 @@
+---
+license: mit
+tags:
+- any-to-any
+- omega
+- omegalabs
+- bittensor
+- agi
+---
+
+This is an Any-to-Any model checkpoint for the OMEGA Labs x Bittensor Any-to-Any subnet.
+
+Check out the [git repo](https://github.com/omegalabsinc/omegalabs-anytoany-bittensor) and find OMEGA on X: [@omegalabsai](https://x.com/omegalabsai).

hotkey.txt

@@ -0,0 +1 @@
+5HEWNpoj22h12CoM4Bue3TyQV9X5ayHcbdG8qwennaqyPw3p

inference.py

@@ -0,0 +1,297 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+import itertools
+import sys
+import time
+from typing import Any, Dict, List
+
+import torch
+from torch import nn
+from omegaconf import DictConfig
+from PIL import Image
+
+from torchtune import config, utils
+from torchtune.utils._generation import sample
+from torchtune.models import convert_weights
+from torchtune.data import Message
+
+from models.tokenizer import START_IMAGE, END_IMAGE, START_AUDIO, END_AUDIO, START_VIDEO, END_VIDEO
+from imagebind.models.imagebind_model import ModalityType
+from diffusers import DiffusionPipeline
+
+from models import add_proj_convert_weights, _BASE_TRAINABLE
+import os
+
+log = utils.get_logger("DEBUG")
+add_proj_convert_weights()
+
+
+class InferenceRecipe:
+    """
+    Recipe for generating tokens from a dense Transformer-based LLM.
+
+    Currently this recipe supports single-GPU generation only. Speculative
+    decoding is not supported.
+
+    For more details on how to use this recipe for generation, please see our
+    tutorial: https://pytorch.org/torchtune/main/tutorials/e2e_flow.html#generation
+
+    For using this recipe with a quantized model, please the following section of
+    the above tutorial:
+    https://pytorch.org/torchtune/main/tutorials/e2e_flow.html#speeding-up-generation-using-quantization
+    """
+
+    def __init__(self, cfg: DictConfig) -> None:
+        self._device = utils.get_device(device=cfg.device)
+        self._dtype = utils.get_dtype(dtype=cfg.dtype)
+        self._quantizer = config.instantiate(cfg.inference.quantizer)
+        self._quantization_mode = utils.get_quantizer_mode(self._quantizer)
+        self.prompt_template = cfg.inference.prompt_template
+        perception_tokens = cfg.model.perception_tokens
+        self._perception_tokens = ("0 " * perception_tokens)[:perception_tokens]
+        utils.set_seed(seed=cfg.seed)
+
+    def setup(self, cfg: DictConfig) -> None:
+        checkpointer = config.instantiate(cfg.checkpointer)
+        if self._quantization_mode is None:
+            ckpt_dict = checkpointer.load_checkpoint()
+        else:
+            # weights_only needs to be False when loading a quantized model
+            # currently loading a quantized model is only supported with the
+            # FullModelTorchTuneCheckpointer
+            ckpt_dict = checkpointer.load_checkpoint(weights_only=False)
+
+        self._model = self._setup_model(
+            model_cfg=cfg.model,
+            model_state_dict=ckpt_dict[utils.MODEL_KEY],
+        )
+        with self._device:
+            self._model.setup_caches(max_batch_size=cfg.batch_size, dtype=self._dtype)
+
+        self._tokenizer = config.instantiate(cfg.tokenizer)
+        self._mm_ids_start = self._tokenizer.encode(START_IMAGE + START_AUDIO + START_VIDEO, add_eos=False, add_bos=False)
+        self._mm_ids_end = self._tokenizer.encode(END_IMAGE + END_AUDIO + END_VIDEO, add_eos=False, add_bos=False)
+        self.use_clip = cfg.model.use_clip
+        if self.use_clip:
+            self._clip_pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-unclip-small", torch_dtype=self._dtype).to(self._device)
+
+    def _setup_model(
+        self,
+        model_cfg: DictConfig,
+        model_state_dict: Dict[str, Any],
+    ) -> nn.Module:
+        with utils.set_default_dtype(self._dtype), self._device:
+            model = config.instantiate(model_cfg)
+
+        if self._quantization_mode is not None:
+            model = self._quantizer.quantize(model)
+            model = model.to(device=self._device, dtype=self._dtype)
+
+        model.load_state_dict(model_state_dict)
+
+        # Validate model was loaded in with the expected dtype.
+        utils.validate_expected_param_dtype(model.named_parameters(), dtype=self._dtype)
+        log.debug(f"Model is initialized with precision {self._dtype}.")
+
+        return model
+
+    def mm_process_prompt(self, prompt):
+        return (
+            prompt
+                .replace("{image}", f"{START_IMAGE}{self._perception_tokens}{END_IMAGE}")
+                .replace("{audio}", f"{START_AUDIO}{self._perception_tokens}{END_AUDIO}")
+                .replace("{video}", f"{START_VIDEO}{self._perception_tokens}{END_VIDEO}")
+            )
+
+    def extract_mm_context(self, video_ib_embed, tokens):
+        context = {}
+        in_mm_embed = False
+        for idx, tok in enumerate(tokens):
+            in_mm_embed = in_mm_embed and not tok in self._mm_ids_end
+            if in_mm_embed:
+                #tokens[idx] # to support multiple embeds: get the value, match it up with the sample embed
+                context[idx] = {
+                    "ib_embed": video_ib_embed.to(dtype=self._dtype, device=self._device),
+                }
+            in_mm_embed = in_mm_embed or tok in self._mm_ids_start
+        return context
+
+    @torch.no_grad()
+    def generate(self, cfg: DictConfig, video_ib_embed: List[float]):
+        messages = [
+            Message(
+                role="user",
+                content=self.mm_process_prompt(self.prompt_template),
+            ),
+            Message(
+                role="assistant",
+                content="",
+            )
+        ]
+        tokens, mask = self._tokenizer.tokenize_messages(messages)
+        tokens = tokens[:-2] # strip eot and eos
+        mm_context = [self.extract_mm_context(video_ib_embed, tokens)] # context should be a list, batch-id indexed
+        prompt = torch.tensor(tokens, dtype=torch.int, device=self._device)
+
+        self._model.tok_embeddings.set_context(mm_context)
+        self._model.output.set_context(mm_context)
+
+        bos_id = self._tokenizer.tt_model.encode("<|begin_of_text|>", allowed_special="all")[0]
+        allowed_id = self._tokenizer.tt_model.encode(f"<|eot_id|>{START_IMAGE}{END_IMAGE}{START_AUDIO}{END_AUDIO}{START_VIDEO}{END_VIDEO}", allowed_special="all")
+        disallowed_tokens = list(set(range(bos_id, bos_id + 256)) - set(allowed_id))
+        # self._model.output.weight.data[disallowed_tokens, :] = 0
+
+        def custom_generate_next_token(model, input_pos, x, temperature=1.0, top_k=None):
+            model.tok_embeddings.set_context([])
+            model.output.set_context([])
+            # x: [1, s]
+            # input_pos: [s]
+            logits = model(x, input_pos=input_pos)
+            # logits: [1, s, v] where v is vocab_size
+            # for sampling we extract the logits for the
+            # last token and convert to shape: [v]
+            logits = logits[0, -1]
+            # logits[disallowed_tokens] = float("-inf")
+            # sample the next token
+            token = sample(logits, temperature, top_k)
+            if token in disallowed_tokens:
+                return torch.tensor([self._tokenizer.eos_id]).to(x)
+            return token
+
+        # since quantized model uses torch.compile to get speedup, it needs a warm up / prefill run
+        # to get the accurate performance measurement
+        if self._quantization_mode is not None:
+            log.info("Starting compilation to improve generation performance ...")
+            custom_generate_next_token = torch.compile(
+                custom_generate_next_token, mode="max-autotune", fullgraph=True
+            )
+            t0 = time.perf_counter()
+            _ = utils.generate(
+                model=self._model,
+                prompt=prompt,
+                max_generated_tokens=2,
+                temperature=cfg.temperature,
+                top_k=cfg.top_k,
+                eos_id=self._tokenizer.eos_id,
+                custom_generate_next_token=custom_generate_next_token,
+            )
+            t = time.perf_counter() - t0
+            log.info(f"Warmup run for quantized model takes: {t:.02f} sec")
+
+        t0 = time.perf_counter()
+        generated_tokens = utils.generate(
+            model=self._model,
+            prompt=prompt,
+            max_generated_tokens=cfg.max_new_tokens,
+            temperature=cfg.temperature,
+            top_k=cfg.top_k,
+            eos_id=self._tokenizer.eos_id,
+            custom_generate_next_token=custom_generate_next_token,
+        )
+        t = time.perf_counter() - t0
+
+        cleaned_tokens = [t for t in generated_tokens[len(prompt):] if t not in disallowed_tokens + allowed_id]
+        caption = self._tokenizer.decode(cleaned_tokens)
+
+        # log.debug(f"Generated caption: {caption} in {t:.02f} sec")
+
+        return caption
+
+
+    @torch.no_grad()
+    def generate_batch(self, cfg: DictConfig, video_ib_embed: torch.Tensor):
+        log.info(f"inside generate_batch, video_ib_embed shape: {video_ib_embed.shape}")
+        batch_dim = video_ib_embed.size(0)
+        messages = [
+            Message(
+                role="user",
+                content=self.mm_process_prompt(self.prompt_template),
+            ),
+            Message(role="assistant", content="")
+        ]
+        tokens, mask = self._tokenizer.tokenize_messages(messages)
+        tokens = tokens[:-2] # strip eot and eos
+        mm_context = [self.extract_mm_context(e, tokens) for e in video_ib_embed] # context should be a list, batch-id indexed
+        prompt = torch.tensor(tokens, dtype=torch.int, device=self._device).expand(batch_dim, -1).clone()
+        prompt_length = prompt.size(1)
+
+        self._model.tok_embeddings.set_context(mm_context)
+        self._model.output.set_context(mm_context)
+
+        bos_id = self._tokenizer.tt_model.encode("<|begin_of_text|>", allowed_special="all")[0]
+        allowed_id = self._tokenizer.tt_model.encode(f"<|eot_id|>{START_IMAGE}{END_IMAGE}{START_AUDIO}{END_AUDIO}{START_VIDEO}{END_VIDEO}", allowed_special="all")
+        disallowed_tokens = list(set(range(bos_id, bos_id + 256)) - set(allowed_id))
+
+        def generate_next_token(model, input_pos, x, temperature=1.0, top_k=None):
+            # x: [B, s]
+            # input_pos: [s]
+            # logits: [B, s, v] where v is vocab_size
+            logits = model(x, input_pos=input_pos)[:, -1]
+            tokens = sample(logits, temperature, top_k)
+            return torch.tensor([
+                [self._tokenizer.eos_id if t in disallowed_tokens else t for t in toks]
+                for toks in tokens
+            ]).to(x.device)
+
+        generated_tokens = prompt.clone()
+        # keeps track at a high level if we've already hit a stop token in a sequence so we can early stop
+        stop_token_reached = torch.zeros(batch_dim, dtype=torch.bool, device=prompt.device)
+
+        # generate the first tokens conditioned on the prompt
+        tokens = generate_next_token(
+            self._model,
+            input_pos=torch.arange(0, prompt_length, device=prompt.device),
+            x=prompt,
+            temperature=cfg.temperature,
+            top_k=cfg.top_k,
+        )
+        eot_reached_b = tokens == self._tokenizer.eot_id
+        generated_tokens = torch.cat([generated_tokens, tokens], dim=-1)
+
+        self._model.tok_embeddings.set_context([])
+        self._model.output.set_context([])
+
+        input_pos = torch.tensor([prompt_length], device=prompt.device)
+        for _ in range(cfg.max_new_tokens - 1):
+            tokens = generate_next_token(
+                self._model, input_pos=input_pos, x=tokens, temperature=cfg.temperature, top_k=cfg.top_k
+            )
+            eot_reached_b |= tokens == self._tokenizer.eot_id
+            tokens *= ~eot_reached_b
+            generated_tokens = torch.cat([generated_tokens, tokens], dim=-1)
+            if eot_reached_b.all():
+                print('eot_reached_b.all()')
+                break
+            input_pos += 1
+
+        captions = []
+        for caption_tokens in generated_tokens.tolist():
+            captions.append(self._tokenizer.decode(caption_tokens[prompt.size(1):]))
+        return captions
+
+
+@config.parse
+def main(cfg: DictConfig) -> None:
+    config.log_config(recipe_name="InferenceRecipe", cfg=cfg)
+    cfg.model = DictConfig({
+        "_component_": "models.mmllama3_8b",
+        "use_clip": False,
+        "perception_tokens": cfg.model.perception_tokens,
+    })
+    cfg.batch_size = 4
+    cfg.checkpointer.checkpoint_dir = os.path.dirname("/home/salman/tezuesh/omegalabs-anytoany-bittensor/sandboxing/cache/xzistance_omega-a2a-hotkey/meta_model_0.pth")
+    
+    cfg.checkpointer.checkpoint_files = ["models/meta_model_0.pt"]
+    cfg.inference.max_new_tokens = 300
+    cfg.tokenizer.path = "./models/tokenizer.model"
+    inference_recipe = InferenceRecipe(cfg)
+    inference_recipe.setup(cfg=cfg)
+    captions = inference_recipe.generate_batch(cfg=cfg, video_ib_embed=torch.randn(4,1024))
+    print(captions)
+
+
+if __name__ == "__main__":
+    sys.exit(main())

models/__init__.py

@@ -0,0 +1,34 @@
+from torchtune.models import convert_weights
+
+from models.tokenizer import a2a_tokenizer
+from models.mmllama3 import lora_mmllama3_8b, mmllama3_8b, imagebind_huge
+
+__all__ = [
+    "a2a_tokenizer",
+    "lora_mmllama3_8b",
+    "mmllama3_8b",
+    "imagebind_huge",
+
+]
+
+_BASE_TRAINABLE = [
+    "tok_embeddings.proj_to_llama.0.weight",
+    "tok_embeddings.proj_to_llama.0.bias",
+    "tok_embeddings.proj_to_llama.2.weight",
+    "tok_embeddings.proj_to_llama.2.bias",
+    "tok_embeddings.proj_to_llama.3.weight",
+    "tok_embeddings.proj_to_llama.3.bias",
+    "output.proj_from_llama.0.weight",
+    "output.proj_from_llama.0.bias",
+    "output.proj_from_llama.2.weight",
+    "output.proj_from_llama.2.bias",
+    "output.proj_from_llama.3.weight",
+    "output.proj_from_llama.3.bias",
+]
+
+def add_proj_convert_weights():
+    # extend _FROM_META torchtune -> meta mapping with new parameter names
+    # allow existing ckpt-save code to work without changes
+    convert_weights._FROM_META.update({a: a for a in _BASE_TRAINABLE})
+
+

models/bpe_simple_vocab_16e6.txt.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:924691ac288e54409236115652ad4aa250f48203de50a9e4722a6ecd48d6804a
+size 1356917

models/imagebind_wrapper.py

@@ -0,0 +1,221 @@
+import os
+import torch
+import numpy as np
+from typing import BinaryIO, List
+
+from imagebind import imagebind_model
+from imagebind.models.imagebind_model import ModalityType
+from imagebind.models.multimodal_preprocessors import SimpleTokenizer, TextPreprocessor
+
+
+V2_URL = "https://huggingface.co/jondurbin/videobind-v0.2/resolve/main/videobind.pth"
+V2_PATH = "./.checkpoints/videobind-v0.2.pth"
+BPE_PATH = "./models/bpe_simple_vocab_16e6.txt.gz"
+TOKENIZER = SimpleTokenizer(bpe_path=BPE_PATH)
+LENGTH_TOKENIZER = SimpleTokenizer(bpe_path=BPE_PATH, context_length=1024)
+TOKEN_CHUNK_SIZE = 74
+
+def get_imagebind_v2(path: str=V2_PATH):
+    if not os.path.isfile(path):
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        torch.hub.download_url_to_file(V2_URL, path, progress=True)
+    imagebind_model = torch.load(path)
+    return imagebind_model
+
+
+def load_and_transform_text(text, device):
+    if text is None:
+        return None
+    tokens = [TOKENIZER(t).unsqueeze(0).to(device) for t in text]
+    tokens = torch.cat(tokens, dim=0)
+    return tokens
+
+def split_text_by_token_limit(text, tokenizer, max_tokens=TOKEN_CHUNK_SIZE):
+    def fits_in_token_limit(text_segment):
+        tokens = tokenizer(text_segment)
+        tokens = tokens[tokens != 0][1:-1].tolist()
+        return len(tokens) <= max_tokens
+
+    def recursive_split(text, delimiters):
+        if fits_in_token_limit(text):
+            return [text]
+        if not delimiters:
+            return split_by_tokens(text)
+        delimiter = delimiters[0]
+        parts = text.split(delimiter)
+        result = []
+        current_segment = ""
+        for part in parts:
+            candidate_segment = current_segment + (delimiter if current_segment else '') + part
+            if fits_in_token_limit(candidate_segment):
+                current_segment = candidate_segment
+            else:
+                if current_segment:
+                    result.append(current_segment)
+                current_segment = part
+        if current_segment:
+            result.append(current_segment)
+        final_result = []
+        for segment in result:
+            if fits_in_token_limit(segment):
+                final_result.append(segment)
+            else:
+                final_result.extend(recursive_split(segment, delimiters[1:]))
+        return final_result
+
+    def split_by_tokens(text):
+        tokens = tokenizer(text)
+        tokens = tokens[tokens != 0][1:-1].tolist()
+        chunks = np.array_split(tokens, int(len(tokens) / max_tokens) or 1)
+        return [
+            tokenizer.decode(segment_tokens)
+            for segment_tokens in chunks
+        ]
+
+    return recursive_split(text, ['\n', '.', '!', '?', ',', ' '])
+
+def load_and_transform_text_chunks(text, device):
+    if not text:
+        return []
+    all_tokens = LENGTH_TOKENIZER(text)
+    all_tokens = all_tokens[all_tokens != 0][1:-1].tolist()
+
+    return [
+        load_and_transform_text([segment], device)
+        for segment in split_text_by_token_limit(text, LENGTH_TOKENIZER)
+    ]
+
+
+class ImageBind:
+    def __init__(self, device="cuda:0", v2=False):
+        self.device = device
+        self.v2 = v2
+        if v2:
+            if not os.path.exists(V2_PATH):
+                os.makedirs(os.path.dirname(V2_PATH), exist_ok=True)
+                torch.hub.download_url_to_file(
+                    V2_URL,
+                    V2_PATH,
+                    progress=True,
+                )
+            self.imagebind = torch.load(V2_PATH)
+        else:
+            self.imagebind = imagebind_model.imagebind_huge(pretrained=True)
+        self.imagebind.eval()
+        self.imagebind.to(self.device)
+
+    def generate_text_embeddings(self, text: str):
+        if not self.v2:
+            return self.imagebind({
+                ModalityType.TEXT: load_and_transform_text([text], self.device)
+            })[ModalityType.TEXT]
+        chunks = load_and_transform_text_chunks(text, self.device)
+        embeddings = [
+            self.imagebind({ModalityType.TEXT: chunk})[ModalityType.TEXT]
+            for chunk in chunks
+        ]
+        return torch.mean(torch.stack(embeddings), dim=0)
+
+    """ Deactivating full embeddings as they are not used in the current implementation
+    def get_inputs(self, video_file: BinaryIO) -> dict:
+        audio_file = video_utils.copy_audio(video_file.name)
+        try:
+            duration = video_utils.get_video_duration(video_file.name)
+            video_data = data.load_and_transform_video_data(
+                [video_file.name],
+                self.device,
+            )
+            audio_data = data.load_and_transform_audio_data(
+                [audio_file.name],
+                self.device,
+            )
+            inputs = {
+                ModalityType.VISION: video_data,
+                ModalityType.AUDIO: audio_data,
+            }
+            return inputs
+        finally:
+            audio_file.close()
+
+    @torch.no_grad()
+    def embed(self, descriptions: List[str], video_files: List[BinaryIO]) -> Embeddings:
+        return_value = None
+        for idx in range(len(descriptions)):
+            inputs = self.get_inputs(video_files[idx])
+            embeddings = self.imagebind(inputs)
+            text_embeddings = self.generate_text_embeddings(descriptions[idx])
+            if not return_value:
+                return_value = Embeddings(
+                    video=embeddings[ModalityType.VISION],
+                    audio=embeddings[ModalityType.AUDIO],
+                    description=text_embeddings,
+                )
+            else:
+                return_value.video = torch.cat((return_value.video, embeddings[ModalityType.VISION]))
+                return_value.audio = torch.cat((return_value.audio, embeddings[ModalityType.AUDIO]))
+                return_value.description = torch.cat((return_value.description, text_embeddings))
+        return return_value
+
+    @torch.no_grad()
+    def embed_only_video(self, video_files: List[BinaryIO]) -> Embeddings:
+        video_filepaths = [video_file.name for video_file in video_files]
+        durations = [video_utils.get_video_duration(f.name) for f in video_files]
+        embeddings = self.imagebind({
+            ModalityType.VISION: [
+                data.load_and_transform_video_data(
+                    [video_filepaths[idx]],
+                    self.device,
+                )[0]
+                for idx in range(len(video_filepaths))
+            ]
+        })
+        return Embeddings(
+            video=embeddings[ModalityType.VISION],
+        )
+
+    @torch.no_grad()
+    def embed_video_and_text(self, video_files: List[BinaryIO], descriptions: List[str]) -> Embeddings:
+        video_filepaths = [video_file.name for video_file in video_files]
+        durations = [video_utils.get_video_duration(f.name) for f in video_files]
+        embeddings = self.imagebind({
+            ModalityType.VISION: [
+                data.load_and_transform_video_data(
+                    [video_filepaths[idx]],
+                    self.device,
+                )[0]
+                for idx in range(len(video_filepaths))
+            ],
+        })
+        description_embeddings = torch.stack([
+            self.generate_text_embeddings(description)
+            for description in descriptions
+        ])
+        return Embeddings(
+            video=embeddings[ModalityType.VISION],
+            description=description_embeddings,
+        )
+
+    @torch.no_grad()
+    def embed_text(self, texts: List[str]) -> torch.Tensor:
+        return_value = None
+        for text in texts:
+            emb = self.generate_text_embeddings(text)
+            if not return_value:
+                return_value = emb
+            else:
+                return_value = torch.cat((return_value, emb))
+        return return_value
+    """
+    
+    @torch.no_grad()
+    def embed_text(self, texts: List[str]) -> torch.Tensor:
+        embeddings = []
+        for text in texts:
+            emb = self.generate_text_embeddings(text)
+            embeddings.append(emb)
+        
+        if not embeddings:
+            return None
+        
+        # Stack all embeddings along dimension 0
+        return torch.stack(embeddings, dim=0)

models/meta_model_5.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a40e2a0f8c070103237a1c2c147fef1ba7cac02cd53b9d320370e3fbfee7ad84
+size 16219158403

models/mmllama3.py

@@ -0,0 +1,153 @@
+from typing import List
+import warnings
+
+import torch
+from torch import nn, Tensor
+from torchvision import transforms
+
+from torchtune.models.llama3 import lora_llama3_8b, llama3_8b
+from torchtune.modules.peft import LORA_ATTN_MODULES, LoRALinear
+from torchtune.modules import TransformerDecoder
+
+with warnings.catch_warnings():
+    warnings.simplefilter("ignore", UserWarning)
+    from imagebind.models import imagebind_model
+    from models.imagebind_wrapper import get_imagebind_v2, V2_PATH
+    from models.imagebind_wrapper import ImageBind
+
+IMAGEBIND_DIM = 1024
+CLIP_DIM = 768
+
+
+class MMEmbedding(nn.Embedding):
+    def __init__(self, e, perception_tokens=1, use_clip=False):
+        super().__init__(
+            num_embeddings=e.num_embeddings,
+            embedding_dim=e.embedding_dim,
+            padding_idx=e.padding_idx,
+            max_norm=e.max_norm,
+            norm_type=e.norm_type,
+            scale_grad_by_freq=e.scale_grad_by_freq,
+            sparse=e.sparse,
+        )
+        self._perception_tokens = perception_tokens
+        self._context = []
+        self._use_clip = use_clip
+
+        dim_in = IMAGEBIND_DIM + (CLIP_DIM if use_clip else 0)
+        dim_out = e.embedding_dim * perception_tokens
+
+        self.proj_to_llama = nn.Sequential(
+            nn.Linear(dim_in, dim_out),
+            nn.GELU(),
+            nn.LayerNorm(dim_out),
+            nn.Linear(dim_out, dim_out),
+        )
+
+    def set_context(self, context):
+        self._context = context
+
+    def forward(self, input: Tensor) -> Tensor:
+        r = super().forward(input)
+        # self._context is first indexed by batch idx
+        for b, context_dict in enumerate(self._context):
+            # then by sequence idx
+            for s, embed in context_dict.items():
+                # and then must be transformed from imagebind dim -> llama3 dim
+                if self._use_clip:
+                    llama_embed = self.proj_to_llama(torch.cat([embed["ib_embed"], embed["clip_embed"]]))
+                else:
+                    llama_embed = self.proj_to_llama(torch.cat([embed["ib_embed"]]))
+                r[b, s:s+self._perception_tokens] = llama_embed.view(self._perception_tokens, -1)
+        return r
+
+
+class MMLinear(nn.Linear):
+    def __init__(self, o):
+        super().__init__(
+            in_features=o.in_features,
+            out_features=o.out_features,
+            bias=(o.bias != None)
+        )
+        self._context = []
+
+        dim_out = CLIP_DIM
+        dim_in = o.in_features
+        self.proj_from_llama = nn.Sequential(
+            nn.Linear(dim_in, dim_out),
+            nn.GELU(),
+            nn.LayerNorm(dim_out),
+            nn.Linear(dim_out, dim_out),
+        )
+
+    def set_context(self, context):
+        self._context = context
+
+    def forward(self, input_bsd: Tensor) -> Tensor:
+        # self._context has the indexes of image llama tokens: process these with proj_from_llama
+        self._clip_projections = []
+        # # self._context is first indexed by batch idx
+        # for b, context_dict in enumerate(self._context):
+        #     # then by sequence idx
+        #     for s, embed in context_dict.items():
+        #         # and then must be transformed from llama3 dim -> clip dim
+        #         self._clip_projections.append((
+        #             self.proj_from_llama(input_bsd[b, s]),
+        #             (embed["clip_embed"] if "clip_embed" in embed else None) # terrible
+        #         ))
+        r = super().forward(input_bsd)
+        return r
+
+
+
+def lora_mmllama3_8b(
+    lora_attn_modules: List[LORA_ATTN_MODULES],
+    apply_lora_to_mlp: bool = False,
+    apply_lora_to_output: bool = False,
+    lora_rank: int = 8,
+    lora_alpha: float = 16,
+    quantize_base: bool = False,
+    perception_tokens: int = 2,
+    use_clip: bool = False
+) -> TransformerDecoder:
+    llama3 = lora_llama3_8b(
+        lora_attn_modules,
+        apply_lora_to_mlp,
+        apply_lora_to_output,
+        lora_rank,
+        lora_alpha,
+        quantize_base,
+    )
+    llama3.tok_embeddings = MMEmbedding(llama3.tok_embeddings, perception_tokens, use_clip)
+    llama3.output = MMLinear(llama3.output)
+    return llama3
+
+
+def mmllama3_8b(
+    perception_tokens: int = 2,
+    use_clip: bool = False
+) -> TransformerDecoder:
+    llama3 = llama3_8b()
+    llama3.tok_embeddings = MMEmbedding(llama3.tok_embeddings, perception_tokens, use_clip)
+    llama3.output = MMLinear(llama3.output)
+    return llama3
+
+
+def imagebind_huge(use_v2: bool=True):
+    if use_v2:
+        imagebind = ImageBind(v2=True)
+    else:
+        imagebind = imagebind_model.imagebind_huge(pretrained=True)
+    imagebind.transform_from_pil = transforms.Compose([
+        transforms.Resize(
+            224, interpolation=transforms.InterpolationMode.BICUBIC
+        ),
+        transforms.CenterCrop(224),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=(0.48145466, 0.4578275, 0.40821073),
+            std=(0.26862954, 0.26130258, 0.27577711),
+        ),
+    ])
+    return imagebind
+

models/tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:82e9d31979e92ab929cd544440f129d9ecd797b69e327f80f17e1c50d5551b55
+size 2183982

models/tokenizer.py

@@ -0,0 +1,85 @@
+from typing import List
+
+from torchtune.modules.tokenizers import TikTokenTokenizer
+from torchtune.modules.tokenizers._utils import _split_long_repetitions
+from torchtune.modules.tokenizers._tiktoken import (
+    MAX_ENCODE_CHARS,
+    MAX_NO_WHITESPACE_CHARS,
+    ALL_SPECIAL_TOKENS,
+)
+  
+
+# use special tokens from TikTokenTokenizer, add some for MM delimiters
+START_IMAGE = "<|start_image|>"
+END_IMAGE = "<|end_image|>"
+START_VIDEO = "<|start_video|>"
+END_VIDEO = "<|end_video|>"
+START_AUDIO = "<|start_audio|>"
+END_AUDIO = "<|end_audio|>"
+
+A2A_SPECIAL_TOKENS = ALL_SPECIAL_TOKENS[:-2] + [
+    START_IMAGE,
+    END_IMAGE,
+    START_VIDEO,
+    END_VIDEO,
+    START_AUDIO,
+    END_AUDIO,
+] + ALL_SPECIAL_TOKENS[-2:]
+
+# override to allow START_IMAGE, END_IMAGE to be encoded
+class A2ATokenizer(TikTokenTokenizer):
+    def encode(
+        self,
+        text: str,
+        add_bos: bool,
+        add_eos: bool,
+    ) -> List[int]:
+        """
+        Encode a string into a list of token ids. Assumes that the string
+        contains no special tokens.
+
+        Args:
+            text (str): The string to encode.
+            add_bos (bool): Whether to add the beginning of sequence token.
+            add_eos (bool): Whether to add the end of sequence token.
+
+        Returns:
+            List[int]: The list of token ids.
+        """
+        substrs: List[str] = []
+        tokens = []
+        for i in range(0, len(text), MAX_ENCODE_CHARS):
+            substr = text[i : i + MAX_ENCODE_CHARS]
+            # See https://github.com/openai/tiktoken/issues/195
+            sliced_substr = _split_long_repetitions(substr, MAX_NO_WHITESPACE_CHARS)
+            substrs.extend(sliced_substr)
+        for substr in substrs:
+            # allowed_special and disallowed_special are used by tiktoken to define
+            # how special tokens are encoded. Our setting here is to encode any
+            # special token as regular text and prevent tiktoken from raising errors.
+            # This means we should only call encode on strings not containing special tokens.
+            tokens.extend(
+                self.tt_model.encode(
+                    substr,
+                    allowed_special=set([
+                        START_IMAGE,
+                        END_IMAGE,
+                        START_VIDEO,
+                        END_VIDEO,
+                        START_AUDIO,
+                        END_AUDIO,
+                    ]),
+                    disallowed_special=(),
+                )
+            )
+        if add_bos:
+            tokens.insert(0, self.bos_id)
+        if add_eos:
+            tokens.append(self.eos_id)
+        return tokens
+
+
+def a2a_tokenizer(path: str) -> TikTokenTokenizer:
+    tiktoken = A2ATokenizer(path, all_special_tokens=A2A_SPECIAL_TOKENS)
+    tiktoken.pad_id = 0
+    return tiktoken

models/training_config.yml

@@ -0,0 +1,83 @@
+model:
+  _component_: models.lora_mmllama3_8b
+  lora_attn_modules:
+  - q_proj
+  - v_proj
+  apply_lora_to_mlp: false
+  apply_lora_to_output: false
+  lora_rank: 8
+  lora_alpha: 16
+  perception_tokens: 2
+  use_clip: false
+tokenizer:
+  _component_: models.a2a_tokenizer
+  path: models/tokenizer.model
+checkpointer:
+  _component_: torchtune.utils.FullModelMetaCheckpointer
+  checkpoint_dir: /workspace/omega_a2a/training
+  checkpoint_files:
+  - consolidated.00.pth
+  adapter_checkpoint: null
+  recipe_checkpoint: null
+  output_dir: /workspace/omega_a2a/checkpoints
+  model_type: LLAMA3
+resume_from_checkpoint: false
+interim_checkpoint_steps: 5000
+interim_gen_steps: null
+max_new_tokens: 170
+temperature: 0.8
+top_k: 200
+dataset:
+  _component_: ds.EvenBatcher
+  buffer_size: 36
+  dataset:
+    _component_: ds.RoundRobinDataset
+    datasets:
+    - _component_: ds.OmegaVideoCaptionDataset
+      length: 500000
+    - _component_: ds.LlavaInstructDataset
+      dataset_path: ds/coco_llava_instruct/output.parquet
+      train_on_input: false
+    - _component_: ds.LlavaInstructDataset
+      dataset_path: ds/vision_flan/output.parquet
+      train_on_input: false
+    - _component_: ds.CaptionInstructDataset
+      dataset_path: ds/sam_llava/output.parquet
+      train_on_input: false
+seed: null
+shuffle: true
+batch_size: 4
+optimizer:
+  _component_: torch.optim.AdamW
+  weight_decay: 0.0001
+  lr: 3.0e-05
+lr_scheduler:
+  _component_: torchtune.modules.get_cosine_schedule_with_warmup
+  num_warmup_steps: 100
+loss:
+  _component_: torch.nn.CrossEntropyLoss
+epochs: 6
+max_steps_per_epoch: null
+gradient_accumulation_steps: 64
+compile: false
+output_dir: /tmp/lora_finetune_output
+metric_logger:
+  _component_: torchtune.utils.metric_logging.DiskLogger
+  log_dir: ${output_dir}
+log_every_n_steps: null
+device: cuda
+dtype: bf16
+enable_activation_checkpointing: false
+profiler:
+  _component_: torchtune.utils.profiler
+  enabled: false
+inference:
+  prompt_template: 'Video:
+
+    {video}
+
+    Caption the previous video.'
+  max_new_tokens: 300
+  temperature: 0.6
+  top_k: 5
+  quantizer: null

requirements.txt

@@ -0,0 +1,18 @@
+torch==2.6.0
+sentencepiece==0.2.0
+tiktoken==0.4.0
+torchtune @ git+https://github.com/pytorch/torchtune.git@8f59c2fecd722691271eecca630a526719a32f76#egg=torchtune
+lm_eval==0.4
+torchvision==0.21.0
+diffusers==0.27.2
+imagebind @ git+https://github.com/omegalabsinc/ImageBind.git@c3c3b2e1ce6fd850ff42ce0375823fe22880a7cc#egg=imagebind
+llama3 @ git+https://github.com/meta-llama/llama3.git@af6eedf7042fb51d00b2b26d8ef1ceaab73e1670
+pytorchvideo @ git+https://github.com/facebookresearch/pytorchvideo.git@28fe037d212663c6a24f373b94cc5d478c8c1a1d
+wandb==0.17.1
+numpy==1.26.4
+huggingface-hub==0.24.0
+omegaconf==2.3.0
+uvicorn==0.25.0
+fastapi==0.104.1
+pydantic==2.5.2
+torchaudio==2.6.0

server.py

@@ -0,0 +1,146 @@
+from fastapi import FastAPI, HTTPException
+import numpy as np
+import torch
+from pydantic import BaseModel
+from typing import List
+import base64
+import io
+import os
+import logging
+from pathlib import Path
+from inference import InferenceRecipe
+from fastapi.middleware.cors import CORSMiddleware
+
+from omegaconf import OmegaConf, DictConfig
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+app = FastAPI()
+
+# Add CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+class EmbeddingRequest(BaseModel):
+    embedding: List[float]
+
+class TextResponse(BaseModel):
+    texts: List[str] = []
+
+# Model initialization status
+INITIALIZATION_STATUS = {
+    "model_loaded": False,
+    "error": None
+}
+
+# Global model instance
+inference_recipe = None
+cfg = None
+
+
+def initialize_model():
+    """Initialize the model with correct path resolution"""
+    global inference_recipe, INITIALIZATION_STATUS, cfg
+    try:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        logger.info(f"Initializing model on device: {device}")
+        
+        # Critical: Use absolute path for model loading
+        model_path = os.path.abspath(os.path.join('/app', 'models'))
+        logger.info(f"Loading models from: {model_path}")
+        
+        if not os.path.exists(model_path):
+            raise RuntimeError(f"Model path {model_path} does not exist")
+            
+        # Log available model files for debugging
+        model_files = os.listdir(model_path)
+        logger.info(f"Available model files: {model_files}")
+        
+        cfg = OmegaConf.load(os.path.join('/app', 'training_config.yml'))
+        cfg.model = DictConfig({
+            "_component_": "models.mmllama3_8b",
+            "use_clip": False,
+            "perception_tokens": cfg.model.perception_tokens,
+        })
+        cfg.checkpointer.checkpoint_dir = model_path
+        cfg.checkpointer.checkpoint_files = ["meta_model_5.pt"]
+        cfg.inference.max_new_tokens = 300
+        cfg.tokenizer.path = os.path.join(model_path, "tokenizer.model")
+        inference_recipe = InferenceRecipe(cfg)
+        inference_recipe.setup(cfg=cfg)
+        INITIALIZATION_STATUS["model_loaded"] = True
+        logger.info("Model initialized successfully")
+        return True
+    except Exception as e:
+        INITIALIZATION_STATUS["error"] = str(e)
+        logger.error(f"Failed to initialize model: {e}")
+        return False
+        
+@app.on_event("startup")
+async def startup_event():
+    """Initialize model on startup"""
+    initialize_model()
+
+@app.get("/api/v1/health")
+def health_check():
+    """Health check endpoint"""
+    status = {
+        "status": "healthy" if INITIALIZATION_STATUS["model_loaded"] else "initializing",
+        "initialization_status": INITIALIZATION_STATUS
+    }
+    
+    if inference_recipe is not None:
+        status.update({
+            "device": str(inference_recipe._device),
+            "dtype": str(inference_recipe._dtype)
+        })
+        
+    return status
+
+@app.post("/api/v1/inference")
+async def inference(request: EmbeddingRequest) -> TextResponse:
+    """Run inference with enhanced error handling and logging"""
+    if not INITIALIZATION_STATUS["model_loaded"]:
+        raise HTTPException(
+            status_code=503,
+            detail=f"Model not ready. Status: {INITIALIZATION_STATUS}"
+        )
+        
+    try:
+        # Log input validation
+        logger.info("Received inference request")
+        
+        # Convert embedding to tensor
+        embedding = request.embedding  # generate() expects List[float]
+        embedding = torch.tensor(embedding)
+        embedding = embedding.unsqueeze(0)  # Add batch dimension
+        embedding = embedding.reshape(-1, 1024)
+        logger.info(f"Converted embedding to tensor with shape: {embedding.shape}")
+        
+        # Run inference
+        results = inference_recipe.generate_batch(cfg=cfg, video_ib_embed=embedding)
+        logger.info("Generation complete")
+        
+        # Convert results to list if it's not already
+        if isinstance(results, str):
+            results = [results]
+        
+        return TextResponse(texts=results)
+        
+    except Exception as e:
+        logger.error(f"Inference failed: {str(e)}", exc_info=True)
+        raise HTTPException(
+            status_code=500, 
+            detail=str(e)
+        )
+        
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)
+

test.sh

@@ -0,0 +1,60 @@
+#!/bin/bash
+
+# Configure bash error handling
+set -euo pipefail
+
+# Configuration
+API_HOST="localhost"
+API_PORT="8000"
+API_VERSION="v1"
+BASE_URL="http://${API_HOST}:${API_PORT}/api/${API_VERSION}"
+
+# Function to generate test embedding data
+generate_test_embedding() {
+    python3 - <<EOF
+import numpy as np
+import json
+
+# Generate a 4096-dimensional embedding vector (correct dimension for model)
+embedding = np.random.randn(4096).astype(np.float32)
+# Normalize the embedding
+embedding = embedding / np.linalg.norm(embedding)
+print(json.dumps(embedding.tolist()), end="")
+EOF
+}
+
+# Function to test health endpoint
+test_health() {
+    echo "Testing health endpoint..."
+    curl -s "${BASE_URL}/health" || {
+        echo "Health check failed"
+        exit 1
+    }
+}
+
+# Function to test inference endpoint
+test_inference() {
+    echo
+    start_time=$(date +%s)
+    echo "Testing inference endpoint..."
+    local embedding_data=$(generate_test_embedding)
+    
+    curl -X POST "${BASE_URL}/inference" \
+        -H "Content-Type: application/json" \
+        -d "{
+            \"embedding\": ${embedding_data}
+        }" || {
+            echo "Inference request failed"
+            exit 1
+        }
+    end_time=$(date +%s)
+    duration=$((end_time - start_time))
+    echo "Inference request completed in ${duration} seconds"
+}
+
+main() {
+    test_health
+    test_inference
+}
+
+main "$@"

training_config.yml

@@ -0,0 +1,76 @@
+identity_token: 0 1 2
+model:
+  _component_: models.lora_mmllama3_8b
+  lora_attn_modules:
+  - q_proj
+  - v_proj
+  apply_lora_to_mlp: false
+  apply_lora_to_output: false
+  lora_rank: 8
+  lora_alpha: 16
+  perception_tokens: 2
+  use_clip: false
+tokenizer:
+  _component_: models.a2a_tokenizer
+  path: checkpoints/Meta-Llama-3-8B-Instruct/original/tokenizer.model
+checkpointer:
+  _component_: torchtune.utils.FullModelMetaCheckpointer
+  checkpoint_dir: checkpoints/Meta-Llama-3-8B-Instruct/original/
+  checkpoint_files:
+  - consolidated.00.pth
+  adapter_checkpoint: null
+  recipe_checkpoint: null
+  output_dir: output_checkpoints/experiment_4
+  model_type: LLAMA3
+resume_from_checkpoint: false
+interim_checkpoint_steps: 1500000
+interim_gen_steps: null
+max_new_tokens: 100
+temperature: 0.6
+top_k: 300
+dataset:
+  _component_: ds.EvenBatcher
+  dataset:
+    _component_: ds.RoundRobinDataset
+    datasets:
+    - _component_: ds.IdentityDataset
+      identity: ${identity_token}
+      length: 250000
+      train_on_input: true
+seed: null
+shuffle: true
+batch_size: 4
+optimizer:
+  _component_: torch.optim.AdamW
+  weight_decay: 0.01
+  lr: 0.0003
+lr_scheduler:
+  _component_: torchtune.modules.get_cosine_schedule_with_warmup
+  num_warmup_steps: 100
+loss:
+  _component_: torch.nn.CrossEntropyLoss
+epochs: 1
+max_steps_per_epoch: null
+gradient_accumulation_steps: 64
+compile: false
+output_dir: /tmp/lora_finetune_output
+metric_logger:
+  _component_: torchtune.utils.metric_logging.DiskLogger
+  log_dir: ${output_dir}
+log_every_n_steps: null
+device: cuda
+dtype: bf16
+enable_activation_checkpointing: false
+profiler:
+  _component_: torchtune.utils.profiler
+  enabled: false
+inference:
+  prompt_template: 'Video:
+
+    {video}
+
+    Caption the previous video.'
+  max_new_tokens: 300
+  temperature: 0.6
+  top_k: 300
+  quantizer: null