dinov2 - read images from disk and compute the class probabilities (#503)

* Load the image from disk and convert it to a tensor.

* Tweak the function name.

candle-examples/examples/dinov2/main.rs

@@ -85,7 +85,7 @@ impl LayerScale {
 
 impl Module for LayerScale {
     fn forward(&self, xs: &Tensor) -> Result<Tensor> {
-        xs * &self.gamma
+        xs.broadcast_mul(&self.gamma)
     }
 }
 
@@ -306,10 +306,17 @@ pub fn main() -> anyhow::Result<()> {
 
     let device = candle_examples::device(args.cpu)?;
 
+    // TODO: apply imagenet normalization.
+    let image = candle_examples::load_image(args.image)?;
+    println!("loaded image {image:?}");
+
     let weights = unsafe { candle::safetensors::MmapedFile::new(args.model)? };
     let weights = weights.deserialize()?;
     let vb = VarBuilder::from_safetensors(vec![weights], DType::F32, &device);
-    let _model = vit_small(vb)?;
+    let model = vit_small(vb)?;
     println!("model built");
+    let logits = model.forward(&image.unsqueeze(0)?)?;
+    let prs = candle_nn::ops::softmax(&logits, D::Minus1)?;
+    println!("{prs}");
     Ok(())
 }

candle-examples/examples/stable-diffusion/main.rs

@@ -332,7 +332,7 @@ fn run(args: Args) -> Result<()> {
                 let image = (image * 255.)?.to_dtype(DType::U8)?.i(0)?;
                 let image_filename =
                     output_filename(&final_image, idx + 1, num_samples, Some(timestep_index + 1));
-                crate::utils::save_image(&image, image_filename)?
+                candle_examples::save_image(&image, image_filename)?
             }
         }
 
@@ -346,7 +346,7 @@ fn run(args: Args) -> Result<()> {
         let image = ((image / 2.)? + 0.5)?.to_device(&Device::Cpu)?;
         let image = (image * 255.)?.to_dtype(DType::U8)?.i(0)?;
         let image_filename = output_filename(&final_image, idx + 1, num_samples, None);
-        crate::utils::save_image(&image, image_filename)?
+        candle_examples::save_image(&image, image_filename)?
     }
     Ok(())
 }

candle-examples/examples/stable-diffusion/utils.rs

@@ -12,25 +12,6 @@ pub fn linspace(start: f64, stop: f64, steps: usize) -> Result<Tensor> {
     Tensor::from_vec(vs, steps, &Device::Cpu)
 }
 
-/// Saves an image to disk using the image crate, this expects an input with shape
-/// (c, width, height).
-pub fn save_image<P: AsRef<std::path::Path>>(img: &Tensor, p: P) -> Result<()> {
-    let p = p.as_ref();
-    let (channel, width, height) = img.dims3()?;
-    if channel != 3 {
-        candle::bail!("save_image expects an input of shape (3, width, height)")
-    }
-    let img = img.transpose(0, 1)?.t()?.flatten_all()?;
-    let pixels = img.to_vec1::<u8>()?;
-    let image: image::ImageBuffer<image::Rgb<u8>, Vec<u8>> =
-        match image::ImageBuffer::from_raw(width as u32, height as u32, pixels) {
-            Some(image) => image,
-            None => candle::bail!("error saving image {p:?}"),
-        };
-    image.save(p).map_err(candle::Error::wrap)?;
-    Ok(())
-}
-
 // Wrap the conv2d op to provide some tracing.
 #[derive(Debug)]
 pub struct Conv2d {