resolve merge conflicts

atoma-network · Sep 2, 2024 · d496cb9 · d496cb9
2 parents cc8a37d + 78555e6
commit d496cb9
Show file tree

Hide file tree

Showing 8 changed files with 468 additions and 11 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -46,7 +46,7 @@ candle-nn = { git = "https://github.com/huggingface/candle", package = "candle-n
 candle-transformers = { git = "https://github.com/huggingface/candle", package = "candle-transformers", branch = "main" }
 clap = "4.5.4"
 config = "0.14.0"
-cudarc = { version = "0.11.6", features = [
+cudarc = { version = "0.12.0", features = [
     "std",
     "cublas",
     "cublaslt",

diff --git a/atoma-inference/src/model_thread.rs b/atoma-inference/src/model_thread.rs
@@ -16,9 +16,9 @@ use crate::models::candle::mixtral_nccl::MixtralNcclModel;
 
 use crate::models::{
     candle::{
-        falcon::FalconModel, llama::LlamaModel, mamba::MambaModel, mistral::MistralModel,
-        mixtral::MixtralModel, phi3::Phi3Model, quantized::QuantizedModel, qwen::QwenModel,
-        stable_diffusion::StableDiffusion,
+        falcon::FalconModel, flux::Flux, llama::LlamaModel, mamba::MambaModel,
+        mistral::MistralModel, mixtral::MixtralModel, phi3::Phi3Model, quantized::QuantizedModel,
+        qwen::QwenModel, stable_diffusion::StableDiffusion,
     },
     config::{ModelConfig, ModelsConfig},
     types::{LlmOutput, ModelType},
@@ -297,6 +297,14 @@ pub(crate) fn dispatch_model_thread(
                     stream_tx,
                 )
             }
+            ModelType::FluxSchnell | ModelType::FluxDev => spawn_model_thread::<Flux>(
+                model_name,
+                api_key.clone(),
+                cache_dir.clone(),
+                model_config,
+                model_receiver,
+                stream_tx,
+            ),
             ModelType::LlamaV1
             | ModelType::LlamaV2
             | ModelType::LlamaTinyLlama1_1BChat