246 lines
9.4 KiB
Python
246 lines
9.4 KiB
Python
import requests
|
|
import base64
|
|
import asyncio
|
|
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 = await asyncio.to_thread(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")
|
|
await asyncio.to_thread(
|
|
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")
|
|
await asyncio.to_thread(
|
|
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.
|
|
"""
|
|
# 1. REQUIRED ARGUMENTS
|
|
prompt = args.get("prompt")
|
|
if not prompt:
|
|
raise ToolError("The 'prompt' argument is required.")
|
|
|
|
model_name = args.get("model_name")
|
|
if not model_name:
|
|
raise ToolError("The 'model_name' argument is required. Use get_model_info to see available models.")
|
|
|
|
res_preset_name = args.get("resolution_preset")
|
|
if not res_preset_name:
|
|
raise ToolError("The 'resolution_preset' argument is required. Use get_model_info to see available resolutions for the chosen model.")
|
|
|
|
# 2. FETCH CURRENT SERVER DEFAULTS & MAP LABELS
|
|
try:
|
|
info_resp = await asyncio.to_thread(requests.get, f"{config.SD_URL}/info", timeout=10)
|
|
info_resp.raise_for_status()
|
|
info_data = info_resp.json()
|
|
params_info = info_data["named_endpoints"]["/txt2img"]["parameters"]
|
|
except Exception as e:
|
|
raise ToolError(f"Failed to connect to Stable Diffusion server: {str(e)}")
|
|
|
|
# Build the label-to-index map and a sparse payload
|
|
label_map = {}
|
|
label_counts = {}
|
|
|
|
# Find the maximum param index to determine payload size
|
|
max_param_idx = 0
|
|
for p in params_info:
|
|
name = p.get("parameter_name", "")
|
|
if name.startswith("param_"):
|
|
try:
|
|
idx = int(name.replace("param_", ""))
|
|
max_param_idx = max(max_param_idx, idx)
|
|
except ValueError:
|
|
pass
|
|
|
|
# Initialize payload with Nones (size is max_idx + 1)
|
|
payload = [None] * (max_param_idx + 1)
|
|
|
|
for idx_in_list, p in enumerate(params_info):
|
|
name = p.get("parameter_name", "")
|
|
label = p.get("label")
|
|
default = p.get("parameter_default")
|
|
|
|
# Determine the absolute index in the payload
|
|
if name == "id_task":
|
|
abs_idx = 0
|
|
elif name.startswith("param_"):
|
|
try:
|
|
abs_idx = int(name.replace("param_", ""))
|
|
except ValueError:
|
|
continue
|
|
else:
|
|
# Fallback for unexpected names, though unlikely
|
|
continue
|
|
|
|
# Set the default value at the absolute index
|
|
payload[abs_idx] = default
|
|
|
|
# Handle label mapping for AI overrides
|
|
if not label or label.startswith("parameter_"):
|
|
label = name
|
|
|
|
if label in label_counts:
|
|
label_counts[label] += 1
|
|
mapped_label = f"{label} [{label_counts[label]}]"
|
|
else:
|
|
label_counts[label] = 0
|
|
mapped_label = label
|
|
|
|
label_map[mapped_label] = abs_idx
|
|
|
|
# 3. LOAD CONFIGS
|
|
models_cfg = load_toml(config.MODEL_PRESETS_PATH)
|
|
res_cfg = load_toml(config.RES_PRESETS_PATH)
|
|
|
|
if model_name not in models_cfg:
|
|
raise ToolError(f"Model '{model_name}' not found in presets. Please call get_model_info to see available models and their correct names.")
|
|
|
|
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"]:
|
|
continue
|
|
if key in label_map:
|
|
payload[label_map[key]] = value
|
|
elif key.startswith("param_"):
|
|
try:
|
|
idx = int(key.replace("param_", ""))
|
|
if 0 <= idx < len(payload):
|
|
payload[idx] = value
|
|
except ValueError:
|
|
pass
|
|
|
|
res_preset_name = args.get("resolution_preset")
|
|
if res_preset_name:
|
|
res_set_name = model_preset.get("Resolution Set")
|
|
if res_set_name and res_set_name in res_cfg:
|
|
res_set = res_cfg[res_set_name]
|
|
if res_preset_name in res_set:
|
|
res_val_str = res_set[res_preset_name]
|
|
try:
|
|
w, h = map(int, res_val_str.split('x'))
|
|
if "Width" in label_map:
|
|
payload[label_map["Width"]] = w
|
|
if "Height" in label_map:
|
|
payload[label_map["Height"]] = h
|
|
except (ValueError, AttributeError):
|
|
raise ToolError(f"Invalid resolution format for preset '{res_preset_name}': {res_val_str}. Expected 'WidthxHeight'.")
|
|
else:
|
|
raise ToolError(f"Resolution preset '{res_preset_name}' not found for this model. Available: {', '.join(res_set.keys())}")
|
|
else:
|
|
raise ToolError(f"No resolution set configured for model '{model_name}'.")
|
|
|
|
overrides = {
|
|
"prompt": "Prompt",
|
|
"negative_prompt": "Negative Prompt",
|
|
"cfg_scale": "CFG Scale",
|
|
}
|
|
|
|
for arg_key, label in overrides.items():
|
|
if arg_key in args and label in label_map:
|
|
payload[label_map[label]] = args[arg_key]
|
|
|
|
if "Prompt" in label_map:
|
|
payload[label_map["Prompt"]] = prompt
|
|
|
|
if "Seed" in label_map:
|
|
payload[label_map["Seed"]] = -1
|
|
|
|
# The "Magic Number" splice is no longer needed as gaps are
|
|
# automatically filled by the sparse-to-dense mapping.
|
|
|
|
# 6. EXECUTE GENERATION
|
|
try:
|
|
gen_payload = {"data": payload}
|
|
gen_resp = await asyncio.to_thread(requests.post, f"{config.SD_URL}/api/txt2img", json=gen_payload, timeout=300)
|
|
gen_resp.raise_for_status()
|
|
res_data = gen_resp.json()
|
|
|
|
gallery_data = res_data["data"][0]["value"]
|
|
if not gallery_data:
|
|
raise ToolError("Server returned successfully but no image was generated.")
|
|
|
|
sd_img_url = gallery_data[0]["image"]["url"]
|
|
|
|
# Extract raw path from SD URL (e.g. http://.../file=/tmp/gradio/abc.png -> /tmp/gradio/abc.png)
|
|
if "/file=" in sd_img_url:
|
|
raw_path = sd_img_url.split("/file=")[1]
|
|
else:
|
|
# Fallback if the URL format changes
|
|
raise ToolError(f"Could not extract file path from SD URL: {sd_img_url}")
|
|
|
|
# Construct proxy URL through our MCP server
|
|
# Use the captured request host to ensure links work across the network
|
|
host = config.request_host.get()
|
|
proxy_url = f"http://{host}/file{raw_path}"
|
|
|
|
# Download the image and convert to base64 for the AI's vision
|
|
img_resp = await asyncio.to_thread(requests.get, sd_img_url, timeout=30)
|
|
img_resp.raise_for_status()
|
|
b64_data = base64.b64encode(img_resp.content).decode('utf-8')
|
|
|
|
return [
|
|
{
|
|
"type": "text",
|
|
"text": f"Image successfully generated using model '{model_name}'.\n\nLocal path: {raw_path}\nTo show the image to the user, you can include this markdown link in your response: "
|
|
},
|
|
{
|
|
"type": "image",
|
|
"data": b64_data,
|
|
"mimeType": "image/png"
|
|
}
|
|
]
|
|
|
|
except requests.exceptions.HTTPError as e:
|
|
raise ToolError(f"Server error during generation: {str(e)}")
|
|
except Exception as e:
|
|
raise ToolError(f"Unexpected error during generation: {str(e)}")
|