From 7ce7a3e8cdc912db4c1cefbffa5ff040b42f09a8 Mon Sep 17 00:00:00 2001
From: NielsRogge <48327001+NielsRogge@users.noreply.github.com>
Date: Mon, 9 Sep 2024 16:34:42 +0200
Subject: [PATCH] Add VFI-Mamba as library (#885)

This PR is linked with https://github.com/MCG-NJU/VFIMamba/pull/2 to
ensure download stats are working for the VFI-Mamba repo.

It also adds a "How to use this model" button with a code snippet.

---------

Co-authored-by: Lucain <lucainp@gmail.com>
---
 packages/tasks/src/model-libraries-snippets.ts | 6 ++++++
 packages/tasks/src/model-libraries.ts          | 7 +++++++
 2 files changed, 13 insertions(+)

diff --git a/packages/tasks/src/model-libraries-snippets.ts b/packages/tasks/src/model-libraries-snippets.ts
index 3247f2e2d..879ce4f06 100644
--- a/packages/tasks/src/model-libraries-snippets.ts
+++ b/packages/tasks/src/model-libraries-snippets.ts
@@ -871,6 +871,12 @@ IWorker engine = WorkerFactory.CreateWorker(BackendType.GPUCompute, model);
 `,
 ];
 
+export const vfimamba = (model: ModelData): string[] => [
+	`from Trainer_finetune import Model
+
+model = Model.from_pretrained("${model.id}")`,
+];
+
 export const voicecraft = (model: ModelData): string[] => [
 	`from voicecraft import VoiceCraft
 
diff --git a/packages/tasks/src/model-libraries.ts b/packages/tasks/src/model-libraries.ts
index 2c163e3de..73c88ce58 100644
--- a/packages/tasks/src/model-libraries.ts
+++ b/packages/tasks/src/model-libraries.ts
@@ -650,6 +650,13 @@ export const MODEL_LIBRARIES_UI_ELEMENTS = {
 		filter: true,
 		countDownloads: `path_extension:"sentis"`,
 	},
+	"vfi-mamba": {
+		prettyLabel: "VFIMamba",
+		repoName: "VFIMamba",
+		repoUrl: "https://github.com/MCG-NJU/VFIMamba",
+		countDownloads: `path_extension:"pkl"`,
+		snippets: snippets.vfimamba,
+	},
 	voicecraft: {
 		prettyLabel: "VoiceCraft",
 		repoName: "VoiceCraft",