1pub mod audio;
2pub mod bs1770;
3pub mod chat_template;
4pub mod coco_classes;
5pub mod hub;
6pub mod imagenet;
7pub mod token_output_stream;
8pub mod wav;
9use candle::utils::{cuda_is_available, metal_is_available};
10use candle::{Device, Result, Tensor};
11
12pub fn device(cpu: bool) -> Result<Device> {
13 if cpu {
14 Ok(Device::Cpu)
15 } else if cuda_is_available() {
16 Ok(Device::new_cuda(0)?)
17 } else if metal_is_available() {
18 Ok(Device::new_metal(0)?)
19 } else {
20 #[cfg(all(target_os = "macos", target_arch = "aarch64"))]
21 {
22 println!(
23 "Running on CPU, to run on GPU(metal), build this example with `--features metal`"
24 );
25 }
26 #[cfg(not(all(target_os = "macos", target_arch = "aarch64")))]
27 {
28 println!("Running on CPU, to run on GPU, build this example with `--features cuda`");
29 }
30 Ok(Device::Cpu)
31 }
32}
33
34pub fn load_image<P: AsRef<std::path::Path>>(
35 p: P,
36 resize_longest: Option<usize>,
37) -> Result<(Tensor, usize, usize)> {
38 let img = image::ImageReader::open(p)?
39 .decode()
40 .map_err(candle::Error::wrap)?;
41 let (initial_h, initial_w) = (img.height() as usize, img.width() as usize);
42 let img = match resize_longest {
43 None => img,
44 Some(resize_longest) => {
45 let (height, width) = (img.height(), img.width());
46 let resize_longest = resize_longest as u32;
47 let (height, width) = if height < width {
48 let h = (resize_longest * height) / width;
49 (h, resize_longest)
50 } else {
51 let w = (resize_longest * width) / height;
52 (resize_longest, w)
53 };
54 img.resize_exact(width, height, image::imageops::FilterType::CatmullRom)
55 }
56 };
57 let (height, width) = (img.height() as usize, img.width() as usize);
58 let img = img.to_rgb8();
59 let data = img.into_raw();
60 let data = Tensor::from_vec(data, (height, width, 3), &Device::Cpu)?.permute((2, 0, 1))?;
61 Ok((data, initial_h, initial_w))
62}
63
64pub fn load_image_and_resize<P: AsRef<std::path::Path>>(
65 p: P,
66 width: usize,
67 height: usize,
68) -> Result<Tensor> {
69 let img = image::ImageReader::open(p)?
70 .decode()
71 .map_err(candle::Error::wrap)?
72 .resize_to_fill(
73 width as u32,
74 height as u32,
75 image::imageops::FilterType::Triangle,
76 );
77 let img = img.to_rgb8();
78 let data = img.into_raw();
79 Tensor::from_vec(data, (width, height, 3), &Device::Cpu)?.permute((2, 0, 1))
80}
81
82pub fn save_image<P: AsRef<std::path::Path>>(img: &Tensor, p: P) -> Result<()> {
85 let p = p.as_ref();
86 let (channel, height, width) = img.dims3()?;
87 if channel != 3 {
88 candle::bail!("save_image expects an input of shape (3, height, width)")
89 }
90 let img = img.permute((1, 2, 0))?.flatten_all()?;
91 let pixels = img.to_vec1::<u8>()?;
92 let image: image::ImageBuffer<image::Rgb<u8>, Vec<u8>> =
93 match image::ImageBuffer::from_raw(width as u32, height as u32, pixels) {
94 Some(image) => image,
95 None => candle::bail!("error saving image {p:?}"),
96 };
97 image.save(p).map_err(candle::Error::wrap)?;
98 Ok(())
99}
100
101pub fn save_image_resize<P: AsRef<std::path::Path>>(
102 img: &Tensor,
103 p: P,
104 h: usize,
105 w: usize,
106) -> Result<()> {
107 let p = p.as_ref();
108 let (channel, height, width) = img.dims3()?;
109 if channel != 3 {
110 candle::bail!("save_image expects an input of shape (3, height, width)")
111 }
112 let img = img.permute((1, 2, 0))?.flatten_all()?;
113 let pixels = img.to_vec1::<u8>()?;
114 let image: image::ImageBuffer<image::Rgb<u8>, Vec<u8>> =
115 match image::ImageBuffer::from_raw(width as u32, height as u32, pixels) {
116 Some(image) => image,
117 None => candle::bail!("error saving image {p:?}"),
118 };
119 let image = image::DynamicImage::from(image);
120 let image = image.resize_to_fill(w as u32, h as u32, image::imageops::FilterType::CatmullRom);
121 image.save(p).map_err(candle::Error::wrap)?;
122 Ok(())
123}
124
125pub fn hub_load_safetensors<T: hf_hub::RepoType>(
127 repo: &crate::hub::Repo<T>,
128 json_file: &str,
129) -> Result<Vec<std::path::PathBuf>> {
130 let json_file = repo.get(json_file).map_err(candle::Error::wrap)?;
131 let json_file = std::fs::File::open(json_file)?;
132 let json: serde_json::Value =
133 serde_json::from_reader(&json_file).map_err(candle::Error::wrap)?;
134 let weight_map = match json.get("weight_map") {
135 None => candle::bail!("no weight map in {json_file:?}"),
136 Some(serde_json::Value::Object(map)) => map,
137 Some(_) => candle::bail!("weight map in {json_file:?} is not a map"),
138 };
139 let mut safetensors_files = std::collections::HashSet::new();
140 for value in weight_map.values() {
141 if let Some(file) = value.as_str() {
142 safetensors_files.insert(file.to_string());
143 }
144 }
145 let safetensors_files = safetensors_files
146 .iter()
147 .map(|v| repo.get(v).map_err(candle::Error::wrap))
148 .collect::<Result<Vec<_>>>()?;
149 Ok(safetensors_files)
150}
151
152pub fn hub_load_local_safetensors<P: AsRef<std::path::Path>>(
153 path: P,
154 json_file: &str,
155) -> Result<Vec<std::path::PathBuf>> {
156 let path = path.as_ref();
157 let jsfile = std::fs::File::open(path.join(json_file))?;
158 let json: serde_json::Value = serde_json::from_reader(&jsfile).map_err(candle::Error::wrap)?;
159 let weight_map = match json.get("weight_map") {
160 None => candle::bail!("no weight map in {json_file:?}"),
161 Some(serde_json::Value::Object(map)) => map,
162 Some(_) => candle::bail!("weight map in {json_file:?} is not a map"),
163 };
164 let mut safetensors_files = std::collections::HashSet::new();
165 for value in weight_map.values() {
166 if let Some(file) = value.as_str() {
167 safetensors_files.insert(file);
168 }
169 }
170 let safetensors_files: Vec<_> = safetensors_files
171 .into_iter()
172 .map(|v| path.join(v))
173 .collect();
174 Ok(safetensors_files)
175}