huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-16 10:38:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

Use the hub to retrieve dinov2 model weights. (#507)

Browse Source
This commit is contained in:
Laurent Mazare
2023-08-18 18:27:31 +01:00
committed by GitHub
parent e5dd5fd1b3
commit b64e782c2d
2 changed files with 16 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
candle-examples/examples/dinov2/main.rs
View File

@ -291,7 +291,7 @@ pub fn vit_small(vb: VarBuilder) -> Result<DinoVisionTransformer> {
#[derive(Parser)]
struct Args {
#[arg(long)]
model: String,
model: Option<String>,
#[arg(long)]
image: String,
@ -309,7 +309,15 @@ pub fn main() -> anyhow::Result<()> {
let image = candle_examples::load_image224(args.image)?;
println!("loaded image {image:?}");
let weights = unsafe { candle::safetensors::MmapedFile::new(args.model)? };
let model_file = match args.model {
None => {
let api = hf_hub::api::sync::Api::new()?;
let api = api.model("lmz/candle-dino-v2".into());
api.get("dinov2_vits14.safetensors")?
}
Some(model) => model.into(),
};
let weights = unsafe { candle::safetensors::MmapedFile::new(model_file)? };
let weights = weights.deserialize()?;
let vb = VarBuilder::from_safetensors(vec![weights], DType::F32, &device);
let model = vit_small(vb)?;