perf: prefer bf16

koharu-ml/src/aot_inpainting/mod.rs

@@ -47,6 +47,7 @@ pub struct AotInpainting {
     model: AotGenerator,
     config: AotInpaintingConfig,
     device: Device,
+    dtype: DType,
 }
 
 #[derive(Debug, Clone, Deserialize)]
@@ -113,17 +114,22 @@ impl AotInpainting {
         cpu: bool,
     ) -> Result<Self> {
         let device = device(cpu)?;
+        let dtype = loading::model_dtype(&device);
         let config = loading::read_json::<AotInpaintingConfig>(config_path.as_ref())
             .with_context(|| format!("failed to parse {}", config_path.as_ref().display()))?;
         config.validate()?;
-        let model = loading::load_mmaped_safetensors_path(weights_path.as_ref(), &device, |vb| {
-            AotGenerator::load(&vb, &config.spec())
-        })?;
+        let model = loading::load_mmaped_safetensors_path_with_dtype(
+            weights_path.as_ref(),
+            &device,
+            dtype,
+            |vb| AotGenerator::load(&vb, &config.spec()),
+        )?;
 
         Ok(Self {
             model,
             config,
             device,
+            dtype,
         })
     }
 
@@ -199,7 +205,7 @@ impl AotInpainting {
             &self.device,
         )?
         .permute((0, 3, 1, 2))?
-        .to_dtype(DType::F32)?
+        .to_dtype(self.dtype)?
             / 127.5)?;
         let image_tensor = (image_tensor - 1.0)?;
 
@@ -209,7 +215,7 @@ impl AotInpainting {
             &self.device,
         )?
         .permute((0, 3, 1, 2))?
-        .to_dtype(DType::F32)?;
+        .to_dtype(self.dtype)?;
         let mask_tensor = (mask_tensor / 255.0)?;
         let mask_inv = (Tensor::ones_like(&mask_tensor)? - &mask_tensor)?;
         let mask_inv_rgb = mask_inv.broadcast_as((1, 3, h as usize, w as usize))?;
@@ -220,7 +226,10 @@ impl AotInpainting {
     }
 
     fn postprocess(&self, output: &Tensor) -> Result<RgbImage> {
-        let output = output.to_device(&Device::Cpu)?.squeeze(0)?;
+        let output = output
+            .to_dtype(DType::F32)?
+            .to_device(&Device::Cpu)?
+            .squeeze(0)?;
         let (channels, height, width) = output.dims3()?;
         if channels != 3 {
             bail!("expected 3 output channels, got {channels}");

koharu-ml/src/aot_inpainting/model.rs

@@ -262,6 +262,8 @@ fn relu_nf(xs: &Tensor) -> candle_core::Result<Tensor> {
 }
 
 fn my_layer_norm(xs: &Tensor) -> candle_core::Result<Tensor> {
+    let dtype = xs.dtype();
+    let xs = xs.to_dtype(candle_core::DType::F32)?;
     let (batch, channels, height, width) = xs.dims4()?;
     let flat = xs.flatten_from(2)?;
     let mean = flat.mean_keepdim(2)?;
@@ -269,7 +271,9 @@ fn my_layer_norm(xs: &Tensor) -> candle_core::Result<Tensor> {
     let normalized = ((flat.broadcast_sub(&mean)? * 2.0)?)
         .broadcast_div(&std)?
         .broadcast_sub(&Tensor::ones_like(&flat)?)?;
-    (normalized * 5.0)?.reshape((batch, channels, height, width))
+    (normalized * 5.0)?
+        .reshape((batch, channels, height, width))?
+        .to_dtype(dtype)
 }
 
 fn load_plain_conv2d(
@@ -342,6 +346,9 @@ fn load_scaled_ws_transpose_conv2d(
 }
 
 fn standardize_conv2d_weight(weight: Tensor, gain: Tensor) -> candle_core::Result<Tensor> {
+    let dtype = weight.dtype();
+    let weight = weight.to_dtype(candle_core::DType::F32)?;
+    let gain = gain.to_dtype(candle_core::DType::F32)?;
     let (out_channels, in_channels, kernel_h, kernel_w) = weight.dims4()?;
     let flat = weight.flatten_from(1)?;
     let fan_in = flat.dim(1)? as f64;
@@ -356,18 +363,19 @@ fn standardize_conv2d_weight(weight: Tensor, gain: Tensor) -> candle_core::Resul
     let scale = variance.maximum(&eps)?.sqrt()?.recip()?;
     let scale = scale.broadcast_mul(&gain.reshape((out_channels, 1))?)?;
     let shift = mean.broadcast_mul(&scale)?;
-    flat.broadcast_mul(&scale)?.broadcast_sub(&shift)?.reshape((
-        out_channels,
-        in_channels,
-        kernel_h,
-        kernel_w,
-    ))
+    flat.broadcast_mul(&scale)?
+        .broadcast_sub(&shift)?
+        .reshape((out_channels, in_channels, kernel_h, kernel_w))?
+        .to_dtype(dtype)
 }
 
 fn standardize_transpose_conv2d_weight(
     weight: Tensor,
     gain: Tensor,
 ) -> candle_core::Result<Tensor> {
+    let dtype = weight.dtype();
+    let weight = weight.to_dtype(candle_core::DType::F32)?;
+    let gain = gain.to_dtype(candle_core::DType::F32)?;
     let (in_channels, out_channels, kernel_h, kernel_w) = weight.dims4()?;
     let flat = weight.flatten_from(1)?;
     let fan_in = flat.dim(1)? as f64;
@@ -382,12 +390,10 @@ fn standardize_transpose_conv2d_weight(
     let scale = variance.maximum(&eps)?.sqrt()?.recip()?;
     let scale = scale.broadcast_mul(&gain.reshape((in_channels, 1))?)?;
     let shift = mean.broadcast_mul(&scale)?;
-    flat.broadcast_mul(&scale)?.broadcast_sub(&shift)?.reshape((
-        in_channels,
-        out_channels,
-        kernel_h,
-        kernel_w,
-    ))
+    flat.broadcast_mul(&scale)?
+        .broadcast_sub(&shift)?
+        .reshape((in_channels, out_channels, kernel_h, kernel_w))?
+        .to_dtype(dtype)
 }
 
 fn reflect_pad2d(xs: &Tensor, pad: usize) -> candle_core::Result<Tensor> {

koharu-ml/src/comic_text_bubble_detector/mod.rs

@@ -3,6 +3,7 @@ mod model;
 use std::{collections::BTreeMap, time::Instant};
 
 use anyhow::{Context, Result, bail};
+use candle_core::DType;
 use image::{DynamicImage, GenericImageView, imageops::FilterType};
 use koharu_runtime::RuntimeManager;
 use serde::{Deserialize, Serialize};
@@ -44,12 +45,14 @@ pub struct ComicTextBubbleDetector {
     config: RTDetrV2Config,
     preprocessor: RTDetrImageProcessorConfig,
     device: Device,
+    dtype: DType,
     slicer: ImageSlicer,
 }
 
 impl ComicTextBubbleDetector {
     pub async fn load(runtime: &RuntimeManager, cpu: bool) -> Result<Self> {
         let device = device(cpu)?;
+        let dtype = loading::model_dtype(&device);
         let downloads = runtime.downloads();
         let config_path = downloads.huggingface_model(HF_REPO, "config.json").await?;
         let preprocessor_path = downloads
@@ -64,15 +67,19 @@ impl ComicTextBubbleDetector {
         config.validate()?;
         let preprocessor = loading::read_json::<RTDetrImageProcessorConfig>(&preprocessor_path)
             .with_context(|| format!("failed to parse {}", preprocessor_path.display()))?;
-        let model = loading::load_mmaped_safetensors_path(&weights_path, &device, |vb| {
-            RTDetrV2ForObjectDetection::load(vb, &config)
-        })?;
+        let model = loading::load_mmaped_safetensors_path_with_dtype(
+            &weights_path,
+            &device,
+            dtype,
+            |vb| RTDetrV2ForObjectDetection::load(vb, &config),
+        )?;
 
         Ok(Self {
             model,
             config,
             preprocessor,
             device,
+            dtype,
             slicer: ImageSlicer::default(),
         })
     }
@@ -121,7 +128,7 @@ impl ComicTextBubbleDetector {
         image: &DynamicImage,
         threshold: f32,
     ) -> Result<Vec<ComicTextBubbleRegion>> {
-        let pixel_values = preprocess_image(image, &self.preprocessor, &self.device)?;
+        let pixel_values = preprocess_image(image, &self.preprocessor, &self.device, self.dtype)?;
         let outputs = self.model.forward(&pixel_values)?;
         post_process_object_detection(&self.config, &outputs, image.dimensions(), threshold)
     }
@@ -365,6 +372,7 @@ fn preprocess_image(
     image: &DynamicImage,
     preprocessor: &RTDetrImageProcessorConfig,
     device: &Device,
+    dtype: DType,
 ) -> Result<candle_core::Tensor> {
     let target_h = preprocessor.size.height;
     let target_w = preprocessor.size.width;
@@ -378,15 +386,17 @@ fn preprocess_image(
         candle_core::Tensor::from_vec(rgb.into_raw(), (1, target_h, target_w, 3), &Device::Cpu)?
             .to_device(device)?
             .permute((0, 3, 1, 2))?
-            .to_dtype(candle_core::DType::F32)?;
+            .to_dtype(dtype)?;
     let tensor = if preprocessor.do_rescale {
         tensor.affine(preprocessor.rescale_factor as f64, 0.0)?
     } else {
         tensor
     };
     if preprocessor.do_normalize {
-        let mean = candle_core::Tensor::from_slice(&preprocessor.image_mean, (1, 3, 1, 1), device)?;
-        let std = candle_core::Tensor::from_slice(&preprocessor.image_std, (1, 3, 1, 1), device)?;
+        let mean = candle_core::Tensor::from_slice(&preprocessor.image_mean, (1, 3, 1, 1), device)?
+            .to_dtype(dtype)?;
+        let std = candle_core::Tensor::from_slice(&preprocessor.image_std, (1, 3, 1, 1), device)?
+            .to_dtype(dtype)?;
         Ok(tensor.broadcast_sub(&mean)?.broadcast_div(&std)?)
     } else {
         Ok(tensor)
@@ -399,8 +409,14 @@ fn post_process_object_detection(
     target_size: (u32, u32),
     threshold: f32,
 ) -> Result<Vec<ComicTextBubbleRegion>> {
-    let logits = outputs.logits.to_device(&Device::Cpu)?;
-    let pred_boxes = outputs.pred_boxes.to_device(&Device::Cpu)?;
+    let logits = outputs
+        .logits
+        .to_dtype(DType::F32)?
+        .to_device(&Device::Cpu)?;
+    let pred_boxes = outputs
+        .pred_boxes
+        .to_dtype(DType::F32)?
+        .to_device(&Device::Cpu)?;
     let (batch_size, num_queries, num_classes) = logits.dims3()?;
     if batch_size != 1 {
         bail!("only single-image inference is supported, got batch_size={batch_size}");

koharu-ml/src/comic_text_bubble_detector/model.rs

@@ -95,6 +95,15 @@ fn load_layer_norm(vb: VarBuilder, hidden_size: usize, eps: f64) -> Result<Layer
     Ok(layer_norm(hidden_size, eps, vb)?)
 }
 
+fn softmax_f32(xs: &Tensor, dim: D) -> Result<Tensor> {
+    let dtype = xs.dtype();
+    if dtype == DType::F32 {
+        Ok(softmax(xs, dim)?)
+    } else {
+        Ok(softmax(&xs.to_dtype(DType::F32)?, dim)?.to_dtype(dtype)?)
+    }
+}
+
 fn pad_all_sides_one(xs: &Tensor) -> candle_core::Result<Tensor> {
     xs.pad_with_zeros(2, 1, 1)?.pad_with_zeros(3, 1, 1)
 }
@@ -539,7 +548,8 @@ impl RTDetrV2MultiheadAttention {
     ) -> Result<Tensor> {
         let (batch_size, sequence_length, hidden_size) = hidden_states.dims3()?;
         let query_key_input = match position_embeddings {
-            Some(position_embeddings) => hidden_states.broadcast_add(position_embeddings)?,
+            Some(position_embeddings) => hidden_states
+                .broadcast_add(&position_embeddings.to_dtype(hidden_states.dtype())?)?,
             None => hidden_states.clone(),
         };
         let shape = (
@@ -570,7 +580,7 @@ impl RTDetrV2MultiheadAttention {
         let mut attention_scores =
             query_states.matmul(&key_states.transpose(2, 3)?.contiguous()?)?;
         attention_scores = (attention_scores * self.scaling)?;
-        let attention_probs = softmax(&attention_scores, D::Minus1)?;
+        let attention_probs = softmax_f32(&attention_scores, D::Minus1)?;
         let context = attention_probs.matmul(&value_states)?;
         let context = context.transpose(1, 2)?.contiguous()?.reshape((
             batch_size,
@@ -985,11 +995,10 @@ impl RTDetrV2HybridEncoder {
             let src_flatten = feature_maps[*encoder_index]
                 .flatten_from(2)?
                 .transpose(1, 2)?;
-            let pos_embed = self.position_embedding.forward(
-                width,
-                height,
-                feature_maps[*encoder_index].device(),
-            )?;
+            let pos_embed = self
+                .position_embedding
+                .forward(width, height, feature_maps[*encoder_index].device())?
+                .to_dtype(src_flatten.dtype())?;
             let encoded = self.encoder[index].forward(&src_flatten, &pos_embed)?;
             feature_maps[*encoder_index] = encoded.transpose(1, 2)?.reshape((
                 batch_size,
@@ -1088,7 +1097,8 @@ impl RTDetrV2MultiscaleDeformableAttention {
         spatial_shapes_list: &[(usize, usize)],
     ) -> Result<Tensor> {
         let hidden_states = match position_embeddings {
-            Some(position_embeddings) => hidden_states.broadcast_add(position_embeddings)?,
+            Some(position_embeddings) => hidden_states
+                .broadcast_add(&position_embeddings.to_dtype(hidden_states.dtype())?)?,
             None => hidden_states.clone(),
         };
         let (batch_size, num_queries, _) = hidden_states.dims3()?;
@@ -1119,7 +1129,7 @@ impl RTDetrV2MultiscaleDeformableAttention {
             self.n_points,
             2,
         ))?;
-        let attention_weights = softmax(
+        let attention_weights = softmax_f32(
             &self.attention_weights.forward(&hidden_states)?.reshape((
                 batch_size,
                 num_queries,
@@ -1346,9 +1356,10 @@ impl RTDetrV2Decoder {
             )?;
 
             let delta = self.bbox_embed[index].forward(&hidden_states)?;
-            reference_points = inverse_sigmoid_to_sigmoid(
-                &delta.broadcast_add(&inverse_sigmoid_tensor(&reference_points)?)?,
-            )?;
+            let reference_points_unact =
+                inverse_sigmoid_tensor(&reference_points)?.to_dtype(delta.dtype())?;
+            reference_points =
+                inverse_sigmoid_to_sigmoid(&delta.broadcast_add(&reference_points_unact)?)?;
             logits = Some(self.class_embed[index].forward(&hidden_states)?);
             pred_boxes = Some(reference_points.clone());
         }
@@ -1619,8 +1630,12 @@ fn batch_gather_rows(tensor: &Tensor, indices: &Tensor) -> Result<Tensor> {
 }
 
 fn inverse_sigmoid_tensor(tensor: &Tensor) -> Result<Tensor> {
+    let dtype = tensor.dtype();
     let tensor = tensor.to_dtype(DType::F32)?.clamp(1e-5, 1.0 - 1e-5)?;
-    Ok(tensor.broadcast_div(&tensor.affine(-1.0, 1.0)?)?.log()?)
+    Ok(tensor
+        .broadcast_div(&tensor.affine(-1.0, 1.0)?)?
+        .log()?
+        .to_dtype(dtype)?)
 }
 
 fn inverse_sigmoid_to_sigmoid(tensor: &Tensor) -> Result<Tensor> {
@@ -1637,6 +1652,9 @@ fn multi_scale_deformable_attention(
     offset_scale: f64,
 ) -> Result<Tensor> {
     let (batch_size, sequence_length, num_heads, head_dim) = value.dims4()?;
+    let reference_points = reference_points.to_dtype(DType::F32)?;
+    let sampling_offsets = sampling_offsets.to_dtype(DType::F32)?;
+    let attention_weights = attention_weights.to_dtype(DType::F32)?;
     let [_, num_queries, _, num_levels, num_points, _] =
         <[usize; 6]>::try_from(sampling_offsets.dims().to_vec()).map_err(|_| {
             anyhow::anyhow!(

koharu-ml/src/comic_text_detector/mod.rs

@@ -57,6 +57,7 @@ pub struct ComicTextDetector {
     unet: unet::UNet,
     dbnet: Option<dbnet::DbNet>,
     device: Device,
+    dtype: DType,
 }
 
 impl ComicTextDetector {
@@ -77,24 +78,32 @@ impl ComicTextDetector {
         load_dbnet: bool,
     ) -> anyhow::Result<Self> {
         let device = device(cpu)?;
+        let dtype = loading::model_dtype(&device);
         let downloads = runtime.downloads();
         let yolo_path = downloads
             .huggingface_model(HF_REPO, "yolo-v5.safetensors")
             .await?;
-        let yolo = loading::load_mmaped_safetensors_path(&yolo_path, &device, |vb| {
-            yolo_v5::YoloV5::load(vb, 2, 3)
-        })?;
+        let yolo =
+            loading::load_mmaped_safetensors_path_with_dtype(&yolo_path, &device, dtype, |vb| {
+                yolo_v5::YoloV5::load(vb, 2, 3)
+            })?;
         let unet_path = downloads
             .huggingface_model(HF_REPO, "unet.safetensors")
             .await?;
-        let unet = loading::load_mmaped_safetensors_path(&unet_path, &device, unet::UNet::load)?;
+        let unet = loading::load_mmaped_safetensors_path_with_dtype(
+            &unet_path,
+            &device,
+            dtype,
+            unet::UNet::load,
+        )?;
         let dbnet = if load_dbnet {
             let dbnet_path = downloads
                 .huggingface_model(HF_REPO, "dbnet.safetensors")
                 .await?;
-            Some(loading::load_mmaped_safetensors_path(
+            Some(loading::load_mmaped_safetensors_path_with_dtype(
                 &dbnet_path,
                 &device,
+                dtype,
                 dbnet::DbNet::load,
             )?)
         } else {
@@ -106,13 +115,14 @@ impl ComicTextDetector {
             unet,
             dbnet,
             device,
+            dtype,
         })
     }
 
     #[instrument(level = "debug", skip_all)]
     pub fn inference(&self, image: &DynamicImage) -> anyhow::Result<ComicTextDetection> {
         let original_dimensions = image.dimensions();
-        let (image_tensor, resized_dimensions) = preprocess(image, &self.device)?;
+        let (image_tensor, resized_dimensions) = preprocess(image, &self.device, self.dtype)?;
         let (predictions, mask, shrink_threshold) = self.forward(&image_tensor)?;
 
         let bboxes = postprocess_yolo(&predictions, original_dimensions, resized_dimensions)?;
@@ -145,7 +155,7 @@ impl ComicTextDetector {
     #[instrument(level = "debug", skip_all)]
     pub fn inference_segmentation(&self, image: &DynamicImage) -> anyhow::Result<GrayImage> {
         let original_dimensions = image.dimensions();
-        let (image_tensor, resized_dimensions) = preprocess(image, &self.device)?;
+        let (image_tensor, resized_dimensions) = preprocess(image, &self.device, self.dtype)?;
         let mask = self.forward_mask(&image_tensor)?;
         postprocess_unet_mask(&mask, original_dimensions, resized_dimensions)
     }
@@ -184,7 +194,11 @@ impl ComicTextDetector {
 }
 
 #[instrument(level = "debug", skip_all)]
-fn preprocess(image: &DynamicImage, device: &Device) -> anyhow::Result<(Tensor, (u32, u32))> {
+fn preprocess(
+    image: &DynamicImage,
+    device: &Device,
+    dtype: DType,
+) -> anyhow::Result<(Tensor, (u32, u32))> {
     let (orig_w, orig_h) = image.dimensions();
     let image_size = match device {
         Device::Cpu => CPU_DETECT_SIZE,
@@ -208,7 +222,7 @@ fn preprocess(image: &DynamicImage, device: &Device) -> anyhow::Result<(Tensor,
     .pad_with_zeros(3, 0, image_size as usize - w)?
         * (1. / 255.))?;
 
-    Ok((tensor, (width, height)))
+    Ok((tensor.to_dtype(dtype)?, (width, height)))
 }
 
 #[instrument(level = "debug", skip(predictions))]
@@ -217,7 +231,7 @@ fn postprocess_yolo(
     original_dimensions: (u32, u32),
     resized_dimensions: (u32, u32),
 ) -> anyhow::Result<Vec<Bbox<usize>>> {
-    let predictions = predictions.squeeze(0)?;
+    let predictions = predictions.squeeze(0)?.to_dtype(DType::F32)?;
     let (_, num_outputs) = predictions.dims2()?;
     if num_outputs < 6 {
         bail!("invalid prediction shape: expected at least 6 outputs, got {num_outputs}");
@@ -272,9 +286,9 @@ fn postprocess_mask(
     resized_dimensions: (u32, u32),
 ) -> anyhow::Result<GrayImage> {
     let shrink_and_thresh = shrink_thresh.squeeze(0)?;
-    let shrink = shrink_and_thresh.i(0)?;
-    let thresh = shrink_and_thresh.i(1)?;
-    let unet_mask = mask.squeeze(0)?;
+    let shrink = shrink_and_thresh.i(0)?.to_dtype(DType::F32)?;
+    let thresh = shrink_and_thresh.i(1)?.to_dtype(DType::F32)?;
+    let unet_mask = mask.squeeze(0)?.to_dtype(DType::F32)?;
 
     let (_, h_db, w_db) = shrink_and_thresh.dims3()?;
     let (_, h_unet, w_unet) = unet_mask.dims3()?;
@@ -334,7 +348,11 @@ fn tensor_channel_to_gray_resized(
     height: u32,
 ) -> anyhow::Result<GrayImage> {
     let (th, tw) = tensor.dims2()?;
-    let values: Vec<f32> = tensor.to_device(&Device::Cpu)?.flatten_all()?.to_vec1()?;
+    let values: Vec<f32> = tensor
+        .to_dtype(DType::F32)?
+        .to_device(&Device::Cpu)?
+        .flatten_all()?
+        .to_vec1()?;
     let pixels: Vec<u8> = values
         .iter()
         .map(|&v| (v.clamp(0.0, 1.0) * 255.0).round() as u8)

koharu-ml/src/comic_text_detector/yolo_v5.rs

@@ -1,4 +1,4 @@
-use candle_core::{DType, Device, IndexOp, Result, Tensor};
+use candle_core::{Device, IndexOp, Result, Tensor};
 use candle_nn::{
     BatchNorm, Conv2d, Conv2dConfig, Module, ModuleT, VarBuilder, batch_norm, conv2d,
     conv2d_no_bias,
@@ -338,8 +338,9 @@ impl YoloV3Head {
         ny: usize,
         dev: &Device,
     ) -> Result<(Tensor, Tensor)> {
-        let gx = Tensor::arange(0, nx as u32, dev)?.to_dtype(DType::F32)?;
-        let gy = Tensor::arange(0, ny as u32, dev)?.to_dtype(DType::F32)?;
+        let dtype = self.anchors.dtype();
+        let gx = Tensor::arange(0, nx as u32, dev)?.to_dtype(dtype)?;
+        let gy = Tensor::arange(0, ny as u32, dev)?.to_dtype(dtype)?;
 
         let gx = gx.reshape((1, 1, 1, nx))?.repeat((1, 1, ny, 1))?;
         let gy = gy.reshape((1, 1, ny, 1))?.repeat((1, 1, 1, nx))?;

koharu-ml/src/flux2_klein/mod.rs

@@ -147,16 +147,18 @@ impl Flux2Klein {
             );
         }
         let prompt_embedder = Flux2PromptEmbedder::new(&model_device);
-        let vae =
-            loading::load_mmaped_safetensors_path(&paths.vae_safetensors, &model_device, |vb| {
-                Flux2Vae::new(vb)
-            })
-            .with_context(|| {
-                format!(
-                    "failed to load Flux2 VAE from {}",
-                    paths.vae_safetensors.display()
-                )
-            })?;
+        let vae = loading::load_mmaped_safetensors_path_with_dtype(
+            &paths.vae_safetensors,
+            &model_device,
+            vae_dtype(&model_device),
+            |vb| Flux2Vae::new(vb),
+        )
+        .with_context(|| {
+            format!(
+                "failed to load Flux2 VAE from {}",
+                paths.vae_safetensors.display()
+            )
+        })?;
 
         Ok(Self {
             transformer,
@@ -454,7 +456,8 @@ impl Flux2Klein {
     }
 
     fn encode_image_latents(&self, image: &RgbImage) -> Result<Tensor> {
-        let image_tensor = image_to_tensor(image, &self.device)?;
+        let image_tensor =
+            image_to_tensor(image, &self.device)?.to_dtype(vae_dtype(&self.device))?;
         let latents = self.vae.encode_patchified_normalized(&image_tensor)?;
         drop(image_tensor);
         Ok(latents.to_dtype(DType::F32)?)
@@ -466,7 +469,8 @@ impl Flux2Klein {
         packed_h: usize,
         packed_w: usize,
     ) -> Result<RgbImage> {
-        let patchified = unpack_latents(&packed_latents, packed_h, packed_w)?;
+        let patchified = unpack_latents(&packed_latents, packed_h, packed_w)?
+            .to_dtype(vae_dtype(&self.device))?;
         drop(packed_latents);
         let decoded = self.vae.decode_patchified_normalized(&patchified)?;
         drop(patchified);
@@ -500,6 +504,14 @@ fn transformer_dtype(device: &Device) -> DType {
     DType::F32
 }
 
+fn vae_dtype(device: &Device) -> DType {
+    if device.is_cuda() {
+        return DType::BF16;
+    }
+
+    DType::F32
+}
+
 fn release_cuda_temporary_memory(device: &Device) -> Result<()> {
     device.synchronize()?;
 

koharu-ml/src/flux2_klein/vae.rs

@@ -54,7 +54,11 @@ fn scaled_dot_product_attention(q: &Tensor, k: &Tensor, v: &Tensor) -> Result<Te
         let len = chunk_size.min(seq_len - start);
         let q_chunk = q.narrow(2, start, len)?;
         let attn_weights = (q_chunk.matmul(&k_t)? * scale)?;
-        chunks.push(candle_nn::ops::softmax_last_dim(&attn_weights)?.matmul(&v)?);
+        let attn_dtype = attn_weights.dtype();
+        let attn_weights =
+            candle_nn::ops::softmax_last_dim(&attn_weights.to_dtype(candle_core::DType::F32)?)?
+                .to_dtype(attn_dtype)?;
+        chunks.push(attn_weights.matmul(&v)?);
     }
     Tensor::cat(&chunks, 2)
 }

koharu-ml/src/font_detector/mod.rs

@@ -37,6 +37,7 @@ pub struct FontDetector {
     model: models::Model,
     labels: FontLabels,
     device: Device,
+    dtype: DType,
 }
 
 impl FontDetector {
@@ -50,18 +51,23 @@ impl FontDetector {
         kind: ModelKind,
     ) -> Result<Self> {
         let device = device(cpu)?;
+        let dtype = loading::model_dtype(&device);
         let downloads = runtime.downloads();
         let weights_path = downloads
             .huggingface_model(HF_REPO, "yuzumarker-font-detection.safetensors")
             .await?;
-        let model = loading::load_mmaped_safetensors_path(&weights_path, &device, move |vb| {
-            models::Model::load(vb.pp("model._orig_mod.model"), kind)
-        })?;
+        let model = loading::load_mmaped_safetensors_path_with_dtype(
+            &weights_path,
+            &device,
+            dtype,
+            move |vb| models::Model::load(vb.pp("model._orig_mod.model"), kind),
+        )?;
         let labels = FontLabels::load(runtime).await?;
 
         Ok(Self {
             model,
             device,
+            dtype,
             labels,
         })
     }
@@ -84,11 +90,17 @@ impl FontDetector {
             .collect::<Vec<_>>()
             .into_iter()
             .collect::<Result<Vec<_>>>()?;
-        let batch = Tensor::stack(&processed, 0)?.to_device(&self.device)?;
+        let batch = Tensor::stack(&processed, 0)?
+            .to_device(&self.device)?
+            .to_dtype(self.dtype)?;
         let preprocess_elapsed = preprocess_started.elapsed();
 
         let forward_started = Instant::now();
-        let logits = self.model.forward(&batch, false)?.to_device(&Device::Cpu)?;
+        let logits = self
+            .model
+            .forward(&batch, false)?
+            .to_dtype(DType::F32)?
+            .to_device(&Device::Cpu)?;
         let forward_elapsed = forward_started.elapsed();
 
         let postprocess_started = Instant::now();

koharu-ml/src/font_detector/models.rs

@@ -1,5 +1,5 @@
 use anyhow::Result;
-use candle_core::{DType, Module, ModuleT, Tensor};
+use candle_core::{Module, ModuleT, Tensor};
 use candle_nn::{BatchNorm, Conv2d, Conv2dConfig, Linear, VarBuilder};
 use clap::ValueEnum;
 
@@ -68,7 +68,8 @@ impl Model {
         // For models that only output font logits (e.g., DeepFont), pad zeros for direction/regression.
         if dim == FONT_COUNT {
             let device = logits.device();
-            let zeros = Tensor::zeros((logits.dim(0)?, REGRESSION_DIM + 2), DType::F32, device)?;
+            let zeros =
+                Tensor::zeros((logits.dim(0)?, REGRESSION_DIM + 2), logits.dtype(), device)?;
             return Tensor::cat(&[logits, zeros], 1);
         }
 

koharu-ml/src/loading.rs

@@ -6,6 +6,14 @@ use candle_core::{DType, Device};
 use candle_nn::VarBuilder;
 use serde::de::DeserializeOwned;
 
+pub fn model_dtype(device: &Device) -> DType {
+    if device.is_cuda() {
+        DType::BF16
+    } else {
+        DType::F32
+    }
+}
+
 pub async fn resolve_manifest_path<F>(manifest: F) -> Result<PathBuf>
 where
     F: Future<Output = Result<PathBuf>>,
@@ -36,7 +44,20 @@ where
     Build: FnOnce(VarBuilder) -> std::result::Result<T, E>,
     E: Into<anyhow::Error>,
 {
-    let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[weights], DType::F32, device)? };
+    load_mmaped_safetensors_path_with_dtype(weights, device, DType::F32, build)
+}
+
+pub fn load_mmaped_safetensors_path_with_dtype<T, Build, E>(
+    weights: &Path,
+    device: &Device,
+    dtype: DType,
+    build: Build,
+) -> Result<T>
+where
+    Build: FnOnce(VarBuilder) -> std::result::Result<T, E>,
+    E: Into<anyhow::Error>,
+{
+    let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[weights], dtype, device)? };
     build(vb).map_err(Into::into)
 }
 
@@ -59,13 +80,26 @@ pub fn load_buffered_safetensors_path<T, Build, E>(
     device: &Device,
     build: Build,
 ) -> Result<T>
+where
+    Build: FnOnce(VarBuilder) -> std::result::Result<T, E>,
+    E: Into<anyhow::Error>,
+{
+    load_buffered_safetensors_path_with_dtype(weights, device, DType::F32, build)
+}
+
+pub fn load_buffered_safetensors_path_with_dtype<T, Build, E>(
+    weights: &Path,
+    device: &Device,
+    dtype: DType,
+    build: Build,
+) -> Result<T>
 where
     Build: FnOnce(VarBuilder) -> std::result::Result<T, E>,
     E: Into<anyhow::Error>,
 {
     let data =
         std::fs::read(weights).with_context(|| format!("failed to read {}", weights.display()))?;
-    let vb = VarBuilder::from_buffered_safetensors(data, DType::F32, device)?;
+    let vb = VarBuilder::from_buffered_safetensors(data, dtype, device)?;
     build(vb).map_err(Into::into)
 }
 

koharu-ml/src/manga_ocr/bert.rs

@@ -175,6 +175,15 @@ impl Module for Dropout {
     }
 }
 
+fn softmax_f32(xs: &Tensor, dim: D) -> candle_core::Result<Tensor> {
+    let dtype = xs.dtype();
+    if dtype == DType::F32 {
+        candle_nn::ops::softmax(xs, dim)
+    } else {
+        candle_nn::ops::softmax(&xs.to_dtype(DType::F32)?, dim)?.to_dtype(dtype)
+    }
+}
+
 struct BertSelfAttention {
     query: Linear,
     key: Linear,
@@ -233,7 +242,7 @@ impl BertSelfAttention {
         if let Some(mask) = attention_mask {
             attention_scores = attention_scores.broadcast_add(mask)?;
         }
-        let attention_probs = candle_nn::ops::softmax(&attention_scores, D::Minus1)?;
+        let attention_probs = softmax_f32(&attention_scores, D::Minus1)?;
         let attention_probs = self.dropout.forward(&attention_probs)?;
         let context_layer = attention_probs.matmul(&value)?;
         context_layer.transpose(1, 2)?.contiguous()?.reshape((

koharu-ml/src/manga_ocr/mod.rs

@@ -27,11 +27,23 @@ pub struct MangaOcr {
     tokenizer: Tokenizer,
     preprocessor: PreprocessorConfig,
     device: Device,
+    dtype: DType,
+}
+
+struct ImagePreprocessOptions<'a> {
+    image_size: u32,
+    image_mean: &'a [f32; 3],
+    image_std: &'a [f32; 3],
+    do_resize: bool,
+    do_normalize: bool,
+    device: &'a Device,
+    dtype: DType,
 }
 
 impl MangaOcr {
     pub async fn load(runtime: &RuntimeManager, cpu: bool) -> Result<Self> {
         let device = device(cpu)?;
+        let dtype = loading::model_dtype(&device);
         let hf = runtime.downloads();
         let config_path = hf.huggingface_model(HF_REPO, "config.json").await?;
         let preprocessor_path = hf
@@ -49,15 +61,19 @@ impl MangaOcr {
         let tokenizer = load_tokenizer(None, &vocab_path, &special_tokens_path)?;
         let model_device = device.clone();
         let weights = hf.huggingface_model(HF_REPO, "model.safetensors").await?;
-        let model = loading::load_mmaped_safetensors_path(&weights, &device, move |vb| {
-            VisionEncoderDecoder::from_config(config, vb, model_device.clone())
-        })?;
+        let model = loading::load_mmaped_safetensors_path_with_dtype(
+            &weights,
+            &device,
+            dtype,
+            move |vb| VisionEncoderDecoder::from_config(config, vb, model_device.clone()),
+        )?;
 
         Ok(Self {
             model,
             tokenizer,
             preprocessor,
             device,
+            dtype,
         })
     }
 
@@ -67,15 +83,16 @@ impl MangaOcr {
             return Ok(Vec::new());
         }
 
-        let pixel_values = preprocess_images(
-            images,
-            self.preprocessor.size,
-            &self.preprocessor.image_mean,
-            &self.preprocessor.image_std,
-            self.preprocessor.do_resize,
-            self.preprocessor.do_normalize,
-            &self.device,
-        )?;
+        let options = ImagePreprocessOptions {
+            image_size: self.preprocessor.size,
+            image_mean: &self.preprocessor.image_mean,
+            image_std: &self.preprocessor.image_std,
+            do_resize: self.preprocessor.do_resize,
+            do_normalize: self.preprocessor.do_normalize,
+            device: &self.device,
+            dtype: self.dtype,
+        };
+        let pixel_values = preprocess_images(images, &options)?;
         let token_ids = self.forward(&pixel_values)?;
         let texts = token_ids
             .into_iter()
@@ -96,24 +113,11 @@ impl MangaOcr {
 #[instrument(level = "debug", skip_all)]
 fn preprocess_images(
     images: &[image::DynamicImage],
-    image_size: u32,
-    image_mean: &[f32; 3],
-    image_std: &[f32; 3],
-    do_resize: bool,
-    do_normalize: bool,
-    device: &Device,
+    options: &ImagePreprocessOptions<'_>,
 ) -> Result<Tensor> {
     let mut batch = Vec::with_capacity(images.len());
     for image in images {
-        let processed = preprocess_single_image(
-            image,
-            image_size,
-            image_mean,
-            image_std,
-            do_resize,
-            do_normalize,
-            device,
-        )?;
+        let processed = preprocess_single_image(image, options)?;
         batch.push(processed);
     }
 
@@ -123,16 +127,11 @@ fn preprocess_images(
 #[instrument(level = "debug", skip_all)]
 fn preprocess_single_image(
     image: &image::DynamicImage,
-    image_size: u32,
-    image_mean: &[f32; 3],
-    image_std: &[f32; 3],
-    do_resize: bool,
-    do_normalize: bool,
-    device: &Device,
+    options: &ImagePreprocessOptions<'_>,
 ) -> Result<Tensor> {
     let (orig_w, orig_h) = image.dimensions();
-    let (width, height) = if do_resize {
-        (image_size as usize, image_size as usize)
+    let (width, height) = if options.do_resize {
+        (options.image_size as usize, options.image_size as usize)
     } else {
         (orig_w as usize, orig_h as usize)
     };
@@ -140,44 +139,44 @@ fn preprocess_single_image(
     let tensor = Tensor::from_vec(
         image.grayscale().to_rgb8().into_raw(),
         (1, orig_h as usize, orig_w as usize, 3),
-        device,
+        options.device,
     )?
     .permute((0, 3, 1, 2))?
     .to_dtype(DType::F32)?;
 
-    let tensor = if do_resize {
+    let tensor = if options.do_resize {
         tensor.interpolate2d(height, width)?
     } else {
         tensor
     };
 
     let tensor = (tensor * (1.0 / 255.0))?;
-    let tensor = if do_normalize {
+    let tensor = if options.do_normalize {
         let std = [
-            if image_std[0] == 0.0 {
+            if options.image_std[0] == 0.0 {
                 1.0
             } else {
-                image_std[0]
+                options.image_std[0]
             },
-            if image_std[1] == 0.0 {
+            if options.image_std[1] == 0.0 {
                 1.0
             } else {
-                image_std[1]
+                options.image_std[1]
             },
-            if image_std[2] == 0.0 {
+            if options.image_std[2] == 0.0 {
                 1.0
             } else {
-                image_std[2]
+                options.image_std[2]
             },
         ];
-        let mean_t = Tensor::from_slice(image_mean, (1, 3, 1, 1), device)?;
-        let std_t = Tensor::from_slice(&std, (1, 3, 1, 1), device)?;
+        let mean_t = Tensor::from_slice(options.image_mean, (1, 3, 1, 1), options.device)?;
+        let std_t = Tensor::from_slice(&std, (1, 3, 1, 1), options.device)?;
         tensor.broadcast_sub(&mean_t)?.broadcast_div(&std_t)?
     } else {
         tensor
     };
 
-    Ok(tensor)
+    Ok(tensor.to_dtype(options.dtype)?)
 }
 
 #[instrument(level = "debug", skip_all)]

koharu-ml/src/manga_ocr/model.rs

@@ -106,7 +106,7 @@ impl VisionEncoderDecoder {
                 }
 
                 let last_idx = seq_lengths[batch_idx].saturating_sub(1);
-                let last_logits = logits.i((batch_idx, last_idx, ..))?;
+                let last_logits = logits.i((batch_idx, last_idx, ..))?.to_dtype(DType::F32)?;
                 let next_id = sampler.sample(&last_logits)?;
                 seq.push(next_id);
                 if next_id == self.eos_token_id {

koharu-ml/src/manga_text_segmentation_2025/mod.rs

@@ -24,6 +24,7 @@ const CPU_MAX_PIXELS: u64 = 1_280 * 1_280;
 pub struct MangaTextSegmentation {
     model: model::MangaTextSegmentationModel,
     device: Device,
+    dtype: DType,
     mean: Tensor,
     std: Tensor,
 }
@@ -44,17 +45,19 @@ impl MangaTextSegmentation {
 
     pub fn load_from_path(path: impl AsRef<Path>, cpu: bool) -> Result<Self> {
         let device = device(cpu)?;
-        let model = loading::load_mmaped_safetensors_path(
+        let dtype = loading::model_dtype(&device);
+        let model = loading::load_mmaped_safetensors_path_with_dtype(
             path.as_ref(),
             &device,
+            dtype,
             model::MangaTextSegmentationModel::load,
         )?;
-        let mean =
-            Tensor::from_slice(&IMAGENET_MEAN, (1, 3, 1, 1), &device)?.to_dtype(DType::F32)?;
-        let std = Tensor::from_slice(&IMAGENET_STD, (1, 3, 1, 1), &device)?.to_dtype(DType::F32)?;
+        let mean = Tensor::from_slice(&IMAGENET_MEAN, (1, 3, 1, 1), &device)?.to_dtype(dtype)?;
+        let std = Tensor::from_slice(&IMAGENET_STD, (1, 3, 1, 1), &device)?.to_dtype(dtype)?;
         Ok(Self {
             model,
             device,
+            dtype,
             mean,
             std,
         })
@@ -92,6 +95,7 @@ impl MangaTextSegmentation {
         } else {
             probabilities
         }
+        .to_dtype(DType::F32)?
         .to_device(&Device::Cpu)?;
         let values = probabilities.flatten_all()?.to_vec1::<f32>()?;
         tracing::info!(
@@ -138,7 +142,7 @@ impl MangaTextSegmentation {
             &self.device,
         )?
         .permute((0, 3, 1, 2))?
-        .to_dtype(DType::F32)?;
+        .to_dtype(self.dtype)?;
         let tensor = (tensor * (1.0 / 255.0))?
             .broadcast_sub(&self.mean)?
             .broadcast_div(&self.std)?;

koharu-ml/src/mit48px_ocr/mod.rs

@@ -75,6 +75,7 @@ pub struct Mit48pxOcr {
     config: Mit48pxConfig,
     dictionary: Vec<String>,
     device: Device,
+    dtype: DType,
 }
 
 impl Mit48pxOcr {
@@ -102,18 +103,20 @@ impl Mit48pxOcr {
 
     fn load_from_files(files: ModelFiles, cpu: bool) -> Result<Self> {
         let device = device(cpu)?;
+        let dtype = loading::model_dtype(&device);
         let config: Mit48pxConfig =
             loading::read_json(&files.config).context("failed to parse mit48px config")?;
         let dictionary = read_dictionary(&files.dictionary)?;
         let data = std::fs::read(&files.weights)
             .with_context(|| format!("failed to read {}", files.weights.display()))?;
-        let vb = VarBuilder::from_buffered_safetensors(data, DType::F32, &device)?;
+        let vb = VarBuilder::from_buffered_safetensors(data, dtype, &device)?;
         let model = Mit48pxModel::new(config.clone(), dictionary.len(), vb, device.clone())?;
         Ok(Self {
             model,
             config,
             dictionary,
             device,
+            dtype,
         })
     }
 
@@ -125,7 +128,7 @@ impl Mit48pxOcr {
 
         let mut predictions = Vec::with_capacity(regions.len());
         for chunk in regions.chunks(OCR_CHUNK_SIZE) {
-            let batch = preprocess_regions(chunk, &self.config, &self.device)?;
+            let batch = preprocess_regions(chunk, &self.config, &self.device, self.dtype)?;
             let raw = self.model.infer_batch(&batch.tensor, &batch.widths)?;
             for prediction in raw {
                 predictions.push(self.decode_prediction(prediction));
@@ -278,6 +281,7 @@ fn preprocess_regions(
     regions: &[DynamicImage],
     config: &Mit48pxConfig,
     device: &Device,
+    dtype: DType,
 ) -> Result<PreparedBatch> {
     let mut resized = Vec::<RgbImage>::with_capacity(regions.len());
     let mut widths = Vec::with_capacity(regions.len());
@@ -310,8 +314,9 @@ fn preprocess_regions(
         }
     }
 
-    let tensor =
-        Tensor::from_vec(flat, (resized.len(), height, width, 3), device)?.permute((0, 3, 1, 2))?;
+    let tensor = Tensor::from_vec(flat, (resized.len(), height, width, 3), device)?
+        .permute((0, 3, 1, 2))?
+        .to_dtype(dtype)?;
     Ok(PreparedBatch { tensor, widths })
 }
 
@@ -397,7 +402,12 @@ mod tests {
     fn preprocessing_resizes_to_48px_and_matches_ballonstranslator_width_padding()
     -> anyhow::Result<()> {
         let image = DynamicImage::ImageRgb8(RgbImage::from_pixel(25, 10, image::Rgb([255, 0, 0])));
-        let batch = preprocess_regions(&[image], &test_config(), &candle_core::Device::Cpu)?;
+        let batch = preprocess_regions(
+            &[image],
+            &test_config(),
+            &candle_core::Device::Cpu,
+            candle_core::DType::F32,
+        )?;
         assert_eq!(batch.widths, vec![120]);
         assert_eq!(batch.tensor.dims(), &[1, 3, 48, 127]);
         Ok(())

koharu-ml/src/mit48px_ocr/model.rs

@@ -35,6 +35,7 @@ pub(crate) struct Mit48pxModel {
     color_pred_fg_ind: Linear,
     color_pred_bg_ind: Linear,
     device: Device,
+    dtype: DType,
 }
 
 #[derive(Clone)]
@@ -46,7 +47,7 @@ struct Hypothesis {
 }
 
 fn topk_last_dim(tensor: &Tensor, topk: usize) -> Result<TopkOutput> {
-    let rows = tensor.to_vec2::<f32>()?;
+    let rows = tensor.to_dtype(DType::F32)?.to_vec2::<f32>()?;
     let mut values = Vec::with_capacity(rows.len());
     let mut indices = Vec::with_capacity(rows.len());
 
@@ -93,10 +94,11 @@ impl Hypothesis {
         decoder_layers: usize,
         embd_dim: usize,
         device: &Device,
+        dtype: DType,
     ) -> Result<Self> {
         let mut cached_activations = Vec::with_capacity(decoder_layers + 1);
         for _ in 0..=decoder_layers {
-            cached_activations.push(Tensor::zeros((1, 0, embd_dim), DType::F32, device)?);
+            cached_activations.push(Tensor::zeros((1, 0, embd_dim), dtype, device)?);
         }
         Ok(Self {
             sample_index,
@@ -160,6 +162,7 @@ impl Mit48pxModel {
         vb: VarBuilder,
         device: Device,
     ) -> Result<Self> {
+        let dtype = vb.dtype();
         let backbone = ConvNextFeatureExtractor::new(vb.pp("backbone"))?;
         let encoders = (0..config.encoder_layers)
             .map(|index| TransformerEncoderLayer::new(vb.pp(format!("encoders.{index}"))))
@@ -190,6 +193,7 @@ impl Mit48pxModel {
             color_pred_fg_ind,
             color_pred_bg_ind,
             device,
+            dtype,
         })
     }
 
@@ -216,6 +220,7 @@ impl Mit48pxModel {
                     self.decoders.len(),
                     self.config.embd_dim,
                     &self.device,
+                    self.dtype,
                 )
             })
             .collect::<Result<Vec<_>>>()?;
@@ -400,7 +405,7 @@ impl Mit48pxModel {
     fn next_token_candidates(&self, decoded: &Tensor, beam_size: usize) -> Result<TopkOutput> {
         let pred_feats = self.pred1.forward(decoded)?.gelu_erf()?;
         let logits = self.pred.forward(&pred_feats)?;
-        let log_probs = candle_nn::ops::log_softmax(&logits, D::Minus1)?;
+        let log_probs = candle_nn::ops::log_softmax(&logits.to_dtype(DType::F32)?, D::Minus1)?;
         topk_last_dim(&log_probs, beam_size)
     }
 
@@ -410,6 +415,7 @@ impl Mit48pxModel {
         let fg_colors = self
             .color_pred_fg
             .forward(&color_feats)?
+            .to_dtype(DType::F32)?
             .squeeze(0)?
             .to_vec2::<f32>()?
             .into_iter()
@@ -418,6 +424,7 @@ impl Mit48pxModel {
         let bg_colors = self
             .color_pred_bg
             .forward(&color_feats)?
+            .to_dtype(DType::F32)?
             .squeeze(0)?
             .to_vec2::<f32>()?
             .into_iter()
@@ -426,6 +433,7 @@ impl Mit48pxModel {
         let fg_indicators = self
             .color_pred_fg_ind
             .forward(&color_feats)?
+            .to_dtype(DType::F32)?
             .squeeze(0)?
             .to_vec2::<f32>()?
             .into_iter()
@@ -434,6 +442,7 @@ impl Mit48pxModel {
         let bg_indicators = self
             .color_pred_bg_ind
             .forward(&color_feats)?
+            .to_dtype(DType::F32)?
             .squeeze(0)?
             .to_vec2::<f32>()?
             .into_iter()
@@ -871,7 +880,8 @@ impl XposMultiheadAttention {
                 f32::NEG_INFINITY,
                 attn_weights_4d.shape().dims(),
                 attn_weights_4d.device(),
-            )?;
+            )?
+            .to_dtype(attn_weights_4d.dtype())?;
             attn_weights = mask.where_cond(&neg_inf, &attn_weights_4d)?.reshape((
                 batch * self.num_heads,
                 tgt_len,
@@ -879,7 +889,9 @@ impl XposMultiheadAttention {
             ))?;
         }
 
-        let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?;
+        let attn_dtype = attn_weights.dtype();
+        let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights.to_dtype(DType::F32)?)?
+            .to_dtype(attn_dtype)?;
         let attn = attn_weights
             .matmul(&v)?
             .reshape((batch, self.num_heads, tgt_len, self.head_dim))?

koharu-ml/src/pp_doclayout_v3/mod.rs

@@ -46,6 +46,7 @@ pub struct PPDocLayoutV3 {
     config: PPDocLayoutV3Config,
     preprocessor: PPDocLayoutV3PreprocessorConfig,
     device: Device,
+    dtype: DType,
     mean: Tensor,
     std: Tensor,
 }
@@ -53,6 +54,7 @@ pub struct PPDocLayoutV3 {
 impl PPDocLayoutV3 {
     pub async fn load(runtime: &RuntimeManager, cpu: bool) -> Result<Self> {
         let device = device(cpu)?;
+        let dtype = loading::model_dtype(&device);
         let downloads = runtime.downloads();
         let config_path = downloads.huggingface_model(HF_REPO, "config.json").await?;
         let preprocessor_path = downloads
@@ -63,22 +65,26 @@ impl PPDocLayoutV3 {
         let preprocessor =
             loading::read_json::<PPDocLayoutV3PreprocessorConfig>(&preprocessor_path)
                 .with_context(|| format!("failed to load {}", preprocessor_path.display()))?;
-        let mean = Tensor::from_slice(&preprocessor.image_mean, (1, 3, 1, 1), &device)?
-            .to_dtype(DType::F32)?;
-        let std = Tensor::from_slice(&preprocessor.image_std, (1, 3, 1, 1), &device)?
-            .to_dtype(DType::F32)?;
+        let mean =
+            Tensor::from_slice(&preprocessor.image_mean, (1, 3, 1, 1), &device)?.to_dtype(dtype)?;
+        let std =
+            Tensor::from_slice(&preprocessor.image_std, (1, 3, 1, 1), &device)?.to_dtype(dtype)?;
         let weights_path = downloads
             .huggingface_model(HF_REPO, "model.safetensors")
             .await?;
-        let model = loading::load_mmaped_safetensors_path(&weights_path, &device, |vb| {
-            PPDocLayoutV3ForObjectDetection::load(vb, &config, &device)
-        })?;
+        let model = loading::load_mmaped_safetensors_path_with_dtype(
+            &weights_path,
+            &device,
+            dtype,
+            |vb| PPDocLayoutV3ForObjectDetection::load(vb, &config, &device),
+        )?;
 
         Ok(Self {
             model,
             config,
             preprocessor,
             device,
+            dtype,
             mean,
             std,
         })
@@ -116,6 +122,7 @@ impl PPDocLayoutV3 {
             images,
             &self.preprocessor,
             &self.device,
+            self.dtype,
             &self.mean,
             &self.std,
         )?;
@@ -376,6 +383,7 @@ fn preprocess_images(
     images: &[DynamicImage],
     preprocessor: &PPDocLayoutV3PreprocessorConfig,
     device: &Device,
+    dtype: DType,
     mean: &Tensor,
     std: &Tensor,
 ) -> Result<Tensor> {
@@ -395,7 +403,7 @@ fn preprocess_images(
     let tensor = Tensor::from_vec(batch, (images.len(), target_h, target_w, 3), &Device::Cpu)?
         .to_device(device)?
         .permute((0, 3, 1, 2))?
-        .to_dtype(DType::F32)?;
+        .to_dtype(dtype)?;
     let tensor = if preprocessor.do_rescale {
         tensor.affine(preprocessor.rescale_factor as f64, 0.0)?
     } else {
@@ -416,9 +424,18 @@ fn post_process_outputs(
     threshold: f32,
     include_polygons: bool,
 ) -> Result<Vec<LayoutDetectionResult>> {
-    let logits = outputs.logits.to_device(&Device::Cpu)?;
-    let boxes = outputs.pred_boxes.to_device(&Device::Cpu)?;
-    let order_logits = outputs.order_logits.to_device(&Device::Cpu)?;
+    let logits = outputs
+        .logits
+        .to_dtype(DType::F32)?
+        .to_device(&Device::Cpu)?;
+    let boxes = outputs
+        .pred_boxes
+        .to_dtype(DType::F32)?
+        .to_device(&Device::Cpu)?;
+    let order_logits = outputs
+        .order_logits
+        .to_dtype(DType::F32)?
+        .to_device(&Device::Cpu)?;
 
     let (batch_size, num_queries, num_classes) = logits.dims3()?;
     if batch_size != images.len() {
@@ -429,7 +446,10 @@ fn post_process_outputs(
     let boxes = boxes.flatten_all()?.to_vec1::<f32>()?;
     let order_logits = order_logits.flatten_all()?.to_vec1::<f32>()?;
     let (masks, mask_h, mask_w) = if include_polygons {
-        let masks = outputs.out_masks.to_device(&Device::Cpu)?;
+        let masks = outputs
+            .out_masks
+            .to_dtype(DType::F32)?
+            .to_device(&Device::Cpu)?;
         let (_, _, mask_h, mask_w) = masks.dims4()?;
         (Some(masks.flatten_all()?.to_vec1::<f32>()?), mask_h, mask_w)
     } else {

koharu-ml/src/pp_doclayout_v3/model.rs

@@ -117,6 +117,15 @@ fn load_layer_norm(vb: VarBuilder, hidden_size: usize, eps: f64) -> Result<Layer
     Ok(layer_norm(hidden_size, eps, vb)?)
 }
 
+fn softmax_f32(xs: &Tensor, dim: D) -> Result<Tensor> {
+    let dtype = xs.dtype();
+    if dtype == DType::F32 {
+        Ok(softmax(xs, dim)?)
+    } else {
+        Ok(softmax(&xs.to_dtype(DType::F32)?, dim)?.to_dtype(dtype)?)
+    }
+}
+
 fn pad_bottom_right_one(xs: &Tensor) -> candle_core::Result<Tensor> {
     xs.pad_with_zeros(2, 0, 1)?.pad_with_zeros(3, 0, 1)
 }
@@ -693,7 +702,8 @@ impl PPDocLayoutV3SelfAttention {
     ) -> Result<Tensor> {
         let (batch_size, sequence_length, hidden_size) = hidden_states.dims3()?;
         let query_key_input = match position_embeddings {
-            Some(position_embeddings) => hidden_states.broadcast_add(position_embeddings)?,
+            Some(position_embeddings) => hidden_states
+                .broadcast_add(&position_embeddings.to_dtype(hidden_states.dtype())?)?,
             None => hidden_states.clone(),
         };
 
@@ -726,9 +736,10 @@ impl PPDocLayoutV3SelfAttention {
         let mut attention_scores = query_states.matmul(&key_states_t)?;
         attention_scores = (attention_scores * self.scaling)?;
         if let Some(attention_mask) = attention_mask {
-            attention_scores = attention_scores.broadcast_add(attention_mask)?;
+            attention_scores = attention_scores
+                .broadcast_add(&attention_mask.to_dtype(attention_scores.dtype())?)?;
         }
-        let attention_probs = softmax(&attention_scores, D::Minus1)?;
+        let attention_probs = softmax_f32(&attention_scores, D::Minus1)?;
         let context = attention_probs.matmul(&value_states)?;
         let context = context.transpose(1, 2)?.contiguous()?.reshape((
             batch_size,
@@ -1084,7 +1095,8 @@ impl PPDocLayoutV3AIFILayer {
         let mut hidden = xs.flatten_from(2)?.transpose(1, 2)?;
         let position_embedding = self
             .position_embedding
-            .forward(width, height, xs.device())?;
+            .forward(width, height, xs.device())?
+            .to_dtype(hidden.dtype())?;
         for layer in &self.layers {
             hidden = layer.forward(&hidden, Some(&position_embedding))?;
         }
@@ -1486,7 +1498,8 @@ impl PPDocLayoutV3MultiscaleDeformableAttention {
         spatial_shapes_list: &[(usize, usize)],
     ) -> Result<Tensor> {
         let hidden_states = match position_embeddings {
-            Some(position_embeddings) => hidden_states.broadcast_add(position_embeddings)?,
+            Some(position_embeddings) => hidden_states
+                .broadcast_add(&position_embeddings.to_dtype(hidden_states.dtype())?)?,
             None => hidden_states.clone(),
         };
         let (batch_size, num_queries, _) = hidden_states.dims3()?;
@@ -1517,7 +1530,7 @@ impl PPDocLayoutV3MultiscaleDeformableAttention {
             self.n_points,
             2,
         ))?;
-        let attention_weights = softmax(
+        let attention_weights = softmax_f32(
             &self.attention_weights.forward(&hidden_states)?.reshape((
                 batch_size,
                 num_queries,
@@ -1748,9 +1761,9 @@ impl PPDocLayoutV3Decoder {
             )?;
 
             let predicted_corners = bbox_head.forward(&hidden_states)?;
-            reference_points = sigmoid(
-                &predicted_corners.broadcast_add(&inverse_sigmoid_tensor(&reference_points)?)?,
-            )?;
+            let reference_points_unact =
+                inverse_sigmoid_tensor(&reference_points)?.to_dtype(predicted_corners.dtype())?;
+            reference_points = sigmoid(&predicted_corners.broadcast_add(&reference_points_unact)?)?;
 
             let out_query = norm.forward(&hidden_states)?;
             let mask_query_embed = mask_query_head.forward(&out_query)?;
@@ -1964,7 +1977,8 @@ impl PPDocLayoutV3ForObjectDetection {
             let enc_out_masks =
                 batched_mask_projection(&mask_query_embed, &encoder_outputs.mask_feat)?;
             reference_points_unact =
-                inverse_sigmoid_tensor(&mask_to_box_coordinate(&enc_out_masks, 0.0)?)?;
+                inverse_sigmoid_tensor(&mask_to_box_coordinate(&enc_out_masks, 0.0)?)?
+                    .to_dtype(target.dtype())?;
         }
 
         let decoder_outputs = self.decoder.forward(
@@ -2097,8 +2111,12 @@ fn batched_mask_projection(mask_query_embed: &Tensor, mask_feat: &Tensor) -> Res
 }
 
 fn inverse_sigmoid_tensor(tensor: &Tensor) -> Result<Tensor> {
+    let dtype = tensor.dtype();
     let tensor = tensor.to_dtype(DType::F32)?.clamp(1e-5, 1.0 - 1e-5)?;
-    Ok(tensor.broadcast_div(&tensor.affine(-1.0, 1.0)?)?.log()?)
+    Ok(tensor
+        .broadcast_div(&tensor.affine(-1.0, 1.0)?)?
+        .log()?
+        .to_dtype(dtype)?)
 }
 
 fn mask_to_box_coordinate(mask_logits: &Tensor, threshold: f32) -> Result<Tensor> {
@@ -2199,6 +2217,9 @@ fn multi_scale_deformable_attention(
     n_points: usize,
 ) -> Result<Tensor> {
     let (batch_size, sequence_length, num_heads, head_dim) = value.dims4()?;
+    let reference_points = reference_points.to_dtype(DType::F32)?;
+    let sampling_offsets = sampling_offsets.to_dtype(DType::F32)?;
+    let attention_weights = attention_weights.to_dtype(DType::F32)?;
     let sampling_dims = sampling_offsets.dims().to_vec();
     let (num_queries, num_levels, num_points) = match sampling_dims.as_slice() {
         [_, queries, _, levels, points, _] => (*queries, *levels, *points),

koharu-ml/src/speech_bubble_segmentation/mod.rs

@@ -45,6 +45,7 @@ pub struct SpeechBubbleSegmentation {
     model: YoloV8Seg,
     config: SpeechBubbleSegmentationConfig,
     device: Device,
+    dtype: DType,
 }
 
 #[derive(Debug, Clone)]
@@ -148,25 +149,32 @@ impl SpeechBubbleSegmentation {
         cpu: bool,
     ) -> Result<Self> {
         let device = device(cpu)?;
+        let dtype = loading::model_dtype(&device);
         let config = loading::read_json::<SpeechBubbleSegmentationConfig>(config_path.as_ref())
             .with_context(|| format!("failed to parse {}", config_path.as_ref().display()))?;
         config.validate()?;
         let multiples = variant_multiples(&config)?;
-        let model = loading::load_mmaped_safetensors_path(weights_path.as_ref(), &device, |vb| {
-            YoloV8Seg::load(
-                vb,
-                multiples,
-                config.num_classes,
-                config.num_masks,
-                config.num_prototypes,
-                config.reg_max,
-            )
-        })?;
+        let model = loading::load_mmaped_safetensors_path_with_dtype(
+            weights_path.as_ref(),
+            &device,
+            dtype,
+            |vb| {
+                YoloV8Seg::load(
+                    vb,
+                    multiples,
+                    config.num_classes,
+                    config.num_masks,
+                    config.num_prototypes,
+                    config.reg_max,
+                )
+            },
+        )?;
 
         Ok(Self {
             model,
             config,
             device,
+            dtype,
         })
     }
 
@@ -258,7 +266,7 @@ impl SpeechBubbleSegmentation {
             &self.device,
         )?
         .permute((0, 3, 1, 2))?
-        .to_dtype(DType::F32)?;
+        .to_dtype(self.dtype)?;
         let pixel_values = (pixel_values * (1.0 / 255.0))?;
 
         Ok(PreparedInput {
@@ -306,8 +314,16 @@ fn postprocess(
     confidence_threshold: f32,
     nms_threshold: f32,
 ) -> Result<SpeechBubbleSegmentationResult> {
-    let pred = outputs.pred.to_device(&Device::Cpu)?.i(0)?;
-    let proto = outputs.proto.to_device(&Device::Cpu)?.i(0)?;
+    let pred = outputs
+        .pred
+        .to_dtype(DType::F32)?
+        .to_device(&Device::Cpu)?
+        .i(0)?;
+    let proto = outputs
+        .proto
+        .to_dtype(DType::F32)?
+        .to_device(&Device::Cpu)?
+        .i(0)?;
     let raw_regions =
         extract_regions(&pred, prepared, config, confidence_threshold, nms_threshold)?;
     let mut probability_map =

koharu-ml/src/speech_bubble_segmentation/model.rs

@@ -1,4 +1,4 @@
-use candle_core::{D, DType, IndexOp, Result, Tensor};
+use candle_core::{D, IndexOp, Result, Tensor};
 use candle_nn::{
     BatchNorm, Conv2d, Conv2dConfig, ConvTranspose2d, ConvTranspose2dConfig, Module, VarBuilder,
     batch_norm, conv_transpose2d, conv2d, conv2d_no_bias,
@@ -399,12 +399,13 @@ fn make_anchors(
     grid_cell_offset: f64,
 ) -> Result<(Tensor, Tensor)> {
     let device = xs0.device();
+    let dtype = xs0.dtype();
     let mut anchor_points = Vec::new();
     let mut stride_tensors = Vec::new();
     for (xs, stride) in [(xs0, strides.0), (xs1, strides.1), (xs2, strides.2)] {
         let (_, _, h, w) = xs.dims4()?;
-        let sx = (Tensor::arange(0, w as u32, device)?.to_dtype(DType::F32)? + grid_cell_offset)?;
-        let sy = (Tensor::arange(0, h as u32, device)?.to_dtype(DType::F32)? + grid_cell_offset)?;
+        let sx = (Tensor::arange(0, w as u32, device)?.to_dtype(dtype)? + grid_cell_offset)?;
+        let sy = (Tensor::arange(0, h as u32, device)?.to_dtype(dtype)? + grid_cell_offset)?;
         let sx = sx
             .reshape((1, sx.elem_count()))?
             .repeat((h, 1))?
@@ -414,7 +415,7 @@ fn make_anchors(
             .repeat((1, w))?
             .flatten_all()?;
         anchor_points.push(Tensor::stack(&[&sx, &sy], D::Minus1)?);
-        stride_tensors.push((Tensor::ones(h * w, DType::F32, device)? * stride as f64)?);
+        stride_tensors.push((Tensor::ones(h * w, dtype, device)? * stride as f64)?);
     }
     let anchor_points = Tensor::cat(anchor_points.as_slice(), 0)?;
     let stride_tensor = Tensor::cat(stride_tensors.as_slice(), 0)?.unsqueeze(1)?;