Add Pruna AI library snippets and integration #1684

davidberenstein1957 · 2025-08-12T06:46:45Z

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

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

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")

- 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. This update enhances support for Pruna AI, providing users with tailored code snippets for model integration.

sharpenb · 2025-08-12T10:38:09Z

This is super cool :) It looks good from the Pruna side. Loading models with this snippet would facilitate their compression to make them more efficient.

merveenoyan

thanks a lot, it's nice to see you here! interestingly this adds a lot of newlines, would be great to run linter

I think no need for a lot of comments to keep the code minimal :)

Vaibhavs10

Thanks for the PR @davidberenstein1957 - let's wait for @Wauplin to give it a review, since this PR is a bit more involved!

- Added a new function `pruna_pro` for generating tailored code snippets for Pruna AI models. - Updated the model libraries to include Pruna AI with appropriate repository details and snippet references. - Improved code readability by standardizing function syntax and cleaning up whitespace. This update further supports the integration of Pruna AI, providing users with enhanced code generation capabilities.

- Introduced a new entry for PXIA in the model-libraries, including repository details and snippet references. - Enhanced the library structure to support additional model integrations. This update expands the capabilities of the model-libraries, providing users with access to PXIA resources.

davidberenstein1957 · 2025-08-19T06:54:09Z

@Vaibhavs10 @merveenoyan, thanks for the review. :) I've resolved the formatting. I realised that the formatting commands in the Readme were merely checking and not applying the formatting directly. @Wauplin, feel free to give it a go. I can remove all comments if needed, but given that it is a bit more involved, I think they help with clarity around the formatting steps. Let me know what you prefer :)

pcuenca · 2025-08-19T08:56:52Z

packages/tasks/src/model-libraries.ts

+	"pruna_pro-ai": {
+		prettyLabel: "Pruna AI",
+		repoName: "Pruna AI",
+		repoUrl: "https://github.com/Pruna-AI/pruna-ai",
+		snippets: snippets.pruna_pro,
+		docsUrl: "https://docs.pruna.ai",
+	},


It's a bit weird to me that two different "libraries" would show up with the same name. In addition, there are only two test models on the Hub that use library pruna_pro-ai.

pcuenca · 2025-08-19T08:59:28Z

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

+	return pruna_default(model);
+};
+
+export const pruna_diffusers = (model: ModelData): string[] => {


Will the snippets always be forward-compatible with new pipelines? As an alternative, we could provide some specific snippets here instead of doing replacements.

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

pcuenca · 2025-08-19T09:05:54Z

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

-	
+


Personally, not a fan of unrelated formatting changes, but I'll let @Wauplin or @hanouticelina comment on best practices in this regard.

Thanks for the review. Happy to revert these changes if you feel it is needed.

Best would be to revert these changes from this PR and open a new one with only the formatting changes 🙏 Would make the discussion cleaner to follow :)

packages/tasks/src/model-libraries.ts

Co-authored-by: Pedro Cuenca <pedro@huggingface.co>

davidberenstein1957 · 2025-08-26T02:44:42Z

@Wauplin @hanouticelina , ping :)

julien-c · 2025-08-27T14:49:37Z

packages/tasks/src/model-libraries.ts

+	"pruna_pro-ai": {
+		prettyLabel: "Pruna AI (pro)",
+		repoName: "Pruna AI (pro)",
+		repoUrl: "https://github.com/Pruna-AI/pruna-ai",
+		snippets: snippets.pruna_pro,
+		docsUrl: "https://docs.pruna.ai",
+	},


not a fan of having two ≠ libraries. IMO this should always be Pruna (and you can prompt users to upgrade to a "pro" version from inside your library/error-messages/documentation/etc)

Same here. Since https://github.com/PrunaAI/pruna is a single library, it should be linked to a single tag and then handle the "pro" models separately.

Note that in the pruna snippet above you receive ModelData as input:

export const pruna = (model: ModelData): string[] => {

this means you can return a different snippet based on model name or tags. For instance, you could choose that all "pro" pruna models have a name ending by -pro and return a different snippet in that case.

Wauplin · 2025-08-28T09:48:01Z

packages/tasks/src/model-libraries.ts

+	"pruna-ai": {
+		prettyLabel: "Pruna AI",
+		repoName: "Pruna AI",
+		repoUrl: "https://github.com/Pruna-AI/pruna-ai",


I'm getting a 404 on this link. Shouldn't it be https://github.com/PrunaAI/pruna instead?

Suggested change

repoUrl: "https://github.com/Pruna-AI/pruna-ai",

repoUrl: "https://github.com/PrunaAI/pruna",

davidberenstein1957 requested review from SBrandeis, gary149, Wauplin, julien-c, pcuenca and ngxson as code owners August 12, 2025 06:46

merveenoyan reviewed Aug 14, 2025

View reviewed changes

Vaibhavs10 reviewed Aug 15, 2025

View reviewed changes

davidberenstein1957 added 2 commits August 19, 2025 08:45

pcuenca reviewed Aug 19, 2025

View reviewed changes

davidberenstein1957 commented Aug 19, 2025

View reviewed changes

packages/tasks/src/model-libraries.ts Outdated Show resolved Hide resolved

Apply suggestions from code review

76ea407

Co-authored-by: Pedro Cuenca <pedro@huggingface.co>

julien-c reviewed Aug 27, 2025

View reviewed changes

Wauplin reviewed Aug 28, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add Pruna AI library snippets and integration #1684

Add Pruna AI library snippets and integration #1684

Uh oh!

davidberenstein1957 commented Aug 12, 2025 •

edited

Loading

Uh oh!

sharpenb commented Aug 12, 2025

Uh oh!

merveenoyan left a comment

Uh oh!

Vaibhavs10 left a comment

Uh oh!

davidberenstein1957 commented Aug 19, 2025

Uh oh!

pcuenca Aug 19, 2025

Uh oh!

pcuenca Aug 19, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

pcuenca Aug 19, 2025

Uh oh!

davidberenstein1957 Aug 19, 2025

Uh oh!

Wauplin Aug 28, 2025

Uh oh!

Uh oh!

davidberenstein1957 commented Aug 26, 2025

Uh oh!

julien-c Aug 27, 2025

Uh oh!

Wauplin Aug 28, 2025

Uh oh!

Wauplin Aug 28, 2025

Uh oh!

Uh oh!

	repoUrl: "https://github.com/Pruna-AI/pruna-ai",
	repoUrl: "https://github.com/PrunaAI/pruna",

Add Pruna AI library snippets and integration #1684

Are you sure you want to change the base?

Add Pruna AI library snippets and integration #1684

Uh oh!

Conversation

davidberenstein1957 commented Aug 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

sharpenb commented Aug 12, 2025

Uh oh!

merveenoyan left a comment

Choose a reason for hiding this comment

Uh oh!

Vaibhavs10 left a comment

Choose a reason for hiding this comment

Uh oh!

davidberenstein1957 commented Aug 19, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

davidberenstein1957 commented Aug 26, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

davidberenstein1957 commented Aug 12, 2025 •

edited

Loading