From f676a07a71e3ed4058a9cd7174e3bcbb7d6f6aae Mon Sep 17 00:00:00 2001 From: moosecrap Date: Sun, 26 Jul 2026 20:06:22 -0700 Subject: [PATCH] Model switch but not quite working yet --- model_presets.toml | 19 ++++++++++------ tools/generate_image.py | 50 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 62 insertions(+), 7 deletions(-) diff --git a/model_presets.toml b/model_presets.toml index eaf87e8..095a826 100644 --- a/model_presets.toml +++ b/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"] -description = "Anime-style model, very stylistic and not aesthetic-tuned. Can do any NSFW." +["NoobAI XL"] +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)>, , , , 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"] -description = "Fast turbo model for high-speed generation." -guide = "Follow the standard prompting guide for the model architecture." +["Z-Image-Turbo"] +description = "General image generation model. Aesthetic tuned, gets good results first try. NSFW is quite limited." +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 diff --git a/tools/generate_image.py b/tools/generate_image.py index 21d3e28..8fd56e6 100644 --- a/tools/generate_image.py +++ b/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"]: