Skip to content

Commit ba42e0b

Browse files
authored
Make ultralytics code snippet generic (huggingface#1101)
Follow-up PR after huggingface#1093 and especially this thread huggingface#1093 (comment). I agree with @NielsRogge's point that `ultralytics` _is_ the library name. Yolov8 or Yolov10 are "just" model names. --- _originally from @merveenoyan (huggingface#1093 (comment) > there's no way of knowing which yolo version would it be because there's no config file. The solution I suggest in this PR is to use the repo tags instead. I have checked and all 258 [ultralytics models](https://huggingface.co/models?other=ultralytics) have a "yolovX" tag. ~In parallel I suggest to remove the `yolov10` library since code snippet is now correctly generated for them as well. I have checked and [29 out of 63 yolov10 models](https://huggingface.co/models?other=yolov10) don't have the `ultranalytics` tag so they'll loose the code snippet. If we agree on this solution I'll open PRs on these repos and everything should be good.~ See below: in the end, I kept `yolov10` but unified the code snippet generation. @merveenoyan @NielsRogge @pcuenca let me know what you think. Will this definitely close the yolo/ultranalytics tags matter ? **Note:** some models have multiple yolo tags (e.g. https://huggingface.co/Ultralytics/YOLOv8 has v10, v8, v3, v5, v9). In that case I suggest to only take the first one (v10 in this case). But to be honest, this looks more like something to fix in the model's metadata [here](https://huggingface.co/Ultralytics/YOLOv8/blob/main/README.md?code=true#L13).
1 parent 51413a4 commit ba42e0b

File tree

2 files changed

+32
-30
lines changed

2 files changed

+32
-30
lines changed

packages/tasks/src/model-libraries-snippets.ts

Lines changed: 14 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1028,23 +1028,24 @@ wavs = chat.infer(texts, )
10281028
torchaudio.save("output1.wav", torch.from_numpy(wavs[0]), 24000)`,
10291029
];
10301030

1031-
export const ultralytics = (model: ModelData): string[] => [
1032-
`from ultralytics import YOLOv8 # modify the yolo version here
1031+
export const ultralytics = (model: ModelData): string[] => {
1032+
// ultralytics models must have a version tag (e.g. `yolov8`)
1033+
const versionTag = model.tags.find((tag) => tag.match(/^yolov\d+$/));
10331034

1034-
model = YOLOv8.from_pretrained("${model.id}")
1035-
source = 'http://images.cocodataset.org/val2017/000000039769.jpg'
1036-
model.predict(source=source, save=True)
1037-
`,
1038-
];
1035+
const className = versionTag ? `YOLOv${versionTag.slice(4)}` : "YOLOvXX";
1036+
const prefix = versionTag
1037+
? ""
1038+
: `# Couldn't find a valid YOLO version tag.\n# Replace XX with the correct version.\n`;
10391039

1040-
export const yolov10 = (model: ModelData): string[] => [
1041-
`from ultralytics import YOLOv10
1040+
return [
1041+
prefix +
1042+
`from ultralytics import ${className}
10421043
1043-
model = YOLOv10.from_pretrained("${model.id}")
1044+
model = ${className}.from_pretrained("${model.id}")
10441045
source = 'http://images.cocodataset.org/val2017/000000039769.jpg'
1045-
model.predict(source=source, save=True)
1046-
`,
1047-
];
1046+
model.predict(source=source, save=True)`,
1047+
];
1048+
};
10481049

10491050
export const birefnet = (model: ModelData): string[] => [
10501051
`# Option 1: use with transformers

packages/tasks/src/model-libraries.ts

Lines changed: 18 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -819,6 +819,15 @@ export const MODEL_LIBRARIES_UI_ELEMENTS = {
819819
repoUrl: "https://github.com/microsoft/TRELLIS",
820820
countDownloads: `path_extension:"safetensors"`,
821821
},
822+
ultralytics: {
823+
prettyLabel: "ultralytics",
824+
repoName: "ultralytics",
825+
repoUrl: "https://github.com/ultralytics/ultralytics",
826+
docsUrl: "https://github.com/ultralytics/ultralytics",
827+
filter: false,
828+
countDownloads: `path_extension:"pt"`,
829+
snippets: snippets.ultralytics,
830+
},
822831
"unity-sentis": {
823832
prettyLabel: "unity-sentis",
824833
repoName: "unity-sentis",
@@ -848,23 +857,6 @@ export const MODEL_LIBRARIES_UI_ELEMENTS = {
848857
docsUrl: "https://github.com/jasonppy/VoiceCraft",
849858
snippets: snippets.voicecraft,
850859
},
851-
ultralytics: {
852-
prettyLabel: "ultralytics",
853-
repoName: "ultralytics",
854-
repoUrl: "https://github.com/ultralytics/ultralytics",
855-
docsUrl: "https://github.com/ultralytics/ultralytics",
856-
filter: false,
857-
countDownloads: `path_extension:"pt"`,
858-
snippets: snippets.ultralytics,
859-
},
860-
yolov10: {
861-
prettyLabel: "YOLOv10",
862-
repoName: "yolov10",
863-
repoUrl: "https://github.com/THU-MIG/yolov10",
864-
docsUrl: "https://github.com/THU-MIG/yolov10",
865-
countDownloads: `path_extension:"pt"`,
866-
snippets: snippets.yolov10,
867-
},
868860
whisperkit: {
869861
prettyLabel: "WhisperKit",
870862
repoName: "WhisperKit",
@@ -873,6 +865,15 @@ export const MODEL_LIBRARIES_UI_ELEMENTS = {
873865
snippets: snippets.whisperkit,
874866
countDownloads: `path_filename:"model" AND path_extension:"mil" AND _exists_:"path_prefix"`,
875867
},
868+
yolov10: {
869+
// YOLOv10 is a fork of ultraLytics. Code snippets and download count are the same but the repo is different.
870+
prettyLabel: "YOLOv10",
871+
repoName: "YOLOv10",
872+
repoUrl: "https://github.com/THU-MIG/yolov10",
873+
docsUrl: "https://github.com/THU-MIG/yolov10",
874+
countDownloads: `path_extension:"pt"`,
875+
snippets: snippets.ultralytics,
876+
},
876877
"3dtopia-xl": {
877878
prettyLabel: "3DTopia-XL",
878879
repoName: "3DTopia-XL",

0 commit comments

Comments
 (0)