reorg + update docs with new paths

integration/utils/src/lib.rs

@@ -0,0 +1,156 @@
+extern crate candle_core;
+extern crate candle_transformers;
+extern crate tokenizers;
+
+pub mod audio;
+pub mod bs1770;
+pub mod coco_classes;
+pub mod imagenet;
+pub mod token_output_stream;
+pub mod wav;
+use candle_core::{Device, Tensor, utils::{cuda_is_available, metal_is_available}};
+
+
+pub fn device(cpu: bool) -> Result<Device, anyhow::Error> {
+    if cpu {
+        Ok(Device::Cpu)
+    } else if cuda_is_available() {
+        Ok(Device::new_cuda(0)?)
+    } else if metal_is_available() {
+        Ok(Device::new_metal(0)?)
+    } else {
+        #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
+        {
+            println!(
+                "Running on CPU, to run on GPU(metal), build this example with `--features metal`"
+            );
+        }
+        #[cfg(not(all(target_os = "macos", target_arch = "aarch64")))]
+        {
+            println!("Running on CPU, to run on GPU, build this example with `--features cuda`");
+        }
+        Ok(Device::Cpu)
+    }
+}
+
+pub fn load_image<P: AsRef<std::path::Path>>(
+    p: P,
+    resize_longest: Option<usize>,
+) -> Result<(Tensor, usize, usize), anyhow::Error> {
+    let img = image::ImageReader::open(p)?
+        .decode()
+        .map_err(candle_core::Error::wrap)?;
+    let (initial_h, initial_w) = (img.height() as usize, img.width() as usize);
+    let img = match resize_longest {
+        None => img,
+        Some(resize_longest) => {
+            let (height, width) = (img.height(), img.width());
+            let resize_longest = resize_longest as u32;
+            let (height, width) = if height < width {
+                let h = (resize_longest * height) / width;
+                (h, resize_longest)
+            } else {
+                let w = (resize_longest * width) / height;
+                (resize_longest, w)
+            };
+            img.resize_exact(width, height, image::imageops::FilterType::CatmullRom)
+        }
+    };
+    let (height, width) = (img.height() as usize, img.width() as usize);
+    let img = img.to_rgb8();
+    let data = img.into_raw();
+    let data = Tensor::from_vec(data, (height, width, 3), &Device::Cpu)?.permute((2, 0, 1))?;
+    Ok((data, initial_h, initial_w))
+}
+
+pub fn load_image_and_resize<P: AsRef<std::path::Path>>(
+    p: P,
+    width: usize,
+    height: usize,
+) -> candle_core::Result<Tensor> {
+    let img = image::ImageReader::open(p)?
+        .decode()
+        .map_err(candle_core::Error::wrap)?
+        .resize_to_fill(
+            width as u32,
+            height as u32,
+            image::imageops::FilterType::Triangle,
+        );
+    let img = img.to_rgb8();
+    let data = img.into_raw();
+    Tensor::from_vec(data, (width, height, 3), &Device::Cpu)?.permute((2, 0, 1))
+}
+
+/// Saves an image to disk using the image crate, this expects an input with shape
+/// (c, height, width).
+pub fn save_image<P: AsRef<std::path::Path>>(img: &Tensor, p: P) -> Result<(), anyhow::Error> {
+    let p = p.as_ref();
+    let (channel, height, width) = img.dims3()?;
+    if channel != 3 {
+        anyhow::bail!("save_image expects an input of shape (3, height, width)")
+    }
+    let img = img.permute((1, 2, 0))?.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 => anyhow::bail!("error saving image {p:?}"),
+        };
+    image.save(p).map_err(candle_core::Error::wrap)?;
+    Ok(())
+}
+
+/// Loads the safetensors files for a model from the hub based on a json index file.
+pub fn hub_load_safetensors(
+    repo: &hf_hub::api::sync::ApiRepo,
+    json_file: &str,
+) -> Result<Vec<std::path::PathBuf>, anyhow::Error> {
+    let json_file = repo.get(json_file).map_err(candle_core::Error::wrap)?;
+    let json_file = std::fs::File::open(json_file)?;
+    let json: serde_json::Value =
+        serde_json::from_reader(&json_file).map_err(candle_core::Error::wrap)?;
+    let weight_map = match json.get("weight_map") {
+        None => anyhow::bail!("no weight map in {json_file:?}"),
+        Some(serde_json::Value::Object(map)) => map,
+        Some(_) => anyhow::bail!("weight map in {json_file:?} is not a map"),
+    };
+    let mut safetensors_files = std::collections::HashSet::new();
+    for value in weight_map.values() {
+        if let Some(file) = value.as_str() {
+            safetensors_files.insert(file.to_string());
+        }
+    }
+    let safetensors_files = safetensors_files
+        .iter()
+        .map(|v| {
+            repo.get(v)
+                .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))
+        })
+        .collect::<Result<Vec<_>, std::io::Error, >>()?;
+    Ok(safetensors_files)
+}
+
+pub fn hub_load_local_safetensors<P: AsRef<std::path::Path>>(
+    path: P,
+    json_file: &str,
+) -> Result<Vec<std::path::PathBuf>, anyhow::Error> {
+    let path = path.as_ref();
+    let jsfile = std::fs::File::open(path.join(json_file))?;
+    let json: serde_json::Value = serde_json::from_reader(&jsfile).map_err(candle_core::Error::wrap)?;
+    let weight_map = match json.get("weight_map") {
+        None => anyhow::bail!("no weight map in {json_file:?}"),
+        Some(serde_json::Value::Object(map)) => map,
+        Some(_) => anyhow::bail!("weight map in {json_file:?} is not a map"),
+    };
+    let mut safetensors_files = std::collections::HashSet::new();
+    for value in weight_map.values() {
+        if let Some(file) = value.as_str() {
+            safetensors_files.insert(file);
+        }
+    }
+    let safetensors_files: Vec<_> = safetensors_files
+        .into_iter()
+        .map(|v| path.join(v))
+        .collect();
+    Ok(safetensors_files)
+}