153 lines
5.6 KiB
Python
153 lines
5.6 KiB
Python
import requests
|
|
import base64
|
|
import config
|
|
from typing import Any, Dict, List
|
|
from tools.utils import ToolError, load_toml
|
|
|
|
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.")
|
|
|
|
# 2. FETCH CURRENT SERVER DEFAULTS & MAP LABELS
|
|
try:
|
|
info_resp = 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
|
|
label_map = {}
|
|
label_counts = {}
|
|
|
|
payload = [p["parameter_default"] for p in params_info]
|
|
|
|
for idx, p in enumerate(params_info):
|
|
label = p["label"]
|
|
if not label or label.startswith("parameter_"):
|
|
label = p["parameter_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] = 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. Available: {', '.join(models_cfg.keys())}")
|
|
|
|
model_preset = models_cfg[model_name]
|
|
|
|
# 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_vals = res_set[res_preset_name]
|
|
if "Width" in label_map:
|
|
payload[label_map["Width"]] = res_vals["width"]
|
|
if "Height" in label_map:
|
|
payload[label_map["Height"]] = res_vals["height"]
|
|
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",
|
|
"steps": "Sampling Steps",
|
|
"cfg_scale": "CFG Scale",
|
|
"seed": "Seed",
|
|
"sampler": "Sampling Method"
|
|
}
|
|
|
|
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
|
|
|
|
for _ in range(7):
|
|
payload.insert(39, None)
|
|
|
|
# 6. EXECUTE GENERATION
|
|
try:
|
|
gen_payload = {"data": payload}
|
|
gen_resp = requests.post(f"{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
|
|
# We strip leading slash from raw_path to avoid double slashes in the proxy URL if we want,
|
|
# but the current endpoint handles it.
|
|
proxy_url = f"http://localhost:{config.PORT}/file{raw_path}"
|
|
|
|
# Download the image and convert to base64 for the AI's vision
|
|
img_resp = 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\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)}")
|