Model switch but not quite working yet

model_presets.toml

@@ -3,12 +3,12 @@
 # Description and Guide are mandatory. 
 # Other keys should be the Labels found in the /info endpoint.
 
-["noobaiXLNAIXL_vPred10Version"]
+["NoobAI XL"]
-description = "Anime-style model, very stylistic and not aesthetic-tuned. Can do any NSFW."
+description = "Anime-style model, very good at specific artist styles and character knowledge. Not aesthetic-tuned: needs precise, explicit prompting for best results. Can do any sort of NSFW."
 guide = """
-This model is not aesthetic tuned, it must be given explicit tags for everything. Omitted parts of the prompt will not default to something 'good'. This model is very fickle, you will almost always need to iterate on the prompt or resubmit to roll the best picture.
+This model is not aesthetic tuned, it must be given explicit tags for everything. Omitted parts of the prompt will not default to something 'good'. This model shows high variance, you can often get a very different image by just resubmitting the same prompt, so do not hesitate to try again.
 Has very excellent understanding of characters and artists down to extremely niche. Unless prompting an original character, their name is enough to decribe their appearance completely except for clothing.
-Accepts a list of comma-separated booru-style tags. Use spaces, not underscores, for tags.
+Accepts a list of comma-separated booru-style tags. Use spaces, not underscores, for tags. **Use the `search_tags` tool to verify your tags**.
 Prompts MUST follow this format: <1girl/1boy/1other/solo/couple/(can use multiple)>, <character(s)>, <series>, <artist>, <tags>
 Every prompt MUST include every one of the above sections.
 Quality tags such as "masterpiece", "best quality", "very awa", are a LAST RESORT, they override the artist tags. If absolutely required they should be prepended.
@@ -19,6 +19,7 @@ Has the SDXL problem with hands, works best if hand posture is explicitly prompt
 No default background, so prompts should include "outdoors", "indoors", or something like "patterned background" etc.
 CFG Scale from 3.0 - 5.5 but the default 4.0 is usually fine.
 """
+filename = "noobaiXLNAIXL_vPred10Version.safetensors"
 "Resolution Set" = "sdxl"
 preset = "xl"
 "CFG Scale" = 4.0
@@ -33,10 +34,14 @@ preset = "xl"
 "Schedule Type" = "Uniform"
 "Rescale CFG" = 0.3
 
-["z_image_turbo_bf16"]
+["Z-Image-Turbo"]
-description = "Fast turbo model for high-speed generation."
+description = "General image generation model. Aesthetic tuned, gets good results first try. NSFW is quite limited."
-guide = "Follow the standard prompting guide for the model architecture."
+guide = """
+This model is aesthetic tuned, regenerating with the same prompt will yield essentially the same image. Change the prompt before resubmitting.
+Understands natural language very well. Characters can be described by naming them and using this name later in the prompt. Longer, detailed prompts work better.
+"""
 preset = "zit"
+filename = "z_image_turbo_bf16.safetensors"
 modules = ["ae.safetensors", "qwen_3_4b_abliterated.safetensors"]
 "Resolution Set" = "2k"
 "CFG Scale" = 1.0

tools/generate_image.py

@@ -4,6 +4,53 @@ import config
 from typing import Any, Dict, List
 from tools.utils import ToolError, load_toml
 
+async def ensure_model_state(model_name: str, model_preset: Dict[str, Any]):
+    """
+    Checks the current server state and switches model/modules if they differ from the preset.
+    """
+    try:
+        config_resp = requests.get(f"{config.SD_URL}/config", timeout=10)
+        config_resp.raise_for_status()
+        cfg_data = config_resp.json()
+        components = cfg_data.get("components", [])
+        
+        # Extract current state from components
+        current_state = {}
+        for comp in components:
+            elem_id = comp.get("props", {}).get("elem_id")
+            if elem_id:
+                current_state[elem_id] = comp.get("props", {}).get("value")
+        
+        # 1. Check and change Checkpoint
+        active_ckpt = current_state.get("setting_sd_model_checkpoint")
+        # Use 'filename' from TOML if available, otherwise fall back to the model_name key
+        target_ckpt = model_preset.get("filename", model_name)
+        
+        if active_ckpt != target_ckpt:
+            preset = model_preset.get("preset", "xl")
+            requests.post(
+                f"{config.SD_URL}/api/predict/checkpoint_change", 
+                json={"data": [target_ckpt, preset]}, 
+                timeout=30
+            )
+        
+        # 2. Check and change VAE / Text Encoders
+        active_modules = current_state.get("setting_sd_modules", [])
+        target_modules = model_preset.get("modules", [])
+        
+        if active_modules != target_modules:
+            preset = model_preset.get("preset", "xl")
+            requests.post(
+                f"{config.SD_URL}/api/predict/modules_change", 
+                json={"data": [target_modules, preset]}, 
+                timeout=30
+            )
+            
+    except Exception as e:
+        # We log this but don't necessarily raise a ToolError unless the generation itself fails,
+        # as the server might still be able to generate if it's just a state-check failure.
+        print(f"Warning: Failed to sync model state: {e}")
+
 async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
     """
     Generates an image using the specified model and parameters.
@@ -55,6 +102,9 @@ async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
     
     model_preset = models_cfg[model_name]
 
+    # Ensure server state (Checkpoint, VAE, etc.) matches the preset before generating
+    await ensure_model_state(model_name, model_preset)
+
     # 4. MERGE PIPELINE
     for key, value in model_preset.items():
         if key in ["description", "guide", "Resolution Set"]: