Add Pruna AI library snippets (no formatting changes) (#1733)

This update enhances support for Pruna AI, providing users with tailored
code snippets for model integrations with Transformers and Diffusers.

- Introduced a new library entry for Pruna AI in model-libraries.
- Added main entry point and specific snippet generation functions for
diffusers and transformers models.
- Cleaned up whitespace inconsistencies in existing snippets.

TLDR: Pruna API normally mimics the Transformers and Diffusers API, so
we can use `PrunaModel.from_pretrained` on top of pipelines or specific
models. We re-use the underlying snippets for both the library and do
some greedy replacements of certain part of the code snippets.

example
```python
import torch
from diffusers import FluxFillPipeline
from diffusers.utils import load_image

image = load_image("https://huggingface.co/datasets/diffusers/diffusers-images-docs/resolve/main/cup.png")
mask = load_image("https://huggingface.co/datasets/diffusers/diffusers-images-docs/resolve/main/cup_mask.png")

pipe = FluxFillPipeline.from_pretrained("black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16).to("cuda")
image = pipe(
    prompt="a white paper cup",
    image=image,
    mask_image=mask,
    height=1632,
    width=1232,
    guidance_scale=30,
    num_inference_steps=50,
    max_sequence_length=512,
    generator=torch.Generator("cpu").manual_seed(0)
).images[0]
image.save(f"flux-fill-dev.png")
```

becomes
```python
import torch
from pruna import PrunaModel
from diffusers.utils import load_image

image = load_image("https://huggingface.co/datasets/diffusers/diffusers-images-docs/resolve/main/cup.png")
mask = load_image("https://huggingface.co/datasets/diffusers/diffusers-images-docs/resolve/main/cup_mask.png")

pipe = PrunaModel.from_pretrained("black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16).to("cuda")
image = pipe(
    prompt="a white paper cup",
    image=image,
    mask_image=mask,
    height=1632,
    width=1232,
    guidance_scale=30,
    num_inference_steps=50,
    max_sequence_length=512,
    generator=torch.Generator("cpu").manual_seed(0)
).images[0]
image.save(f"flux-fill-dev.png")
```

---------

Co-authored-by: Lucain <[email protected]>

Author: davidberenstein1957

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

@@ -1905,6 +1905,87 @@ export const model2vec = (model: ModelData): string[] => [
 model = StaticModel.from_pretrained("${model.id}")`,
 ];
 
+export const pruna = (model: ModelData): string[] => {
+	let snippets: string[];
+
+	if (model.tags.includes("diffusers")) {
+		snippets = pruna_diffusers(model);
+	} else if (model.tags.includes("transformers")) {
+		snippets = pruna_transformers(model);
+	} else {
+		snippets = pruna_default(model);
+	}
+
+	const ensurePrunaModelImport = (snippet: string): string => {
+		if (!/^from pruna import PrunaModel/m.test(snippet)) {
+			return `from pruna import PrunaModel\n${snippet}`;
+		}
+		return snippet;
+	};
+	snippets = snippets.map(ensurePrunaModelImport);
+
+	if (model.tags.includes("pruna_pro-ai")) {
+		return snippets.map((snippet) =>
+			snippet.replace(/\bpruna\b/g, "pruna_pro").replace(/\bPrunaModel\b/g, "PrunaProModel")
+		);
+	}
+
+	return snippets;
+};
+
+const pruna_diffusers = (model: ModelData): string[] => {
+	const diffusersSnippets = diffusers(model);
+
+	return diffusersSnippets.map((snippet) =>
+		snippet
+			// Replace pipeline classes with PrunaModel
+			.replace(/\b\w*Pipeline\w*\b/g, "PrunaModel")
+			// Clean up diffusers imports containing PrunaModel
+			.replace(/from diffusers import ([^,\n]*PrunaModel[^,\n]*)/g, "")
+			.replace(/from diffusers import ([^,\n]+),?\s*([^,\n]*PrunaModel[^,\n]*)/g, "from diffusers import $1")
+			.replace(/from diffusers import\s*(\n|$)/g, "")
+			// Fix PrunaModel imports
+			.replace(/from diffusers import PrunaModel/g, "from pruna import PrunaModel")
+			.replace(/from diffusers import ([^,\n]+), PrunaModel/g, "from diffusers import $1")
+			.replace(/from diffusers import PrunaModel, ([^,\n]+)/g, "from diffusers import $1")
+			// Clean up whitespace
+			.replace(/\n\n+/g, "\n")
+			.trim()
+	);
+};
+
+const pruna_transformers = (model: ModelData): string[] => {
+	const info = model.transformersInfo;
+	const transformersSnippets = transformers(model);
+
+	// Replace pipeline with PrunaModel
+	let processedSnippets = transformersSnippets.map((snippet) =>
+		snippet
+			.replace(/from transformers import pipeline/g, "from pruna import PrunaModel")
+			.replace(/pipeline\([^)]*\)/g, `PrunaModel.from_pretrained("${model.id}")`)
+	);
+
+	// Additional cleanup if auto_model info is available
+	if (info?.auto_model) {
+		processedSnippets = processedSnippets.map((snippet) =>
+			snippet
+				.replace(new RegExp(`from transformers import ${info.auto_model}\n?`, "g"), "")
+				.replace(new RegExp(`${info.auto_model}.from_pretrained`, "g"), "PrunaModel.from_pretrained")
+				.replace(new RegExp(`^.*from.*import.*(, *${info.auto_model})+.*$`, "gm"), (line) =>
+					line.replace(new RegExp(`, *${info.auto_model}`, "g"), "")
+				)
+		);
+	}
+
+	return processedSnippets;
+};
+
+const pruna_default = (model: ModelData): string[] => [
+	`from pruna import PrunaModel
+model = PrunaModel.from_pretrained("${model.id}")
+`,
+];
+
 export const nemo = (model: ModelData): string[] => {
 	let command: string[] | undefined = undefined;
 	// Resolve the tag to a nemo domain/sub-domain

packages/tasks/src/model-libraries.ts

@@ -801,6 +801,13 @@ export const MODEL_LIBRARIES_UI_ELEMENTS = {
 		filter: false,
 		countDownloads: `path_extension:"pth"`,
 	},
+	"pruna-ai": {
+		prettyLabel: "Pruna AI",
+		repoName: "Pruna AI",
+		repoUrl: "https://github.com/PrunaAI/pruna",
+		snippets: snippets.pruna,
+		docsUrl: "https://docs.pruna.ai",
+	},
 	pxia: {
 		prettyLabel: "pxia",
 		repoName: "pxia",