Detect and use z-image

Author: cmdr2

include/image_generator.h

@@ -101,6 +101,7 @@ class ImageGenerator {
     std::string current_clip_l_path_;
     std::string current_clip_g_path_;
     std::string current_t5xxl_path_;
+    std::string current_llm_path_;
     std::string current_taesd_path_;
     std::string current_lora_model_dir_;
     std::string current_embeddings_dir_;

include/model_detection.h

@@ -13,12 +13,12 @@ std::string detectModelFormat(const std::string& model_path);
 std::vector<std::string> extractTensorKeys(const std::string& model_path, const std::string& format);
 
 // Determine model type from tensor keys
-// Returns: "clip_l", "clip_g", "t5xxl", or "vae" (default)
+// Returns: "clip_l", "clip_g", "t5xxl", "llm", or "vae" (default)
 std::string inferModelTypeFromTensorKeys(const std::vector<std::string>& tensor_keys);
 
 // Inspect a model file and determine its type
 // Supports both GGUF and safetensors formats
-// Returns: "vae", "clip_l", "clip_g", "t5xxl", or "unknown"
+// Returns: "vae", "clip_l", "clip_g", "t5xxl", "llm", or "unknown"
 std::string inspectModelType(const std::string& model_path);
 
 #endif  // MODEL_DETECTION_H

src/image_generator.cpp

@@ -400,6 +400,7 @@ bool ImageGenerator::ensureModelLoaded(const std::string& controlnet_model) {
     std::string clip_l_path_str;
     std::string clip_g_path_str;
     std::string t5xxl_path_str;
+    std::string llm_path_str;
 
     if (options.has("forge_additional_modules")) {
         auto modules = options["forge_additional_modules"];
@@ -423,6 +424,8 @@ bool ImageGenerator::ensureModelLoaded(const std::string& controlnet_model) {
                     clip_g_path_str = module_full_path;
                 } else if (model_type == "t5xxl") {
                     t5xxl_path_str = module_full_path;
+                } else if (model_type == "llm") {
+                    llm_path_str = module_full_path;
                 } else {
                     LOG_WARNING("Unknown model type for: %s (detected as: %s)", module_full_path.c_str(),
                                 model_type.c_str());
@@ -437,7 +440,7 @@ bool ImageGenerator::ensureModelLoaded(const std::string& controlnet_model) {
     bool needs_reload = !initialized_ || !sd_ctx_ || model_path != current_model_path_ ||
                         vae_path_str != current_vae_path_ || clip_l_path_str != current_clip_l_path_ ||
                         clip_g_path_str != current_clip_g_path_ || t5xxl_path_str != current_t5xxl_path_ ||
-                        controlnet_path_str != current_controlnet_path_;
+                        llm_path_str != current_llm_path_ || controlnet_path_str != current_controlnet_path_;
 
     if (!needs_reload) {
         LOG_DEBUG("Model already loaded: %s", model_path.c_str());
@@ -471,7 +474,8 @@ bool ImageGenerator::ensureModelLoaded(const std::string& controlnet_model) {
     params.free_params_immediately = false;
 
     // Check if we have additional modules (VAE, CLIP, etc.)
-    bool has_additional_modules = !clip_l_path_str.empty() || !clip_g_path_str.empty() || !t5xxl_path_str.empty();
+    bool has_additional_modules =
+        !clip_l_path_str.empty() || !clip_g_path_str.empty() || !t5xxl_path_str.empty() || !llm_path_str.empty();
 
     if (has_additional_modules) {
         LOG_INFO("Using additional modules - loading diffusion model separately");
@@ -488,6 +492,7 @@ bool ImageGenerator::ensureModelLoaded(const std::string& controlnet_model) {
     params.clip_l_path = clip_l_path_str.empty() ? nullptr : clip_l_path_str.c_str();
     params.clip_g_path = clip_g_path_str.empty() ? nullptr : clip_g_path_str.c_str();
     params.t5xxl_path = t5xxl_path_str.empty() ? nullptr : t5xxl_path_str.c_str();
+    params.llm_path = llm_path_str.empty() ? nullptr : llm_path_str.c_str();
     params.taesd_path = nullptr;
     params.control_net_path = controlnet_path_str.empty() ? nullptr : controlnet_path_str.c_str();
     params.lora_model_dir = lora_dir_str.empty() ? nullptr : lora_dir_str.c_str();
@@ -507,6 +512,9 @@ bool ImageGenerator::ensureModelLoaded(const std::string& controlnet_model) {
     if (!t5xxl_path_str.empty()) {
         LOG_INFO("Loading T5XXL model: %s", t5xxl_path_str.c_str());
     }
+    if (!llm_path_str.empty()) {
+        LOG_INFO("Loading LLM model: %s", llm_path_str.c_str());
+    }
     if (!controlnet_path_str.empty()) {
         LOG_INFO("Loading ControlNet model: %s", controlnet_path_str.c_str());
     }
@@ -530,6 +538,7 @@ bool ImageGenerator::ensureModelLoaded(const std::string& controlnet_model) {
     current_clip_l_path_ = clip_l_path_str;
     current_clip_g_path_ = clip_g_path_str;
     current_t5xxl_path_ = t5xxl_path_str;
+    current_llm_path_ = llm_path_str;
     current_taesd_path_ = "";
     current_lora_model_dir_ = lora_dir_str;
     current_embeddings_dir_ = "";

src/model_detection.cpp

@@ -75,12 +75,21 @@ static std::vector<std::string> getGGUFTensorKeys(const std::string& model_path)
 }
 
 // Helper function to determine model type from tensor keys
-// Returns: "clip_l", "clip_g", "t5xxl", or "vae" (default)
+// Returns: "clip_l", "clip_g", "t5xxl", "llm", or "vae" (default)
 std::string inferModelTypeFromTensorKeys(const std::vector<std::string>& tensor_keys) {
     if (tensor_keys.empty()) {
         return "vae";  // Default to VAE if we can't determine
     }
 
+    // Check for LLM model indicators
+    for (const std::string& name : tensor_keys) {
+        if (name.find("blk.35.attn_k.weight") != std::string::npos ||
+            name.find("model.layers.35.post_attention_layernorm.weight") != std::string::npos) {
+            LOG_DEBUG("Detected LLM model");
+            return "llm";
+        }
+    }
+
     bool has_text_model = false;
     bool has_text_projection = false;
     bool has_position_ids = false;
@@ -209,7 +218,7 @@ std::vector<std::string> extractTensorKeys(const std::string& model_path, const
 
 // Helper function to inspect a model file and determine its type
 // Supports both GGUF and safetensors formats
-// Returns: "vae", "clip_l", "clip_g", "t5xxl", or "unknown"
+// Returns: "vae", "clip_l", "clip_g", "t5xxl", "llm", or "unknown"
 std::string inspectModelType(const std::string& model_path) {
     // First, detect the file type
     std::string format = detectModelFormat(model_path);