Model switch but not quite working yet

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"]: