Add Sana .pth files into download counting in huggingface repo (#1061)

Co-authored-by: vb <[email protected]> Co-authored-by: Lucain <[email protected]> Co-authored-by: Lucain <[email protected]>
huggingface · Dec 5, 2024 · 690e143 · 690e143
1 parent 32365c0
commit 690e143
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 0 deletions.
diff --git a/packages/tasks/src/model-libraries-snippets.ts b/packages/tasks/src/model-libraries-snippets.ts
@@ -982,6 +982,26 @@ IWorker engine = WorkerFactory.CreateWorker(BackendType.GPUCompute, model);
 `,
 ];
 
+export const sana = (model: ModelData): string[] => [
+	`
+# Load the model and infer image from text
+import torch
+from app.sana_pipeline import SanaPipeline
+from torchvision.utils import save_image
+
+sana = SanaPipeline("configs/sana_config/1024ms/Sana_1600M_img1024.yaml")
+sana.from_pretrained("hf://${model.id}")
+
+image = sana(
+    prompt='a cyberpunk cat with a neon sign that says "Sana"',
+    height=1024,
+    width=1024,
+    guidance_scale=5.0,
+    pag_guidance_scale=2.0,
+    num_inference_steps=18,
+) `,
+];
+
 export const vfimamba = (model: ModelData): string[] => [
 	`from Trainer_finetune import Model
 

diff --git a/packages/tasks/src/model-libraries.ts b/packages/tasks/src/model-libraries.ts
@@ -787,6 +787,13 @@ export const MODEL_LIBRARIES_UI_ELEMENTS = {
 		filter: true,
 		countDownloads: `path_extension:"sentis"`,
 	},
+	sana: {
+		prettyLabel: "Sana",
+		repoName: "Sana",
+		repoUrl: "https://github.com/NVlabs/Sana",
+		countDownloads: `path_extension:"pth"`,
+		snippets: snippets.sana,
+	},
 	"vfi-mamba": {
 		prettyLabel: "VFIMamba",
 		repoName: "VFIMamba",