diff --git a/packages/tasks/src/model-libraries-snippets.ts b/packages/tasks/src/model-libraries-snippets.ts index ad4c771ac..d2c19e4b4 100644 --- a/packages/tasks/src/model-libraries-snippets.ts +++ b/packages/tasks/src/model-libraries-snippets.ts @@ -139,6 +139,32 @@ depth = model.infer_image(raw_img) # HxW raw depth map in numpy ]; }; +export const depth_pro = (model: ModelData): string[] => { + const installSnippet = `# Download checkpoint +pip install huggingface-hub +huggingface-cli download --local-dir checkpoints ${model.id}`; + + const inferenceSnippet = `import depth_pro + +# Load model and preprocessing transform +model, transform = depth_pro.create_model_and_transforms() +model.eval() + +# Load and preprocess an image. +image, _, f_px = depth_pro.load_rgb("example.png") +image = transform(image) + +# Run inference. +prediction = model.infer(image, f_px=f_px) + +# Results: 1. Depth in meters +depth = prediction["depth"] +# Results: 2. Focal length in pixels +focallength_px = prediction["focallength_px"]`; + + return [installSnippet, inferenceSnippet]; +}; + const diffusersDefaultPrompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"; const diffusers_default = (model: ModelData) => [ diff --git a/packages/tasks/src/model-libraries.ts b/packages/tasks/src/model-libraries.ts index 8066b5121..488d83e4d 100644 --- a/packages/tasks/src/model-libraries.ts +++ b/packages/tasks/src/model-libraries.ts @@ -165,6 +165,14 @@ export const MODEL_LIBRARIES_UI_ELEMENTS = { filter: false, countDownloads: `path_extension:"pth"`, }, + "depth-pro": { + prettyLabel: "Depth Pro", + repoName: "Depth Pro", + repoUrl: "https://github.com/apple/ml-depth-pro", + countDownloads: `path_extension:"pt"`, + snippets: snippets.depth_pro, + filter: false, + }, diffree: { prettyLabel: "Diffree", repoName: "Diffree",