From 2182d4dde1d520d09a2c35250e190a710f1bef50 Mon Sep 17 00:00:00 2001
From: Brandon Thomas <echelon@gmail.com>
Date: Wed, 26 Mar 2025 20:32:05 -0400
Subject: [PATCH] lots of new models in R2

---
 .../src/weights_registry/flux_weights.rs      | 81 +++++++++++++++++++
 .../src/weights_registry/mod.rs               |  3 +-
 .../src/weights_registry/weight_descriptor.rs |  2 +
 3 files changed, 85 insertions(+), 1 deletion(-)
 create mode 100644 crates/desktop/ml_weights_registry/src/weights_registry/flux_weights.rs

diff --git a/crates/desktop/ml_weights_registry/src/weights_registry/flux_weights.rs b/crates/desktop/ml_weights_registry/src/weights_registry/flux_weights.rs
new file mode 100644
index 0000000000..33ce0047dc
--- /dev/null
+++ b/crates/desktop/ml_weights_registry/src/weights_registry/flux_weights.rs
@@ -0,0 +1,81 @@
+use crate::weights_registry::weight_descriptor::{WeightDescriptor, WeightFunction};
+use crate::weights_registry::weight_descriptor_builder::weight;
+
+
+/// https://huggingface.co/black-forest-labs/FLUX.1-dev
+pub const FLUX_DEV : WeightDescriptor = weight!(
+  "Flux-1 Dev",
+  "flux1-dev.safetensors",
+  WeightFunction::ImageGeneration
+);
+
+// TODO(bt): I think the autoencoder for Dev and Schnell are the same.
+//  They have the same checksum unless I mixed up the files when downloading.
+
+/// https://huggingface.co/black-forest-labs/FLUX.1-dev
+pub const FLUX_DEV_AUTOENCODER : WeightDescriptor = weight!(
+  "Flux-1 Dev (Autoencoder)",
+  "flux1-dev.ae.safetensors",
+  WeightFunction::ImageGeneration
+);
+
+/// https://huggingface.co/black-forest-labs/FLUX.1-schnell
+pub const FLUX_SCHNELL : WeightDescriptor = weight!(
+  "Flux-1 Schnell",
+  "flux1-schnell.safetensors",
+  WeightFunction::ImageGeneration
+);
+
+/// https://huggingface.co/black-forest-labs/FLUX.1-schnell
+pub const FLUX_SCHNELL_AUTOENCODER : WeightDescriptor = weight!(
+  "Flux-1 Schnell (Autoencoder)",
+  "flux1-schnell.ae.safetensors",
+  WeightFunction::ImageGeneration
+);
+
+/// https://huggingface.co/lmz/candle-flux
+pub const LMZ_CANDLE_FLUX_SCHNELL_QUANTIZED_GGUF : WeightDescriptor = weight!(
+  "LMZ Flux Schnell Quantized",
+  "lmz_candle_flux_flux1-schnell.gguf",
+  WeightFunction::ImageGeneration
+);
+
+// TODO: Still uploading
+/// https://huggingface.co/google/t5-v1_1-xxl/tree/main
+/// NB: I couldn't find the safetensors file here
+/// NB: This should live under "refs/pr/2"
+pub const GOOGLE_T5_V1_1_XXL_MODEL : WeightDescriptor = weight!(
+  "Google T5 v1.1 (XXL)",
+  "google_t5_v1_1_xxl_model.safetensors",
+  WeightFunction::TextEncoder
+);
+
+/// https://huggingface.co/google/t5-v1_1-xxl/tree/main
+/// NB: I couldn't find the safetensors file here
+/// NB: This should live under "refs/pr/2"
+pub const GOOGLE_T5_V1_1_XXL_CONFIG : WeightDescriptor = weight!(
+  "Google T5 v1.1 (XXL) Config",
+  "google_t5_v1_1_xxl_config.json",
+  WeightFunction::TextEncoder
+);
+
+/// https://huggingface.co/lmz/mt5-tokenizers
+pub const LMZ_T5_TOKENIZER_JSON : WeightDescriptor = weight!(
+  "LMZ T5 Tokenizer",
+  "lmt_mt5_tokenizers_t5-v1_1-xxl.tokenizer.json",
+  WeightFunction::TextTokenizer
+);
+
+/// https://huggingface.co/openai/clip-vit-large-patch14/tree/main
+pub const OPENAI_CLIP_VIT_P14_MODEL : WeightDescriptor = weight!(
+  "OpenAI CLIP VIT (patch 14) model",
+  "openai_clip-vit-large-patch14.model.safetensors",
+  WeightFunction::TextTokenizer
+);
+
+/// https://huggingface.co/openai/clip-vit-large-patch14/tree/main
+pub const OPENAI_CLIP_VIT_P14_TOKENIZER_JSON: WeightDescriptor = weight!(
+  "OpenAI CLIP VIT (patch 14) tokenizer",
+  "openai_clip-vit-large-patch14.tokenizer.json",
+  WeightFunction::TextTokenizer
+);
diff --git a/crates/desktop/ml_weights_registry/src/weights_registry/mod.rs b/crates/desktop/ml_weights_registry/src/weights_registry/mod.rs
index 23e8434c4b..a49023702b 100644
--- a/crates/desktop/ml_weights_registry/src/weights_registry/mod.rs
+++ b/crates/desktop/ml_weights_registry/src/weights_registry/mod.rs
@@ -1,3 +1,4 @@
 mod weight_descriptor_builder;
 pub mod weight_descriptor;
-pub mod weights;
\ No newline at end of file
+pub mod weights;
+mod flux_weights;
\ No newline at end of file
diff --git a/crates/desktop/ml_weights_registry/src/weights_registry/weight_descriptor.rs b/crates/desktop/ml_weights_registry/src/weights_registry/weight_descriptor.rs
index dca41c0099..d26d125451 100644
--- a/crates/desktop/ml_weights_registry/src/weights_registry/weight_descriptor.rs
+++ b/crates/desktop/ml_weights_registry/src/weights_registry/weight_descriptor.rs
@@ -5,6 +5,7 @@
 
 #[derive(Clone, Copy)]
 pub enum WeightFunction {
+  ImageGeneration,
   ImageSegmentation,
   TextEncoder,
   TextTokenizer,
@@ -29,6 +30,7 @@ pub struct WeightDescriptor {
 impl WeightFunction {
   pub fn as_descriptive_str(&self) -> &'static str {
     match self {
+      WeightFunction::ImageGeneration => "image generation",
       WeightFunction::ImageSegmentation => "image segmentation",
       WeightFunction::TextEncoder => "text encoder",
       WeightFunction::TextTokenizer => "text tokenizer",