Bringup tt-torch models in forge

env/core_requirements.txt

@@ -51,3 +51,4 @@ pytorch_forecasting==1.0.0
 patool
 openpyxl==3.1.5
 GitPython==3.1.44
+kornia==0.8.0

forge/test/models/pytorch/text/bloom/test_bloom.py

@@ -0,0 +1,45 @@
+# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+import torch
+
+import forge
+from forge.verify.verify import verify
+
+from test.models.pytorch.text.bloom.utils.utils import load_input, load_model
+from test.models.utils import Framework, Source, Task, build_module_name
+
+
+def test_bloom(record_forge_property):
+
+    # Wrapper to get around past key values
+    class Wrapper(torch.nn.Module):
+        def __init__(self, model):
+            super().__init__()
+            self.model = model
+
+        def forward(self, input_ids, attention_mask):
+            output = self.model(input_ids, None, attention_mask)
+            return output
+
+    # Build Module Name
+    module_name = build_module_name(
+        framework=Framework.PYTORCH,
+        model="bloom",
+        source=Source.HUGGINGFACE,
+        task=Task.CAUSAL_LM,
+    )
+
+    # Record Forge Property
+    record_forge_property("model_name", module_name)
+
+    # Load model and input
+    model = load_model()
+    framework_model = Wrapper(model)
+    inputs = load_input()
+
+    # Forge compile framework model
+    compiled_model = forge.compile(framework_model, sample_inputs=inputs, module_name=module_name)
+
+    # Model Verification
+    verify(inputs, framework_model, compiled_model)

forge/test/models/pytorch/text/bloom/utils/utils.py

@@ -0,0 +1,25 @@
+# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+def load_model():
+    model = AutoModelForCausalLM.from_pretrained("bigscience/bloom-1b1")
+    model.config.use_cache = False
+    model.eval()
+    return model
+
+
+def load_input():
+    test_input = "This is a sample text from "
+    tokenizer = AutoTokenizer.from_pretrained("bigscience/bloom-1b1", padding_side="left")
+    inputs = tokenizer.encode_plus(
+        test_input,
+        return_tensors="pt",
+        max_length=32,
+        padding="max_length",
+        add_special_tokens=True,
+        truncation=True,
+    )
+    return [inputs["input_ids"], inputs["attention_mask"]]

forge/test/models/pytorch/text/mamba/test_mamba.py

@@ -0,0 +1,53 @@
+# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+# Reference: https://huggingface.co/state-spaces/mamba-2.8b-hf
+
+import pytest
+import torch
+
+import forge
+from forge.verify.verify import verify
+
+from test.models.pytorch.text.mamba.utils.utils import load_input, load_model
+from test.models.utils import Framework, Source, Task, build_module_name
+
+variants = [
+    "state-spaces/mamba-790m-hf",
+    "state-spaces/mamba-2.8b-hf",
+    "state-spaces/mamba-1.4b-hf",
+    "state-spaces/mamba-370m-hf",
+]
+
+
+@pytest.mark.parametrize("variant", variants)
+def test_mamba(record_forge_property, variant):
+
+    # Build Module Name
+    module_name = build_module_name(
+        framework=Framework.PYTORCH, model="mamba", variant=variant, task=Task.CAUSAL_LM, source=Source.HUGGINGFACE
+    )
+
+    # Record Forge Property
+    record_forge_property("model_name", module_name)
+
+    # Wrapper to return only the output tensor, excluding cache or additional outputs
+    class Wrapper(torch.nn.Module):
+        def __init__(self, model):
+            super().__init__()
+            self.model = model
+
+        def forward(self, input_ids):
+            output = self.model(input_ids)
+            return output[0]
+
+    # Load model and input
+    model = load_model(variant)
+    framework_model = Wrapper(model)
+    inputs = load_input(variant)
+
+    # Forge compile framework model
+    compiled_model = forge.compile(framework_model, sample_inputs=inputs)
+
+    # Model Verification
+    verify(inputs, framework_model, compiled_model, module_name=module_name)

forge/test/models/pytorch/text/mamba/utils/utils.py

@@ -0,0 +1,19 @@
+# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+# Reference: https://huggingface.co/state-spaces/mamba-2.8b-hf
+
+from transformers import AutoTokenizer, MambaForCausalLM
+
+
+def load_model(variant):
+    model = MambaForCausalLM.from_pretrained(variant)
+    model.eval()
+    return model
+
+
+def load_input(variant):
+    prompt = "Hey how are you doing?"
+    tokenizer = AutoTokenizer.from_pretrained(variant)
+    input_ids = tokenizer(prompt, return_tensors="pt")["input_ids"]
+    return [input_ids]

forge/test/models/pytorch/vision/beit/test_beit.py

@@ -0,0 +1,38 @@
+# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+import pytest
+
+import forge
+from forge.verify.verify import verify
+
+from test.models.pytorch.vision.beit.utils.utils import load_input, load_model
+from test.models.utils import Framework, Source, Task, build_module_name
+
+variants = ["microsoft/beit-base-patch16-224", "microsoft/beit-large-patch16-224"]
+
+
+@pytest.mark.parametrize("variant", variants)
+def test_beit_image_classification(record_forge_property, variant):
+
+    # Build Module Name
+    module_name = build_module_name(
+        framework=Framework.PYTORCH,
+        model="beit",
+        variant=variant,
+        source=Source.HUGGINGFACE,
+        task=Task.IMAGE_CLASSIFICATION,
+    )
+
+    # Record Forge Property
+    record_forge_property("model_name", module_name)
+
+    # Load model and input
+    framework_model = load_model(variant)
+    inputs = load_input(variant)
+
+    # Forge compile framework model
+    compiled_model = forge.compile(framework_model, sample_inputs=inputs, module_name=module_name)
+
+    # Model Verification
+    verify(inputs, framework_model, compiled_model)

forge/test/models/pytorch/vision/beit/utils/utils.py

@@ -0,0 +1,20 @@
+# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+import requests
+from PIL import Image
+from transformers import BeitForImageClassification, BeitImageProcessor
+
+
+def load_model(variant):
+    model = BeitForImageClassification.from_pretrained(variant)
+    model.eval()
+    return model
+
+
+def load_input(variant):
+    url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+    image = Image.open(requests.get(url, stream=True).raw)
+    processor = BeitImageProcessor.from_pretrained(variant)
+    inputs = processor(images=image, return_tensors="pt")
+    return [inputs["pixel_values"]]

forge/test/models/pytorch/vision/glpn_kitti/test_glpn_kitti.py

@@ -0,0 +1,32 @@
+# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+import forge
+from forge.verify.verify import verify
+
+from test.models.pytorch.vision.glpn_kitti.utils.utils import load_input, load_model
+from test.models.utils import Framework, Source, Task, build_module_name
+
+
+def test_glpn_kitti(record_forge_property):
+
+    # Build Module Name
+    module_name = build_module_name(
+        framework=Framework.PYTORCH,
+        model="glpn_kitti",
+        source=Source.HUGGINGFACE,
+        task=Task.DEPTH_ESTIMATION,
+    )
+
+    # Record Forge Property
+    record_forge_property("model_name", module_name)
+
+    # Load model and input
+    framework_model = load_model()
+    inputs = load_input()
+
+    # Forge compile framework model
+    compiled_model = forge.compile(framework_model, sample_inputs=inputs, module_name=module_name)
+
+    # Model Verification
+    verify(inputs, framework_model, compiled_model)

forge/test/models/pytorch/vision/glpn_kitti/utils/utils.py

@@ -0,0 +1,20 @@
+# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+import requests
+from PIL import Image
+from transformers import GLPNForDepthEstimation, GLPNImageProcessor
+
+
+def load_model():
+    model = GLPNForDepthEstimation.from_pretrained("vinvino02/glpn-kitti")
+    model.eval()
+    return model
+
+
+def load_input():
+    url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+    image = Image.open(requests.get(url, stream=True).raw)
+    processor = GLPNImageProcessor.from_pretrained("vinvino02/glpn-kitti")
+    inputs = processor(images=image, return_tensors="pt")
+    return [inputs["pixel_values"]]

forge/test/models/pytorch/vision/rmbg/test_rmbg.py

@@ -0,0 +1,33 @@
+# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+
+import forge
+from forge.verify.verify import verify
+
+from test.models.pytorch.vision.rmbg.utils.utils import load_input, load_model
+from test.models.utils import Framework, Source, Task, build_module_name
+
+
+def test_RMBG(record_forge_property):
+
+    # Build Module Name
+    module_name = build_module_name(
+        framework=Framework.PYTORCH,
+        model="rmbg_2_0",
+        source=Source.HUGGINGFACE,
+        task=Task.IMAGE_SEGMENTATION,
+    )
+
+    # Record Forge Property
+    record_forge_property("model_name", module_name)
+
+    # Load model and input
+    framework_model = load_model()
+    inputs = load_input()
+
+    # Forge compile framework model
+    compiled_model = forge.compile(framework_model, sample_inputs=inputs, module_name=module_name)
+
+    # Model Verification
+    verify(inputs, framework_model, compiled_model)

forge/test/models/pytorch/vision/rmbg/utils/utils.py

@@ -0,0 +1,28 @@
+# SPDX-FileCopyrightText: (c) 2024 Tenstorrent AI ULC
+
+# SPDX-License-Identifier: Apache-2.0
+import requests
+from PIL import Image
+from torchvision import transforms
+from transformers import AutoModelForImageSegmentation
+
+
+def load_model():
+    model = AutoModelForImageSegmentation.from_pretrained("briaai/RMBG-2.0", trust_remote_code=True)
+    model.eval()
+    return model
+
+
+def load_input():
+    url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+    image = Image.open(requests.get(url, stream=True).raw)
+    image_size = (1024, 1024)
+    transform_image = transforms.Compose(
+        [
+            transforms.Resize(image_size),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ]
+    )
+    inputs = transform_image(image).unsqueeze(0)
+    return [inputs]

forge/test/models/utils.py

@@ -36,6 +36,7 @@ class Task(StrEnum):
     CONDITIONAL_GENERATION = "cond_gen"
     IMAGE_ENCODING = "img_enc"
     VISUAL_BACKBONE = "visual_bb"
+    DEPTH_ESTIMATION = "depth_estimation"
 
 
 class Source(StrEnum):