Repository: tenstorrent/tt-metal
Pull request #19899: Add dependency for yolo models for downloading model weights
Author: login: ssinghalTT
Created: 2025-03-31T07:59:50Z | Merged: 2025-04-06T10:24:46Z

### Problem description
Yolo models download weights from Ultralytics. To use the libs, we need the dependencies.

### What's changed
Added the dependency for python_env and system dependency.

### Checklist
- [x] [All post commit](https://github.com/tenstorrent/tt-metal/actions/workflows/all-post-commit-workflows.yaml) CI passes
- [ ] [Blackhole Post commit](https://github.com/tenstorrent/tt-metal/actions/workflows/blackhole-post-commit.yaml) CI passes (if applicable)
- [x] [Model regression](https://github.com/tenstorrent/tt-metal/actions/workflows/perf-models.yaml) CI passes (if applicable)
- [x] [Device performance regression](https://github.com/tenstorrent/tt-metal/actions/workflows/perf-device-models.yaml) CI passes (if applicable)
- [ ] **(For models and ops writers)** Full [new models tests](https://github.com/tenstorrent/tt-metal/actions/workflows/full-new-models-suite.yaml) CI passes (if applicable)
- [ ] New/Existing tests provide coverage for changes


Diff:
diff --git a/dockerfile/Dockerfile b/dockerfile/Dockerfile
index 1cb31eb27573..e59f36a9aba3 100644
--- a/dockerfile/Dockerfile
+++ b/dockerfile/Dockerfile
@@ -284,6 +284,7 @@ RUN apt-get remove -y gdb || true \
     && rm -rf /tmp/gdb-build
 
 # Install dev deps
+# libgl1-mesa-glx is needed for Yolo models https://github.com/tenstorrent/tt-metal/pull/19899
 RUN apt-get update && apt-get install -y --no-install-recommends \
     acl \
     emacs \
@@ -291,6 +292,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     nano \
     openssh-server \
     vim \
+    libgl1-mesa-glx \
     && apt-get clean && rm -rf /var/lib/apt/lists/*
 
 #############################################################
diff --git a/tests/ttnn/integration_tests/yolov8x/test_yolov8x.py b/tests/ttnn/integration_tests/yolov8x/test_yolov8x.py
index a365ff812546..b373bac6fcf4 100644
--- a/tests/ttnn/integration_tests/yolov8x/test_yolov8x.py
+++ b/tests/ttnn/integration_tests/yolov8x/test_yolov8x.py
@@ -10,9 +10,9 @@
 from pathlib import Path
 import torch.nn as nn
 from loguru import logger
+from ultralytics import YOLO
 from tests.ttnn.utils_for_testing import assert_with_pcc
 from models.utility_functions import disable_persistent_kernel_cache
-from models.experimental.functional_yolov8x.reference import yolov8x_utils
 from models.experimental.functional_yolov8x.tt.ttnn_yolov8x import YOLOv8xModel, Conv, C2f, SPPF, DFL
 from models.experimental.functional_yolov8x.tt.ttnn_yolov8x_utils import (
     ttnn_decode_bboxes,
@@ -20,18 +20,7 @@
 )
 from models.experimental.functional_yolov8x.reference import yolov8x
 
-try:
-    sys.modules["ultralytics"] = yolov8x_utils
-    sys.modules["ultralytics.nn.tasks"] = yolov8x_utils
-    sys.modules["ultralytics.nn.modules.conv"] = yolov8x_utils
-    sys.modules["ultralytics.nn.modules.block"] = yolov8x_utils
-    sys.modules["ultralytics.nn.modules.head"] = yolov8x_utils
 
-except KeyError:
-    logger.info("models.experimental.functional_yolov8x.reference.yolov8x_utils not found.")
-
-
-# For testing reference
 def decode_bboxes(distance, anchor_points, xywh=True, dim=1):
     lt, rb = distance.chunk(2, dim)
     x1y1 = anchor_points - lt
@@ -59,77 +48,7 @@ def make_anchors(feats, strides, grid_cell_offset=0.5):
     return torch.cat(anchor_points), torch.cat(stride_tensor)
 
 
-class Ensemble(nn.ModuleList):
-    def __init__(self):
-        super(Ensemble, self).__init__()
-
-    def forward(self, x, augment=False):
-        y = []
-        for module in self:
-            y.append(module(x, augment)[0])
-        y = torch.cat(y, 1)
-        return y, None
-
-
-def attempt_download(file, repo="ultralytics/assets"):
-    tests = Path(__file__).parent.parent / "yolov8x"
-    file_path = tests / Path(str(file).strip().replace("'", "").lower())
-
-    if not file_path.exists():
-        name = "yolov8x.pt"
-        msg = f"{file_path} missing, try downloading from https://github.com/{repo}/releases/"
-        try:
-            url = f"https://github.com/{repo}/releases/download/v8.3.0/{name}"
-            logger.info(f"Downloading {url} to {file_path}...")
-            torch.hub.download_url_to_file(url, file_path)
-
-            assert file_path.exists() and file_path.stat().st_size > 1e6, f"Download failed for {name}"
-        except Exception as e:
-            logger.info(f"Error downloading from GitHub: {e}. Trying secondary source...")
-
-            url = f"https://storage.googleapis.com/{repo}/ckpt/{name}"
-            logger.info(f"Downloading {url} to {file_path}...")
-            os.system(f"curl -L {url} -o {file_path}")
-
-            if not file_path.exists() or file_path.stat().st_size < 1e6:
-                file_path.unlink(missing_ok=True)
-                logger.info(f"ERROR: Download failure for {msg}")
-            else:
-                logger.info(f"Download succeeded from secondary source!")
-    return file_path
-
-
-def attempt_load(weights, map_location=None):
-    model = Ensemble()
-    for w in weights if isinstance(weights, list) else [weights]:
-        weight_path = attempt_download(w)
-        logger.info("Loading weights from:", weight_path)
-        ckpt = torch.load(weight_path, map_location=map_location)
-        model.append(ckpt["ema" if ckpt.get("ema") else "model"].float().eval())
-    for m in model.modules():
-        if isinstance(m, (nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU)):
-            m.inplace = True
-        elif isinstance(m, nn.Upsample):
-            m.recompute_scale_factor = None
-    if len(model) == 1:
-        return model[-1]
-    else:
-        for k in ["names", "stride"]:
-            setattr(model, k, getattr(model[-1], k))
-        return model
-
-
-def run_submodule(x, submodule):
-    y = []
-    for m in submodule:
-        if m.f != -1:
-            x = y[m.f] if isinstance(m.f, int) else [x if j == -1 else y[j] for j in m.f]
-        x = m(x)
-        y.append(x)
-    return x
-
-
-@pytest.mark.parametrize("device_params", [{"l1_small_size": 32768}], indirect=True)
+@pytest.mark.parametrize("device_params", [{"l1_small_size": 32768}], indirect=True, ids=["0"])
 @pytest.mark.parametrize(
     "input_tensor",
     [torch.rand((1, 3, 640, 640))],
@@ -151,10 +70,13 @@ def test_yolov8x_640(device, input_tensor, use_pretrained_weight):
 
     inp_h, inp_w = input_tensor.shape[2], input_tensor.shape[3]
     if use_pretrained_weight:
-        torch_model = attempt_load("yolov8x.pt", map_location="cpu")
+        torch_model = YOLO("yolov8x.pt")
+        torch_model = torch_model.model
+        torch_model.eval()
         state_dict = torch_model.state_dict()
     else:
         torch_model = yolov8x.DetectionModel()
+        torch_model.eval()
         state_dict = torch_model.state_dict()
     parameters = custom_preprocessor(device, state_dict, inp_h, inp_w)
     ttnn_model = YOLOv8xModel(device=device, parameters=parameters)
@@ -180,7 +102,8 @@ def test_yolov8x_640(device, input_tensor, use_pretrained_weight):
 def test_Conv(device, input_tensor):
     disable_persistent_kernel_cache()
 
-    torch_model = attempt_load("yolov8x.pt", map_location="cpu")
+    torch_model = YOLO("yolov8x.pt")
+    torch_model = torch_model.model
     torch_model.eval()
 
     ttnn_input = ttnn.from_torch(input_tensor, dtype=ttnn.bfloat16, layout=ttnn.ROW_MAJOR_LAYOUT, device=device)
@@ -219,7 +142,8 @@ def test_Conv(device, input_tensor):
 def test_C2f(device, input_tensor, reset_seeds):
     disable_persistent_kernel_cache()
 
-    torch_model = attempt_load("yolov8x.pt", map_location="cpu")
+    torch_model = YOLO("yolov8x.pt")
+    torch_model = torch_model.model
     torch_model.eval()
 
     ttnn_input = ttnn.from_torch(input_tensor, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device)
@@ -264,7 +188,8 @@ def test_C2f(device, input_tensor, reset_seeds):
 def test_SPPF(device, input_tensor, reset_seeds):
     disable_persistent_kernel_cache()
 
-    torch_model = attempt_load("yolov8x.pt", map_location="cpu")
+    torch_model = YOLO("yolov8x.pt")
+    torch_model = torch_model.model
     torch_model.eval()
 
     ttnn_input = ttnn.from_torch(input_tensor, dtype=ttnn.bfloat16, layout=ttnn.ROW_MAJOR_LAYOUT, device=device)
@@ -299,7 +224,8 @@ def test_SPPF(device, input_tensor, reset_seeds):
 def test_DFL(device, input_tensor, reset_seeds):
     disable_persistent_kernel_cache()
 
-    torch_model = attempt_load("yolov8x.pt", map_location="cpu")
+    torch_model = YOLO("yolov8x.pt")
+    torch_model = torch_model.model
     torch_model.eval()
 
     ttnn_input = ttnn.from_torch(input_tensor, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=device)
diff --git a/tests/ttnn/integration_tests/yolov9c/test_ttnn_yolov9c.py b/tests/ttnn/integration_tests/yolov9c/test_ttnn_yolov9c.py
index 3c4a0ce59141..0959de20fe9b 100644
--- a/tests/ttnn/integration_tests/yolov9c/test_ttnn_yolov9c.py
+++ b/tests/ttnn/integration_tests/yolov9c/test_ttnn_yolov9c.py
@@ -19,14 +19,13 @@
 from models.experimental.functional_yolov9c.tt import ttnn_yolov9c
 from models.experimental.functional_yolov9c.reference import yolov9c
 from models.experimental.functional_yolov9c.demo.demo_utils import attempt_load
+from ultralytics import YOLO
 
 
 @run_for_wormhole_b0()
 @pytest.mark.parametrize(
     "use_pretrained_weight",
-    [
-        False,
-    ],
+    [True],
 )
 @pytest.mark.parametrize("device_params", [{"l1_small_size": 79104}], indirect=True)
 def test_yolov9c(use_pretrained_weight, device, use_program_cache, reset_seeds):
@@ -34,7 +33,7 @@ def test_yolov9c(use_pretrained_weight, device, use_program_cache, reset_seeds):
     state_dict = None
 
     if use_pretrained_weight:
-        torch_model = attempt_load("yolov9c.pt", map_location="cpu")
+        torch_model = YOLO("yolov9c.pt")
         state_dict = torch_model.state_dict()
 
     torch_model = yolov9c.YoloV9()
diff --git a/tt_metal/python_env/requirements-dev.txt b/tt_metal/python_env/requirements-dev.txt
index 540b919e0fa6..db66f7a2d553 100644
--- a/tt_metal/python_env/requirements-dev.txt
+++ b/tt_metal/python_env/requirements-dev.txt
@@ -49,6 +49,7 @@ sentencepiece==0.1.97
 numba>=0.58.1
 librosa==0.10.0
 timm==0.6.13
+ultralytics==8.2
 opencv-python-headless==4.8.1.78
 diffusers==0.32.2
 accelerate==0.27.2


Reviews:
user: login: tt-rkim
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-01T21:40:48Z
user: login: ssinghalTT
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-02T02:26:13Z
user: login: tt-rkim
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-02T11:53:46Z
user: login: ssinghalTT
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-02T22:55:45Z
user: login: dvartaniansTT
body: LGTM as long as you have tested all the available yolo variants and it runs smoothly
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-02T23:04:15Z
user: login: bbradelTT
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-02T23:11:59Z
user: login: tt-rkim
body: Please leave some comments in the requirements file and Dockerfile explaining what the new deps are for like you see for some other dependencies. You should link the PR link for this PR as well

Ensure they build and pass pipelines
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-03T12:44:15Z

Review comments:
diff_hunk: @@ -214,6 +214,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     libtbb-dev \
     libcapstone-dev \
     libgtest-dev \
+    libgl1-mesa-glx \
path: dockerfile/Dockerfile
user: login: tt-rkim
body: Why not in `ci-test` ?
created_at: 2025-04-01T21:40:47Z
author_association: COLLABORATOR
diff_hunk: @@ -214,6 +214,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     libtbb-dev \
     libcapstone-dev \
     libgtest-dev \
+    libgl1-mesa-glx \
path: dockerfile/Dockerfile
user: login: ssinghalTT
body: @tt-rkim  was recommended here:
https://github.com/tenstorrent/tt-metal/issues/19812
created_at: 2025-04-02T02:26:13Z
author_association: CONTRIBUTOR
diff_hunk: @@ -214,6 +214,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     libtbb-dev \
     libcapstone-dev \
     libgtest-dev \
+    libgl1-mesa-glx \
path: dockerfile/Dockerfile
user: login: tt-rkim
body: Was that a mistake? @blozano-tt 
I think we prefer to have it in `ci-test` rather than `ci-build` as it's a python dep
created_at: 2025-04-02T11:53:42Z
author_association: COLLABORATOR
diff_hunk: @@ -214,6 +214,7 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     libtbb-dev \
     libcapstone-dev \
     libgtest-dev \
+    libgl1-mesa-glx \
path: dockerfile/Dockerfile
user: login: ssinghalTT
body: fixed
created_at: 2025-04-02T22:55:45Z
author_association: CONTRIBUTOR

PR comments:
user: login: ssinghalTT
created_at: 2025-03-31T19:16:17Z
body: @blozano-tt @bkeith-TT 

I have the PR here: https://github.com/tenstorrent/tt-metal/pull/19899 with the fix. Can you take a look
author_association: CONTRIBUTOR
user: login: ssinghalTT
created_at: 2025-04-02T08:09:51Z
body: CIs:
https://github.com/tenstorrent/tt-metal/actions/runs/14209377918
https://github.com/tenstorrent/tt-metal/actions/runs/14213489248
https://github.com/tenstorrent/tt-metal/actions/runs/14209195867
https://github.com/tenstorrent/tt-metal/actions/runs/14209380912
author_association: CONTRIBUTOR
user: login: dvartaniansTT
created_at: 2025-04-03T01:25:59Z
body: @ssinghalTT please rebase and fix any potential conflicts (hopefully none) :) and we can merge this asap. 
author_association: CONTRIBUTOR
user: login: dvartaniansTT
created_at: 2025-04-03T01:27:10Z
body: @tt-rkim could you please review and if ok approve this? cc: @ssinghalTT 
author_association: CONTRIBUTOR