Add VSCode settings and update dependencies; refactor model prediction logic to include modality and targets

.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+  "python.analysis.typeCheckingMode": "standard"
+}

inference_utils/inference.py

@@ -15,13 +15,15 @@ import numpy as np
 import torch.nn.functional as F
 from PIL import Image
 from torchvision import transforms
-#from utils.visualizer import Visualizer
+
+# from utils.visualizer import Visualizer
 # from detectron2.utils.colormap import random_color
 # from detectron2.data import MetadataCatalog
 # from detectron2.structures import BitMasks
 from modeling.language.loss import vl_similarity
 from utilities.constants import BIOMED_CLASSES
-#from detectron2.data.datasets.builtin_meta import COCO_CATEGORIES
+
+# from detectron2.data.datasets.builtin_meta import COCO_CATEGORIES
 
 # import cv2
 # import os
@@ -33,62 +35,70 @@ import random
 t = []
 t.append(transforms.Resize((1024, 1024), interpolation=Image.BICUBIC))
 transform = transforms.Compose(t)
-#metadata = MetadataCatalog.get('coco_2017_train_panoptic')
-all_classes = ['background'] + [name.replace('-other','').replace('-merged','') 
-                                for name in BIOMED_CLASSES] + ["others"]
+# metadata = MetadataCatalog.get('coco_2017_train_panoptic')
+all_classes = (
+    ["background"]
+    + [name.replace("-other", "").replace("-merged", "") for name in BIOMED_CLASSES]
+    + ["others"]
+)
 # colors_list = [(np.array(color['color'])/255).tolist() for color in COCO_CATEGORIES] + [[1, 1, 1]]
 
 # use color list from matplotlib
 import matplotlib.colors as mcolors
+
 colors = dict(mcolors.TABLEAU_COLORS, **mcolors.BASE_COLORS)
-colors_list = [list(colors.values())[i] for i in range(16)] 
+colors_list = [list(colors.values())[i] for i in range(16)]
 
 from .output_processing import mask_stats, combine_masks
-    
+
 
 @torch.no_grad()
-def interactive_infer_image(model, image, prompts):
+def interactive_infer_image(model, image, prompts) -> np.ndarray:
 
     image_resize = transform(image)
     width = image.size[0]
     height = image.size[1]
     image_resize = np.asarray(image_resize)
-    image = torch.from_numpy(image_resize.copy()).permute(2,0,1).cuda()
+    image = torch.from_numpy(image_resize.copy()).permute(2, 0, 1).cuda()
 
-    data = {"image": image, 'text': prompts, "height": height, "width": width}
-    
-    # inistalize task
-    model.model.task_switch['spatial'] = False
-    model.model.task_switch['visual'] = False
-    model.model.task_switch['grounding'] = True
-    model.model.task_switch['audio'] = False
-    model.model.task_switch['grounding'] = True
+    data = {"image": image, "text": prompts, "height": height, "width": width}
 
+    # inistalize task
+    model.model.task_switch["spatial"] = False
+    model.model.task_switch["visual"] = False
+    model.model.task_switch["grounding"] = True
+    model.model.task_switch["audio"] = False
+    model.model.task_switch["grounding"] = True
 
     batch_inputs = [data]
-    results,image_size,extra = model.model.evaluate_demo(batch_inputs)
+    results, image_size, extra = model.model.evaluate_demo(batch_inputs)
 
-    pred_masks = results['pred_masks'][0]
-    v_emb = results['pred_captions'][0]
-    t_emb = extra['grounding_class']
+    pred_masks = results["pred_masks"][0]
+    v_emb = results["pred_captions"][0]
+    t_emb = extra["grounding_class"]
 
     t_emb = t_emb / (t_emb.norm(dim=-1, keepdim=True) + 1e-7)
     v_emb = v_emb / (v_emb.norm(dim=-1, keepdim=True) + 1e-7)
 
     temperature = model.model.sem_seg_head.predictor.lang_encoder.logit_scale
     out_prob = vl_similarity(v_emb, t_emb, temperature=temperature)
-    
+
     matched_id = out_prob.max(0)[1]
-    pred_masks_pos = pred_masks[matched_id,:,:]
-    pred_class = results['pred_logits'][0][matched_id].max(dim=-1)[1]
+    pred_masks_pos = pred_masks[matched_id, :, :]
+    pred_class = results["pred_logits"][0][matched_id].max(dim=-1)[1]
 
     # interpolate mask to ori size
-    pred_mask_prob = F.interpolate(pred_masks_pos[None,], (data['height'], data['width']), 
-                                   mode='bilinear')[0,:,:data['height'],:data['width']].sigmoid().cpu().numpy()
-    pred_masks_pos = (1*(pred_mask_prob > 0.5)).astype(np.uint8)
-    
-    return pred_mask_prob
+    pred_mask_prob = (
+        F.interpolate(
+            pred_masks_pos[None,], (data["height"], data["width"]), mode="bilinear"
+        )[0, :, : data["height"], : data["width"]]
+        .sigmoid()
+        .cpu()
+        .numpy()
+    )
+    pred_masks_pos = (1 * (pred_mask_prob > 0.5)).astype(np.uint8)
 
+    return pred_mask_prob
 
 
 # def interactive_infer_panoptic_biomedseg(model, image, tasks, reftxt=None):
@@ -103,7 +113,7 @@ def interactive_infer_image(model, image, prompts):
 #     data = {"image": images, "height": height, "width": width}
 #     if len(tasks) == 0:
 #         tasks = ["Panoptic"]
-    
+
 #     # inistalize task
 #     model.model.task_switch['spatial'] = False
 #     model.model.task_switch['visual'] = False
@@ -114,7 +124,7 @@ def interactive_infer_image(model, image, prompts):
 #     assert isinstance(reftxt, list), f"reftxt should be a list of strings, but got {type(reftxt)}"
 #     model.model.task_switch['grounding'] = True
 #     predicts = {}
-#     for i, txt in enumerate(reftxt): 
+#     for i, txt in enumerate(reftxt):
 #         data['text'] = txt
 #         batch_inputs = [data]
 
@@ -129,7 +139,7 @@ def interactive_infer_image(model, image, prompts):
 
 #         temperature = model.model.sem_seg_head.predictor.lang_encoder.logit_scale
 #         out_prob = vl_similarity(v_emb, t_emb, temperature=temperature)
-        
+
 #         matched_id = out_prob.max(0)[1]
 #         pred_masks_pos = pred_masks[matched_id,:,:]
 #         pred_class = results['pred_logits'][0][matched_id].max(dim=-1)[1]
@@ -144,18 +154,17 @@ def interactive_infer_image(model, image, prompts):
 #         pred_mask_prob = F.interpolate(pred_masks_pos[None,], image_size[-2:], mode='bilinear')[0,:,:data['height'],:data['width']].sigmoid().cpu().numpy()
 #         #pred_masks_pos = 1*(pred_mask_prob > 0.5)
 #         predicts[txt] = pred_mask_prob[0]
-        
+
 #     masks = combine_masks(predicts)
-        
+
 #     predict_mask_stats = {}
 #     print(masks.keys())
 #     for i, txt in enumerate(masks):
 #         mask = masks[txt]
 #         demo = visual.draw_binary_mask(mask, color=colors_list[i], text=txt)
 #         predict_mask_stats[txt] = mask_stats((predicts[txt]*255), image_ori)
-        
+
 #     res = demo.get_image()
 #     torch.cuda.empty_cache()
 #     # return Image.fromarray(res), stroke_inimg, stroke_refimg
 #     return Image.fromarray(res), None, predict_mask_stats
-

inference_utils/model.py

@@ -10,14 +10,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from dataclasses import dataclass
 import os
+from typing import Tuple
+
 from PIL import Image
 from huggingface_hub import hf_hub_download
 import matplotlib.pyplot as plt
 import numpy as np
 import torch
-from inference_utils.inference import interactive_infer_image
 
+from inference_utils.inference import interactive_infer_image
+from inference_utils.output_processing import check_mask_stats
 from modeling import build_model
 from modeling.BaseModel import BaseModel
 from utilities.arguments import load_opt_from_config_files
@@ -25,15 +29,35 @@ from utilities.constants import BIOMED_CLASSES
 from utilities.distributed import init_distributed
 
 
+zero_tensor = torch.zeros(1, 1, 1)
+
+
+@dataclass
+class PredictionTarget:
+    target: str
+    pred_mask: torch.Tensor = zero_tensor
+    adjusted_p_value: float = -1.0
+
+
 class Model:
     def init(self):
         self._model = init_model()
 
-    def predict(self, image: Image, modality_type: str, targets: list[str]) -> Image:
-        return predict(self._model, image, targets)
+    def predict(
+        self, image: Image.Image, modality_type: str, targets: list[str]
+    ) -> Tuple[Image.Image, str]:
+        image_annotated, prediction_targets_not_found = predict(
+            self._model, image, modality_type, targets
+        )
+        targets_not_found_str = (
+            "\n".join(t.target for t in prediction_targets_not_found)
+            if prediction_targets_not_found
+            else "All targets were found!"
+        )
+        return image_annotated, targets_not_found_str
 
 
-def init_model():
+def init_model() -> BaseModel:
     # Download model
     model_file = hf_hub_download(
         repo_id="microsoft/BiomedParse",
@@ -55,33 +79,77 @@ def init_model():
     return model
 
 
-def predict(model, image: Image, prompts: list[str]):
-    if not prompts:
-        return None
+def predict(
+    model: BaseModel, image: Image.Image, modality_type: str, targets: list[str]
+) -> Tuple[Image.Image, list[PredictionTarget]]:
+    assert len(targets) > 0, "At least one target is required"
+
+    prediction_tasks = [PredictionTarget(target=target) for target in targets]
 
     # Convert to RGB if needed
     if image.mode != "RGB":
         image = image.convert("RGB")
 
     # Get predictions
-    pred_mask = interactive_infer_image(model, image, prompts)
+    pred_mask = interactive_infer_image(model, image, targets)
 
-    # Generate visualization
-    colors = generate_colors(len(prompts))
-    pred_overlay = overlay_masks(
-        image, [1 * (pred_mask[i] > 0.5) for i in range(len(prompts))], colors
+    for i, pt in enumerate(prediction_tasks):
+        pt.pred_mask = pred_mask[i]
+
+    image_np = np.array(image)
+
+    for pt in prediction_tasks:
+        adj_p_value = check_mask_stats(
+            image_np, pt.pred_mask * 255, modality_type, pt.target
+        )
+        pt.adjusted_p_value = float(adj_p_value)
+
+    pred_tasks_found, pred_tasks_not_found = segregate_prediction_tasks(
+        prediction_tasks, 0.05
     )
 
-    return pred_overlay
+    # Generate visualization
+    colors = generate_colors(len(pred_tasks_found))
+    masks = [1 * (pred_mask[i] > 0.5) for i in range(len(pred_tasks_found))]
+    pred_overlay = overlay_masks(image, masks, colors)
+
+    return pred_overlay, pred_tasks_not_found
+
+
+def segregate_prediction_tasks(
+    prediction_tasks: list[PredictionTarget], p_value_threshold: float
+) -> tuple[list[PredictionTarget], list[PredictionTarget]]:
+    """Segregates Prediction Tasks by p-value
+
+    Prediction tasks with a p-value higher than p_value_threshold go into the targets_found list.
+    Otherwise, they go into the targets_not_found list.
+    """
+
+    targets_found = []
+    targets_not_found = []
+    for pt in prediction_tasks:
+        if pt.adjusted_p_value > p_value_threshold:
+            targets_found.append(pt)
+        else:
+            targets_not_found.append(pt)
+
+    return targets_found, targets_not_found
 
 
-def generate_colors(n):
+def generate_colors(n: int) -> list[Tuple[int, int, int]]:
     cmap = plt.get_cmap("tab10")
-    colors = [tuple(int(255 * val) for val in cmap(i)[:3]) for i in range(n)]
+    colors = [
+        (int(255 * cmap(i)[0]), int(255 * cmap(i)[1]), int(255 * cmap(i)[2]))
+        for i in range(n)
+    ]
     return colors
 
 
-def overlay_masks(image, masks, colors):
+def overlay_masks(
+    image: Image.Image,
+    masks: list[np.ndarray],
+    colors: list[Tuple[int, int, int]],
+) -> Image.Image:
     overlay = image.copy()
     overlay = np.array(overlay, dtype=np.uint8)
     for mask, color in zip(masks, colors):

inference_utils/model_mock.py

@@ -12,9 +12,7 @@
 
 
 from typing import Tuple
-from PIL import ImageDraw, ImageFont
-from PIL.Image import Image
-import gradio as gr
+from PIL import ImageDraw, ImageFont, Image
 import random
 
 
@@ -23,8 +21,8 @@ class Model:
         pass
 
     def predict(
-        self, image: Image, modality_type: str, targets: list[str]
-    ) -> Tuple[Image, str]:
+        self, image: Image.Image, modality_type: str, targets: list[str]
+    ) -> Tuple[Image.Image, str]:
         # Randomly split targets into found and not found
         targets_found = random.sample(targets, k=len(targets) // 2)
         targets_not_found = [t for t in targets if t not in targets_found]

pyproject.toml

@@ -7,6 +7,7 @@ requires-python = ">=3.9"
 dependencies = [
     "gradio==4.44.1",
     "pytest>=8.3.4",
+    "torch>=2.5.1",
 ]
 
 [tool.pytest.ini_options]

uv.lock

@@ -46,12 +46,14 @@ source = { virtual = "." }
 dependencies = [
     { name = "gradio" },
     { name = "pytest" },
+    { name = "torch" },
 ]
 
 [package.metadata]
 requires-dist = [
     { name = "gradio", specifier = "==4.44.1" },
     { name = "pytest", specifier = ">=8.3.4" },
+    { name = "torch", specifier = ">=2.5.1" },
 ]
 
 [[package]]
@@ -722,6 +724,24 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979 },
 ]
 
+[[package]]
+name = "mpmath"
+version = "1.3.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/e0/47/dd32fa426cc72114383ac549964eecb20ecfd886d1e5ccf5340b55b02f57/mpmath-1.3.0.tar.gz", hash = "sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f", size = 508106 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198 },
+]
+
+[[package]]
+name = "networkx"
+version = "3.2.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/c4/80/a84676339aaae2f1cfdf9f418701dd634aef9cc76f708ef55c36ff39c3ca/networkx-3.2.1.tar.gz", hash = "sha256:9f1bb5cf3409bf324e0a722c20bdb4c20ee39bf1c30ce8ae499c8502b0b5e0c6", size = 2073928 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/d5/f0/8fbc882ca80cf077f1b246c0e3c3465f7f415439bdea6b899f6b19f61f70/networkx-3.2.1-py3-none-any.whl", hash = "sha256:f18c69adc97877c42332c170849c96cefa91881c99a7cb3e95b7c659ebdc1ec2", size = 1647772 },
+]
+
 [[package]]
 name = "numpy"
 version = "2.0.2"
@@ -774,6 +794,126 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/cc/dc/d330a6faefd92b446ec0f0dfea4c3207bb1fef3c4771d19cf4543efd2c78/numpy-2.0.2-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:a46288ec55ebbd58947d31d72be2c63cbf839f0a63b49cb755022310792a3385", size = 15828784 },
 ]
 
+[[package]]
+name = "nvidia-cublas-cu12"
+version = "12.4.5.8"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/7f/7f/7fbae15a3982dc9595e49ce0f19332423b260045d0a6afe93cdbe2f1f624/nvidia_cublas_cu12-12.4.5.8-py3-none-manylinux2014_aarch64.whl", hash = "sha256:0f8aa1706812e00b9f19dfe0cdb3999b092ccb8ca168c0db5b8ea712456fd9b3", size = 363333771 },
+    { url = "https://files.pythonhosted.org/packages/ae/71/1c91302526c45ab494c23f61c7a84aa568b8c1f9d196efa5993957faf906/nvidia_cublas_cu12-12.4.5.8-py3-none-manylinux2014_x86_64.whl", hash = "sha256:2fc8da60df463fdefa81e323eef2e36489e1c94335b5358bcb38360adf75ac9b", size = 363438805 },
+]
+
+[[package]]
+name = "nvidia-cuda-cupti-cu12"
+version = "12.4.127"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/93/b5/9fb3d00386d3361b03874246190dfec7b206fd74e6e287b26a8fcb359d95/nvidia_cuda_cupti_cu12-12.4.127-py3-none-manylinux2014_aarch64.whl", hash = "sha256:79279b35cf6f91da114182a5ce1864997fd52294a87a16179ce275773799458a", size = 12354556 },
+    { url = "https://files.pythonhosted.org/packages/67/42/f4f60238e8194a3106d06a058d494b18e006c10bb2b915655bd9f6ea4cb1/nvidia_cuda_cupti_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:9dec60f5ac126f7bb551c055072b69d85392b13311fcc1bcda2202d172df30fb", size = 13813957 },
+]
+
+[[package]]
+name = "nvidia-cuda-nvrtc-cu12"
+version = "12.4.127"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/77/aa/083b01c427e963ad0b314040565ea396f914349914c298556484f799e61b/nvidia_cuda_nvrtc_cu12-12.4.127-py3-none-manylinux2014_aarch64.whl", hash = "sha256:0eedf14185e04b76aa05b1fea04133e59f465b6f960c0cbf4e37c3cb6b0ea198", size = 24133372 },
+    { url = "https://files.pythonhosted.org/packages/2c/14/91ae57cd4db3f9ef7aa99f4019cfa8d54cb4caa7e00975df6467e9725a9f/nvidia_cuda_nvrtc_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:a178759ebb095827bd30ef56598ec182b85547f1508941a3d560eb7ea1fbf338", size = 24640306 },
+]
+
+[[package]]
+name = "nvidia-cuda-runtime-cu12"
+version = "12.4.127"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/a1/aa/b656d755f474e2084971e9a297def515938d56b466ab39624012070cb773/nvidia_cuda_runtime_cu12-12.4.127-py3-none-manylinux2014_aarch64.whl", hash = "sha256:961fe0e2e716a2a1d967aab7caee97512f71767f852f67432d572e36cb3a11f3", size = 894177 },
+    { url = "https://files.pythonhosted.org/packages/ea/27/1795d86fe88ef397885f2e580ac37628ed058a92ed2c39dc8eac3adf0619/nvidia_cuda_runtime_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:64403288fa2136ee8e467cdc9c9427e0434110899d07c779f25b5c068934faa5", size = 883737 },
+]
+
+[[package]]
+name = "nvidia-cudnn-cu12"
+version = "9.1.0.70"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "nvidia-cublas-cu12" },
+]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/9f/fd/713452cd72343f682b1c7b9321e23829f00b842ceaedcda96e742ea0b0b3/nvidia_cudnn_cu12-9.1.0.70-py3-none-manylinux2014_x86_64.whl", hash = "sha256:165764f44ef8c61fcdfdfdbe769d687e06374059fbb388b6c89ecb0e28793a6f", size = 664752741 },
+]
+
+[[package]]
+name = "nvidia-cufft-cu12"
+version = "11.2.1.3"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "nvidia-nvjitlink-cu12" },
+]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/7a/8a/0e728f749baca3fbeffad762738276e5df60851958be7783af121a7221e7/nvidia_cufft_cu12-11.2.1.3-py3-none-manylinux2014_aarch64.whl", hash = "sha256:5dad8008fc7f92f5ddfa2101430917ce2ffacd86824914c82e28990ad7f00399", size = 211422548 },
+    { url = "https://files.pythonhosted.org/packages/27/94/3266821f65b92b3138631e9c8e7fe1fb513804ac934485a8d05776e1dd43/nvidia_cufft_cu12-11.2.1.3-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f083fc24912aa410be21fa16d157fed2055dab1cc4b6934a0e03cba69eb242b9", size = 211459117 },
+]
+
+[[package]]
+name = "nvidia-curand-cu12"
+version = "10.3.5.147"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/80/9c/a79180e4d70995fdf030c6946991d0171555c6edf95c265c6b2bf7011112/nvidia_curand_cu12-10.3.5.147-py3-none-manylinux2014_aarch64.whl", hash = "sha256:1f173f09e3e3c76ab084aba0de819c49e56614feae5c12f69883f4ae9bb5fad9", size = 56314811 },
+    { url = "https://files.pythonhosted.org/packages/8a/6d/44ad094874c6f1b9c654f8ed939590bdc408349f137f9b98a3a23ccec411/nvidia_curand_cu12-10.3.5.147-py3-none-manylinux2014_x86_64.whl", hash = "sha256:a88f583d4e0bb643c49743469964103aa59f7f708d862c3ddb0fc07f851e3b8b", size = 56305206 },
+]
+
+[[package]]
+name = "nvidia-cusolver-cu12"
+version = "11.6.1.9"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "nvidia-cublas-cu12" },
+    { name = "nvidia-cusparse-cu12" },
+    { name = "nvidia-nvjitlink-cu12" },
+]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/46/6b/a5c33cf16af09166845345275c34ad2190944bcc6026797a39f8e0a282e0/nvidia_cusolver_cu12-11.6.1.9-py3-none-manylinux2014_aarch64.whl", hash = "sha256:d338f155f174f90724bbde3758b7ac375a70ce8e706d70b018dd3375545fc84e", size = 127634111 },
+    { url = "https://files.pythonhosted.org/packages/3a/e1/5b9089a4b2a4790dfdea8b3a006052cfecff58139d5a4e34cb1a51df8d6f/nvidia_cusolver_cu12-11.6.1.9-py3-none-manylinux2014_x86_64.whl", hash = "sha256:19e33fa442bcfd085b3086c4ebf7e8debc07cfe01e11513cc6d332fd918ac260", size = 127936057 },
+]
+
+[[package]]
+name = "nvidia-cusparse-cu12"
+version = "12.3.1.170"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "nvidia-nvjitlink-cu12" },
+]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/96/a9/c0d2f83a53d40a4a41be14cea6a0bf9e668ffcf8b004bd65633f433050c0/nvidia_cusparse_cu12-12.3.1.170-py3-none-manylinux2014_aarch64.whl", hash = "sha256:9d32f62896231ebe0480efd8a7f702e143c98cfaa0e8a76df3386c1ba2b54df3", size = 207381987 },
+    { url = "https://files.pythonhosted.org/packages/db/f7/97a9ea26ed4bbbfc2d470994b8b4f338ef663be97b8f677519ac195e113d/nvidia_cusparse_cu12-12.3.1.170-py3-none-manylinux2014_x86_64.whl", hash = "sha256:ea4f11a2904e2a8dc4b1833cc1b5181cde564edd0d5cd33e3c168eff2d1863f1", size = 207454763 },
+]
+
+[[package]]
+name = "nvidia-nccl-cu12"
+version = "2.21.5"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/df/99/12cd266d6233f47d00daf3a72739872bdc10267d0383508b0b9c84a18bb6/nvidia_nccl_cu12-2.21.5-py3-none-manylinux2014_x86_64.whl", hash = "sha256:8579076d30a8c24988834445f8d633c697d42397e92ffc3f63fa26766d25e0a0", size = 188654414 },
+]
+
+[[package]]
+name = "nvidia-nvjitlink-cu12"
+version = "12.4.127"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/02/45/239d52c05074898a80a900f49b1615d81c07fceadd5ad6c4f86a987c0bc4/nvidia_nvjitlink_cu12-12.4.127-py3-none-manylinux2014_aarch64.whl", hash = "sha256:4abe7fef64914ccfa909bc2ba39739670ecc9e820c83ccc7a6ed414122599b83", size = 20552510 },
+    { url = "https://files.pythonhosted.org/packages/ff/ff/847841bacfbefc97a00036e0fce5a0f086b640756dc38caea5e1bb002655/nvidia_nvjitlink_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:06b3b9b25bf3f8af351d664978ca26a16d2c5127dbd53c0497e28d1fb9611d57", size = 21066810 },
+]
+
+[[package]]
+name = "nvidia-nvtx-cu12"
+version = "12.4.127"
+source = { registry = "https://pypi.org/simple" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/06/39/471f581edbb7804b39e8063d92fc8305bdc7a80ae5c07dbe6ea5c50d14a5/nvidia_nvtx_cu12-12.4.127-py3-none-manylinux2014_aarch64.whl", hash = "sha256:7959ad635db13edf4fc65c06a6e9f9e55fc2f92596db928d169c0bb031e88ef3", size = 100417 },
+    { url = "https://files.pythonhosted.org/packages/87/20/199b8713428322a2f22b722c62b8cc278cc53dffa9705d744484b5035ee9/nvidia_nvtx_cu12-12.4.127-py3-none-manylinux2014_x86_64.whl", hash = "sha256:781e950d9b9f60d8241ccea575b32f5105a5baf4c2351cab5256a24869f12a1a", size = 99144 },
+]
+
 [[package]]
 name = "orjson"
 version = "3.10.13"
@@ -1294,6 +1434,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/6a/23/8146aad7d88f4fcb3a6218f41a60f6c2d4e3a72de72da1825dc7c8f7877c/semantic_version-2.10.0-py2.py3-none-any.whl", hash = "sha256:de78a3b8e0feda74cabc54aab2da702113e33ac9d9eb9d2389bcf1f58b7d9177", size = 15552 },
 ]
 
+[[package]]
+name = "setuptools"
+version = "75.6.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/43/54/292f26c208734e9a7f067aea4a7e282c080750c4546559b58e2e45413ca0/setuptools-75.6.0.tar.gz", hash = "sha256:8199222558df7c86216af4f84c30e9b34a61d8ba19366cc914424cdbd28252f6", size = 1337429 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/55/21/47d163f615df1d30c094f6c8bbb353619274edccf0327b185cc2493c2c33/setuptools-75.6.0-py3-none-any.whl", hash = "sha256:ce74b49e8f7110f9bf04883b730f4765b774ef3ef28f722cce7c273d253aaf7d", size = 1224032 },
+]
+
 [[package]]
 name = "shellingham"
 version = "1.5.4"
@@ -1334,6 +1483,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/96/00/2b325970b3060c7cecebab6d295afe763365822b1306a12eeab198f74323/starlette-0.41.3-py3-none-any.whl", hash = "sha256:44cedb2b7c77a9de33a8b74b2b90e9f50d11fcf25d8270ea525ad71a25374ff7", size = 73225 },
 ]
 
+[[package]]
+name = "sympy"
+version = "1.13.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "mpmath" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/ca/99/5a5b6f19ff9f083671ddf7b9632028436167cd3d33e11015754e41b249a4/sympy-1.13.1.tar.gz", hash = "sha256:9cebf7e04ff162015ce31c9c6c9144daa34a93bd082f54fd8f12deca4f47515f", size = 7533040 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/b2/fe/81695a1aa331a842b582453b605175f419fe8540355886031328089d840a/sympy-1.13.1-py3-none-any.whl", hash = "sha256:db36cdc64bf61b9b24578b6f7bab1ecdd2452cf008f34faa33776680c26d66f8", size = 6189177 },
+]
+
 [[package]]
 name = "tomli"
 version = "2.2.1"
@@ -1382,6 +1543,52 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/68/4f/12207897848a653d03ebbf6775a29d949408ded5f99b2d87198bc5c93508/tomlkit-0.12.0-py3-none-any.whl", hash = "sha256:926f1f37a1587c7a4f6c7484dae538f1345d96d793d9adab5d3675957b1d0766", size = 37334 },
 ]
 
+[[package]]
+name = "torch"
+version = "2.5.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "filelock" },
+    { name = "fsspec" },
+    { name = "jinja2" },
+    { name = "networkx" },
+    { name = "nvidia-cublas-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" },
+    { name = "nvidia-cuda-cupti-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" },
+    { name = "nvidia-cuda-nvrtc-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" },
+    { name = "nvidia-cuda-runtime-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" },
+    { name = "nvidia-cudnn-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" },
+    { name = "nvidia-cufft-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" },
+    { name = "nvidia-curand-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" },
+    { name = "nvidia-cusolver-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" },
+    { name = "nvidia-cusparse-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" },
+    { name = "nvidia-nccl-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" },
+    { name = "nvidia-nvjitlink-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" },
+    { name = "nvidia-nvtx-cu12", marker = "platform_machine == 'x86_64' and platform_system == 'Linux'" },
+    { name = "setuptools", marker = "python_full_version >= '3.12'" },
+    { name = "sympy" },
+    { name = "triton", marker = "python_full_version < '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'" },
+    { name = "typing-extensions" },
+]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/2a/ef/834af4a885b31a0b32fff2d80e1e40f771e1566ea8ded55347502440786a/torch-2.5.1-cp310-cp310-manylinux1_x86_64.whl", hash = "sha256:71328e1bbe39d213b8721678f9dcac30dfc452a46d586f1d514a6aa0a99d4744", size = 906446312 },
+    { url = "https://files.pythonhosted.org/packages/69/f0/46e74e0d145f43fa506cb336eaefb2d240547e4ce1f496e442711093ab25/torch-2.5.1-cp310-cp310-manylinux2014_aarch64.whl", hash = "sha256:34bfa1a852e5714cbfa17f27c49d8ce35e1b7af5608c4bc6e81392c352dbc601", size = 91919522 },
+    { url = "https://files.pythonhosted.org/packages/a5/13/1eb674c8efbd04d71e4a157ceba991904f633e009a584dd65dccbafbb648/torch-2.5.1-cp310-cp310-win_amd64.whl", hash = "sha256:32a037bd98a241df6c93e4c789b683335da76a2ac142c0973675b715102dc5fa", size = 203088048 },
+    { url = "https://files.pythonhosted.org/packages/a9/9d/e0860474ee0ff8f6ef2c50ec8f71a250f38d78a9b9df9fd241ad3397a65b/torch-2.5.1-cp310-none-macosx_11_0_arm64.whl", hash = "sha256:23d062bf70776a3d04dbe74db950db2a5245e1ba4f27208a87f0d743b0d06e86", size = 63877046 },
+    { url = "https://files.pythonhosted.org/packages/d1/35/e8b2daf02ce933e4518e6f5682c72fd0ed66c15910ea1fb4168f442b71c4/torch-2.5.1-cp311-cp311-manylinux1_x86_64.whl", hash = "sha256:de5b7d6740c4b636ef4db92be922f0edc425b65ed78c5076c43c42d362a45457", size = 906474467 },
+    { url = "https://files.pythonhosted.org/packages/40/04/bd91593a4ca178ece93ca55f27e2783aa524aaccbfda66831d59a054c31e/torch-2.5.1-cp311-cp311-manylinux2014_aarch64.whl", hash = "sha256:340ce0432cad0d37f5a31be666896e16788f1adf8ad7be481196b503dad675b9", size = 91919450 },
+    { url = "https://files.pythonhosted.org/packages/0d/4a/e51420d46cfc90562e85af2fee912237c662ab31140ab179e49bd69401d6/torch-2.5.1-cp311-cp311-win_amd64.whl", hash = "sha256:603c52d2fe06433c18b747d25f5c333f9c1d58615620578c326d66f258686f9a", size = 203098237 },
+    { url = "https://files.pythonhosted.org/packages/d0/db/5d9cbfbc7968d79c5c09a0bc0bc3735da079f2fd07cc10498a62b320a480/torch-2.5.1-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:31f8c39660962f9ae4eeec995e3049b5492eb7360dd4f07377658ef4d728fa4c", size = 63884466 },
+    { url = "https://files.pythonhosted.org/packages/8b/5c/36c114d120bfe10f9323ed35061bc5878cc74f3f594003854b0ea298942f/torch-2.5.1-cp312-cp312-manylinux1_x86_64.whl", hash = "sha256:ed231a4b3a5952177fafb661213d690a72caaad97d5824dd4fc17ab9e15cec03", size = 906389343 },
+    { url = "https://files.pythonhosted.org/packages/6d/69/d8ada8b6e0a4257556d5b4ddeb4345ea8eeaaef3c98b60d1cca197c7ad8e/torch-2.5.1-cp312-cp312-manylinux2014_aarch64.whl", hash = "sha256:3f4b7f10a247e0dcd7ea97dc2d3bfbfc90302ed36d7f3952b0008d0df264e697", size = 91811673 },
+    { url = "https://files.pythonhosted.org/packages/5f/ba/607d013b55b9fd805db2a5c2662ec7551f1910b4eef39653eeaba182c5b2/torch-2.5.1-cp312-cp312-win_amd64.whl", hash = "sha256:73e58e78f7d220917c5dbfad1a40e09df9929d3b95d25e57d9f8558f84c9a11c", size = 203046841 },
+    { url = "https://files.pythonhosted.org/packages/57/6c/bf52ff061da33deb9f94f4121fde7ff3058812cb7d2036c97bc167793bd1/torch-2.5.1-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:8c712df61101964eb11910a846514011f0b6f5920c55dbf567bff8a34163d5b1", size = 63858109 },
+    { url = "https://files.pythonhosted.org/packages/69/72/20cb30f3b39a9face296491a86adb6ff8f1a47a897e4d14667e6cf89d5c3/torch-2.5.1-cp313-cp313-manylinux1_x86_64.whl", hash = "sha256:9b61edf3b4f6e3b0e0adda8b3960266b9009d02b37555971f4d1c8f7a05afed7", size = 906393265 },
+    { url = "https://files.pythonhosted.org/packages/a9/18/81c399e8f4f1580d34bf99d827cb5fb5cf7a18a266bb5d30ca3ec2e89ba6/torch-2.5.1-cp39-cp39-manylinux1_x86_64.whl", hash = "sha256:1f3b7fb3cf7ab97fae52161423f81be8c6b8afac8d9760823fd623994581e1a3", size = 906479005 },
+    { url = "https://files.pythonhosted.org/packages/5d/86/1c4b168d52cddb8d17952a7b5b25f69ef0da1fc34de1223d73d0d9db1801/torch-2.5.1-cp39-cp39-manylinux2014_aarch64.whl", hash = "sha256:7974e3dce28b5a21fb554b73e1bc9072c25dde873fa00d54280861e7a009d7dc", size = 91846074 },
+    { url = "https://files.pythonhosted.org/packages/76/49/4a0a8b19ce8f9bf32fcab4e863c7e2366f519f9826c84ca250567b11a014/torch-2.5.1-cp39-cp39-win_amd64.whl", hash = "sha256:46c817d3ea33696ad3b9df5e774dba2257e9a4cd3c4a3afbf92f6bb13ac5ce2d", size = 203000888 },
+    { url = "https://files.pythonhosted.org/packages/25/07/3548a7cfcf69d0eccec2ee79ee3913f1cdaadb27b36946774db86729ee47/torch-2.5.1-cp39-none-macosx_11_0_arm64.whl", hash = "sha256:8046768b7f6d35b85d101b4b38cba8aa2f3cd51952bc4c06a49580f2ce682291", size = 63876023 },
+]
+
 [[package]]
 name = "tqdm"
 version = "4.67.1"
@@ -1394,6 +1601,20 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/d0/30/dc54f88dd4a2b5dc8a0279bdd7270e735851848b762aeb1c1184ed1f6b14/tqdm-4.67.1-py3-none-any.whl", hash = "sha256:26445eca388f82e72884e0d580d5464cd801a3ea01e63e5601bdff9ba6a48de2", size = 78540 },
 ]
 
+[[package]]
+name = "triton"
+version = "3.1.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "filelock" },
+]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/98/29/69aa56dc0b2eb2602b553881e34243475ea2afd9699be042316842788ff5/triton-3.1.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6b0dd10a925263abbe9fa37dcde67a5e9b2383fc269fdf59f5657cac38c5d1d8", size = 209460013 },
+    { url = "https://files.pythonhosted.org/packages/86/17/d9a5cf4fcf46291856d1e90762e36cbabd2a56c7265da0d1d9508c8e3943/triton-3.1.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0f34f6e7885d1bf0eaaf7ba875a5f0ce6f3c13ba98f9503651c1e6dc6757ed5c", size = 209506424 },
+    { url = "https://files.pythonhosted.org/packages/78/eb/65f5ba83c2a123f6498a3097746607e5b2f16add29e36765305e4ac7fdd8/triton-3.1.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c8182f42fd8080a7d39d666814fa36c5e30cc00ea7eeeb1a2983dbb4c99a0fdc", size = 209551444 },
+    { url = "https://files.pythonhosted.org/packages/c4/69/57e0fed438d547524e08bfedc587078314176ad1c15c8be904d3f03149ec/triton-3.1.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:aafa9a20cd0d9fee523cd4504aa7131807a864cd77dcf6efe7e981f18b8c6c11", size = 209460480 },
+]
+
 [[package]]
 name = "typer"
 version = "0.15.1"