Add ability to use any CLIPModel with model-id and revision

candle-examples/examples/bert/main.rs

@@ -76,7 +76,7 @@ impl Args {
  };
  (config, tokenizer, weights)
  };
- let config = std::fs::read_to_string(config_filename)?;
+ let config: String = std::fs::read_to_string(config_filename)?;
  let mut config: Config = serde_json::from_str(&config)?;
  let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?;
 

candle-examples/examples/clip/README.md

@@ -13,19 +13,34 @@ https://github.com/huggingface/transformers/tree/f6fa0f0bf0796ac66f201f23bdb8585
 $ cargo run --example clip --release -- --images "candle-examples/examples/stable-diffusion/assets/stable-diffusion-xl.jpg","candle-examples/examples/yolo-v8/assets/bike.jpg" --cpu --sequences "a cycling race","a photo of two cats","a robot holding a candle"
 
 
-Results for image: candle-examples/examples/stable-diffusion/assets/stable-diffusion-xl.jpg
+> Results for image: candle-examples/examples/stable-diffusion/assets/stable-diffusion-xl.jpg
 
-INFO clip: Probability: 0.0000% Text: a cycling race
-INFO clip: Probability: 0.0000% Text: a photo of two cats
-INFO clip: Probability: 100.0000% Text: a robot holding a candle
+> INFO clip: Probability: 0.0000% Text: a cycling race
+> INFO clip: Probability: 0.0000% Text: a photo of two cats
+> INFO clip: Probability: 100.0000% Text: a robot holding a candle
 
-Results for image: candle-examples/examples/yolo-v8/assets/bike.jpg
+> Results for image: candle-examples/examples/yolo-v8/assets/bike.jpg
 
-INFO clip: Probability: 99.9999% Text: a cycling race
-INFO clip: Probability: 0.0001% Text: a photo of two cats
-INFO clip: Probability: 0.0000% Text: a robot holding a candle
+> INFO clip: Probability: 99.9999% Text: a cycling race
+> INFO clip: Probability: 0.0001% Text: a photo of two cats
+> INFO clip: Probability: 0.0000% Text: a robot holding a candle
 ```
 
+### Arguments
+- `--model`: local path to the model. If not provided will download from huggingface
+- `--tokenizer`: local path to the tokenizer.json file. If not provided will download from huggingface
+- `--images`: list of images to use. 
+
+Example: `--images candle-examples/examples/stable-diffusion/assets/stable-diffusion-xl.jpg,candle-examples/examples/yolo-v8/assets/bike.jpg`
+- `sequences`: list of text sequences to use. 
+
+Example: `--sequences "a cycling race","bike"`
+
+- `--model-id`: model id to use from huggingface. Example: `--model-id openai/clip-vit-large-patch14`
+- `--revision`: revision to use from huggingface. Example: `--revision refs/pr/4`
+- `--use-pth`: Use the pytorch weights rather than the safetensors ones. Default: true
+- `--cpu`: Use cpu. Use `--cpu false` for gpu but requires gpu support with `--features cuda` 
+
 ## Running on an example with metal feature (mac)
 
 ```

candle-examples/examples/clip/main.rs

@@ -9,26 +9,42 @@ use clap::Parser;
 
 use candle::{DType, Device, Tensor};
 use candle_nn::{ops::softmax, VarBuilder};
-use candle_transformers::models::clip;
+use candle_transformers::models::clip::{self, ClipConfig};
 
+use hf_hub::{api::sync::Api, Repo, RepoType};
 use tokenizers::Tokenizer;
 
 #[derive(Parser)]
 struct Args {
+ /// model to use from local file. If not provided will download from huggingface
  #[arg(long)]
  model: Option<String>,
 
  #[arg(long)]
  tokenizer: Option<String>,
 
+ /// The images to use
  #[arg(long, use_value_delimiter = true)]
  images: Option<Vec<String>>,
 
  #[arg(long)]
  cpu: bool,
 
+ /// The sequences to use
  #[arg(long, use_value_delimiter = true)]
  sequences: Option<Vec<String>>,
+
+ /// The model_id to use from huggingface, check out available models:https://huggingface.co/openai
+ #[arg(long)]
+ model_id: Option<String>,
+
+ /// The revision to use from huggingface
+ #[arg(long)]
+ revision: Option<String>,
+
+ /// Use the pytorch weights rather than the safetensors ones
+ #[arg(long, default_value = "false")]
+ use_pth: bool,
 }
 
 fn load_image<T: AsRef<std::path::Path>>(path: T, image_size: usize) -> anyhow::Result<Tensor> {
@@ -63,33 +79,60 @@ fn load_images<T: AsRef<std::path::Path>>(
 
 pub fn main() -> anyhow::Result<()> {
  let args = Args::parse();
- let model_file = match args.model {
- None => {
- let api = hf_hub::api::sync::Api::new()?;
 
- let api = api.repo(hf_hub::Repo::with_revision(
- "openai/clip-vit-base-patch32".to_string(),
- hf_hub::RepoType::Model,
- "refs/pr/15".to_string(),
- ));
+ tracing_subscriber::fmt::init();
 
- api.get("model.safetensors")?
- }
- Some(model) => model.into(),
- };
- let tokenizer = get_tokenizer(args.tokenizer)?;
- let config = clip::ClipConfig::vit_base_patch32();
  let device = candle_examples::device(args.cpu)?;
+ let default_model = "openai/clip-vit-base-patch16".to_string();
+ let default_revision = "refs/pr/4".to_string();
+ let (model_id, revision) = match (args.model_id.to_owned(), args.revision.to_owned()) {
+ (Some(model_id), Some(revision)) => (model_id, revision),
+ (Some(model_id), None) => (model_id, "main".to_string()),
+ (None, Some(revision)) => (default_model, revision),
+ (None, None) => (default_model, default_revision),
+ };
+
+ let repo = Repo::with_revision(model_id, RepoType::Model, revision);
+ let (config_filename, tokenizer_filename, weights_filename) = {
+ let api = Api::new()?;
+ let api = api.repo(repo);
+ let config = api.get("config.json")?;
+ let tokenizer = match args.tokenizer {
+ Some(tokenizer) => tokenizer.into(),
+ None => api.get("tokenizer.json")?,
+ };
+ let weights = match args.model {
+ Some(model) => model.into(),
+ None => {
+ if args.use_pth {
+ api.get("pytorch_model.bin")?
+ } else {
+ api.get("model.safetensors")?
+ }
+ }
+ };
+
+ (config, tokenizer, weights)
+ };
+
+ let config: String = std::fs::read_to_string(config_filename)?;
+ let config: ClipConfig = serde_json::from_str(&config)?;
+ let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?;
+
  let vec_imgs = match args.images {
  Some(imgs) => imgs,
  None => vec![
  "candle-examples/examples/stable-diffusion/assets/stable-diffusion-xl.jpg".to_string(),
  "candle-examples/examples/yolo-v8/assets/bike.jpg".to_string(),
  ],
  };
- let images = load_images(&vec_imgs, config.image_size)?.to_device(&device)?;
- let vb =
- unsafe { VarBuilder::from_mmaped_safetensors(&[model_file.clone()], DType::F32, &device)? };
+
+ let images = load_images(&vec_imgs, config.vision_config.image_size)?.to_device(&device)?;
+
+ let vb = unsafe {
+ VarBuilder::from_mmaped_safetensors(&[weights_filename.clone()], DType::F32, &device)?
+ };
+
  let model = clip::ClipModel::new(vb, &config)?;
  let (input_ids, vec_seq) = tokenize_sequences(args.sequences, &tokenizer, &device)?;
  let (_logits_per_text, logits_per_image) = model.forward(&images, &input_ids)?;
@@ -113,22 +156,6 @@ pub fn main() -> anyhow::Result<()> {
  Ok(())
 }
 
-pub fn get_tokenizer(tokenizer: Option<String>) -> anyhow::Result<Tokenizer> {
- let tokenizer = match tokenizer {
- None => {
- let api = hf_hub::api::sync::Api::new()?;
- let api = api.repo(hf_hub::Repo::with_revision(
- "openai/clip-vit-base-patch32".to_string(),
- hf_hub::RepoType::Model,
- "refs/pr/15".to_string(),
- ));
- api.get("tokenizer.json")?
- }
- Some(file) => file.into(),
- };
- Tokenizer::from_file(tokenizer).map_err(E::msg)
-}
-
 pub fn tokenize_sequences(
  sequences: Option<Vec<String>>,
  tokenizer: &Tokenizer,

candle-examples/examples/flux/main.rs

@@ -133,9 +133,9 @@ fn run(args: Args) -> Result<()> {
  let config = clip::text_model::ClipTextConfig {
  vocab_size: 49408,
  projection_dim: 768,
- activation: clip::text_model::Activation::QuickGelu,
+ hidden_act: clip::text_model::Activation::QuickGelu,
  intermediate_size: 3072,
- embed_dim: 768,
+ hidden_size: 768,
  max_position_embeddings: 77,
  pad_with: None,
  num_hidden_layers: 12,

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

@@ -10,6 +10,7 @@ use self::{
  vision_model::ClipVisionTransformer,
 };
 use candle::{Result, Tensor, D};
+use serde::Deserialize;
 
 pub mod text_model;
 pub mod vision_model;
@@ -32,8 +33,8 @@ pub enum EncoderConfig {
 impl EncoderConfig {
  pub fn embed_dim(&self) -> usize {
  match self {
- Self::Text(c) => c.embed_dim,
- Self::Vision(c) => c.embed_dim,
+ Self::Text(c) => c.hidden_size,
+ Self::Vision(c) => c.hidden_size,
  }
  }
 
@@ -61,17 +62,16 @@ impl EncoderConfig {
  pub fn activation(&self) -> Activation {
  match self {
  Self::Text(_c) => Activation::QuickGelu,
- Self::Vision(c) => c.activation,
+ Self::Vision(c) => c.hidden_act,
  }
  }
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Deserialize)]
 pub struct ClipConfig {
  pub text_config: text_model::ClipTextConfig,
  pub vision_config: vision_model::ClipVisionConfig,
  pub logit_scale_init_value: f32,
- pub image_size: usize,
 }
 
 impl ClipConfig {
@@ -84,7 +84,6 @@ impl ClipConfig {
  text_config,
  vision_config,
  logit_scale_init_value: 2.6592,
- image_size: 224,
  }
  }
 }
@@ -94,12 +93,12 @@ impl ClipModel {
  let text_model = ClipTextTransformer::new(vs.pp("text_model"), &c.text_config)?;
  let vision_model = ClipVisionTransformer::new(vs.pp("vision_model"), &c.vision_config)?;
  let visual_projection = candle_nn::linear_no_bias(
- c.vision_config.embed_dim,
+ c.vision_config.hidden_size,
  c.vision_config.projection_dim,
  vs.pp("visual_projection"),
  )?;
  let text_projection = candle_nn::linear_no_bias(
- c.text_config.embed_dim,
+ c.text_config.hidden_size,
  c.text_config.projection_dim,
  vs.pp("text_projection"),
  )?;

candle-transformers/src/models/clip/text_model.rs

@@ -9,10 +9,12 @@
 use candle::{DType, Device, IndexOp, Result, Tensor, D};
 use candle_nn as nn;
 use candle_nn::Module;
+use serde::Deserialize;
 
 use super::EncoderConfig;
 
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, Clone, Copy, Deserialize)]
+#[serde(rename_all = "snake_case")]
 pub enum Activation {
  QuickGelu,
 }
@@ -25,11 +27,11 @@ impl Module for Activation {
  }
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Deserialize)]
 pub struct ClipTextConfig {
  pub vocab_size: usize,
- pub embed_dim: usize,
- pub activation: Activation,
+ pub hidden_size: usize,
+ pub hidden_act: Activation,
  pub intermediate_size: usize,
  pub max_position_embeddings: usize,
  pub pad_with: Option<String>,
@@ -45,14 +47,14 @@ impl ClipTextConfig {
  pub fn vit_base_patch32() -> Self {
  Self {
  vocab_size: 49408,
- embed_dim: 512,
+ hidden_size: 512,
  intermediate_size: 2048,
  max_position_embeddings: 77,
  pad_with: None,
  num_hidden_layers: 12,
  num_attention_heads: 8,
  projection_dim: 512,
- activation: Activation::QuickGelu,
+ hidden_act: Activation::QuickGelu,
  }
  }
 }
@@ -69,10 +71,10 @@ struct ClipTextEmbeddings {
 impl ClipTextEmbeddings {
  fn new(vs: candle_nn::VarBuilder, c: &ClipTextConfig) -> Result<Self> {
  let token_embedding =
- candle_nn::embedding(c.vocab_size, c.embed_dim, vs.pp("token_embedding"))?;
+ candle_nn::embedding(c.vocab_size, c.hidden_size, vs.pp("token_embedding"))?;
  let position_embedding: nn::Embedding = candle_nn::embedding(
  c.max_position_embeddings,
- c.embed_dim,
+ c.hidden_size,
  vs.pp("position_embedding"),
  )?;
  let position_ids =
@@ -108,13 +110,13 @@ struct ClipAttention {
 
 impl ClipAttention {
  fn new(vs: candle_nn::VarBuilder, c: &EncoderConfig) -> Result<Self> {
- let embed_dim = c.embed_dim();
+ let hidden_size = c.embed_dim();
  let num_attention_heads = c.num_attention_heads();
- let k_proj = candle_nn::linear(embed_dim, embed_dim, vs.pp("k_proj"))?;
- let v_proj = candle_nn::linear(embed_dim, embed_dim, vs.pp("v_proj"))?;
- let q_proj = candle_nn::linear(embed_dim, embed_dim, vs.pp("q_proj"))?;
- let out_proj = candle_nn::linear(embed_dim, embed_dim, vs.pp("out_proj"))?;
- let head_dim = embed_dim / num_attention_heads;
+ let k_proj = candle_nn::linear(hidden_size, hidden_size, vs.pp("k_proj"))?;
+ let v_proj = candle_nn::linear(hidden_size, hidden_size, vs.pp("v_proj"))?;
+ let q_proj = candle_nn::linear(hidden_size, hidden_size, vs.pp("q_proj"))?;
+ let out_proj = candle_nn::linear(hidden_size, hidden_size, vs.pp("out_proj"))?;
+ let head_dim = hidden_size / num_attention_heads;
  let scale = (head_dim as f64).powf(-0.5);
 
  Ok(ClipAttention {
@@ -136,7 +138,7 @@ impl ClipAttention {
 
  fn forward(&self, xs: &Tensor, causal_attention_mask: Option<&Tensor>) -> Result<Tensor> {
  let in_dtype = xs.dtype();
- let (bsz, seq_len, embed_dim) = xs.dims3()?;
+ let (bsz, seq_len, hidden_size) = xs.dims3()?;
 
  let query_states = (self.q_proj.forward(xs)? * self.scale)?;
  let proj_shape = (bsz * self.num_attention_heads, seq_len, self.head_dim);
@@ -171,7 +173,7 @@ impl ClipAttention {
  let attn_output = attn_output
  .reshape((bsz, self.num_attention_heads, seq_len, self.head_dim))?
  .transpose(1, 2)?
- .reshape((bsz, seq_len, embed_dim))?;
+ .reshape((bsz, seq_len, hidden_size))?;
  self.out_proj.forward(&attn_output)
  }
 }
@@ -290,7 +292,8 @@ impl ClipTextTransformer {
  pub fn new(vs: candle_nn::VarBuilder, c: &ClipTextConfig) -> Result<Self> {
  let embeddings = ClipTextEmbeddings::new(vs.pp("embeddings"), c)?;
  let encoder = ClipEncoder::new(vs.pp("encoder"), &EncoderConfig::Text(c.clone()))?;
- let final_layer_norm = candle_nn::layer_norm(c.embed_dim, 1e-5, vs.pp("final_layer_norm"))?;
+ let final_layer_norm =
+ candle_nn::layer_norm(c.hidden_size, 1e-5, vs.pp("final_layer_norm"))?;
  Ok(ClipTextTransformer {
  embeddings,
  encoder,