Add quantized CLIP

candle-transformers/src/models/mod.rs

@@ -33,6 +33,7 @@ pub mod phi;
 pub mod phi3;
 pub mod quantized_blip;
 pub mod quantized_blip_text;
+pub mod quantized_clip;
 pub mod quantized_llama;
 pub mod quantized_llama2_c;
 pub mod quantized_metavoice;

candle-transformers/src/models/quantized_clip/mod.rs

@@ -0,0 +1,126 @@
+use candle::{Result, Tensor, D};
+
+use crate::models::clip;
+use crate::models::clip::text_model::{ClipTextConfig, Activation};
+use crate::models::clip::vision_model::ClipVisionConfig;
+use crate::quantized_nn as quantized_nn;
+use crate::quantized_var_builder::VarBuilder;
+
+pub mod text_model;
+pub mod vision_model;
+
+#[derive(Clone, Debug)]
+pub enum EncoderConfig {
+ Text(ClipTextConfig),
+ Vision(ClipVisionConfig),
+}
+
+impl EncoderConfig {
+ pub fn embed_dim(&self) -> usize {
+ match self {
+ Self::Text(c) => c.embed_dim,
+ Self::Vision(c) => c.embed_dim,
+ }
+ }
+
+ pub fn num_attention_heads(&self) -> usize {
+ match self {
+ Self::Text(c) => c.num_attention_heads,
+ Self::Vision(c) => c.num_attention_heads,
+ }
+ }
+
+ pub fn intermediate_size(&self) -> usize {
+ match self {
+ Self::Text(c) => c.intermediate_size,
+ Self::Vision(c) => c.intermediate_size,
+ }
+ }
+
+ pub fn num_hidden_layers(&self) -> usize {
+ match self {
+ Self::Text(c) => c.num_hidden_layers,
+ Self::Vision(c) => c.num_hidden_layers,
+ }
+ }
+
+ pub fn activation(&self) -> Activation {
+ match self {
+ Self::Text(_c) => Activation::QuickGelu,
+ Self::Vision(c) => c.activation,
+ }
+ }
+}
+
+#[derive(Clone, Debug)]
+pub struct ClipModel {
+ text_model: text_model::ClipTextTransformer,
+ vision_model: vision_model::ClipVisionTransformer,
+ visual_projection: quantized_nn::Linear,
+ text_projection: quantized_nn::Linear,
+ logit_scale: Tensor,
+}
+
+impl ClipModel {
+ pub fn new(vs: VarBuilder, c: &clip::ClipConfig) -> Result<Self> {
+ let text_model = text_model::ClipTextTransformer::new(vs.pp("text_model"), &c.text_config)?;
+
+ let vision_model = vision_model::ClipVisionTransformer::new(vs.pp("vision_model"), &c.vision_config)?;
+
+ let visual_projection = quantized_nn::linear_no_bias(
+ c.vision_config.embed_dim,
+ c.vision_config.projection_dim,
+ vs.pp("visual_projection"),
+ )?;
+
+ let text_projection = quantized_nn::linear_no_bias(
+ c.text_config.embed_dim,
+ c.text_config.projection_dim,
+ vs.pp("text_projection"),
+ )?;
+
+ // originally nn.Parameter
+ let logit_scale = if vs.contains_key("logit_scale") {
+ vs.get(&[], "logit_scale")?.dequantize(vs.device())?
+ } else {
+ Tensor::new(&[c.logit_scale_init_value], vs.device())?
+ };
+
+ Ok(Self {
+ text_model,
+ vision_model,
+ visual_projection,
+ text_projection,
+ logit_scale,
+ })
+ }
+
+ pub fn get_text_features(&self, input_ids: &Tensor) -> Result<Tensor> {
+ input_ids
+ .apply(&self.text_model)?
+ .apply(&self.text_projection)
+ }
+
+ pub fn get_image_features(&self, pixel_values: &Tensor) -> Result<Tensor> {
+ pixel_values
+ .apply(&self.vision_model)?
+ .apply(&self.visual_projection)
+ }
+
+ pub fn forward(&self, pixel_values: &Tensor, input_ids: &Tensor) -> Result<(Tensor, Tensor)> {
+ let image_features = self.get_image_features(pixel_values)?;
+ let text_features = self.get_text_features(input_ids)?;
+ let image_features_normalized = div_l2_norm(&image_features)?;
+ let text_features_normalized = div_l2_norm(&text_features)?;
+ let logits_per_text = text_features_normalized.matmul(&image_features_normalized.t()?)?;
+ let logit_scale = self.logit_scale.exp()?;
+ let logits_per_text = logits_per_text.broadcast_mul(&logit_scale)?;
+ let logits_per_image = logits_per_text.t()?;
+ Ok((logits_per_text, logits_per_image))
+ }
+}
+
+pub fn div_l2_norm(v: &Tensor) -> Result<Tensor> {
+ let l2_norm = v.sqr()?.sum_keepdim(D::Minus1)?.sqrt()?;
+ v.broadcast_div(&l2_norm)
+}