Skip to content

Add Argmax DiffusionKit Snippet #869

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
Aug 30, 2024
Merged

Add Argmax DiffusionKit Snippet #869

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
7c47f27
add diffusionkit snippet and add diffusionkit to MODEL_LIBRARIES_UI_E…
ardaatahan Aug 26, 2024
44a5f8b
remove last 3 prompts and add a new one
ardaatahan Aug 26, 2024
57c93ca
move hf login instructions down
ardaatahan Aug 26, 2024
031b28d
fix snippet content and move it up to preserve alphabetical order
ardaatahan Aug 27, 2024
6055b69
show relevant pipeline according to tag
ardaatahan Aug 28, 2024
9d55cdd
Merge branch 'huggingface:main' into add-diffusionkit-snippet
ardaatahan Aug 29, 2024
49b810c
dynamically change num_steps and cfg_weight in diffusionkit snippet a…
ardaatahan Aug 29, 2024
6b983bb
Linter
pcuenca Aug 30, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 42 additions & 0 deletions packages/tasks/src/model-libraries-snippets.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -170,6 +170,48 @@ export const diffusers = (model: ModelData): string[] => {
}
};

export const diffusionkit = (model: ModelData): string[] => {
const sd3Snippet = `# Pipeline for Stable Diffusion 3
from diffusionkit.mlx import DiffusionPipeline

pipeline = DiffusionPipeline(
shift=3.0,
use_t5=False,
model_version=${model.id},
low_memory_mode=True,
a16=True,
w16=True,
)`;

const fluxSnippet = `# Pipeline for Flux
from diffusionkit.mlx import FluxPipeline

pipeline = FluxPipeline(
shift=1.0,
model_version=${model.id},
low_memory_mode=True,
a16=True,
w16=True,
)`;

const generateSnippet = `# Image Generation
HEIGHT = 512
WIDTH = 512
NUM_STEPS = ${model.tags.includes("flux") ? 4 : 50}
CFG_WEIGHT = ${model.tags.includes("flux") ? 0 : 5}

image, _ = pipeline.generate_image(
"a photo of a cat",
cfg_weight=CFG_WEIGHT,
num_steps=NUM_STEPS,
latent_size=(HEIGHT // 8, WIDTH // 8),
)`;

const pipelineSnippet = model.tags.includes("flux") ? fluxSnippet : sd3Snippet;

return [pipelineSnippet, generateSnippet];
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We usually show a single snippet per task (in this case it would be pipeline preparation followed by generation), but I'm not opposed to using two to differentiate instantiation from generation.

};

export const cartesia_pytorch = (model: ModelData): string[] => [
`# pip install --no-binary :all: cartesia-pytorch
from cartesia_pytorch import ReneLMHeadModel
Expand Down
6 changes: 6 additions & 0 deletions packages/tasks/src/model-libraries.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -181,6 +181,12 @@ export const MODEL_LIBRARIES_UI_ELEMENTS = {
filter: true,
/// diffusers has its own more complex "countDownloads" query
},
diffusionkit: {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would be nice to have some models listed in https://huggingface.co/models?other=diffusionkit before merging this PR. You can do that by tagging models as diffusionkit in the model card metadata.

I opened a first PR to do that: https://huggingface.co/argmaxinc/mlx-FLUX.1-schnell/discussions/2

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Merged 👍

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I noticed that the SD3 models have the DiffusionKit tag: https://huggingface.co/models?other=DiffusionKit. The correct tag should be diffusionkit (lowercased). Make sure to update them to benefit from this PR.

(the correct key is the key defined here, so diffusionkit in this case. Below we are defining prettyLabel and repoName but that's just for aesthetic in the UI)

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We generated these models with diffusionkit much before we had an MLX implementation, and in fact don't yet have a public CoreML inference implementation of DiffusionKit (in the works). These models are actually more directly usable by the huggingface swift-coreml-diffusers app, although they were generated with diffusionkit. So in this case we might want to just remove the lowercase diffusionkit tag? @pcuenca what do you think about a local-app deeplink for Diffusers.app? We had discussed adding arbitrary repos in the past.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated our tags and model cards: https://huggingface.co/models?other=diffusionkit

prettyLabel: "DiffusionKit",
repoName: "DiffusionKit",
repoUrl: "https://github.com/argmaxinc/DiffusionKit",
snippets: snippets.diffusionkit,
},
doctr: {
prettyLabel: "docTR",
repoName: "doctr",
Expand Down
Loading